module PGF.TrieMap
         ( TrieMap

         , empty
         , singleton

         , lookup
         
         , null
         , compose
         , decompose
         
         , insertWith

         , union,  unionWith
         , unions, unionsWith
         
         , elems
         , toList
         , fromList, fromListWith

         , map
         , mapWithKey
         ) where

import Prelude hiding (lookup, null, map)
import qualified Data.Map as Map
import Data.List (foldl')

data TrieMap k v = Tr (Maybe v) (Map.Map k (TrieMap k v))

empty :: TrieMap k v
empty = Maybe v -> Map k (TrieMap k v) -> TrieMap k v
forall k v. Maybe v -> Map k (TrieMap k v) -> TrieMap k v
Tr Maybe v
forall a. Maybe a
Nothing Map k (TrieMap k v)
forall k a. Map k a
Map.empty

singleton :: [k] -> a -> TrieMap k a
singleton :: [k] -> a -> TrieMap k a
singleton []     a
v = Maybe a -> Map k (TrieMap k a) -> TrieMap k a
forall k v. Maybe v -> Map k (TrieMap k v) -> TrieMap k v
Tr (a -> Maybe a
forall a. a -> Maybe a
Just a
v) Map k (TrieMap k a)
forall k a. Map k a
Map.empty
singleton (k
k:[k]
ks) a
v = Maybe a -> Map k (TrieMap k a) -> TrieMap k a
forall k v. Maybe v -> Map k (TrieMap k v) -> TrieMap k v
Tr Maybe a
forall a. Maybe a
Nothing  (k -> TrieMap k a -> Map k (TrieMap k a)
forall k a. k -> a -> Map k a
Map.singleton k
k ([k] -> a -> TrieMap k a
forall k a. [k] -> a -> TrieMap k a
singleton [k]
ks a
v))

lookup :: Ord k => [k] -> TrieMap k a -> Maybe a
lookup :: [k] -> TrieMap k a -> Maybe a
lookup []     (Tr Maybe a
mb_v Map k (TrieMap k a)
m) = Maybe a
mb_v
lookup (k
k:[k]
ks) (Tr Maybe a
mb_v Map k (TrieMap k a)
m) = k -> Map k (TrieMap k a) -> Maybe (TrieMap k a)
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup k
k Map k (TrieMap k a)
m Maybe (TrieMap k a) -> (TrieMap k a -> Maybe a) -> Maybe a
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= [k] -> TrieMap k a -> Maybe a
forall k a. Ord k => [k] -> TrieMap k a -> Maybe a
lookup [k]
ks

null :: TrieMap k v -> Bool
null :: TrieMap k v -> Bool
null (Tr Maybe v
Nothing Map k (TrieMap k v)
m) = Map k (TrieMap k v) -> Bool
forall k a. Map k a -> Bool
Map.null Map k (TrieMap k v)
m
null TrieMap k v
_              = Bool
False

compose :: Maybe v -> Map.Map k (TrieMap k v) -> TrieMap k v
compose :: Maybe v -> Map k (TrieMap k v) -> TrieMap k v
compose Maybe v
mb_v Map k (TrieMap k v)
m = Maybe v -> Map k (TrieMap k v) -> TrieMap k v
forall k v. Maybe v -> Map k (TrieMap k v) -> TrieMap k v
Tr Maybe v
mb_v Map k (TrieMap k v)
m

decompose :: TrieMap k v -> (Maybe v, Map.Map k (TrieMap k v))
decompose :: TrieMap k v -> (Maybe v, Map k (TrieMap k v))
decompose (Tr Maybe v
mb_v Map k (TrieMap k v)
m) = (Maybe v
mb_v,Map k (TrieMap k v)
m)

insertWith :: Ord k => (v -> v -> v) -> [k] -> v -> TrieMap k v -> TrieMap k v
insertWith :: (v -> v -> v) -> [k] -> v -> TrieMap k v -> TrieMap k v
insertWith v -> v -> v
f []     v
v0 (Tr Maybe v
mb_v Map k (TrieMap k v)
m) = case Maybe v
mb_v of
                                       Just  v
v -> Maybe v -> Map k (TrieMap k v) -> TrieMap k v
forall k v. Maybe v -> Map k (TrieMap k v) -> TrieMap k v
Tr (v -> Maybe v
forall a. a -> Maybe a
Just (v -> v -> v
f v
v0 v
v)) Map k (TrieMap k v)
m
                                       Maybe v
Nothing -> Maybe v -> Map k (TrieMap k v) -> TrieMap k v
forall k v. Maybe v -> Map k (TrieMap k v) -> TrieMap k v
Tr (v -> Maybe v
forall a. a -> Maybe a
Just v
v0      ) Map k (TrieMap k v)
m
insertWith v -> v -> v
f (k
k:[k]
ks) v
v0 (Tr Maybe v
mb_v Map k (TrieMap k v)
m) = case k -> Map k (TrieMap k v) -> Maybe (TrieMap k v)
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup k
k Map k (TrieMap k v)
m of
                                       Maybe (TrieMap k v)
Nothing -> Maybe v -> Map k (TrieMap k v) -> TrieMap k v
forall k v. Maybe v -> Map k (TrieMap k v) -> TrieMap k v
Tr Maybe v
mb_v (k -> TrieMap k v -> Map k (TrieMap k v) -> Map k (TrieMap k v)
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert k
k ([k] -> v -> TrieMap k v
forall k a. [k] -> a -> TrieMap k a
singleton [k]
ks v
v0) Map k (TrieMap k v)
m)
                                       Just TrieMap k v
tr -> Maybe v -> Map k (TrieMap k v) -> TrieMap k v
forall k v. Maybe v -> Map k (TrieMap k v) -> TrieMap k v
Tr Maybe v
mb_v (k -> TrieMap k v -> Map k (TrieMap k v) -> Map k (TrieMap k v)
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert k
k ((v -> v -> v) -> [k] -> v -> TrieMap k v -> TrieMap k v
forall k v.
Ord k =>
(v -> v -> v) -> [k] -> v -> TrieMap k v -> TrieMap k v
insertWith v -> v -> v
f [k]
ks v
v0 TrieMap k v
tr) Map k (TrieMap k v)
m)

union :: Ord k => TrieMap k v -> TrieMap k v -> TrieMap k v
union :: TrieMap k v -> TrieMap k v -> TrieMap k v
union = (v -> v -> v) -> TrieMap k v -> TrieMap k v -> TrieMap k v
forall k v.
Ord k =>
(v -> v -> v) -> TrieMap k v -> TrieMap k v -> TrieMap k v
unionWith (\v
a v
b -> v
a)

unionWith :: Ord k => (v -> v -> v) -> TrieMap k v -> TrieMap k v -> TrieMap k v
unionWith :: (v -> v -> v) -> TrieMap k v -> TrieMap k v -> TrieMap k v
unionWith v -> v -> v
f (Tr Maybe v
mb_v1 Map k (TrieMap k v)
m1) (Tr Maybe v
mb_v2 Map k (TrieMap k v)
m2) =
  let mb_v :: Maybe v
mb_v = case (Maybe v
mb_v1,Maybe v
mb_v2) of
               (Maybe v
Nothing,Maybe v
Nothing) -> Maybe v
forall a. Maybe a
Nothing
               (Just v
v ,Maybe v
Nothing) -> v -> Maybe v
forall a. a -> Maybe a
Just v
v
               (Maybe v
Nothing,Just v
v ) -> v -> Maybe v
forall a. a -> Maybe a
Just v
v
               (Just v
v1,Just v
v2) -> v -> Maybe v
forall a. a -> Maybe a
Just (v -> v -> v
f v
v1 v
v2)
      m :: Map k (TrieMap k v)
m    = (TrieMap k v -> TrieMap k v -> TrieMap k v)
-> Map k (TrieMap k v)
-> Map k (TrieMap k v)
-> Map k (TrieMap k v)
forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
Map.unionWith ((v -> v -> v) -> TrieMap k v -> TrieMap k v -> TrieMap k v
forall k v.
Ord k =>
(v -> v -> v) -> TrieMap k v -> TrieMap k v -> TrieMap k v
unionWith v -> v -> v
f) Map k (TrieMap k v)
m1 Map k (TrieMap k v)
m2
  in Maybe v -> Map k (TrieMap k v) -> TrieMap k v
forall k v. Maybe v -> Map k (TrieMap k v) -> TrieMap k v
Tr Maybe v
mb_v Map k (TrieMap k v)
m

unions :: Ord k => [TrieMap k v] -> TrieMap k v
unions :: [TrieMap k v] -> TrieMap k v
unions = (TrieMap k v -> TrieMap k v -> TrieMap k v)
-> TrieMap k v -> [TrieMap k v] -> TrieMap k v
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl TrieMap k v -> TrieMap k v -> TrieMap k v
forall k v. Ord k => TrieMap k v -> TrieMap k v -> TrieMap k v
union TrieMap k v
forall k v. TrieMap k v
empty

unionsWith :: Ord k => (v -> v -> v) -> [TrieMap k v] -> TrieMap k v
unionsWith :: (v -> v -> v) -> [TrieMap k v] -> TrieMap k v
unionsWith v -> v -> v
f = (TrieMap k v -> TrieMap k v -> TrieMap k v)
-> TrieMap k v -> [TrieMap k v] -> TrieMap k v
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl ((v -> v -> v) -> TrieMap k v -> TrieMap k v -> TrieMap k v
forall k v.
Ord k =>
(v -> v -> v) -> TrieMap k v -> TrieMap k v -> TrieMap k v
unionWith v -> v -> v
f) TrieMap k v
forall k v. TrieMap k v
empty

elems :: TrieMap k v -> [v]
elems :: TrieMap k v -> [v]
elems TrieMap k v
tr = TrieMap k v -> [v] -> [v]
forall k a. TrieMap k a -> [a] -> [a]
collect TrieMap k v
tr []
  where
    collect :: TrieMap k a -> [a] -> [a]
collect (Tr Maybe a
mb_v Map k (TrieMap k a)
m) [a]
xs = ([a] -> [a]) -> (a -> [a] -> [a]) -> Maybe a -> [a] -> [a]
forall b a. b -> (a -> b) -> Maybe a -> b
maybe [a] -> [a]
forall a. a -> a
id (:) Maybe a
mb_v ((TrieMap k a -> [a] -> [a]) -> [a] -> Map k (TrieMap k a) -> [a]
forall a b k. (a -> b -> b) -> b -> Map k a -> b
Map.foldr TrieMap k a -> [a] -> [a]
collect [a]
xs Map k (TrieMap k a)
m)

toList :: TrieMap k v -> [([k],v)]
toList :: TrieMap k v -> [([k], v)]
toList TrieMap k v
tr = [k] -> TrieMap k v -> [([k], v)] -> [([k], v)]
forall a b. [a] -> TrieMap a b -> [([a], b)] -> [([a], b)]
collect [] TrieMap k v
tr []
  where
    collect :: [a] -> TrieMap a b -> [([a], b)] -> [([a], b)]
collect [a]
ks (Tr Maybe b
mb_v Map a (TrieMap a b)
m) [([a], b)]
xs = ([([a], b)] -> [([a], b)])
-> (b -> [([a], b)] -> [([a], b)])
-> Maybe b
-> [([a], b)]
-> [([a], b)]
forall b a. b -> (a -> b) -> Maybe a -> b
maybe [([a], b)] -> [([a], b)]
forall a. a -> a
id (\b
v -> (:) ([a]
ks,b
v)) Maybe b
mb_v ((a -> TrieMap a b -> [([a], b)] -> [([a], b)])
-> [([a], b)] -> Map a (TrieMap a b) -> [([a], b)]
forall k a b. (k -> a -> b -> b) -> b -> Map k a -> b
Map.foldrWithKey (\a
k -> [a] -> TrieMap a b -> [([a], b)] -> [([a], b)]
collect (a
ka -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
ks)) [([a], b)]
xs Map a (TrieMap a b)
m)

fromListWith :: Ord k => (v -> v -> v) -> [([k],v)] -> TrieMap k v
fromListWith :: (v -> v -> v) -> [([k], v)] -> TrieMap k v
fromListWith v -> v -> v
f [([k], v)]
xs = (TrieMap k v -> ([k], v) -> TrieMap k v)
-> TrieMap k v -> [([k], v)] -> TrieMap k v
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\TrieMap k v
trie ([k]
ks,v
v) -> (v -> v -> v) -> [k] -> v -> TrieMap k v -> TrieMap k v
forall k v.
Ord k =>
(v -> v -> v) -> [k] -> v -> TrieMap k v -> TrieMap k v
insertWith v -> v -> v
f [k]
ks v
v TrieMap k v
trie) TrieMap k v
forall k v. TrieMap k v
empty [([k], v)]
xs

fromList :: Ord k => [([k],v)] -> TrieMap k v
fromList :: [([k], v)] -> TrieMap k v
fromList [([k], v)]
xs = (v -> v -> v) -> [([k], v)] -> TrieMap k v
forall k v. Ord k => (v -> v -> v) -> [([k], v)] -> TrieMap k v
fromListWith v -> v -> v
forall a b. a -> b -> a
const [([k], v)]
xs

map :: (a -> b) -> TrieMap k a -> TrieMap k b
map :: (a -> b) -> TrieMap k a -> TrieMap k b
map a -> b
f (Tr Maybe a
mb_v Map k (TrieMap k a)
m) = Maybe b -> Map k (TrieMap k b) -> TrieMap k b
forall k v. Maybe v -> Map k (TrieMap k v) -> TrieMap k v
Tr ((a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f Maybe a
mb_v) ((TrieMap k a -> TrieMap k b)
-> Map k (TrieMap k a) -> Map k (TrieMap k b)
forall a b k. (a -> b) -> Map k a -> Map k b
Map.map ((a -> b) -> TrieMap k a -> TrieMap k b
forall a b k. (a -> b) -> TrieMap k a -> TrieMap k b
map a -> b
f) Map k (TrieMap k a)
m)

mapWithKey :: ([k] -> a -> b) -> TrieMap k a -> TrieMap k b
mapWithKey :: ([k] -> a -> b) -> TrieMap k a -> TrieMap k b
mapWithKey [k] -> a -> b
f (Tr Maybe a
mb_v Map k (TrieMap k a)
m) = Maybe b -> Map k (TrieMap k b) -> TrieMap k b
forall k v. Maybe v -> Map k (TrieMap k v) -> TrieMap k v
Tr ((a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ([k] -> a -> b
f []) Maybe a
mb_v) ((k -> TrieMap k a -> TrieMap k b)
-> Map k (TrieMap k a) -> Map k (TrieMap k b)
forall k a b. (k -> a -> b) -> Map k a -> Map k b
Map.mapWithKey (\k
k -> ([k] -> a -> b) -> TrieMap k a -> TrieMap k b
forall k a b. ([k] -> a -> b) -> TrieMap k a -> TrieMap k b
mapWithKey ([k] -> a -> b
f ([k] -> a -> b) -> ([k] -> [k]) -> [k] -> a -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (k
kk -> [k] -> [k]
forall a. a -> [a] -> [a]
:))) Map k (TrieMap k a)
m)