module Data.TreeMap.Strict where
import Control.Applicative (Applicative(..))
import Control.DeepSeq (NFData(..))
import Control.Monad (Monad(..))
import Data.Bool
import Data.Data (Data)
import Data.Eq (Eq)
import Data.Foldable (Foldable, foldMap)
import Data.Function (($), (.), const, flip, id)
import Data.Functor (Functor(..), (<$>))
import qualified Data.List
import qualified Data.List.NonEmpty
import Data.List.NonEmpty (NonEmpty(..))
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map
import Data.Maybe (Maybe(..), maybe)
import Data.Monoid (Monoid(..))
import Data.Ord (Ord(..))
import qualified Data.Strict.Maybe as Strict
import Data.Traversable (Traversable(..))
import Data.Typeable (Typeable)
import Prelude (Int, Num(..), seq)
import Text.Show (Show(..))
deriving instance Data x => Data (Strict.Maybe x)
deriving instance Typeable Strict.Maybe
instance Monoid x => Monoid (Strict.Maybe x) where
mempty = Strict.Nothing
mappend (Strict.Just x) (Strict.Just y) = Strict.Just (x `mappend` y)
mappend x Strict.Nothing = x
mappend Strict.Nothing y = y
instance NFData x => NFData (Strict.Maybe x) where
rnf Strict.Nothing = ()
rnf (Strict.Just x) = rnf x
newtype TreeMap k x
= TreeMap (Map k (Node k x))
deriving (Data, Eq, Show, Typeable)
instance (Ord k, Monoid v) => Monoid (TreeMap k v) where
mempty = empty
mappend = union mappend
instance Ord k => Functor (TreeMap k) where
fmap f (TreeMap m) = TreeMap $ fmap (fmap f) m
instance Ord k => Foldable (TreeMap k) where
foldMap f (TreeMap m) = foldMap (foldMap f) m
instance Ord k => Traversable (TreeMap k) where
traverse f (TreeMap m) = TreeMap <$> traverse (traverse f) m
instance (Ord k, NFData k, NFData x) => NFData (TreeMap k x) where
rnf (TreeMap m) = rnf m
type Path = NonEmpty
path :: k -> [k] -> Path k
path = (:|)
list :: Path k -> [k]
list = Data.List.NonEmpty.toList
reverse :: Path k -> Path k
reverse = Data.List.NonEmpty.reverse
data Node k x
= Node
{ node_size :: !Int
, node_value :: !(Strict.Maybe x)
, node_descendants :: !(TreeMap k x)
} deriving (Data, Eq, Show, Typeable)
instance (Ord k, Monoid v) => Monoid (Node k v) where
mempty = node Strict.Nothing (TreeMap mempty)
mappend
Node{node_value=x0, node_descendants=m0}
Node{node_value=x1, node_descendants=m1} =
node (x0 `mappend` x1) (union const m0 m1)
instance Ord k => Functor (Node k) where
fmap f Node{node_value=x, node_descendants=m, node_size} =
Node
{ node_value = fmap f x
, node_descendants = map f m
, node_size
}
instance Ord k => Foldable (Node k) where
foldMap f Node{node_value=Strict.Nothing, node_descendants=TreeMap m} =
foldMap (foldMap f) m
foldMap f Node{node_value=Strict.Just x, node_descendants=TreeMap m} =
f x `mappend` foldMap (foldMap f) m
instance Ord k => Traversable (Node k) where
traverse f Node{node_value=Strict.Nothing, node_descendants=TreeMap m, node_size} =
Node node_size <$> pure Strict.Nothing <*> (TreeMap <$> traverse (traverse f) m)
traverse f Node{node_value=Strict.Just x, node_descendants=TreeMap m, node_size} =
Node node_size <$> (Strict.Just <$> f x) <*> (TreeMap <$> traverse (traverse f) m)
instance (Ord k, NFData k, NFData x) => NFData (Node k x) where
rnf (Node s v d) = rnf s `seq` rnf v `seq` rnf d
node :: Strict.Maybe x -> TreeMap k x -> Node k x
node node_value node_descendants =
Node
{ node_value
, node_size =
size node_descendants +
Strict.maybe 0 (const 1) node_value
, node_descendants
}
node_empty :: Node k x
node_empty = node Strict.Nothing empty
node_find :: Ord k => [k] -> Node k x -> Strict.Maybe (Node k x)
node_find [] n = Strict.Just n
node_find (k:ks) Node{node_descendants=TreeMap m} =
maybe Strict.Nothing (node_find ks) $
Map.lookup k m
empty :: TreeMap k x
empty = TreeMap Map.empty
singleton :: Ord k => Path k -> x -> TreeMap k x
singleton ks x = insert const ks x empty
leaf :: x -> Node k x
leaf x = node (Strict.Just x) empty
insert :: Ord k => (x -> x -> x) -> Path k -> x -> TreeMap k x -> TreeMap k x
insert merge (k:|[]) x (TreeMap m) =
TreeMap $
Map.insertWith (\_ Node{..} -> node
(Strict.maybe (Strict.Just x) (Strict.Just . merge x) node_value)
node_descendants)
k (leaf x) m
insert merge (k:|k':ks) x (TreeMap m) =
TreeMap $
Map.insertWith (\_ Node{..} -> node node_value $
insert merge (path k' ks) x node_descendants)
k
(node Strict.Nothing (insert merge (path k' ks) x empty))
m
from_List :: Ord k => (x -> x -> x) -> [(Path k, x)] -> TreeMap k x
from_List merge = Data.List.foldl (\acc (p, x) -> insert merge p x acc) empty
from_Map :: Ord k => (x -> x -> x) -> Map (Path k) x -> TreeMap k x
from_Map merge = Map.foldlWithKey (\acc p x -> insert merge p x acc) empty
nodes :: TreeMap k x -> Map k (Node k x)
nodes (TreeMap m) = m
null :: TreeMap k x -> Bool
null (TreeMap m) = Map.null m
size :: TreeMap k x -> Int
size = Map.foldr ((+) . node_size) 0 . nodes
find :: Ord k => Path k -> TreeMap k x -> Strict.Maybe x
find (k:|[]) (TreeMap m) = maybe Strict.Nothing node_value $ Map.lookup k m
find (k:|k':ks) (TreeMap m) =
maybe Strict.Nothing (find (path k' ks) . node_descendants) $
Map.lookup k m
find_along :: Ord k => Path k -> TreeMap k x -> [x]
find_along p (TreeMap tm) =
go (list p) tm
where
go :: Ord k => [k] -> Map k (Node k x) -> [x]
go [] _m = []
go (k:ks) m =
case Map.lookup k m of
Nothing -> []
Just nod ->
Strict.maybe id (:) (node_value nod) $
go ks $ nodes (node_descendants nod)
find_node :: Ord k => Path k -> TreeMap k x -> Maybe (Node k x)
find_node (k:|[]) (TreeMap m) = Map.lookup k m
find_node (k:|k':ks) (TreeMap m) =
Map.lookup k m >>=
find_node (path k' ks) . node_descendants
union :: Ord k => (x -> x -> x) -> TreeMap k x -> TreeMap k x -> TreeMap k x
union merge (TreeMap tm0) (TreeMap tm1) =
TreeMap $
Map.unionWith
(\Node{node_value=x0, node_descendants=m0}
Node{node_value=x1, node_descendants=m1} ->
node (Strict.maybe x1 (\x0' -> Strict.maybe (Strict.Just x0') (Strict.Just . merge x0') x1) x0)
(union merge m0 m1))
tm0 tm1
unions :: Ord k => (x -> x -> x) -> [TreeMap k x] -> TreeMap k x
unions merge = Data.List.foldl' (union merge) empty
map :: Ord k => (x -> y) -> TreeMap k x -> TreeMap k y
map f =
TreeMap .
Map.map
(\n@Node{node_value=x, node_descendants=m} ->
n{ node_value = fmap f x
, node_descendants = map f m
}) .
nodes
map_monotonic :: (Ord k, Ord l) => (k -> l) -> (x -> y) -> TreeMap k x -> TreeMap l y
map_monotonic fk fx =
TreeMap .
Map.mapKeysMonotonic fk .
Map.map
(\n@Node{node_value=x, node_descendants=m} ->
n{ node_value = fmap fx x
, node_descendants = map_monotonic fk fx m
}) .
nodes
map_by_depth_first :: Ord k => (TreeMap k y -> Strict.Maybe x -> y) -> TreeMap k x -> TreeMap k y
map_by_depth_first f =
TreeMap .
Map.map
(\Node{node_value, node_descendants} ->
let m = map_by_depth_first f node_descendants in
node (Strict.Just $ f m node_value) m) .
nodes
alterl_path :: Ord k => (Strict.Maybe x -> Strict.Maybe x) -> Path k -> TreeMap k x -> TreeMap k x
alterl_path fct =
go fct . list
where
go :: Ord k
=> (Strict.Maybe x -> Strict.Maybe x) -> [k]
-> TreeMap k x -> TreeMap k x
go _f [] m = m
go f (k:p) (TreeMap m) =
TreeMap $
Map.alter
(\c ->
let (cv, cm) =
case c of
Just Node{node_value=v, node_descendants=d} -> (v, d)
Nothing -> (Strict.Nothing, empty) in
let fx = f cv in
let gm = go f p cm in
case (fx, size gm) of
(Strict.Nothing, 0) -> Nothing
(_, s) -> Just
Node
{ node_value = fx
, node_descendants = gm
, node_size = s + 1
}
) k m
foldl_with_Path :: Ord k => (a -> Path k -> x -> a) -> a -> TreeMap k x -> a
foldl_with_Path =
foldp []
where
foldp :: Ord k
=> [k] -> (a -> Path k -> x -> a)
-> a -> TreeMap k x -> a
foldp p fct a (TreeMap m) =
Map.foldlWithKey
(\acc k Node{..} ->
let acc' = Strict.maybe acc (fct acc (reverse $ path k p)) node_value in
foldp (k:p) fct acc' node_descendants) a m
foldl_with_Path_and_Node :: Ord k => (a -> Node k x -> Path k -> x -> a) -> a -> TreeMap k x -> a
foldl_with_Path_and_Node =
foldp []
where
foldp :: Ord k
=> [k] -> (a -> Node k x -> Path k -> x -> a)
-> a -> TreeMap k x -> a
foldp p fct a (TreeMap m) =
Map.foldlWithKey
(\acc k n@Node{..} ->
let acc' = Strict.maybe acc (fct acc n (reverse $ path k p)) node_value in
foldp (k:p) fct acc' node_descendants) a m
foldr_with_Path :: Ord k => (Path k -> x -> a -> a) -> a -> TreeMap k x -> a
foldr_with_Path =
foldp []
where
foldp :: Ord k
=> [k] -> (Path k -> x -> a -> a)
-> a -> TreeMap k x -> a
foldp p fct a (TreeMap m) =
Map.foldrWithKey
(\k Node{..} acc ->
let acc' = foldp (k:p) fct acc node_descendants in
Strict.maybe acc' (\x -> fct (reverse $ path k p) x acc') node_value) a m
foldr_with_Path_and_Node :: Ord k => (Node k x -> Path k -> x -> a -> a) -> a -> TreeMap k x -> a
foldr_with_Path_and_Node =
foldp []
where
foldp :: Ord k
=> [k] -> (Node k x -> Path k -> x -> a -> a)
-> a -> TreeMap k x -> a
foldp p fct a (TreeMap m) =
Map.foldrWithKey
(\k n@Node{..} acc ->
let acc' = foldp (k:p) fct acc node_descendants in
Strict.maybe acc' (\x -> fct n (reverse $ path k p) x acc') node_value) a m
foldl_path :: Ord k => (Path k -> x -> a -> a) -> Path k -> TreeMap k x -> a -> a
foldl_path fct =
go fct [] . list
where
go :: Ord k
=> (Path k -> x -> a -> a) -> [k] -> [k]
-> TreeMap k x -> a -> a
go _f _ [] _t a = a
go f p (k:n) (TreeMap t) a =
case Map.lookup k t of
Nothing -> a
Just Node{..} ->
case node_value of
Strict.Nothing -> go f (k:p) n node_descendants a
Strict.Just x -> go f (k:p) n node_descendants (f (reverse $ path k p) x a)
foldr_path :: Ord k => (Path k -> x -> a -> a) -> Path k -> TreeMap k x -> a -> a
foldr_path fct =
go fct [] . list
where
go :: Ord k
=> (Path k -> x -> a -> a) -> [k] -> [k]
-> TreeMap k x -> a -> a
go _f _ [] _t a = a
go f p (k:n) (TreeMap t) a =
case Map.lookup k t of
Nothing -> a
Just Node{..} ->
case node_value of
Strict.Nothing -> go f (k:p) n node_descendants a
Strict.Just x -> f (reverse $ path k p) x $ go f (k:p) n node_descendants a
flatten :: Ord k => (x -> y) -> TreeMap k x -> Map (Path k) y
flatten = flatten_with_Path . const
flatten_with_Path :: Ord k => (Path k -> x -> y) -> TreeMap k x -> Map (Path k) y
flatten_with_Path =
flat_map []
where
flat_map :: Ord k
=> [k] -> (Path k -> x -> y)
-> TreeMap k x
-> Map (Path k) y
flat_map p f (TreeMap m) =
Map.unions $
Map.mapKeysMonotonic (reverse . flip path p) (
Map.mapMaybeWithKey (\k Node{node_value} ->
case node_value of
Strict.Nothing -> Nothing
Strict.Just x -> Just $ f (reverse $ path k p) x) m
) :
Map.foldrWithKey
(\k -> (:) . flat_map (k:p) f . node_descendants)
[] m
filter :: Ord k => (x -> Bool) -> TreeMap k x -> TreeMap k x
filter f =
map_Maybe_with_Path
(\_p x -> if f x then Strict.Just x else Strict.Nothing)
filter_with_Path :: Ord k => (Path k -> x -> Bool) -> TreeMap k x -> TreeMap k x
filter_with_Path f =
map_Maybe_with_Path
(\p x -> if f p x then Strict.Just x else Strict.Nothing)
filter_with_Path_and_Node :: Ord k => (Node k x -> Path k -> x -> Bool) -> TreeMap k x -> TreeMap k x
filter_with_Path_and_Node f =
map_Maybe_with_Path_and_Node
(\n p x -> if f n p x then Strict.Just x else Strict.Nothing)
map_Maybe :: Ord k => (x -> Strict.Maybe y) -> TreeMap k x -> TreeMap k y
map_Maybe = map_Maybe_with_Path . const
map_Maybe_with_Path :: Ord k => (Path k -> x -> Strict.Maybe y) -> TreeMap k x -> TreeMap k y
map_Maybe_with_Path = map_Maybe_with_Path_and_Node . const
map_Maybe_with_Path_and_Node :: Ord k => (Node k x -> Path k -> x -> Strict.Maybe y) -> TreeMap k x -> TreeMap k y
map_Maybe_with_Path_and_Node =
go []
where
go :: Ord k
=> [k] -> (Node k x -> Path k -> x -> Strict.Maybe y)
-> TreeMap k x
-> TreeMap k y
go p test (TreeMap m) =
TreeMap $
Map.mapMaybeWithKey
(\k nod@Node{node_value=v, node_descendants=ns} ->
let node_descendants = go (k:p) test ns in
let node_size = size node_descendants in
case v of
Strict.Just x ->
let node_value = test nod (reverse $ path k p) x in
case node_value of
Strict.Nothing | null node_descendants -> Nothing
Strict.Nothing -> Just Node{node_value, node_descendants, node_size=1 + node_size}
Strict.Just _ -> Just Node{node_value, node_descendants, node_size}
_ ->
if null node_descendants
then Nothing
else Just Node{node_value=Strict.Nothing, node_descendants, node_size}
) m