module Data.NMap (
NMap, (.<), (.>), drawNMap, mapKeys, mapWithKey, mapWithKey0,
traverseKeys, traverseWithKey, bitraverse, bisequence, mapKeysM, fromList, elems,
delete, delete0, rootKeys, roots, lookup, (!), member, branch, branch0, insert, insert0
)where
import qualified Data.Map as M
import Data.Bifunctor
import Data.List (intersperse)
import Data.Tuple (swap)
import Prelude hiding (lookup)
data NMap k a =
Branch { roots :: M.Map k (NMap k a) }
| Leaf { leaf :: a }
deriving Eq
infixr 0 .<
k .< xs = (k, Branch $ M.fromList xs)
infixr 0 .>
k .> a = (k, Leaf a)
infixr 0 <:>
m <:> n = m*n
instance (Show a, Show b) => Show (NMap a b) where
show (Leaf a) = concat ["[", show a, "]"]
show (Branch m) = concat $ ["["] ++ (intersperse ", " $ show' <$> (M.toList m)) ++ ["]"]
where show' (k,a) = show k ++ " -> " ++ show a
drawNMap :: (Show k, Show a) => NMap k a -> String
drawNMap mt = unlines $ firstdraw mt where
firstdraw (Leaf a) = concat ["{", show a, "}"] :[]
firstdraw (Branch m) = firstdraw' (M.toList m)
firstdraw' [] = ["()"]
firstdraw' [(x,m)] = shift (show x ++ "─") (show x ++ " " *> " ") (draw m)
firstdraw' ((x,m):xs) = (shift ("┌─" ++ show x ++ "─") ("│ " ++ (show x ++ "-" *> " ")) (draw m)) ++ (firstdraw'' xs)
firstdraw'' [(x,m)] = shift ("└─" ++ show x ++ "─") ("└─" ++ show x ++ "─" *> " ") (draw m)
firstdraw'' ((x,m):xs) = shift ("├─" ++ show x ++ "─") ("│ " ++ (show x ++ " " *> " ")) (draw m) ++ (firstdraw'' xs)
draw (Leaf a) = concat ["──{", show a, "}"] :[]
draw (Branch m) = draw' $ M.toList m
draw' [] = "()":[]
draw' [(x,m)] = shift ("──" ++ show x ++ "─") ("--" ++ show x ++ "-" *> " ") (draw m)
draw' ((x,m):xs) = shift ("┬─" ++ show x ++ "─") ("│ " ++ (show x ++ " " *> " ")) (draw m) ++ (draw'' xs)
draw'' [(x,m)] = shift ("└─" ++ show x ++ "─") (" " ++ show x ++ " " *> " ") (draw m)
draw'' ((x,m):xs) = shift ("├─" ++ show x ++ "─") ("│ " ++ (show x ++ " " *> " ")) (draw m) ++ (draw'' xs)
shift first other = zipWith (++) (first : repeat other)
instance Ord k => Monoid (NMap k b) where
mempty = Branch M.empty
(Branch m1) `mappend` (Branch m2) = Branch $ M.unionWith mappend m1 m2
fst `mappend` _ = fst
instance Functor (NMap k) where
fmap f (Leaf a) = Leaf (f a)
fmap f (Branch m) = Branch $ (fmap f) <$> m
instance Applicative (NMap k) where
pure a = Leaf a
Leaf f <*> t = f <$> t
Branch m <*> Leaf a = ($ a) <$> Branch m
Branch m <*> Branch m' = Branch $ (Branch m <*>) <$> m'
instance Monad (NMap k) where
return = pure
Leaf a >>= f = f a
Branch m >>= f = Branch $ (>>= f) <$> m
mapKeys :: Ord k2 => (k1 -> k2) -> NMap k1 a -> NMap k2 a
mapKeys f (Leaf a) = Leaf a
mapKeys f (Branch m) = Branch $ (mapKeys f) <$> (M.mapKeys f m)
mapWithKey :: (k -> a -> b) -> NMap k a -> NMap k b
mapWithKey _ (Leaf _) = error "can't map without a key"
mapWithKey f (Branch m) = Branch $ M.mapWithKey (\k mt -> mapWithKey0 f k mt) m
mapWithKey0 :: (k -> a -> b) -> k -> NMap k a -> NMap k b
mapWithKey0 f k (Leaf a) = Leaf (f k a)
mapWithKey0 f k (Branch m) = Branch $ M.mapWithKey (\k' mt -> mapWithKey0 f k' mt) m
mapWithKeys :: ([k] -> a -> b) -> NMap k a -> NMap k b
mapWithKeys f m = mapWithKeys0 f [] m
where
mapWithKeys0 f ks (Leaf a) = Leaf $ f ks a
mapWithKeys0 f ks (Branch m) = Branch $ M.mapWithKey (\k mt -> mapWithKeys0 f (ks ++ [k]) mt) m
instance Foldable (NMap k) where
foldr = fold
fold :: (a -> b -> b) -> b -> NMap k a -> b
fold f b (Leaf a) = f a b
fold f b (Branch m) = foldr (flip $ fold f) b m
instance Traversable (NMap k) where
traverse f (Leaf a) = Leaf <$> f a
traverse f (Branch m) = Branch <$> (traverse (traverse f) m)
traverseKeys :: (Ord k2, Applicative t) => (k1 -> t k2) -> NMap k1 a -> t (NMap k2 a)
traverseKeys f = bitraverse f pure
traverseWithKey :: (Ord k, Applicative t) => (k -> a -> t b) -> NMap k a -> t (NMap k b)
traverseWithKey f = sequenceA . (mapWithKey f)
bitraverse :: (Applicative t, Ord k2) => (k1 -> t k2) -> (a1 -> t a2) -> NMap k1 a1 -> t (NMap k2 a2)
bitraverse _ g (Leaf a) = Leaf <$> g a
bitraverse f g (Branch m) = Branch <$> (traverseK2 f (bitraverse f g) m)
traverseK2 :: (Applicative t, Ord k2) => (k1 -> t k2) -> (a -> t b) -> M.Map k1 a -> t (M.Map k2 b)
traverseK2 f g = (M.fromList <$>) . (traverse (\(x,y) -> (,) <$> f x <*> g y)) . M.toList
bisequence :: (Applicative t, Ord k) => NMap (t k) (t a) -> t (NMap k a)
bisequence = bitraverse id id
mapKeysM :: (Ord k2, Monad m) => (k1 -> m k2) -> NMap k1 a -> m (NMap k2 a)
mapKeysM = traverseKeys
fromList :: Ord k => [(k,NMap k a)] -> NMap k a
fromList = Branch . M.fromList
fromKeys :: Ord k => [([k], a)] -> NMap k a
fromKeys xs = foldMap (uncurry branch) xs
toKeys :: Ord k => NMap k a -> [([k], a)]
toKeys = elems . (mapWithKeys (,))
elems :: NMap k a -> [a]
elems = foldr (:) []
delete :: Ord k => [k] -> NMap k a -> NMap k a
delete (k:[]) (Branch m) = Branch $ M.delete k m
delete (k:ks) (Branch m) = Branch $ M.adjust (delete ks) k m
delete [] nm = nm
delete _ l = l
delete0 :: Ord k => k -> NMap k a -> NMap k a
delete0 k (Branch m) = Branch $ M.delete k m
delete0 _ l = l
rootKeys :: NMap k a -> [k]
rootKeys (Leaf _) = error "0-depth trees don't have keys"
rootKeys (Branch m) = M.keys m
lookup :: Ord k => k -> NMap k a -> Maybe (Either (NMap k a) a)
lookup k (Leaf _) = Nothing
lookup k (Branch m) = case M.lookup k m of
Just (Branch m') -> Just (Left $ Branch m')
Just (Leaf a) -> Just (Right a)
Nothing -> Nothing
(!) :: Ord k => NMap k a -> [k] -> a
(Leaf a)![] = a
(Leaf a)!_ = error "too many keys provided"
(Branch m)!(k:ks) = case m M.! k of
(Leaf a') -> a'
b -> b!ks
(Branch m)![] = error "not enough keys"
member :: Ord k => k -> NMap k a -> Bool
member k (Branch m) = M.member k m
member _ _ = False
branch :: Ord k => [k] -> a -> NMap k a
branch ks a = foldr (\a b -> fromList [(a,b)]) (Leaf a) ks
branch0 :: Ord k => k -> a -> NMap k a
branch0 k a = fromList [k .> a]
insert :: Ord k => [k] -> a -> NMap k a -> NMap k a
insert ks a m = (branch ks a) `mappend` m
insert0 :: Ord k => k -> a -> NMap k a -> NMap k a
insert0 k a m = insert [k] a m
pt :: (Show k, Show a) => NMap k a -> IO ()
pt = putStr . drawNMap
test1 = fromList
[ 1 .<
[ 1 .> 'a'
, 2 .<
[ 1 .> 'b'
, 2 .> 'c']
, 3 .<
[ 4 .> 'd' ] ]
, 2 .> 'e' ]
test2 = fromList
[ 1.< [ 2 .> 'a' ]]
test3 = fromList
[ 1.<
[ 2 .> 'a' ]
, 3 .> 'b'
, 4 .< [] ]
test4 = fromList
[ 1 .> (:[]) ]
test5 = fromList
[ 1 .> (:[])
, 2 .> (:"-d")]