Data/TreeMap/Strict/Zipper.hs

   1 {-# LANGUAGE DeriveDataTypeable #-}
   2 {-# LANGUAGE NamedFieldPuns #-}
   3 {-# OPTIONS_GHC -fno-warn-tabs #-}
   4
   5 module Data.TreeMap.Strict.Zipper where
   6
   7 import           Control.Monad (Monad(..), (>=>))
   8 import           Data.Data (Data)
   9 import           Data.Eq (Eq)
  10 import           Data.Function (($), (.))
  11 import qualified Data.List as List
  12 import           Data.List.NonEmpty (NonEmpty(..))
  13 import qualified Data.Map.Strict as Map
  14 import           Data.Maybe (Maybe(..), maybe, maybeToList)
  15 import           Data.Ord (Ord(..))
  16 import           Data.Typeable (Typeable)
  17 import           Text.Show (Show(..))
  18
  19 import           Data.TreeMap.Strict (TreeMap(..))
  20 import qualified Data.TreeMap.Strict as TreeMap
  21
  22 -- * Type 'Zipper'
  23
  24 data Zipper k x
  25  =   Zipper
  26  {   zipper_path :: [Zipper_Step k x]
  27  ,   zipper_curr :: TreeMap k x
  28  } deriving (Data, Eq, Show, Typeable)
  29
  30 zipper :: TreeMap k x -> Zipper k x
  31 zipper = Zipper []
  32
  33 zipper_root :: Ord k => Zipper k x -> TreeMap k x
  34 zipper_root =
  35         zipper_curr . List.last .
  36         zipper_collect zipper_parent
  37
  38 -- * Type 'Zipper_Step'
  39
  40 data Zipper_Step k x
  41  =   Zipper_Step
  42  {   zipper_step_prec :: TreeMap k x
  43  ,   zipper_step_self :: (k, TreeMap.Node k x)
  44  ,   zipper_step_foll :: TreeMap k x
  45  } deriving (Data, Eq, Show, Typeable)
  46
  47 -- * Axis
  48
  49 -- | Collect all 'Zipper's along a given axis,
  50 --   including the first 'Zipper'.
  51 zipper_collect :: (z -> Maybe z) -> z -> [z]
  52 zipper_collect f z = z : maybe [] (zipper_collect f) (f z)
  53
  54 -- | Collect all 'Zipper's along a given axis,
  55 --   excluding the first 'Zipper'.
  56 zipper_collect_without_self :: (z -> Maybe z) -> z -> [z]
  57 zipper_collect_without_self f z = maybe [] (zipper_collect f) (f z)
  58
  59 -- ** Axis self
  60
  61 zipper_self :: Ord k => Zipper k x -> Maybe (k, TreeMap.Node k x)
  62 zipper_self z =
  63         case z of
  64          Zipper{ zipper_path=
  65                  Zipper_Step{zipper_step_self}
  66                  : _ } -> Just zipper_step_self
  67          _ -> Nothing
  68
  69 -- ** Axis child
  70
  71 zipper_child :: Ord k => Zipper k x -> [Zipper k x]
  72 zipper_child z =
  73         maybeToList (zipper_child_first z)
  74         >>= zipper_collect zipper_foll
  75
  76 zipper_child_at :: Ord k => k -> Zipper k x -> Maybe (Zipper k x)
  77 zipper_child_at k (Zipper path (TreeMap m)) =
  78         case Map.splitLookup k m of
  79          (_, Nothing, _) -> Nothing
  80          (ps, Just s, fs) ->
  81                 Just Zipper
  82                  { zipper_path = Zipper_Step (TreeMap ps) (k, s) (TreeMap fs) : path
  83                  , zipper_curr = TreeMap.node_descendants s
  84                  }
  85
  86 zipper_child_first :: Ord k => Zipper k x -> Maybe (Zipper k x)
  87 zipper_child_first (Zipper path (TreeMap m)) =
  88         case Map.minViewWithKey m of
  89          Nothing -> Nothing
  90          Just ((k', s'), fs') ->
  91                 Just Zipper
  92                  { zipper_path = Zipper_Step TreeMap.empty (k', s') (TreeMap fs') : path
  93                  , zipper_curr = TreeMap.node_descendants s'
  94                  }
  95
  96 zipper_child_last :: Ord k => Zipper k x -> Maybe (Zipper k x)
  97 zipper_child_last (Zipper path (TreeMap m)) =
  98         case Map.maxViewWithKey m of
  99          Nothing -> Nothing
 100          Just ((k', s'), ps') ->
 101                 Just Zipper
 102                  { zipper_path = Zipper_Step (TreeMap ps') (k', s') TreeMap.empty : path
 103                  , zipper_curr = TreeMap.node_descendants s'
 104                  }
 105
 106 -- ** Axis ancestor
 107
 108 zipper_ancestor :: Ord k => Zipper k x -> [Zipper k x]
 109 zipper_ancestor = zipper_collect_without_self zipper_parent
 110
 111 zipper_ancestor_or_self :: Ord k => Zipper k x -> [Zipper k x]
 112 zipper_ancestor_or_self = zipper_collect zipper_parent
 113
 114 -- ** Axis descendant
 115
 116 zipper_descendant_or_self :: Ord k => Zipper k x -> [Zipper k x]
 117 zipper_descendant_or_self =
 118         collect_child []
 119         where
 120                 collect_child acc z =
 121                         z : maybe acc
 122                          (collect_foll acc)
 123                          (zipper_child_first z)
 124                 collect_foll  acc z =
 125                         collect_child
 126                          (maybe acc
 127                                  (collect_foll acc)
 128                                  (zipper_foll z)
 129                          ) z
 130
 131 zipper_descendant_or_self_reverse :: Ord k => Zipper k x -> [Zipper k x]
 132 zipper_descendant_or_self_reverse z =
 133         z : List.concatMap
 134          zipper_descendant_or_self_reverse
 135          (List.reverse $ zipper_child z)
 136
 137 zipper_descendant :: Ord k => Zipper k x -> [Zipper k x]
 138 zipper_descendant = List.tail . zipper_descendant_or_self
 139
 140 zipper_descendant_at :: Ord k => TreeMap.Path k -> Zipper k x -> Maybe (Zipper k x)
 141 zipper_descendant_at (k:|ks) =
 142         case ks of
 143          []     -> zipper_child_at k
 144          k':ks' -> zipper_child_at k >=> zipper_descendant_at (k':|ks')
 145
 146 -- ** Axis preceding
 147
 148 zipper_prec :: Ord k => Zipper k x -> Maybe (Zipper k x)
 149 zipper_prec (Zipper path _curr) =
 150         case path of
 151          [] -> Nothing
 152          Zipper_Step (TreeMap ps) (k, s) (TreeMap fs):steps ->
 153                 case Map.maxViewWithKey ps of
 154                  Nothing -> Nothing
 155                  Just ((k', s'), ps') ->
 156                         Just Zipper
 157                          { zipper_path = Zipper_Step (TreeMap ps')
 158                                                      (k', s')
 159                                                      (TreeMap $ Map.insert k s fs)
 160                                          : steps
 161                          , zipper_curr = TreeMap.node_descendants s'
 162                          }
 163
 164 zipper_preceding :: Ord k => Zipper k x -> [Zipper k x]
 165 zipper_preceding =
 166         zipper_ancestor_or_self >=>
 167         zipper_preceding_sibling >=>
 168         zipper_descendant_or_self_reverse
 169
 170 zipper_preceding_sibling :: Ord k => Zipper k x -> [Zipper k x]
 171 zipper_preceding_sibling = zipper_collect_without_self zipper_prec
 172
 173 -- ** Axis following
 174
 175 zipper_foll :: Ord k => Zipper k x -> Maybe (Zipper k x)
 176 zipper_foll (Zipper path _curr) =
 177         case path of
 178          [] -> Nothing
 179          Zipper_Step (TreeMap ps) (k, s) (TreeMap fs):steps ->
 180                 case Map.minViewWithKey fs of
 181                  Nothing -> Nothing
 182                  Just ((k', s'), fs') ->
 183                         Just Zipper
 184                          { zipper_path = Zipper_Step (TreeMap $ Map.insert k s ps)
 185                                                      (k', s')
 186                                                      (TreeMap fs')
 187                                          : steps
 188                          , zipper_curr = TreeMap.node_descendants s'
 189                          }
 190
 191 zipper_following :: Ord k => Zipper k x -> [Zipper k x]
 192 zipper_following =
 193         zipper_ancestor_or_self >=>
 194         zipper_following_sibling >=>
 195         zipper_descendant_or_self
 196
 197 zipper_following_sibling :: Ord k => Zipper k x -> [Zipper k x]
 198 zipper_following_sibling = zipper_collect_without_self zipper_foll
 199
 200 -- ** Axis parent
 201
 202 zipper_parent :: Ord k => Zipper k x -> Maybe (Zipper k x)
 203 zipper_parent (Zipper path curr) =
 204         case path of
 205          [] -> Nothing
 206          Zipper_Step (TreeMap ps) (k, s) (TreeMap fs):steps ->
 207                 let node = TreeMap.Node
 208                          { TreeMap.node_value       = TreeMap.node_value s
 209                          , TreeMap.node_size        = TreeMap.size curr
 210                          , TreeMap.node_descendants = curr
 211                          } in
 212                 Just Zipper
 213                  { zipper_path = steps
 214                  , zipper_curr = TreeMap $ Map.union ps $ Map.insert k node fs
 215                  }