Trie implementation

[Keith Wansbrough]

Hi... I was writing some code yesterday to walk over a directory tree, 
and needed a Trie.  Not seeing one around, I wrote a basic 
implementation of the functions I needed (attached below).

Has anyone else done this?  Should I polish it up and offer it for 
inclusion in Data?

--KW 8-)

-------------- next part --------------
-----------------------------------------------------------------------------
-- |
-- Module      :  Unstable.Org.Lochan.Trie
-- Copyright   :  (c) Keith Wansbrough 2005
-- License     :  BSD-style
-- 
-- Maintainer  :  keith.hlib at lochan.org
-- 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.
--
-----------------------------------------------------------------------------

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

import Prelude hiding (takeWhile)
import Data.FiniteMap
import Data.Maybe
import Control.Monad

-- |A Trie with key elements of type @k@ (keys of type @[k]@) and values of type @v at .
data Trie k v = Trie { value :: Maybe v,
                       children :: FiniteMap 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 :: (FiniteMap k (Trie k v) -> FiniteMap 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 = emptyFM }

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

-- |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 = plusFM_C 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 = plusFM_C (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')
                             (fmToList (children p))
    where getNode p = maybe [] (\ v -> [(ks,v)]) (value p)

-- |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 (mapFM (const (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 (mapFM (const (downwards f))) . f

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

-- |Return the prefix of the trie satisfying @f@ 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 isEmptyFM (children p) then p else value_u (const Nothing) p)