{-# LANGUAGE CPP #-}
{-# LANGUAGE StrictData #-}
module Trie
  ( Trie
  , empty
  , insert
  , alter
  , unfoldTrie
  , fromList
  , matchLongestPrefix
  )
  where

import Control.Monad (foldM)
import qualified Data.IntMap as M
import Data.Bifunctor (first)
import Data.Maybe (fromMaybe)

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
$cshowsPrec :: forall a. Show a => Int -> Trie a -> ShowS
showsPrec :: Int -> Trie a -> ShowS
$cshow :: forall a. Show a => Trie a -> String
show :: Trie a -> String
$cshowList :: forall a. Show a => [Trie a] -> ShowS
showList :: [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
$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
/= :: 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
$ccompare :: forall a. Ord a => Trie a -> Trie a -> Ordering
compare :: Trie a -> Trie a -> Ordering
$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
>= :: Trie a -> Trie a -> Bool
$cmax :: forall a. Ord a => Trie a -> Trie a -> Trie a
max :: Trie a -> Trie a -> Trie a
$cmin :: forall a. Ord a => Trie a -> Trie a -> Trie a
min :: Trie a -> Trie a -> Trie a
Ord)

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 a b. (a -> b -> b) -> b -> [a] -> b
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
(<>)

empty :: Trie a
empty :: forall a. 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 :: forall a. 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 a b c. (a -> b) -> (a, c) -> (b, c)
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 :: forall a. [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))

fromList :: [([Int], a)] -> Trie a
fromList :: forall a. [([Int], a)] -> Trie a
fromList = (([Int], a) -> Trie a -> Trie a)
-> Trie a -> [([Int], a)] -> Trie a
forall a b. (a -> b -> b) -> b -> [a] -> b
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
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) = 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 :: forall (t :: * -> *) a.
Foldable t =>
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 :: 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 Maybe (IntMap (Trie a))
-> (IntMap (Trie a) -> Maybe (Trie a)) -> Maybe (Trie a)
forall a b. Maybe a -> (a -> Maybe b) -> Maybe b
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)) (Maybe (a, Int, Trie a), Int, Trie 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))
_)
               -> (Maybe (a, Int, Trie a), Int, Trie a)
-> Either
     (Maybe (a, Int, Trie a)) (Maybe (a, Int, Trie a), Int, Trie 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))
_)
               -> (Maybe (a, Int, Trie a), Int, Trie a)
-> Either
     (Maybe (a, Int, Trie a)) (Maybe (a, Int, Trie a), Int, Trie 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)