module Data.Trie
    ( Trie
    , singleton, fromList
    , terminal, step
    ) where

import Control.Applicative

import qualified Data.Map.Strict as M

data Trie k v
    = TrieNode !(Maybe v) !(M.Map k (Trie k v))

instance Ord k => Monoid (Trie k v) where
    mempty :: Trie k v
mempty = forall k v. Maybe v -> Map k (Trie k v) -> Trie k v
TrieNode forall a. Maybe a
Nothing forall k a. Map k a
M.empty

instance Ord k => Semigroup (Trie k v) where
    TrieNode Maybe v
v0 Map k (Trie k v)
ys0 <> :: Trie k v -> Trie k v -> Trie k v
<> TrieNode Maybe v
v1 Map k (Trie k v)
ys1 =
        forall k v. Maybe v -> Map k (Trie k v) -> Trie k v
TrieNode (Maybe v
v1 forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> Maybe v
v0) (forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
M.unionWith forall a. Semigroup a => a -> a -> a
(<>) Map k (Trie k v)
ys0 Map k (Trie k v)
ys1)

singleton :: Ord k => [k] -> v -> Trie k v
singleton :: forall k v. Ord k => [k] -> v -> Trie k v
singleton = forall {k} {t}. [k] -> t -> Trie k t
go
  where
    go :: [k] -> t -> Trie k t
go [] t
v = forall k v. Maybe v -> Map k (Trie k v) -> Trie k v
TrieNode (forall a. a -> Maybe a
Just t
v) forall k a. Map k a
M.empty
    go (k
x:[k]
xs) t
v = forall k v. Maybe v -> Map k (Trie k v) -> Trie k v
TrieNode forall a. Maybe a
Nothing (forall k a. k -> a -> Map k a
M.singleton k
x ([k] -> t -> Trie k t
go [k]
xs t
v))

fromList :: Ord k => [([k], v)] -> Trie k v
fromList :: forall k v. Ord k => [([k], v)] -> Trie k v
fromList = forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry forall k v. Ord k => [k] -> v -> Trie k v
singleton)

terminal :: Trie k v -> Maybe v
terminal :: forall k v. Trie k v -> Maybe v
terminal (TrieNode Maybe v
v Map k (Trie k v)
_) = Maybe v
v

step :: Ord k => k -> Trie k v -> Trie k v
step :: forall k v. Ord k => k -> Trie k v -> Trie k v
step k
k (TrieNode Maybe v
_ Map k (Trie k v)
xs) = forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault forall a. Monoid a => a
mempty k
k Map k (Trie k v)
xs