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