Ticket #2037: Dict.hs

{-# OPTIONS_GHC -XTypeOperators #-}
{-# OPTIONS_GHC -XScopedTypeVariables #-}

-----------------------------------------------------------------------------
-- |
-- Module      :  Data.Dict
-- Copyright   :  (c) Bas van Dijk 2008
-- License     :  BSD-style
-- Maintainer  :  v.dijk.bas@gmail.com
-- Stability   :  provisional
-- Portability :  portable
--
-----------------------------------------------------------------------------

module Data.Dict where

import Prelude hiding (lookup)

import qualified Data.Map as M
import Data.Monoid (Monoid(..))

{--------------------------------------------------------------------
  Type
--------------------------------------------------------------------}

newtype Dict k a = D (DMap k a) deriving Show

type DMap k a = M.Map k (Value k a)

data Value k a = NoValue (Dict k a) 
               | AValue  (Dict k a, a)
                 deriving Show


{--------------------------------------------------------------------
  Instances
--------------------------------------------------------------------}

instance (Ord k) => Monoid (Dict k a) where
    mempty  = empty
    mappend = union
    mconcat = unions

{--------------------------------------------------------------------
  Operators
--------------------------------------------------------------------}



{--------------------------------------------------------------------
  Construction
--------------------------------------------------------------------}

empty :: Dict k a
empty = D M.empty

singleton :: forall k a. [k] -> a -> Dict k a
singleton []     _ = empty -- TODO: or error ???
singleton (x:xs) y = go x xs
    where
      go :: k -> [k] -> Dict k a
      go x xs = D $ M.singleton x $ case xs of
                                      []   -> AValue (empty, y) 
                                      x:xs -> NoValue $ go x xs


{--------------------------------------------------------------------
  Query
--------------------------------------------------------------------}

lookup :: forall k a m. (Monad m, Ord k) => [k] -> Dict k a -> m a
lookup []     _     = fail "not found"
lookup (x:xs) (D m) = go x xs m
    where
      go :: k -> [k] -> DMap k a -> m a
      go x xs m = do v <- M.lookup x m
                     case v of
                       NoValue (D m)   -> case xs of
                                            []   -> fail "not found"
                                            x:xs -> go x xs m
                       AValue (D m, y) -> case xs of
                                            []   -> return y
                                            x:xs -> go x xs m


{--------------------------------------------------------------------
  Insertion
--------------------------------------------------------------------}

-- TODO: error when xs is null?
insert :: Ord k => [k] -> a -> Dict k a -> Dict k a
insert xs y d = d `union` (singleton xs y)


{--------------------------------------------------------------------
  Union
--------------------------------------------------------------------}

union :: Ord k => Dict k a -> Dict k a -> Dict k a
union = unionWithKey (\k l r -> l)

-- I use foldl instead of foldr because union is more efficient on (bigset `union` smallset)
-- TODO: Data.Map uses a strict foldl. Investigate if a strict version is also better here...
unions :: Ord k => [Dict k a] -> Dict k a
unions = foldl union empty 

unionWithKey :: forall k a. Ord k => (k -> a -> a -> a) -> Dict k a -> Dict k a -> Dict k a
unionWithKey f dl@(D ml) dr@(D mr) | M.null ml = dr
                                   | M.null mr = dl
                                   | otherwise = D $ M.unionWithKey unify ml mr
    where
      unify :: k ->  Value k a -> Value k a -> Value k a
      unify x (NoValue   dl)      (NoValue dr)      = NoValue(union dl dr)
      unify x (NoValue   dl)      (AValue (dr, yr)) = AValue (union dl dr, yr) 
      unify x (AValue   (dl, yl)) (NoValue dr)      = AValue (union dl dr, yl)
      unify x (AValue   (dl, yl)) (AValue (dr, yr)) = AValue (union dl dr, f x yl yr)