module Data.Map.TernaryMap (
TernaryMap,
insert,
singleton,
member,
size,
) where
import Data.Binary
import Control.Monad
import qualified Data.Set.TernarySet as S
import Prelude hiding (lookup)
data Elem2 a b = C !a
| Val b
deriving (Show)
data TernaryMap a b = TNode !(Elem2 a b) !(TernaryMap a b) !(TernaryMap a b) !(TernaryMap a b)
| TEnd
deriving (Show, Eq)
instance Eq a => Eq (Elem2 a b) where
(Val _) == (Val _) = True
(Val _) == x = False
x == (Val _) = False
(C a) == (C b) = a == b
instance (Ord a) => Ord (Elem2 a b) where
compare (Val _) (Val _) = EQ
compare (Val _) x = LT
compare x (Val _) = GT
compare (C x) (C y) = compare x y
isVal (Val _) = True
isVal _ = False
singleton :: Ord a => [a] -> b -> TernaryMap a b
singleton (x:xs) b = TNode (C x) TEnd (singleton xs b) TEnd
singleton [] b = TNode (Val b) TEnd TEnd TEnd
insert :: Ord a => [a] -> b -> TernaryMap a b -> TernaryMap a b
insert xss@(x:xs) b (TNode ele l e h) =
case compare (C x) ele of
LT -> TNode ele (insert xss b l) e h
EQ -> TNode ele l (insert xs b e) h
GT -> TNode ele l e (insert xss b h)
insert xss@(x:xs) b TEnd =
singleton xss b
insert [] b (TNode ele l e h) =
case compare (Val b) ele of
EQ -> TNode (Val b) l e h
LT -> TNode ele (insert [] b l) e h
insert [] b TEnd =
TNode (Val b) TEnd TEnd TEnd
member :: Ord a => [a] -> TernaryMap a b -> Bool
member _ TEnd = False
member [] (TNode ele l e h) = isVal ele || member [] l
member xss@(x:xs) (TNode ele l e h) =
case compare (C x) ele of
LT -> member xss l
EQ -> member xs e
GT -> member xss h
lookup :: Ord a => [a] -> TernaryMap a b -> Maybe b
lookup _ TEnd = Nothing
lookup [] (TNode (Val b) _ _ _) = Just b
lookup [] (TNode ele l _ _) = lookup [] l
lookup xss@(x:xs) (TNode ele l e h) =
case compare (C x) ele of
LT -> lookup xss l
EQ -> lookup xs e
GT -> lookup xss h
(!) :: Ord a => TernaryMap a b -> [a] -> Maybe b
(!) = flip lookup
treeSize :: TernaryMap a b -> Int
treeSize TEnd = 0
treeSize (TNode (Val _) l e h) = treeSize l + treeSize e + treeSize h
treeSize (TNode _ l e h) = 1 + treeSize l + treeSize e + treeSize h
size :: TernaryMap a b -> Int
size TEnd = 0
size (TNode (Val _) l _ h) = 1 + size l + size h
size (TNode _ l e h) = size l + size e + size h
fromList :: Ord a => [([a],b)] -> TernaryMap a b
fromList = foldl (\tree (as,b) -> insert as b tree) TEnd
empty :: TernaryMap a b
empty = TEnd
elems :: TernaryMap a b -> [b]
elems (TEnd) = []
elems (TNode (Val v) l _ h) = elems l ++ (v : elems h)
elems (TNode _ l e h) = elems l ++ (elems e ++ elems h)
null :: TernaryMap a b -> Bool
null TEnd = True
null _ = False
instance Functor (Elem2 a) where
fmap _ (C x) = C x
fmap f (Val b) = Val . f $ b
instance Functor (TernaryMap a) where
fmap f (TNode ele l e h) = TNode (fmap f ele) (fmap f l) (fmap f e) (fmap f h)
fmap _ TEnd = TEnd
instance (Binary a, Binary b) => Binary (Elem2 a b) where
put (C a) = putWord8 0 >> put a
put (Val b) = putWord8 1 >> put b
get = do
n <- getWord8
case n of
0 -> liftM C get
1 -> liftM Val get
instance (Binary a, Binary b) => Binary (TernaryMap a b) where
put (TNode ch TEnd TEnd TEnd) = do
putWord8 0
put ch
put (TNode ch TEnd TEnd h) = do
putWord8 1
put ch
put h
put (TNode ch TEnd e TEnd) = do
putWord8 2
put ch
put e
put (TNode ch TEnd e h) = do
putWord8 3
put ch
put e
put h
put (TNode ch l TEnd TEnd) = do
putWord8 4
put ch
put l
put (TNode ch l TEnd h) = do
putWord8 5
put ch
put l
put h
put (TNode ch l e TEnd) = do
putWord8 6
put ch
put l
put e
put (TNode ch l e h) = do
putWord8 7
put ch
put l
put e
put h
put TEnd = putWord8 8
get = do
tag <- getWord8
case tag of
8 -> return TEnd
_ -> do
ch <- get
case tag of
0 -> return (TNode ch TEnd TEnd TEnd)
1 -> do
h <- get
return (TNode ch TEnd TEnd h)
2 -> do
e <- get
return (TNode ch TEnd e TEnd)
3 -> do
e <- get
h <- get
return (TNode ch TEnd e h)
4 -> do
l <- get
return (TNode ch l TEnd TEnd)
5 -> do
l <- get
h <- get
return (TNode ch l TEnd h)
6 -> do
l <- get
e <- get
return (TNode ch l e TEnd)
7 -> do
l <- get
e <- get
h <- get
return (TNode ch l e h)