{-# 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
[Trie a] -> ShowS
Trie a -> String
(Int -> Trie a -> ShowS)
-> (Trie a -> String) -> ([Trie a] -> ShowS) -> Show (Trie a)
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
(Trie a -> Trie a -> Bool)
-> (Trie a -> Trie a -> Bool) -> Eq (Trie a)
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, Eq (Trie a)
Eq (Trie a)
-> (Trie a -> Trie a -> Ordering)
-> (Trie a -> Trie a -> Bool)
-> (Trie a -> Trie a -> Bool)
-> (Trie a -> Trie a -> Bool)
-> (Trie a -> Trie a -> Bool)
-> (Trie a -> Trie a -> Trie a)
-> (Trie a -> Trie a -> Trie a)
-> Ord (Trie a)
Trie a -> Trie a -> Bool
Trie a -> Trie a -> Ordering
Trie a -> Trie a -> Trie a
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
$cp1Ord :: forall a. Ord a => Eq (Trie a)
Ord, Trie a -> Q Exp
Trie a -> Q (TExp (Trie a))
(Trie a -> Q Exp) -> (Trie a -> Q (TExp (Trie a))) -> Lift (Trie a)
forall a. Lift a => Trie a -> Q Exp
forall a. Lift a => Trie a -> Q (TExp (Trie a))
forall t. (t -> Q Exp) -> (t -> Q (TExp t)) -> Lift t
liftTyped :: Trie a -> Q (TExp (Trie a))
$cliftTyped :: forall a. Lift a => Trie a -> Q (TExp (Trie a))
lift :: Trie a -> Q Exp
$clift :: forall a. Lift a => Trie a -> Q Exp
Lift, a -> Trie b -> Trie a
(a -> b) -> Trie a -> Trie b
(forall a b. (a -> b) -> Trie a -> Trie b)
-> (forall a b. a -> Trie b -> Trie a) -> Functor Trie
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
<$ :: a -> Trie b -> Trie a
$c<$ :: forall a b. a -> Trie b -> Trie a
fmap :: (a -> b) -> Trie a -> Trie b
$cfmap :: forall a b. (a -> b) -> Trie a -> Trie b
Functor, Trie a -> Bool
(a -> m) -> Trie a -> m
(a -> b -> b) -> b -> Trie a -> b
(forall m. Monoid m => Trie m -> m)
-> (forall m a. Monoid m => (a -> m) -> Trie a -> m)
-> (forall m a. Monoid m => (a -> m) -> Trie a -> m)
-> (forall a b. (a -> b -> b) -> b -> Trie a -> b)
-> (forall a b. (a -> b -> b) -> b -> Trie a -> b)
-> (forall b a. (b -> a -> b) -> b -> Trie a -> b)
-> (forall b a. (b -> a -> b) -> b -> Trie a -> b)
-> (forall a. (a -> a -> a) -> Trie a -> a)
-> (forall a. (a -> a -> a) -> Trie a -> a)
-> (forall a. Trie a -> [a])
-> (forall a. Trie a -> Bool)
-> (forall a. Trie a -> Int)
-> (forall a. Eq a => a -> Trie a -> Bool)
-> (forall a. Ord a => Trie a -> a)
-> (forall a. Ord a => Trie a -> a)
-> (forall a. Num a => Trie a -> a)
-> (forall a. Num a => Trie a -> a)
-> Foldable Trie
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 :: Trie a -> a
$cproduct :: forall a. Num a => Trie a -> a
sum :: Trie a -> a
$csum :: forall a. Num a => Trie a -> a
minimum :: Trie a -> a
$cminimum :: forall a. Ord a => Trie a -> a
maximum :: Trie a -> a
$cmaximum :: forall a. Ord a => Trie a -> a
elem :: a -> Trie a -> Bool
$celem :: forall a. Eq a => a -> Trie a -> Bool
length :: Trie a -> Int
$clength :: forall a. Trie a -> Int
null :: Trie a -> Bool
$cnull :: forall a. Trie a -> Bool
toList :: Trie a -> [a]
$ctoList :: forall a. Trie a -> [a]
foldl1 :: (a -> a -> a) -> Trie a -> a
$cfoldl1 :: forall a. (a -> a -> a) -> Trie a -> a
foldr1 :: (a -> a -> a) -> Trie a -> a
$cfoldr1 :: forall a. (a -> a -> a) -> Trie a -> a
foldl' :: (b -> a -> b) -> b -> Trie a -> b
$cfoldl' :: forall b a. (b -> a -> b) -> b -> Trie a -> b
foldl :: (b -> a -> b) -> b -> Trie a -> b
$cfoldl :: forall b a. (b -> a -> b) -> b -> Trie a -> b
foldr' :: (a -> b -> b) -> b -> Trie a -> b
$cfoldr' :: forall a b. (a -> b -> b) -> b -> Trie a -> b
foldr :: (a -> b -> b) -> b -> Trie a -> b
$cfoldr :: forall a b. (a -> b -> b) -> b -> Trie a -> b
foldMap' :: (a -> m) -> Trie a -> m
$cfoldMap' :: forall m a. Monoid m => (a -> m) -> Trie a -> m
foldMap :: (a -> m) -> Trie a -> m
$cfoldMap :: forall m a. Monoid m => (a -> m) -> Trie a -> m
fold :: Trie m -> m
$cfold :: forall m. Monoid m => Trie m -> m
Foldable, Functor Trie
Foldable Trie
Functor Trie
-> Foldable Trie
-> (forall (f :: * -> *) a b.
    Applicative f =>
    (a -> f b) -> Trie a -> f (Trie b))
-> (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 (m :: * -> *) a. Monad m => Trie (m a) -> m (Trie a))
-> Traversable Trie
(a -> f b) -> Trie a -> f (Trie b)
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 :: Trie (m a) -> m (Trie a)
$csequence :: forall (m :: * -> *) a. Monad m => Trie (m a) -> m (Trie a)
mapM :: (a -> m b) -> Trie a -> m (Trie b)
$cmapM :: forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Trie a -> m (Trie b)
sequenceA :: Trie (f a) -> f (Trie a)
$csequenceA :: forall (f :: * -> *) a. Applicative f => Trie (f a) -> f (Trie a)
traverse :: (a -> f b) -> Trie a -> f (Trie b)
$ctraverse :: forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Trie a -> f (Trie b)
$cp2Traversable :: Foldable Trie
$cp1Traversable :: Functor Trie
Traversable)

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

instance Monoid (Trie a) where
   mempty :: Trie a
mempty = Maybe a -> Maybe (IntMap (Trie a)) -> Trie a
forall a. Maybe a -> Maybe (IntMap (Trie a)) -> Trie a
Trie Maybe a
forall a. Maybe a
Nothing Maybe (IntMap (Trie a))
forall a. Maybe a
Nothing
   mappend :: Trie a -> Trie a -> Trie a
mappend = Trie a -> Trie a -> Trie a
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) = (Maybe a, Maybe (IntMap (Trie a))) -> Put
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) <- Get (Maybe a, Maybe (IntMap (Trie a)))
forall t. Binary t => Get t
get
     Trie a -> Get (Trie a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Trie a -> Get (Trie a)) -> Trie a -> Get (Trie a)
forall a b. (a -> b) -> a -> b
$ Maybe a -> Maybe (IntMap (Trie a)) -> Trie a
forall a. Maybe a -> Maybe (IntMap (Trie a)) -> Trie a
Trie Maybe a
mbv Maybe (IntMap (Trie a))
mbm

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

unfoldTrie :: Trie a -> [([Int], a)]
unfoldTrie :: Trie a -> [([Int], a)]
unfoldTrie  = (([Int], a) -> ([Int], a)) -> [([Int], a)] -> [([Int], a)]
forall a b. (a -> b) -> [a] -> [b]
map (([Int] -> [Int]) -> ([Int], a) -> ([Int], a)
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first [Int] -> [Int]
forall a. [a] -> [a]
reverse) ([([Int], a)] -> [([Int], a)])
-> (Trie a -> [([Int], a)]) -> Trie a -> [([Int], a)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Int] -> Trie a -> [([Int], a)]
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) ([Int], b) -> [([Int], b)] -> [([Int], b)]
forall a. a -> [a] -> [a]
: ((Int, Trie b) -> [([Int], b)]) -> [(Int, Trie b)] -> [([Int], b)]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap ([Int] -> (Int, Trie b) -> [([Int], b)]
gopair [Int]
xs) (IntMap (Trie b) -> [(Int, Trie b)]
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)) =
    ((Int, Trie b) -> [([Int], b)]) -> [(Int, Trie b)] -> [([Int], b)]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap ([Int] -> (Int, Trie b) -> [([Int], b)]
gopair [Int]
xs) (IntMap (Trie b) -> [(Int, Trie b)]
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
iInt -> [Int] -> [Int]
forall a. a -> [a] -> [a]
:[Int]
xs) Trie b
trie

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

alter :: (Maybe a -> Maybe a) -> [Int] -> Trie a -> Trie a
alter :: (Maybe a -> Maybe a) -> [Int] -> Trie a -> Trie a
alter Maybe a -> Maybe a
f [] (Trie Maybe a
mbv Maybe (IntMap (Trie a))
mbm) = Maybe a -> Maybe (IntMap (Trie a)) -> Trie a
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)) =
  Maybe a -> Maybe (IntMap (Trie a)) -> Trie a
forall a. Maybe a -> Maybe (IntMap (Trie a)) -> Trie a
Trie Maybe a
mbv (IntMap (Trie a) -> Maybe (IntMap (Trie a))
forall a. a -> Maybe a
Just (Int -> Trie a -> IntMap (Trie a) -> IntMap (Trie a)
forall a. Int -> a -> IntMap a -> IntMap a
M.insert Int
c ((Maybe a -> Maybe a) -> [Int] -> Trie a -> Trie a
forall a. (Maybe a -> Maybe a) -> [Int] -> Trie a -> Trie a
alter Maybe a -> Maybe a
f [Int]
cs (Trie a -> Trie a) -> Trie a -> Trie a
forall a b. (a -> b) -> a -> b
$ Trie a -> Maybe (Trie a) -> Trie a
forall a. a -> Maybe a -> a
fromMaybe Trie a
forall a. Trie a
empty (Maybe (Trie a) -> Trie a) -> Maybe (Trie a) -> Trie a
forall a b. (a -> b) -> a -> b
$ Int -> IntMap (Trie a) -> Maybe (Trie a)
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) =
  Maybe a -> Maybe (IntMap (Trie a)) -> Trie a
forall a. Maybe a -> Maybe (IntMap (Trie a)) -> Trie a
Trie Maybe a
mbv (IntMap (Trie a) -> Maybe (IntMap (Trie a))
forall a. a -> Maybe a
Just (Int -> Trie a -> IntMap (Trie a) -> IntMap (Trie a)
forall a. Int -> a -> IntMap a -> IntMap a
M.insert Int
c ((Maybe a -> Maybe a) -> [Int] -> Trie a -> Trie a
forall a. (Maybe a -> Maybe a) -> [Int] -> Trie a -> Trie a
alter Maybe a -> Maybe a
f [Int]
cs Trie a
forall a. Trie a
empty) IntMap (Trie a)
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 :: Trie a -> t Int -> Maybe (a, Int, Trie a)
matchLongestPrefix Trie a
trie = (Maybe (a, Int, Trie a) -> Maybe (a, Int, Trie a))
-> ((Maybe (a, Int, Trie a), Int, Trie a)
    -> Maybe (a, Int, Trie a))
-> Either
     (Maybe (a, Int, Trie a)) (Maybe (a, Int, Trie a), Int, Trie a)
-> Maybe (a, Int, Trie a)
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either Maybe (a, Int, Trie a) -> Maybe (a, Int, Trie a)
forall a. a -> a
id (Maybe (a, Int, Trie a), Int, Trie a) -> Maybe (a, Int, Trie a)
forall a b c. (a, b, c) -> a
getBest (Either
   (Maybe (a, Int, Trie a)) (Maybe (a, Int, Trie a), Int, Trie a)
 -> Maybe (a, Int, Trie a))
-> (t Int
    -> Either
         (Maybe (a, Int, Trie a)) (Maybe (a, Int, Trie a), Int, Trie a))
-> t Int
-> Maybe (a, Int, Trie a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Maybe (a, Int, Trie a), Int, Trie a)
 -> Int
 -> Either
      (Maybe (a, Int, Trie a)) (Maybe (a, Int, Trie a), Int, Trie a))
-> (Maybe (a, Int, Trie a), Int, Trie a)
-> t Int
-> Either
     (Maybe (a, Int, Trie a)) (Maybe (a, Int, Trie a), Int, Trie a)
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM (Maybe (a, Int, Trie a), Int, Trie a)
-> Int
-> Either
     (Maybe (a, Int, Trie a)) (Maybe (a, Int, Trie a), Int, Trie a)
forall a.
MatchState a
-> Int -> Either (Maybe (a, Int, Trie a)) (MatchState a)
go (Maybe (a, Int, Trie a)
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 :: 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 Maybe (IntMap (Trie a))
-> (IntMap (Trie a) -> Maybe (Trie a)) -> Maybe (Trie a)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Int -> IntMap (Trie a) -> Maybe (Trie a)
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 -> Maybe (a, Int, Trie a)
-> Either (Maybe (a, Int, Trie a)) (MatchState a)
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))
_)
               -> MatchState a -> Either (Maybe (a, Int, Trie a)) (MatchState a)
forall a b. b -> Either a b
Right ((a, Int, Trie a) -> Maybe (a, Int, Trie a)
forall a. a -> Maybe a
Just (a
x, Int
consumed Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1, Trie a
subtrie), Int
consumed Int -> Int -> Int
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))
_)
               -> MatchState a -> Either (Maybe (a, Int, Trie a)) (MatchState a)
forall a b. b -> Either a b
Right (Maybe (a, Int, Trie a)
best, Int
consumed Int -> Int -> Int
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 :: Trie a -> Int -> Maybe (a, Trie a)
lookupNonEmptyChild (Trie Maybe a
_ Maybe (IntMap (Trie a))
Nothing) Int
_ = Maybe (a, Trie a)
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' <- Int -> IntMap (Trie a) -> Maybe (Trie a)
forall a. Int -> IntMap a -> Maybe a
M.lookup Int
idx IntMap (Trie a)
m
  a
node <- Maybe a
mnode
  (a, Trie a) -> Maybe (a, Trie a)
forall (m :: * -> *) a. Monad m => a -> m a
return (a
node, Maybe a -> Maybe (IntMap (Trie a)) -> Trie a
forall a. Maybe a -> Maybe (IntMap (Trie a)) -> Trie a
Trie Maybe a
forall a. Maybe a
Nothing Maybe (IntMap (Trie a))
m')