module Data.StableTree.Mutate
( insert
, delete
) where
import Data.StableTree.Build ( consume, consumeBranches', consumeMap, merge )
import Data.StableTree.Key ( StableKey )
import Data.StableTree.Properties ( bottomChildren, selectNode )
import Data.StableTree.Types
import qualified Data.Map as Map
import Data.Map ( Map )
insert :: (Ord k, StableKey k)
=> k -> v -> StableTree k v -> StableTree k v
insert k v (StableTree_C c) = (uncurry consume) $ insert' k v c
insert k v (StableTree_I i) = (uncurry consume) $ insert' k v i
delete :: (Ord k, StableKey k)
=> k -> StableTree k v -> StableTree k v
delete k (StableTree_C c) = (uncurry consume) $ delete' k c
delete k (StableTree_I i) = (uncurry consume) $ delete' k i
insert' :: (Ord k, StableKey k)
=> k
-> v
-> Tree d c k v
-> ([Tree d Complete k v], Maybe (Tree d Incomplete k v))
insert' k v = mutateBottom k $ Map.insert k v
delete' :: (Ord k, StableKey k)
=> k
-> Tree d c k v
-> ([Tree d Complete k v], Maybe (Tree d Incomplete k v))
delete' k = mutateBottom k $ Map.delete k
mutateBottom :: (Ord k, StableKey k)
=> k
-> (Map k v -> Map k v)
-> Tree d c k v
-> ([Tree d Complete k v], Maybe (Tree d Incomplete k v))
mutateBottom search_key mut_fn = \case
bottom@(Bottom _ _ _ _) -> consumeMap $ mut_fn $ bottomChildren bottom
bottom@(IBottom0 _) -> consumeMap $ mut_fn $ bottomChildren bottom
bottom@(IBottom1 _ _ _) -> consumeMap $ mut_fn $ bottomChildren bottom
branch@(Branch _ _ _ _ _) -> mutate search_key mut_fn branch
branch@(IBranch0 _ _) -> mutate search_key mut_fn branch
branch@(IBranch1 _ _ _) -> mutate search_key mut_fn branch
branch@(IBranch2 _ _ _ _ _) -> mutate search_key mut_fn branch
where
mutate :: (Ord k, StableKey k)
=> k
-> (Map k v -> Map k v)
-> Tree (S d) c k v
-> ([Tree (S d) Complete k v], Maybe (Tree (S d) Incomplete k v))
mutate key fn b =
case selectNode key b of
(Left (before, incomplete)) ->
let (mut_before, mut_minc) = mutateBottom key fn incomplete
in consumeBranches' (before++mut_before) mut_minc
(Right (before, tree, after, mincomplete)) ->
let (mut_before, mut_minc) = mutateBottom key fn tree
(merged_before, merged_minc) = merge (before++mut_before)
mut_minc
after
mincomplete
in consumeBranches' merged_before merged_minc