{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE NamedFieldPuns #-}
--- | This module implements a 'TreeMap',
+-- | This module implements a strict 'TreeMap',
-- which is like a 'Map'
-- but whose key is now a 'NonEmpty' list of 'Map' keys (a 'Path')
-- enabling the possibility to gather mapped values
-- by 'Path' prefixes (inside a 'Node').
module Hcompta.Lib.TreeMap where
-import Control.Applicative ((<$>), (<*>), pure)
+import Control.DeepSeq (NFData(..))
+-- import Control.Applicative ((<$>), (<*>), pure)
import Data.Data (Data)
-import Data.Foldable (Foldable(..))
+-- import Data.Foldable (Foldable)
import qualified Data.List
import qualified Data.List.NonEmpty
import Data.List.NonEmpty (NonEmpty(..))
import qualified Data.Map.Strict as Data.Map
import Data.Map.Strict (Map)
-import Data.Monoid (Monoid(..))
-import Data.Traversable (Traversable(..))
+-- import Data.Monoid (Monoid(..))
+import qualified Data.Strict.Maybe as Strict
+-- import Data.Traversable (Traversable(..))
import Data.Typeable (Typeable)
import Prelude hiding (filter, null, reverse)
+import qualified Hcompta.Lib.Strict as Strict ()
+
-- * Type 'TreeMap'
newtype TreeMap k x
instance (Ord k, Monoid v) => Monoid (TreeMap k v) where
mempty = empty
- mappend = union const
+ mappend = union mappend
-- mconcat = Data.List.foldr mappend mempty
instance Ord k => Functor (TreeMap k) where
fmap f (TreeMap m) = TreeMap $ fmap (fmap f) m
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'
data Ord k
=> Node k x
= Node
- { node_size :: Int -- ^ The number of non-'Nothing' 'node_value's reachable from this 'Node'.
- , node_value :: Maybe x -- ^ Some value, or 'Nothing' if this 'Node' is intermediary.
- , node_descendants :: TreeMap k x -- ^ Descendants 'Node's.
+ { node_size :: !Int -- ^ The number of non-'Strict.Nothing' 'node_value's reachable from this 'Node'.
+ , node_value :: !(Strict.Maybe x) -- ^ Some value, or 'Strict.Nothing' if this 'Node' is intermediary.
+ , node_descendants :: !(TreeMap k x) -- ^ Descendants 'Node's.
} deriving (Data, Eq, Read, Show, Typeable)
+
instance (Ord k, Monoid v) => Monoid (Node k v) where
mempty =
Node
- { node_value = Nothing
+ { node_value = Strict.Nothing
, node_size = 0
, node_descendants = TreeMap mempty
}
let x = x0 `mappend` x1 in
Node
{ node_value = x
- , node_size = size m + maybe 0 (const 1) x
+ , node_size = size m + Strict.maybe 0 (const 1) x
, node_descendants = union const m0 m1
}
-- mconcat = Data.List.foldr mappend mempty
-
instance Ord k => Functor (Node k) where
fmap f Node{node_value=x, node_descendants=m, node_size} =
Node
, node_descendants = Hcompta.Lib.TreeMap.map f m
, node_size
}
-
instance Ord k => Foldable (Node k) where
- foldMap f Node{node_value=Nothing, node_descendants=TreeMap m} =
+ foldMap f Node{node_value=Strict.Nothing, node_descendants=TreeMap m} =
foldMap (foldMap f) m
- foldMap f Node{node_value=Just x, node_descendants=TreeMap 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=Nothing, node_descendants=TreeMap m, node_size} =
- Node node_size <$> pure Nothing <*> (TreeMap <$> traverse (traverse f) m)
- traverse f Node{node_value=Just x, node_descendants=TreeMap m, node_size} =
- Node node_size <$> (Just <$> f x) <*> (TreeMap <$> traverse (traverse f) m)
+ 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
-- * Construct
leaf :: Ord k => x -> Node k x
leaf x =
Node
- { node_value = Just x
+ { node_value = Strict.Just x
, node_descendants = empty
, node_size = 1
}
-- | Return the given 'TreeMap' associating the given 'Path' with the given value,
-- merging values if the given 'TreeMap' already associates the given 'Path'
--- with a non-'Nothing' 'node_value'.
+-- with a non-'Strict.Nothing' 'node_value'.
insert :: Ord k => (x -> x -> x) -> Path k -> x -> TreeMap k x -> TreeMap k x
insert merge (k:|[]) x (TreeMap m) =
TreeMap $
Data.Map.insertWith
- (\_ Node{node_value=x1, node_descendants=m1, node_size=s1} ->
+ (\_ Node{node_value = x1, node_descendants = m1, node_size = s1} ->
Node
- { node_value = maybe (Just x) (Just . merge x) x1
+ { node_value = Strict.maybe (Strict.Just x) (Strict.Just . merge x) x1
, node_descendants = m1
- , node_size = maybe (s1 + 1) (const s1) x1
+ , node_size = Strict.maybe (s1 + 1) (const s1) x1
})
k (leaf x) m
insert merge (k:|k':ks) x (TreeMap m) =
TreeMap $
Data.Map.insertWith
- (\_ Node{node_value=x1, node_descendants=m1} ->
- let m' = insert merge (path k' ks) x m1 in
- Node{node_value=x1, node_descendants=m', node_size=size m' + maybe 0 (const 1) x1})
+ (\_ Node{node_value = x1, node_descendants = m1} ->
+ let m' = insert merge (path k' ks) x $ m1 in
+ let s' = size m' + Strict.maybe 0 (const 1) x1 in
+ Node{node_value=x1, node_descendants=m', node_size=s'})
k
(Node
- { node_value = Nothing
+ { node_value = Strict.Nothing
, node_descendants = insert merge (path k' ks) x empty
, node_size = 1
})
null :: TreeMap k x -> Bool
null (TreeMap m) = Data.Map.null m
--- | Return the number of non-'Nothing' 'node_value's in the given 'TreeMap'.
+-- | Return the number of non-'Strict.Nothing' 'node_value's in the given 'TreeMap'.
--
-- * Complexity: O(r) where r is the size of the root 'Map'.
size :: Ord k => TreeMap k x -> Int
-- * Find
-- | Return the value (if any) associated with the given 'Path'.
-find :: Ord k => Path k -> TreeMap k x -> Maybe x
-find (k:|[]) (TreeMap m) = maybe Nothing node_value $ Data.Map.lookup k m
+find :: Ord k => Path k -> TreeMap k x -> Strict.Maybe x
+find (k:|[]) (TreeMap m) = maybe Strict.Nothing node_value $ Data.Map.lookup k m
find (k:|k':ks) (TreeMap m) =
- maybe Nothing (find (path k' ks) . node_descendants) $
+ maybe Strict.Nothing (find (path k' ks) . node_descendants) $
Data.Map.lookup k m
-- * Union
-- | Return a 'TreeMap' associating the same 'Path's as both given 'TreeMap's,
-- merging values (in respective order) when a 'Path' leads
--- to a non-'Nothing' 'node_value' in both given 'TreeMap's.
+-- to a non-'Strict.Nothing' 'node_value' in both given 'TreeMap's.
union :: Ord k => (x -> x -> x) -> TreeMap k x -> TreeMap k x -> TreeMap k x
union merge (TreeMap tm0) (TreeMap tm1) =
TreeMap $
(\Node{node_value=x0, node_descendants=m0}
Node{node_value=x1, node_descendants=m1} ->
let m = union merge m0 m1 in
- let x = maybe x1 (\x0' -> maybe (Just x0') (Just . merge x0') x1) x0 in
+ let x = Strict.maybe x1 (\x0' -> Strict.maybe (Strict.Just x0') (Strict.Just . merge x0') x1) x0 in
Node
{ node_value = x
, node_descendants = m
- , node_size = size m + maybe 0 (const 1) x
+ , node_size = size m + Strict.maybe 0 (const 1) x
})
tm0 tm1
--
-- NOTE: use 'Data.List.foldl'' to reduce demand on the control-stack.
unions :: Ord k => (x -> x -> x) -> [TreeMap k x] -> TreeMap k x
-unions merge ts = Data.List.foldl' (union merge) empty ts
+unions merge = Data.List.foldl' (union merge) empty
-- foldl' :: (a -> b -> a) -> a -> [b] -> a
-- foldl' f = go
-- * Map
--- | Return the given 'TreeMap' with each non-'Nothing' 'node_value'
+-- | Return the given 'TreeMap' with each non-'Strict.Nothing' 'node_value'
-- mapped by the given function.
map :: Ord k => (x -> y) -> TreeMap k x -> TreeMap k y
map f =
TreeMap .
Data.Map.map
(\n@Node{node_value=x, node_descendants=m} ->
- n{ node_value=maybe Nothing (Just . f) x
+ n{ node_value=Strict.maybe Strict.Nothing (Strict.Just . f) x
, node_descendants=Hcompta.Lib.TreeMap.map f m
}) .
nodes
-- | Return the given 'TreeMap' with each 'node_value'
-- mapped by the given function supplied with
-- the already mapped 'node_descendants' of the current 'Node'.
-map_by_depth_first :: Ord k => (TreeMap k y -> Maybe x -> y) -> TreeMap k x -> TreeMap k y
+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 .
Data.Map.map
- (\n@Node{node_value, node_descendants} ->
+ (\Node{node_value, node_descendants} ->
let m = map_by_depth_first f node_descendants in
- let x = f m node_value in
- n{ node_value = Just x
+ Node
+ { node_value = Strict.Just $ f m node_value
, node_descendants = m
, node_size = size m + 1
}) .
nodes
+-- * Alter
+
+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 $
+ Data.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
+
-- * Fold
-- | Return the given accumulator folded by the given function
--- applied on non-'Nothing' 'node_value's
+-- applied on non-'Strict.Nothing' 'node_value's
-- from left to right through the given 'TreeMap'.
foldl_with_Path :: Ord k => (a -> Path k -> x -> a) -> a -> TreeMap k x -> a
foldl_with_Path =
foldp p fct a (TreeMap m) =
Data.Map.foldlWithKey
(\acc k Node{node_value, node_descendants} ->
- let acc' = maybe acc (fct acc (reverse $ path k p)) node_value in
+ let acc' = Strict.maybe acc (fct acc (reverse $ path k p)) node_value in
foldp (k:p) fct acc' node_descendants) a m
-- | Return the given accumulator folded by the given function
--- applied on non-'Nothing' 'Node's and 'node_value's
+-- applied on non-'Strict.Nothing' 'Node's and 'node_value's
-- from left to right through the given 'TreeMap'.
-foldl_with_Path_and_Node :: Ord k => (a -> Path k -> Node k x -> x -> a) -> a -> TreeMap k x -> a
+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 -> Path k -> Node k x -> x -> a)
+ => [k] -> (a -> Node k x -> Path k -> x -> a)
-> a -> TreeMap k x -> a
foldp p fct a (TreeMap m) =
Data.Map.foldlWithKey
(\acc k n@Node{node_value, node_descendants} ->
- let acc' = maybe acc (fct acc (reverse $ path k p) n) node_value in
+ let acc' = Strict.maybe acc (fct acc n (reverse $ path k p)) node_value in
foldp (k:p) fct acc' node_descendants) a m
-- | Return the given accumulator folded by the given function
--- applied on non-'Nothing' 'node_value's
+-- applied on non-'Strict.Nothing' 'node_value's
-- from right to left through the given 'TreeMap'.
foldr_with_Path :: Ord k => (Path k -> x -> a -> a) -> a -> TreeMap k x -> a
foldr_with_Path =
Data.Map.foldrWithKey
(\k Node{node_value, node_descendants} acc ->
let acc' = foldp (k:p) fct acc node_descendants in
- maybe acc' (\x -> fct (reverse $ path k p) x acc') node_value) a m
+ Strict.maybe acc' (\x -> fct (reverse $ path k p) x acc') node_value) a m
-- | Return the given accumulator folded by the given function
--- applied on non-'Nothing' 'Node's and 'node_value's
+-- applied on non-'Strict.Nothing' 'Node's and 'node_value's
-- from right to left through the given 'TreeMap'.
-foldr_with_Path_and_Node :: Ord k => (Path k -> Node k x -> x -> a -> a) -> a -> TreeMap k x -> a
+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] -> (Path k -> Node k x -> x -> a -> a)
+ => [k] -> (Node k x -> Path k -> x -> a -> a)
-> a -> TreeMap k x -> a
foldp p fct a (TreeMap m) =
Data.Map.foldrWithKey
(\k n@Node{node_value, node_descendants} acc ->
let acc' = foldp (k:p) fct acc node_descendants in
- maybe acc' (\x -> fct (reverse $ path k p) n x acc') node_value) a m
+ Strict.maybe acc' (\x -> fct n (reverse $ path k p) x acc') node_value) a m
+
+-- | Return the given accumulator folded by the given function
+-- applied on non-'Strict.Nothing' 'node_value's
+-- from left to right along the given 'Path'.
+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 Data.Map.lookup k t of
+ Nothing -> a
+ Just Node{node_value=v, node_descendants=d} ->
+ case v of
+ Strict.Nothing -> go f (k:p) n d a
+ Strict.Just x -> go f (k:p) n d (f (reverse $ path k p) x a)
+
+-- | Return the given accumulator folded by the given function
+-- applied on non-'Strict.Nothing' 'node_value's
+-- from right to left along the given 'Path'.
+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 Data.Map.lookup k t of
+ Nothing -> a
+ Just Node{node_value=v, node_descendants=d} ->
+ case v of
+ Strict.Nothing -> go f (k:p) n d a
+ Strict.Just x -> f (reverse $ path k p) x $ go f (k:p) n d a
-- * Flatten
-- | Return a 'Map' associating each 'Path'
--- leading to a non-'Nothing' 'node_value' in the given 'TreeMap',
+-- leading to a non-'Strict.Nothing' 'node_value' in the given 'TreeMap',
-- with its value mapped by the given function.
flatten :: Ord k => (x -> y) -> TreeMap k x -> Map (Path k) y
-flatten =
+flatten = flatten_with_Path . const
+
+-- | Like 'flatten' but with also the current 'Path' given to the mapping function.
+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] -> (x -> y)
+ => [k] -> (Path k -> x -> y)
-> TreeMap k x
-> Map (Path k) y
flat_map p f (TreeMap m) =
Data.Map.unions $
(
Data.Map.mapKeysMonotonic (reverse . flip path p) $
- Data.Map.mapMaybe (\Node{node_value=x} -> f <$> x) m
+ Data.Map.mapMaybeWithKey (\k Node{node_value} ->
+ case node_value of
+ Strict.Nothing -> Nothing
+ Strict.Just x -> Just $ f (reverse $ path k p) x) m
) :
Data.Map.foldrWithKey
(\k -> (:) . flat_map (k:p) f . node_descendants)
-- * Filter
-- | Return the given 'TreeMap'
--- keeping only its non-'Nothing' 'node_value's
+-- keeping only its non-'Strict.Nothing' 'node_value's
-- passing the given predicate.
filter :: Ord k => (x -> Bool) -> TreeMap k x -> TreeMap k x
filter f =
map_Maybe_with_Path
- (\_p x -> if f x then Just x else Nothing)
+ (\_p x -> if f x then Strict.Just x else Strict.Nothing)
-- | Like 'filter' but with also the current 'Path' given to the predicate.
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 Just x else Nothing)
+ (\p x -> if f p x then Strict.Just x else Strict.Nothing)
+
+-- | Like 'filter_with_Path' but with also the current 'Node' given to the predicate.
+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)
-- | Return the given 'TreeMap'
--- mapping its non-'Nothing' 'node_value's
--- and keeping only the non-'Nothing' results.
-map_Maybe :: Ord k => (x -> Maybe y) -> TreeMap k x -> TreeMap k y
-map_Maybe f = map_Maybe_with_Path (const f)
+-- mapping its non-'Strict.Nothing' 'node_value's
+-- and keeping only the non-'Strict.Nothing' results.
+map_Maybe :: Ord k => (x -> Strict.Maybe y) -> TreeMap k x -> TreeMap k y
+map_Maybe = map_Maybe_with_Path . const
-- | Like 'map_Maybe' but with also the current 'Path' given to the predicate.
-map_Maybe_with_Path :: Ord k => (Path k -> x -> Maybe y) -> TreeMap k x -> TreeMap k y
-map_Maybe_with_Path =
+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
+
+-- | Like 'map_Maybe_with_Path' but with also the current 'Node' given to the predicate.
+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] -> (Path k -> x -> Maybe y)
+ => [k] -> (Node k x -> Path k -> x -> Strict.Maybe y)
-> TreeMap k x
-> TreeMap k y
go p test (TreeMap m) =
TreeMap $
Data.Map.mapMaybeWithKey
- (\k Node{node_value=v, node_descendants=ns} ->
+ (\k node@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
- Just x ->
- let node_value = test (reverse $ path k p) x in
+ Strict.Just x ->
+ let node_value = test node (reverse $ path k p) x in
case node_value of
- Nothing | null node_descendants -> Nothing
- Nothing -> Just Node{node_value, node_descendants, node_size=1 + node_size}
- Just _ -> Just Node{node_value, node_descendants, node_size}
+ 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=Nothing, node_descendants, node_size}
+ else Just Node{node_value=Strict.Nothing, node_descendants, node_size}
) m