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_first :: Alternative f => Zipper k a -> f (Zipper k a)
  97 zipper_child_first (Zipper path (TreeMap m)) =
  98         case Map.minViewWithKey m of
  99          Nothing -> empty
 100          Just ((k', s'), fs') ->
 101                 pure Zipper
 102                  { zipper_path = Zipper_Step TreeMap.empty (k', s') (TreeMap fs') : path
 103                  , zipper_curr = TreeMap.node_descendants s'
 104                  }
 105
 106 zipper_child_last :: Alternative f => Zipper k a -> f (Zipper k a)
 107 zipper_child_last (Zipper path (TreeMap m)) =
 108         case Map.maxViewWithKey m of
 109          Nothing -> empty
 110          Just ((k', s'), ps') ->
 111                 pure Zipper
 112                  { zipper_path = Zipper_Step (TreeMap ps') (k', s') TreeMap.empty : path
 113                  , zipper_curr = TreeMap.node_descendants s'
 114                  }
 115
 116 -- ** Axis ancestor
 117
 118 zipper_ancestor :: Ord k => Zipper k a -> [Zipper k a]
 119 zipper_ancestor = zipper_collect_without_self zipper_parent
 120
 121 zipper_ancestor_or_self :: Ord k => Zipper k a -> [Zipper k a]
 122 zipper_ancestor_or_self = zipper_collect zipper_parent
 123
 124 -- ** Axis descendant
 125
 126 zipper_descendant_or_self :: Ord k => Zipper k a -> [Zipper k a]
 127 zipper_descendant_or_self =
 128         collect_child []
 129         where
 130                 collect_child acc z =
 131                         z : maybe acc
 132                          (collect_foll acc)
 133                          (zipper_child_first z)
 134                 collect_foll  acc z =
 135                         collect_child
 136                          (maybe acc
 137                                  (collect_foll acc)
 138                                  (zipper_foll z)
 139                          ) z
 140
 141 zipper_descendant_or_self_reverse :: Ord k => Zipper k a -> [Zipper k a]
 142 zipper_descendant_or_self_reverse z =
 143         z : List.concatMap
 144          zipper_descendant_or_self_reverse
 145          (List.reverse $ zipper_child z)
 146
 147 zipper_descendant :: Ord k => Zipper k a -> [Zipper k a]
 148 zipper_descendant = List.tail . zipper_descendant_or_self
 149
 150 zipper_descendant_lookup
 151  :: (Ord k, Alternative f, Monad f)
 152  => TreeMap.Path k -> Zipper k a -> f (Zipper k a)
 153 zipper_descendant_lookup (k:|ks) =
 154         case ks of
 155          []     -> zipper_child_lookup k
 156          k':ks' -> zipper_child_lookup k >=> zipper_descendant_lookup (k':|ks')
 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