{- |
Module      :  Trie
Copyright   :  (c) Keith Wansbrough 2005, C. Maeder 2006
License     :  BSD-style

Maintainer  :  none
Stability   :  experimental
Portability :  portable

This module provides a very basic implementation of the Trie data type,
 with no great concern for efficiency, or for completeness of API.

original version
<http://article.gmane.org/gmane.comp.lang.haskell.libraries/2571>

modified using Data.Map and added functions insert and null
-}

module Trie
    (
    -- * Data type
    Trie,
    -- * Constructors
    empty, insert, unit, plus, plus_C,
    -- * Primitive discriminators, accessors and mutators
    null, value, children, value_u, children_u,
    -- * Basic operations
    preOrder, upwards, downwards,
    -- * Derived operations
    takeWhile, takeWhile_V, fringe,
    ) where

import Prelude hiding (takeWhile, null)
import qualified Data.Map as Map
import Data.Maybe
import Control.Monad

-- |A Trie with key elements of type  <at> k <at>
-- (keys of type  <at> [k] <at> ) and values of type  <at> v <at> .
data Trie k v = Trie { value :: Maybe v,
                       children :: Map.Map k (Trie k v)
                     }

-- |Modify the 'children' field of a trie.
value_u :: (Maybe v -> Maybe v) -> Trie k v -> Trie k v
value_u f p = p { value = f (value p) }

-- |Modify the 'children' field of a trie.
children_u :: (Map.Map k (Trie k v) -> Map.Map k (Trie k v))
           -> Trie k v -> Trie k v
children_u f p = p { children = f (children p) }

-- |The empty trie.
empty :: Trie k v
empty = Trie { value = Nothing, children = Map.empty }

-- |Test for the empty trie
null :: Trie k v -> Bool
null t = isNothing (value t) && Map.null (children t)

-- |The singleton trie.
unit :: Ord k => [k] -> v -> Trie k v
unit [] x = Trie { value = Just x, children = Map.empty }
unit (k:ks) x = Trie { value = Nothing
                     , children = Map.singleton k (unit ks x) }

insert :: Ord k => [k] -> (Maybe v -> Maybe v) -> Trie k v -> Trie k v
insert l f t = case l of
    [] -> t { value = f (value t) }
    k : ks -> let cs = children t
                  nt = insert ks f $ Map.findWithDefault empty k cs
              in if null nt then t { children = Map.delete k cs }
                 else t { children = Map.insert k nt cs }

-- |Combining two tries.  The first shadows the second.
plus :: Ord k => Trie k v -> Trie k v -> Trie k v
plus p1 p2 =
    Trie {
          value = mplus (value p1) (value p2),
          children = Map.unionWith plus (children p1) (children p2)
         }

-- |Combining two tries.  If the two define the same key, the
-- specified combining function is used.
plus_C :: Ord k => (v -> v -> v) -> Trie k v -> Trie k v -> Trie k v
plus_C f p1 p2 =
    Trie {
          value =  lift f (value p1) (value p2),
          children = Map.unionWith (plus_C f) (children p1) (children p2)
         }
    where lift _ Nothing y = y
          lift _ x Nothing = x
          lift _ (Just x) (Just y) = Just (f x y)

-- |Enumerate all (key,value) pairs, in preorder.
preOrder :: Ord k => [k] -> Trie k v -> [([k],v)]
preOrder ks p = getNode p
                ++ concatMap (\(k,p') -> preOrder (ks++[k]) p')
                             (Map.toList (children p))
    where getNode q = maybe [] (\ v -> [(ks,v)]) (value q)

-- |An upwards accumulation on the trie.
upwards :: Ord k => (Trie k v -> Trie k v) -> Trie k v -> Trie k v
upwards f = f . children_u (Map.map (upwards f))

-- |A downwards accumulation on the trie.
downwards :: Ord k => (Trie k v -> Trie k v) -> Trie k v -> Trie k v
downwards f = children_u (Map.map (downwards f)) . f

-- |Return the prefix of the trie satisfying  <at> f <at> .
takeWhile :: Ord k => (Trie k v -> Bool) -> Trie k v -> Trie k v
takeWhile f = downwards (children_u (Map.filter f))

-- |Return the prefix of the trie satisfying  <at> f <at>
-- on all values present.
takeWhile_V :: Ord k => (v -> Bool) -> Trie k v -> Trie k v
takeWhile_V f = takeWhile (maybe True f . value)

-- |Return the fringe of the trie (the trie composed of only the leaf nodes).
fringe :: Ord k => Trie k v -> Trie k v
fringe = upwards (\ p -> if Map.null (children p)
                         then p else value_u (const Nothing) p)