-- A strict general trie for types with orderings
module Hope.Util.Trie 
    (

     Trie(..), 

     -- * Construction
     empty, 
     fromList, fromListWith, fromListWith',
     insert, insertWith, insertWith',

     -- * Query
     lookup, member,

     -- * Destruction
     toList, toTree,

     -- * Key manipulation
     addPrefix

    ) where

import Data.List (foldl')
import Data.Map (Map)
import qualified Data.Map as Map
import Data.Maybe (maybe, fromMaybe, isJust)
import Data.Tree (Tree)
import qualified Data.Tree as Tree
import Prelude hiding (lookup)

data Trie a b = Node !(Map a (Trie a b)) !(Maybe b)
-- deriving (Show)

-- | An empty trie
empty :: Trie a b
empty = Node Map.empty Nothing

-- | Construct a trie from a list of (key, value) pairs.
fromList :: Ord a => [([a], b)] -> Trie a b
fromList = fromListWith const

-- | Construct a trie from a list of (key, value) pairs
--   with a combining function for values of equals keys.
--   The order in which the elements are inserted and passed
--   to the combining function is unspecified.
fromListWith :: Ord a => (b -> b -> b) -> [([a], b)] -> Trie a b
fromListWith f = fromListWith' f id

-- | Construct a trie from a list of (key, value) pairs
--   with a combining function for values of equals keys, 
--   and a construction function used the first time a key is inserted.
--   The order in which the elements are inserted and passed
--   to the combining function is unspecified.
fromListWith' :: Ord a => (b -> c -> b) -> (c -> b) -> [([a], c)] -> Trie a b
fromListWith' f g = foldl' (flip (uncurry (insertWith' f g))) empty

-- | Insert an element into a trie.
insert ::  Ord a => [a] -> b -> Trie a b -> Trie a b
insert = insertWith const

-- | Insert an element into a trie with a combining
--   function for the value of an existing key equal to
--   the given one.
insertWith :: Ord a => 
              (b -> b -> b) -- ^ The first argument is the existing value,
	                    --   the second is the new value.
	   -> [a] -> b -> Trie a b -> Trie a b
insertWith f = insertWith' f id


-- | Insert an element into a trie with a combining
--   function for the value of an existing key equal to
--   the given one, and a construction function used the
--   first time a key is inserted.
insertWith' :: Ord a => 
               (b -> c -> b) -- ^ The first argument is the existing value,
	                     --   the second is the new value.
            -> (c -> b) -- ^ Used to transform the input value to a value
                        --   in the map if there is no existing value.
            -> [a] -> c -> Trie a b -> Trie a b
insertWith' f g cs y = insertWith_ cs
    where 
    insertWith_ []     (Node m mv) = Node m (Just $ maybe (g y) (flip f y) mv)
    insertWith_ (c:cs) (Node m mv) = Node (h m) mv
	where h m = Map.insert c (insertWith_ cs t) m
                  where t = fromMaybe empty (Map.lookup c m)

-- | Lookup a key in a trie.
lookup :: (Monad m, Ord a) => [a] -> Trie a b -> m b
lookup []     (Node _ v) = maybe (fail "Not found") return v
lookup (c:cs) (Node m _) = maybe (fail "Not found") (lookup cs) (Map.lookup c m)

-- | Check if a key is in a trie.
member :: Ord a => [a] -> Trie a b -> Bool
member k tr = isJust (lookup k tr)

-- | Get the (key,value) pairs from a trie.
--   The results are ordered by key.
toList :: Trie a b -> [([a],b)]
toList t = collapse t []
    where 
    collapse (Node m mv) xs = maybe [] (\v -> [(reverse xs, v)]) mv ++ rest m xs
    rest m xs = concat [collapse tr (c:xs) | (c,tr) <- Map.toList m]

-- | Convert a 'Trie' to a 'Tree'.
toTree :: ([a] -> Maybe b -> c) -> Trie a b -> Tree c
toTree f = toTree_ []
  where toTree_ p (Node m x) = Tree.Node (f (reverse p) x) cs
            where cs = [toTree_ (k:p) t | (k,t) <- Map.toList m]

-- | Add a prefix to all keys in the map.
addPrefix :: a       -- ^ Prefix to ad to the keys.
          -> Maybe b -- ^ Value for the new @[]@ key.
          -> Trie a b -> Trie a b
addPrefix k mv t = Node (Map.singleton k t) mv