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 = Maybe v -> Map k (Trie k v) -> Trie k v
forall k v. Maybe v -> Map k (Trie k v) -> Trie k v
TrieNode Maybe v
forall a. Maybe a
Nothing Map k (Trie k v)
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 =
        Maybe v -> Map k (Trie k v) -> Trie k v
forall k v. Maybe v -> Map k (Trie k v) -> Trie k v
TrieNode (Maybe v
v1 Maybe v -> Maybe v -> Maybe v
forall a. Maybe a -> Maybe a -> Maybe a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> Maybe v
v0) ((Trie k v -> Trie k v -> Trie k v)
-> Map k (Trie k v) -> Map k (Trie k v) -> Map k (Trie k v)
forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
M.unionWith Trie k v -> Trie k v -> Trie k v
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 = [k] -> v -> Trie k v
forall {k} {t}. [k] -> t -> Trie k t
go
  where
    go :: [k] -> t -> Trie k t
go [] t
v = Maybe t -> Map k (Trie k t) -> Trie k t
forall k v. Maybe v -> Map k (Trie k v) -> Trie k v
TrieNode (t -> Maybe t
forall a. a -> Maybe a
Just t
v) Map k (Trie k t)
forall k a. Map k a
M.empty
    go (k
x:[k]
xs) t
v = Maybe t -> Map k (Trie k t) -> Trie k t
forall k v. Maybe v -> Map k (Trie k v) -> Trie k v
TrieNode Maybe t
forall a. Maybe a
Nothing (k -> Trie k t -> Map k (Trie k t)
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 = (([k], v) -> Trie k v) -> [([k], v)] -> Trie k v
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (([k] -> v -> Trie k v) -> ([k], v) -> Trie k v
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry [k] -> v -> Trie k v
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) = Trie k v -> k -> Map k (Trie k v) -> Trie k v
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault Trie k v
forall a. Monoid a => a
mempty k
k Map k (Trie k v)
xs