{-# LANGUAGE CPP #-}
{-# LANGUAGE StrictData #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE DeriveTraversable #-}
{-# LANGUAGE DeriveLift #-}
module Text.Collate.Trie
  ( Trie
  , empty
  , insert
  , alter
  , unfoldTrie
  , matchLongestPrefix
  , lookupNonEmptyChild
  )
  where

import Control.Monad (foldM)
import qualified Data.IntMap as M
import Data.Bifunctor (first)
import Data.Binary (Binary(..))
import Language.Haskell.TH.Syntax (Lift(..))
import Instances.TH.Lift ()
import Data.Maybe (fromMaybe)
#if MIN_VERSION_base(4,11,0)
#else
import Data.Semigroup (Semigroup(..))
#endif

data Trie a = Trie (Maybe a) (Maybe (M.IntMap (Trie a)))
  deriving (Int -> Trie a -> ShowS
forall a. Show a => Int -> Trie a -> ShowS
forall a. Show a => [Trie a] -> ShowS
forall a. Show a => Trie a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Trie a] -> ShowS
$cshowList :: forall a. Show a => [Trie a] -> ShowS
show :: Trie a -> String
$cshow :: forall a. Show a => Trie a -> String
showsPrec :: Int -> Trie a -> ShowS
$cshowsPrec :: forall a. Show a => Int -> Trie a -> ShowS
Show, Trie a -> Trie a -> Bool
forall a. Eq a => Trie a -> Trie a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Trie a -> Trie a -> Bool
$c/= :: forall a. Eq a => Trie a -> Trie a -> Bool
== :: Trie a -> Trie a -> Bool
$c== :: forall a. Eq a => Trie a -> Trie a -> Bool
Eq, Trie a -> Trie a -> Bool
Trie a -> Trie a -> Ordering
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
forall {a}. Ord a => Eq (Trie a)
forall a. Ord a => Trie a -> Trie a -> Bool
forall a. Ord a => Trie a -> Trie a -> Ordering
forall a. Ord a => Trie a -> Trie a -> Trie a
min :: Trie a -> Trie a -> Trie a
$cmin :: forall a. Ord a => Trie a -> Trie a -> Trie a
max :: Trie a -> Trie a -> Trie a
$cmax :: forall a. Ord a => Trie a -> Trie a -> Trie a
>= :: Trie a -> Trie a -> Bool
$c>= :: forall a. Ord a => Trie a -> Trie a -> Bool
> :: Trie a -> Trie a -> Bool
$c> :: forall a. Ord a => Trie a -> Trie a -> Bool
<= :: Trie a -> Trie a -> Bool
$c<= :: forall a. Ord a => Trie a -> Trie a -> Bool
< :: Trie a -> Trie a -> Bool
$c< :: forall a. Ord a => Trie a -> Trie a -> Bool
compare :: Trie a -> Trie a -> Ordering
$ccompare :: forall a. Ord a => Trie a -> Trie a -> Ordering
Ord, forall a (m :: * -> *). (Lift a, Quote m) => Trie a -> m Exp
forall a (m :: * -> *).
(Lift a, Quote m) =>
Trie a -> Code m (Trie a)
forall t.
(forall (m :: * -> *). Quote m => t -> m Exp)
-> (forall (m :: * -> *). Quote m => t -> Code m t) -> Lift t
forall (m :: * -> *). Quote m => Trie a -> m Exp
forall (m :: * -> *). Quote m => Trie a -> Code m (Trie a)
liftTyped :: forall (m :: * -> *). Quote m => Trie a -> Code m (Trie a)
$cliftTyped :: forall a (m :: * -> *).
(Lift a, Quote m) =>
Trie a -> Code m (Trie a)
lift :: forall (m :: * -> *). Quote m => Trie a -> m Exp
$clift :: forall a (m :: * -> *). (Lift a, Quote m) => Trie a -> m Exp
Lift, forall a b. a -> Trie b -> Trie a
forall a b. (a -> b) -> Trie a -> Trie b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: forall a b. a -> Trie b -> Trie a
$c<$ :: forall a b. a -> Trie b -> Trie a
fmap :: forall a b. (a -> b) -> Trie a -> Trie b
$cfmap :: forall a b. (a -> b) -> Trie a -> Trie b
Functor, forall a. Eq a => a -> Trie a -> Bool
forall a. Num a => Trie a -> a
forall a. Ord a => Trie a -> a
forall m. Monoid m => Trie m -> m
forall a. Trie a -> Bool
forall a. Trie a -> Int
forall a. Trie a -> [a]
forall a. (a -> a -> a) -> Trie a -> a
forall m a. Monoid m => (a -> m) -> Trie a -> m
forall b a. (b -> a -> b) -> b -> Trie a -> b
forall a b. (a -> b -> b) -> b -> Trie a -> b
forall (t :: * -> *).
(forall m. Monoid m => t m -> m)
-> (forall m a. Monoid m => (a -> m) -> t a -> m)
-> (forall m a. Monoid m => (a -> m) -> t a -> m)
-> (forall a b. (a -> b -> b) -> b -> t a -> b)
-> (forall a b. (a -> b -> b) -> b -> t a -> b)
-> (forall b a. (b -> a -> b) -> b -> t a -> b)
-> (forall b a. (b -> a -> b) -> b -> t a -> b)
-> (forall a. (a -> a -> a) -> t a -> a)
-> (forall a. (a -> a -> a) -> t a -> a)
-> (forall a. t a -> [a])
-> (forall a. t a -> Bool)
-> (forall a. t a -> Int)
-> (forall a. Eq a => a -> t a -> Bool)
-> (forall a. Ord a => t a -> a)
-> (forall a. Ord a => t a -> a)
-> (forall a. Num a => t a -> a)
-> (forall a. Num a => t a -> a)
-> Foldable t
product :: forall a. Num a => Trie a -> a
$cproduct :: forall a. Num a => Trie a -> a
sum :: forall a. Num a => Trie a -> a
$csum :: forall a. Num a => Trie a -> a
minimum :: forall a. Ord a => Trie a -> a
$cminimum :: forall a. Ord a => Trie a -> a
maximum :: forall a. Ord a => Trie a -> a
$cmaximum :: forall a. Ord a => Trie a -> a
elem :: forall a. Eq a => a -> Trie a -> Bool
$celem :: forall a. Eq a => a -> Trie a -> Bool
length :: forall a. Trie a -> Int
$clength :: forall a. Trie a -> Int
null :: forall a. Trie a -> Bool
$cnull :: forall a. Trie a -> Bool
toList :: forall a. Trie a -> [a]
$ctoList :: forall a. Trie a -> [a]
foldl1 :: forall a. (a -> a -> a) -> Trie a -> a
$cfoldl1 :: forall a. (a -> a -> a) -> Trie a -> a
foldr1 :: forall a. (a -> a -> a) -> Trie a -> a
$cfoldr1 :: forall a. (a -> a -> a) -> Trie a -> a
foldl' :: forall b a. (b -> a -> b) -> b -> Trie a -> b
$cfoldl' :: forall b a. (b -> a -> b) -> b -> Trie a -> b
foldl :: forall b a. (b -> a -> b) -> b -> Trie a -> b
$cfoldl :: forall b a. (b -> a -> b) -> b -> Trie a -> b
foldr' :: forall a b. (a -> b -> b) -> b -> Trie a -> b
$cfoldr' :: forall a b. (a -> b -> b) -> b -> Trie a -> b
foldr :: forall a b. (a -> b -> b) -> b -> Trie a -> b
$cfoldr :: forall a b. (a -> b -> b) -> b -> Trie a -> b
foldMap' :: forall m a. Monoid m => (a -> m) -> Trie a -> m
$cfoldMap' :: forall m a. Monoid m => (a -> m) -> Trie a -> m
foldMap :: forall m a. Monoid m => (a -> m) -> Trie a -> m
$cfoldMap :: forall m a. Monoid m => (a -> m) -> Trie a -> m
fold :: forall m. Monoid m => Trie m -> m
$cfold :: forall m. Monoid m => Trie m -> m
Foldable, Functor Trie
Foldable Trie
forall (t :: * -> *).
Functor t
-> Foldable t
-> (forall (f :: * -> *) a b.
    Applicative f =>
    (a -> f b) -> t a -> f (t b))
-> (forall (f :: * -> *) a. Applicative f => t (f a) -> f (t a))
-> (forall (m :: * -> *) a b.
    Monad m =>
    (a -> m b) -> t a -> m (t b))
-> (forall (m :: * -> *) a. Monad m => t (m a) -> m (t a))
-> Traversable t
forall (m :: * -> *) a. Monad m => Trie (m a) -> m (Trie a)
forall (f :: * -> *) a. Applicative f => Trie (f a) -> f (Trie a)
forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Trie a -> m (Trie b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Trie a -> f (Trie b)
sequence :: forall (m :: * -> *) a. Monad m => Trie (m a) -> m (Trie a)
$csequence :: forall (m :: * -> *) a. Monad m => Trie (m a) -> m (Trie a)
mapM :: forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Trie a -> m (Trie b)
$cmapM :: forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Trie a -> m (Trie b)
sequenceA :: forall (f :: * -> *) a. Applicative f => Trie (f a) -> f (Trie a)
$csequenceA :: forall (f :: * -> *) a. Applicative f => Trie (f a) -> f (Trie a)
traverse :: forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Trie a -> f (Trie b)
$ctraverse :: forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Trie a -> f (Trie b)
Traversable)

instance Semigroup (Trie a) where
   Trie a
trie1 <> :: Trie a -> Trie a -> Trie a
<> Trie a
trie2 = forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry forall a. [Int] -> a -> Trie a -> Trie a
insert) Trie a
trie1 (forall a. Trie a -> [([Int], a)]
unfoldTrie Trie a
trie2)

instance Monoid (Trie a) where
   mempty :: Trie a
mempty = forall a. Maybe a -> Maybe (IntMap (Trie a)) -> Trie a
Trie forall a. Maybe a
Nothing forall a. Maybe a
Nothing
   mappend :: Trie a -> Trie a -> Trie a
mappend = forall a. Semigroup a => a -> a -> a
(<>)

instance Binary a => Binary (Trie a) where
   put :: Trie a -> Put
put (Trie Maybe a
mbv Maybe (IntMap (Trie a))
mbm) = forall t. Binary t => t -> Put
put (Maybe a
mbv, Maybe (IntMap (Trie a))
mbm)
   get :: Get (Trie a)
get = do
     (Maybe a
mbv,Maybe (IntMap (Trie a))
mbm) <- forall t. Binary t => Get t
get
     forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. Maybe a -> Maybe (IntMap (Trie a)) -> Trie a
Trie Maybe a
mbv Maybe (IntMap (Trie a))
mbm

empty :: Trie a
empty :: forall a. Trie a
empty = forall a. Maybe a -> Maybe (IntMap (Trie a)) -> Trie a
Trie forall a. Maybe a
Nothing forall a. Maybe a
Nothing

unfoldTrie :: Trie a -> [([Int], a)]
unfoldTrie :: forall a. Trie a -> [([Int], a)]
unfoldTrie  = forall a b. (a -> b) -> [a] -> [b]
map (forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first forall a. [a] -> [a]
reverse) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {b}. [Int] -> Trie b -> [([Int], b)]
go []
 where
  go :: [Int] -> Trie b -> [([Int], b)]
go [Int]
xs (Trie (Just b
v) (Just IntMap (Trie b)
m)) =
    ([Int]
xs, b
v) forall a. a -> [a] -> [a]
: forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap ([Int] -> (Int, Trie b) -> [([Int], b)]
gopair [Int]
xs) (forall a. IntMap a -> [(Int, a)]
M.toList IntMap (Trie b)
m)
  go [Int]
xs (Trie (Just b
v) Maybe (IntMap (Trie b))
Nothing) = [([Int]
xs, b
v)]
  go [Int]
xs (Trie Maybe b
Nothing (Just IntMap (Trie b)
m)) =
    forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap ([Int] -> (Int, Trie b) -> [([Int], b)]
gopair [Int]
xs) (forall a. IntMap a -> [(Int, a)]
M.toList IntMap (Trie b)
m)
  go [Int]
_ (Trie Maybe b
Nothing Maybe (IntMap (Trie b))
Nothing) = []
  gopair :: [Int] -> (Int, Trie b) -> [([Int], b)]
gopair [Int]
xs (Int
i, Trie b
trie) = [Int] -> Trie b -> [([Int], b)]
go (Int
iforall a. a -> [a] -> [a]
:[Int]
xs) Trie b
trie

insert :: [Int] -> a -> Trie a -> Trie a
insert :: forall a. [Int] -> a -> Trie a -> Trie a
insert [] a
x (Trie Maybe a
_ Maybe (IntMap (Trie a))
mbm) = forall a. Maybe a -> Maybe (IntMap (Trie a)) -> Trie a
Trie (forall a. a -> Maybe a
Just a
x) Maybe (IntMap (Trie a))
mbm
insert (Int
c:[Int]
cs) a
x (Trie Maybe a
mbv (Just IntMap (Trie a)
m)) =
  case forall a. Int -> IntMap a -> Maybe a
M.lookup Int
c IntMap (Trie a)
m of
    Maybe (Trie a)
Nothing   -> forall a. Maybe a -> Maybe (IntMap (Trie a)) -> Trie a
Trie Maybe a
mbv (forall a. a -> Maybe a
Just (forall a. Int -> a -> IntMap a -> IntMap a
M.insert Int
c (forall a. [Int] -> a -> Trie a -> Trie a
insert [Int]
cs a
x forall a. Trie a
empty) IntMap (Trie a)
m))
    Just Trie a
trie -> forall a. Maybe a -> Maybe (IntMap (Trie a)) -> Trie a
Trie Maybe a
mbv (forall a. a -> Maybe a
Just (forall a. Int -> a -> IntMap a -> IntMap a
M.insert Int
c (forall a. [Int] -> a -> Trie a -> Trie a
insert [Int]
cs a
x Trie a
trie) IntMap (Trie a)
m))
insert (Int
c:[Int]
cs) a
x (Trie Maybe a
mbv Maybe (IntMap (Trie a))
Nothing) =
  forall a. Maybe a -> Maybe (IntMap (Trie a)) -> Trie a
Trie Maybe a
mbv (forall a. a -> Maybe a
Just (forall a. Int -> a -> IntMap a -> IntMap a
M.insert Int
c (forall a. [Int] -> a -> Trie a -> Trie a
insert [Int]
cs a
x forall a. Trie a
empty) forall a. Monoid a => a
mempty))

alter :: (Maybe a -> Maybe a) -> [Int] -> Trie a -> Trie a
alter :: forall a. (Maybe a -> Maybe a) -> [Int] -> Trie a -> Trie a
alter Maybe a -> Maybe a
f [] (Trie Maybe a
mbv Maybe (IntMap (Trie a))
mbm) = forall a. Maybe a -> Maybe (IntMap (Trie a)) -> Trie a
Trie (Maybe a -> Maybe a
f Maybe a
mbv) Maybe (IntMap (Trie a))
mbm
alter Maybe a -> Maybe a
f (Int
c:[Int]
cs) (Trie Maybe a
mbv (Just IntMap (Trie a)
m)) =
  forall a. Maybe a -> Maybe (IntMap (Trie a)) -> Trie a
Trie Maybe a
mbv (forall a. a -> Maybe a
Just (forall a. Int -> a -> IntMap a -> IntMap a
M.insert Int
c (forall a. (Maybe a -> Maybe a) -> [Int] -> Trie a -> Trie a
alter Maybe a -> Maybe a
f [Int]
cs forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a -> a
fromMaybe forall a. Trie a
empty forall a b. (a -> b) -> a -> b
$ forall a. Int -> IntMap a -> Maybe a
M.lookup Int
c IntMap (Trie a)
m) IntMap (Trie a)
m))
alter Maybe a -> Maybe a
f (Int
c:[Int]
cs) (Trie Maybe a
mbv Maybe (IntMap (Trie a))
Nothing) =
  forall a. Maybe a -> Maybe (IntMap (Trie a)) -> Trie a
Trie Maybe a
mbv (forall a. a -> Maybe a
Just (forall a. Int -> a -> IntMap a -> IntMap a
M.insert Int
c (forall a. (Maybe a -> Maybe a) -> [Int] -> Trie a -> Trie a
alter Maybe a -> Maybe a
f [Int]
cs forall a. Trie a
empty) forall a. Monoid a => a
mempty))

type MatchState a = (Maybe (a, Int, Trie a), Int, Trie a)
  -- best match so far, number of code points consumed, current subtrie

{-# SPECIALIZE matchLongestPrefix :: Trie a -> [Int] -> Maybe (a, Int, Trie a) #-}
-- returns Nothing for no match, or:
-- Just (value, number of code points consumed, subtrie)
matchLongestPrefix :: Foldable t => Trie a -> t Int -> Maybe (a, Int, Trie a)
matchLongestPrefix :: forall (t :: * -> *) a.
Foldable t =>
Trie a -> t Int -> Maybe (a, Int, Trie a)
matchLongestPrefix Trie a
trie = forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either forall a. a -> a
id forall {a} {b} {c}. (a, b, c) -> a
getBest forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM forall a.
MatchState a
-> Int -> Either (Maybe (a, Int, Trie a)) (MatchState a)
go (forall a. Maybe a
Nothing, Int
0, Trie a
trie)
 where
   getBest :: (a, b, c) -> a
getBest (a
x,b
_,c
_) = a
x
   -- Left means we've failed, Right means we're still pursuing a match
   go :: MatchState a -> Int -> Either (Maybe (a, Int, Trie a)) (MatchState a)
   go :: forall a.
MatchState a
-> Int -> Either (Maybe (a, Int, Trie a)) (MatchState a)
go (Maybe (a, Int, Trie a)
best, Int
consumed, Trie Maybe a
_ Maybe (IntMap (Trie a))
mbm) Int
c =
     case Maybe (IntMap (Trie a))
mbm forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= forall a. Int -> IntMap a -> Maybe a
M.lookup Int
c of
       -- char not matched: stop processing, return best so far:
       Maybe (Trie a)
Nothing -> forall a b. a -> Either a b
Left Maybe (a, Int, Trie a)
best
       -- char matched, with value: replace best, keep going:
       Just subtrie :: Trie a
subtrie@(Trie (Just a
x) Maybe (IntMap (Trie a))
_)
               -> forall a b. b -> Either a b
Right (forall a. a -> Maybe a
Just (a
x, Int
consumed forall a. Num a => a -> a -> a
+ Int
1, Trie a
subtrie), Int
consumed forall a. Num a => a -> a -> a
+ Int
1, Trie a
subtrie)
       -- char matched, but not value: keep best, keep going:
       Just subtrie :: Trie a
subtrie@(Trie Maybe a
Nothing Maybe (IntMap (Trie a))
_)
               -> forall a b. b -> Either a b
Right (Maybe (a, Int, Trie a)
best, Int
consumed forall a. Num a => a -> a -> a
+ Int
1, Trie a
subtrie)

-- | Return the sub-trie at the given branch if it exists and has a
-- non-empty node
lookupNonEmptyChild :: Trie a -> Int -> Maybe (a, Trie a)
lookupNonEmptyChild :: forall a. Trie a -> Int -> Maybe (a, Trie a)
lookupNonEmptyChild (Trie Maybe a
_ Maybe (IntMap (Trie a))
Nothing) Int
_ = forall a. Maybe a
Nothing
lookupNonEmptyChild (Trie Maybe a
_ (Just IntMap (Trie a)
m)) Int
idx = do
  Trie Maybe a
mnode Maybe (IntMap (Trie a))
m' <- forall a. Int -> IntMap a -> Maybe a
M.lookup Int
idx IntMap (Trie a)
m
  a
node <- Maybe a
mnode
  forall (m :: * -> *) a. Monad m => a -> m a
return (a
node, forall a. Maybe a -> Maybe (IntMap (Trie a)) -> Trie a
Trie forall a. Maybe a
Nothing Maybe (IntMap (Trie a))
m')