module Data.StableTree.Walk
( foldM
, foldr
, foldl
) where
import Data.StableTree.Properties ( bottomChildren, branchChildren )
import Data.StableTree.Types
import qualified Control.Monad as M
import qualified Data.List as List
import qualified Data.Map as Map
import qualified Prelude
import Prelude hiding ( foldr, foldl )
foldM :: (Monad m, Ord k)
=> (a -> k -> v -> m a) -> a -> StableTree k v -> m a
foldM fn a0 tree =
case tree of
StableTree_I i -> foldM' fn a0 i
StableTree_C c -> foldM' fn a0 c
where
foldM' :: (Monad m, Ord k)
=> (a -> k -> v -> m a) -> a -> Tree d c k v -> m a
foldM' f accum t =
case t of
Bottom _ _ _ _ -> bottom f accum t
IBottom0 _ -> bottom f accum t
IBottom1 _ _ _ -> bottom f accum t
Branch _ _ _ _ _ -> branch f accum t
IBranch0 _ _ -> branch f accum t
IBranch1 _ _ _ -> branch f accum t
IBranch2 _ _ _ _ _ -> branch f accum t
bottom :: (Monad m, Ord k)
=> (a -> k -> v -> m a) -> a -> Tree Z c k v -> m a
bottom f accum t =
let children = Map.assocs $ bottomChildren t
f' a (k,v) = f a k v
in M.foldM f' accum children
branch :: (Monad m, Ord k)
=> (a -> k -> v -> m a) -> a -> Tree (S d) c k v -> m a
branch f accum t = do
let (compMap, mi) = branchChildren t
accum' <- M.foldM (foldM' f) accum (map snd $ Map.elems compMap)
case mi of
Nothing -> return accum'
Just i -> foldM' f accum' (third i)
third (_, _, t) = t
foldr :: Ord k => (k -> v -> a -> a) -> a -> StableTree k v -> a
foldr fn a0 tree =
case tree of
StableTree_I i -> foldr' fn a0 i
StableTree_C c -> foldr' fn a0 c
where
foldr' :: Ord k => (k -> v -> a -> a) -> a -> Tree d c k v -> a
foldr' f accum t =
case t of
Bottom _ _ _ _ -> bottom f accum t
IBottom0 _ -> bottom f accum t
IBottom1 _ _ _ -> bottom f accum t
Branch _ _ _ _ _ -> branch f accum t
IBranch0 _ _ -> branch f accum t
IBranch1 _ _ _ -> branch f accum t
IBranch2 _ _ _ _ _ -> branch f accum t
bottom :: Ord k => (k -> v -> a -> a) -> a -> Tree Z c k v -> a
bottom f accum t =
let children = Map.assocs $ bottomChildren t
f' (k,v) a = f k v a
in Prelude.foldr f' accum children
branch :: Ord k => (k -> v -> a -> a) -> a -> Tree (S d) c k v -> a
branch f accum t =
let (compMap, mi) = branchChildren t
elems = map snd $ Map.elems compMap
accum' = case mi of
Nothing -> accum
Just i -> foldr' f accum (third i)
in Prelude.foldr (flip $ foldr' f) accum' elems
third (_, _, t) = t
foldl :: Ord k => (a -> k -> v -> a) -> a -> StableTree k v -> a
foldl fn a0 tree =
case tree of
StableTree_I i -> foldl' fn a0 i
StableTree_C c -> foldl' fn a0 c
where
foldl' :: Ord k => (a -> k -> v -> a) -> a -> Tree d c k v -> a
foldl' f accum t =
case t of
Bottom _ _ _ _ -> bottom f accum t
IBottom0 _ -> bottom f accum t
IBottom1 _ _ _ -> bottom f accum t
Branch _ _ _ _ _ -> branch f accum t
IBranch0 _ _ -> branch f accum t
IBranch1 _ _ _ -> branch f accum t
IBranch2 _ _ _ _ _ -> branch f accum t
bottom :: Ord k => (a -> k -> v -> a) -> a -> Tree Z c k v -> a
bottom f accum t =
let children = Map.assocs $ bottomChildren t
f' a (k,v) = f a k v
in List.foldl' f' accum children
branch :: Ord k => (a -> k -> v -> a) -> a -> Tree (S d) c k v -> a
branch f accum t =
let (compMap, mi) = branchChildren t
elems = map snd $ Map.elems compMap
accum' = List.foldl' (foldl' f) accum elems
in case mi of
Nothing -> accum'
Just i -> foldl' f accum' (third i)
third (_, _, t) = t