-- |
-- Module      :  ELynx.Tree.Zipper
-- Description :  Zippers on rooted rose trees with branch labels
-- Copyright   :  (c) Dominik Schrempf, 2021
-- License     :  GPL-3.0-or-later
--
-- Maintainer  :  dominik.schrempf@gmail.com
-- Stability   :  unstable
-- Portability :  portable
--
-- Creation date: Thu Jul 23 08:42:37 2020.
module ELynx.Tree.Zipper
  ( -- * Data type
    TreePos (..),

    -- * Conversion
    fromTree,
    toTree,

    -- * Movement
    goParent,
    goParentUnsafe,
    goRoot,
    goLeft,
    goRight,
    goChild,
    goChildUnsafe,

    -- * Paths
    Path,
    goPath,
    goPathUnsafe,
    getSubTreeUnsafe,
    isValidPath,
    isLeafPath,

    -- * Modification
    insertTree,
    modifyTree,
    insertBranch,
    insertLabel,
  )
where

import Data.Foldable
import ELynx.Tree.Rooted

-- | Tree zipper. For reference, please see http://hackage.haskell.org/package/rosezipper.
data TreePos e a = Pos
  { -- | The currently selected tree.
    TreePos e a -> Tree e a
current :: Tree e a,
    -- | Forest to the left in reversed order.
    TreePos e a -> Forest e a
before :: Forest e a,
    -- | Forest to the right
    TreePos e a -> Forest e a
after :: Forest e a,
    -- | Finger to the selected tree
    TreePos e a -> [([Tree e a], e, a, [Tree e a])]
parents :: [([Tree e a], e, a, [Tree e a])]
  }
  deriving (Int -> TreePos e a -> ShowS
[TreePos e a] -> ShowS
TreePos e a -> String
(Int -> TreePos e a -> ShowS)
-> (TreePos e a -> String)
-> ([TreePos e a] -> ShowS)
-> Show (TreePos e a)
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall e a. (Show e, Show a) => Int -> TreePos e a -> ShowS
forall e a. (Show e, Show a) => [TreePos e a] -> ShowS
forall e a. (Show e, Show a) => TreePos e a -> String
showList :: [TreePos e a] -> ShowS
$cshowList :: forall e a. (Show e, Show a) => [TreePos e a] -> ShowS
show :: TreePos e a -> String
$cshow :: forall e a. (Show e, Show a) => TreePos e a -> String
showsPrec :: Int -> TreePos e a -> ShowS
$cshowsPrec :: forall e a. (Show e, Show a) => Int -> TreePos e a -> ShowS
Show, TreePos e a -> TreePos e a -> Bool
(TreePos e a -> TreePos e a -> Bool)
-> (TreePos e a -> TreePos e a -> Bool) -> Eq (TreePos e a)
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall e a. (Eq e, Eq a) => TreePos e a -> TreePos e a -> Bool
/= :: TreePos e a -> TreePos e a -> Bool
$c/= :: forall e a. (Eq e, Eq a) => TreePos e a -> TreePos e a -> Bool
== :: TreePos e a -> TreePos e a -> Bool
$c== :: forall e a. (Eq e, Eq a) => TreePos e a -> TreePos e a -> Bool
Eq)

-- | Get a zipper pointing to the root.
fromTree :: Tree e a -> TreePos e a
fromTree :: Tree e a -> TreePos e a
fromTree Tree e a
t = Pos :: forall e a.
Tree e a
-> Forest e a
-> Forest e a
-> [(Forest e a, e, a, Forest e a)]
-> TreePos e a
Pos {current :: Tree e a
current = Tree e a
t, before :: Forest e a
before = [], after :: Forest e a
after = [], parents :: [(Forest e a, e, a, Forest e a)]
parents = []}

-- | Get the complete tree of the zipper.
toTree :: TreePos e a -> Tree e a
toTree :: TreePos e a -> Tree e a
toTree = TreePos e a -> Tree e a
forall e a. TreePos e a -> Tree e a
current (TreePos e a -> Tree e a)
-> (TreePos e a -> TreePos e a) -> TreePos e a -> Tree e a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TreePos e a -> TreePos e a
forall e a. TreePos e a -> TreePos e a
goRoot

getForest :: TreePos e a -> Forest e a
getForest :: TreePos e a -> Forest e a
getForest TreePos e a
pos = (Forest e a -> Tree e a -> Forest e a)
-> Forest e a -> Forest e a -> Forest e a
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl ((Tree e a -> Forest e a -> Forest e a)
-> Forest e a -> Tree e a -> Forest e a
forall a b c. (a -> b -> c) -> b -> a -> c
flip (:)) (TreePos e a -> Tree e a
forall e a. TreePos e a -> Tree e a
current TreePos e a
pos Tree e a -> Forest e a -> Forest e a
forall a. a -> [a] -> [a]
: TreePos e a -> Forest e a
forall e a. TreePos e a -> Forest e a
after TreePos e a
pos) (TreePos e a -> Forest e a
forall e a. TreePos e a -> Forest e a
before TreePos e a
pos)

-- | Go to parent.
goParent :: TreePos e a -> Maybe (TreePos e a)
goParent :: TreePos e a -> Maybe (TreePos e a)
goParent TreePos e a
pos = case TreePos e a -> [([Tree e a], e, a, [Tree e a])]
forall e a. TreePos e a -> [(Forest e a, e, a, Forest e a)]
parents TreePos e a
pos of
  ([Tree e a]
ls, e
br, a
lb, [Tree e a]
rs) : [([Tree e a], e, a, [Tree e a])]
ps ->
    TreePos e a -> Maybe (TreePos e a)
forall a. a -> Maybe a
Just
      Pos :: forall e a.
Tree e a
-> Forest e a
-> Forest e a
-> [(Forest e a, e, a, Forest e a)]
-> TreePos e a
Pos
        { current :: Tree e a
current = e -> a -> [Tree e a] -> Tree e a
forall e a. e -> a -> Forest e a -> Tree e a
Node e
br a
lb ([Tree e a] -> Tree e a) -> [Tree e a] -> Tree e a
forall a b. (a -> b) -> a -> b
$ TreePos e a -> [Tree e a]
forall e a. TreePos e a -> Forest e a
getForest TreePos e a
pos,
          before :: [Tree e a]
before = [Tree e a]
ls,
          after :: [Tree e a]
after = [Tree e a]
rs,
          parents :: [([Tree e a], e, a, [Tree e a])]
parents = [([Tree e a], e, a, [Tree e a])]
ps
        }
  [] -> Maybe (TreePos e a)
forall a. Maybe a
Nothing

-- | Go to parent.
--
-- Call 'error' if no parent is found.
goParentUnsafe :: TreePos e a -> TreePos e a
goParentUnsafe :: TreePos e a -> TreePos e a
goParentUnsafe TreePos e a
pos = case TreePos e a -> [([Tree e a], e, a, [Tree e a])]
forall e a. TreePos e a -> [(Forest e a, e, a, Forest e a)]
parents TreePos e a
pos of
  ([Tree e a]
ls, e
br, a
lb, [Tree e a]
rs) : [([Tree e a], e, a, [Tree e a])]
ps ->
    Pos :: forall e a.
Tree e a
-> Forest e a
-> Forest e a
-> [(Forest e a, e, a, Forest e a)]
-> TreePos e a
Pos
      { current :: Tree e a
current = e -> a -> [Tree e a] -> Tree e a
forall e a. e -> a -> Forest e a -> Tree e a
Node e
br a
lb ([Tree e a] -> Tree e a) -> [Tree e a] -> Tree e a
forall a b. (a -> b) -> a -> b
$ TreePos e a -> [Tree e a]
forall e a. TreePos e a -> Forest e a
getForest TreePos e a
pos,
        before :: [Tree e a]
before = [Tree e a]
ls,
        after :: [Tree e a]
after = [Tree e a]
rs,
        parents :: [([Tree e a], e, a, [Tree e a])]
parents = [([Tree e a], e, a, [Tree e a])]
ps
      }
  [] -> String -> TreePos e a
forall a. HasCallStack => String -> a
error String
"goUpUnsafe: No parent found."

-- | Go to root.
goRoot :: TreePos e a -> TreePos e a
goRoot :: TreePos e a -> TreePos e a
goRoot TreePos e a
pos = TreePos e a
-> (TreePos e a -> TreePos e a)
-> Maybe (TreePos e a)
-> TreePos e a
forall b a. b -> (a -> b) -> Maybe a -> b
maybe TreePos e a
pos TreePos e a -> TreePos e a
forall e a. TreePos e a -> TreePos e a
goRoot (TreePos e a -> Maybe (TreePos e a)
forall e a. TreePos e a -> Maybe (TreePos e a)
goParent TreePos e a
pos)

-- | Go to left sibling in current forest.
goLeft :: TreePos e a -> Maybe (TreePos e a)
goLeft :: TreePos e a -> Maybe (TreePos e a)
goLeft TreePos e a
pos =
  case TreePos e a -> Forest e a
forall e a. TreePos e a -> Forest e a
before TreePos e a
pos of
    Tree e a
t : Forest e a
ts ->
      TreePos e a -> Maybe (TreePos e a)
forall a. a -> Maybe a
Just
        TreePos e a
pos
          { current :: Tree e a
current = Tree e a
t,
            before :: Forest e a
before = Forest e a
ts,
            after :: Forest e a
after = TreePos e a -> Tree e a
forall e a. TreePos e a -> Tree e a
current TreePos e a
pos Tree e a -> Forest e a -> Forest e a
forall a. a -> [a] -> [a]
: TreePos e a -> Forest e a
forall e a. TreePos e a -> Forest e a
after TreePos e a
pos
          }
    [] -> Maybe (TreePos e a)
forall a. Maybe a
Nothing

-- | Go to right sibling in current forest.
goRight :: TreePos e a -> Maybe (TreePos e a)
goRight :: TreePos e a -> Maybe (TreePos e a)
goRight TreePos e a
pos =
  case TreePos e a -> Forest e a
forall e a. TreePos e a -> Forest e a
after TreePos e a
pos of
    Tree e a
t : Forest e a
ts ->
      TreePos e a -> Maybe (TreePos e a)
forall a. a -> Maybe a
Just
        TreePos e a
pos
          { current :: Tree e a
current = Tree e a
t,
            before :: Forest e a
before = TreePos e a -> Tree e a
forall e a. TreePos e a -> Tree e a
current TreePos e a
pos Tree e a -> Forest e a -> Forest e a
forall a. a -> [a] -> [a]
: TreePos e a -> Forest e a
forall e a. TreePos e a -> Forest e a
before TreePos e a
pos,
            after :: Forest e a
after = Forest e a
ts
          }
    [] -> Maybe (TreePos e a)
forall a. Maybe a
Nothing

-- | Go to child with given index in forest.
goChild :: Int -> TreePos e a -> Maybe (TreePos e a)
goChild :: Int -> TreePos e a -> Maybe (TreePos e a)
goChild Int
n TreePos e a
pos = case TreePos e a -> Tree e a
forall e a. TreePos e a -> Tree e a
current TreePos e a
pos of
  (Node e
br a
lb Forest e a
ts)
    | Forest e a -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null Forest e a
ts -> Maybe (TreePos e a)
forall a. Maybe a
Nothing
    | Forest e a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Forest e a
ts Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
n -> Maybe (TreePos e a)
forall a. Maybe a
Nothing
    | Bool
otherwise ->
      TreePos e a -> Maybe (TreePos e a)
forall a. a -> Maybe a
Just (TreePos e a -> Maybe (TreePos e a))
-> TreePos e a -> Maybe (TreePos e a)
forall a b. (a -> b) -> a -> b
$
        Pos :: forall e a.
Tree e a
-> Forest e a
-> Forest e a
-> [(Forest e a, e, a, Forest e a)]
-> TreePos e a
Pos
          { current :: Tree e a
current = Forest e a -> Tree e a
forall a. [a] -> a
head Forest e a
rs',
            before :: Forest e a
before = Forest e a -> Forest e a
forall a. [a] -> [a]
reverse Forest e a
ls',
            after :: Forest e a
after = Forest e a -> Forest e a
forall a. [a] -> [a]
tail Forest e a
rs',
            parents :: [(Forest e a, e, a, Forest e a)]
parents = (TreePos e a -> Forest e a
forall e a. TreePos e a -> Forest e a
before TreePos e a
pos, e
br, a
lb, TreePos e a -> Forest e a
forall e a. TreePos e a -> Forest e a
after TreePos e a
pos) (Forest e a, e, a, Forest e a)
-> [(Forest e a, e, a, Forest e a)]
-> [(Forest e a, e, a, Forest e a)]
forall a. a -> [a] -> [a]
: TreePos e a -> [(Forest e a, e, a, Forest e a)]
forall e a. TreePos e a -> [(Forest e a, e, a, Forest e a)]
parents TreePos e a
pos
          }
    where
      (Forest e a
ls', Forest e a
rs') = Int -> Forest e a -> (Forest e a, Forest e a)
forall a. Int -> [a] -> ([a], [a])
splitAt Int
n Forest e a
ts

-- | Go to child with given index in forest. Call 'error' if child does not
-- exist.
goChildUnsafe :: Int -> TreePos e a -> TreePos e a
goChildUnsafe :: Int -> TreePos e a -> TreePos e a
goChildUnsafe Int
n TreePos e a
pos = case TreePos e a -> Tree e a
forall e a. TreePos e a -> Tree e a
current TreePos e a
pos of
  (Node e
br a
lb Forest e a
ts)
    | Forest e a -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null Forest e a
ts -> String -> TreePos e a
forall a. HasCallStack => String -> a
error String
"goChildUnsafe: Forest is empty."
    | Forest e a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Forest e a
ts Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
n -> String -> TreePos e a
forall a. HasCallStack => String -> a
error String
"goChildUnsafe: Forest is too short."
    | Bool
otherwise ->
      Pos :: forall e a.
Tree e a
-> Forest e a
-> Forest e a
-> [(Forest e a, e, a, Forest e a)]
-> TreePos e a
Pos
        { current :: Tree e a
current = Forest e a -> Tree e a
forall a. [a] -> a
head Forest e a
rs',
          before :: Forest e a
before = Forest e a -> Forest e a
forall a. [a] -> [a]
reverse Forest e a
ls',
          after :: Forest e a
after = Forest e a -> Forest e a
forall a. [a] -> [a]
tail Forest e a
rs',
          parents :: [(Forest e a, e, a, Forest e a)]
parents = (TreePos e a -> Forest e a
forall e a. TreePos e a -> Forest e a
before TreePos e a
pos, e
br, a
lb, TreePos e a -> Forest e a
forall e a. TreePos e a -> Forest e a
after TreePos e a
pos) (Forest e a, e, a, Forest e a)
-> [(Forest e a, e, a, Forest e a)]
-> [(Forest e a, e, a, Forest e a)]
forall a. a -> [a] -> [a]
: TreePos e a -> [(Forest e a, e, a, Forest e a)]
forall e a. TreePos e a -> [(Forest e a, e, a, Forest e a)]
parents TreePos e a
pos
        }
    where
      (Forest e a
ls', Forest e a
rs') = Int -> Forest e a -> (Forest e a, Forest e a)
forall a. Int -> [a] -> ([a], [a])
splitAt Int
n Forest e a
ts

-- | Path from the root of a tree to the node of the tree.
--
-- The position is specific to a tree topology. If the topology changes, the
-- position becomes invalid.
type Path = [Int]

-- | Go to node with given path.
goPath :: Path -> TreePos e a -> Maybe (TreePos e a)
goPath :: Path -> TreePos e a -> Maybe (TreePos e a)
goPath Path
pos TreePos e a
pth = (TreePos e a -> Int -> Maybe (TreePos e a))
-> TreePos e a -> Path -> Maybe (TreePos e a)
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldlM ((Int -> TreePos e a -> Maybe (TreePos e a))
-> TreePos e a -> Int -> Maybe (TreePos e a)
forall a b c. (a -> b -> c) -> b -> a -> c
flip Int -> TreePos e a -> Maybe (TreePos e a)
forall e a. Int -> TreePos e a -> Maybe (TreePos e a)
goChild) TreePos e a
pth Path
pos

-- | Check if a path is valid in that it leads to a node on a tree.
isValidPath :: Tree e a -> Path -> Bool
isValidPath :: Tree e a -> Path -> Bool
isValidPath Tree e a
t Path
p = case Path -> TreePos e a -> Maybe (TreePos e a)
forall e a. Path -> TreePos e a -> Maybe (TreePos e a)
goPath Path
p (Tree e a -> TreePos e a
forall e a. Tree e a -> TreePos e a
fromTree Tree e a
t) of
  Maybe (TreePos e a)
Nothing -> Bool
False
  Just TreePos e a
_ -> Bool
True

-- | Check if a path leads to a leaf.
isLeafPath :: Tree e a -> Path -> Bool
isLeafPath :: Tree e a -> Path -> Bool
isLeafPath Tree e a
t Path
p = case Path -> TreePos e a -> Maybe (TreePos e a)
forall e a. Path -> TreePos e a -> Maybe (TreePos e a)
goPath Path
p (Tree e a -> TreePos e a
forall e a. Tree e a -> TreePos e a
fromTree Tree e a
t) of
  Maybe (TreePos e a)
Nothing -> Bool
False
  Just TreePos e a
pos -> [Tree e a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null ([Tree e a] -> Bool) -> [Tree e a] -> Bool
forall a b. (a -> b) -> a -> b
$ Tree e a -> [Tree e a]
forall e a. Tree e a -> Forest e a
forest (TreePos e a -> Tree e a
forall e a. TreePos e a -> Tree e a
current TreePos e a
pos)

-- | Got to node with given path.
--
-- Call 'error' if path is invalid.
goPathUnsafe :: Path -> TreePos e a -> TreePos e a
goPathUnsafe :: Path -> TreePos e a -> TreePos e a
goPathUnsafe Path
pos TreePos e a
pth =
  {-# SCC "goPathUnsafe" #-}
  (TreePos e a -> Int -> TreePos e a)
-> TreePos e a -> Path -> TreePos e a
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl ((Int -> TreePos e a -> TreePos e a)
-> TreePos e a -> Int -> TreePos e a
forall a b c. (a -> b -> c) -> b -> a -> c
flip Int -> TreePos e a -> TreePos e a
forall e a. Int -> TreePos e a -> TreePos e a
goChildUnsafe) TreePos e a
pth Path
pos

-- | Get the sub tree at path.
--
-- Call 'error' if path is invalid.
getSubTreeUnsafe :: Path -> Tree e a -> Tree e a
getSubTreeUnsafe :: Path -> Tree e a -> Tree e a
getSubTreeUnsafe Path
p = TreePos e a -> Tree e a
forall e a. TreePos e a -> Tree e a
current (TreePos e a -> Tree e a)
-> (Tree e a -> TreePos e a) -> Tree e a -> Tree e a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Path -> TreePos e a -> TreePos e a
forall e a. Path -> TreePos e a -> TreePos e a
goPathUnsafe Path
p (TreePos e a -> TreePos e a)
-> (Tree e a -> TreePos e a) -> Tree e a -> TreePos e a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Tree e a -> TreePos e a
forall e a. Tree e a -> TreePos e a
fromTree

-- | Insert a new tree into the current focus of the zipper.
insertTree :: Tree e a -> TreePos e a -> TreePos e a
insertTree :: Tree e a -> TreePos e a -> TreePos e a
insertTree Tree e a
t TreePos e a
pos = TreePos e a
pos {current :: Tree e a
current = Tree e a
t}

-- | Modify the tree at the current focus of the zipper.
modifyTree :: (Tree e a -> Tree e a) -> TreePos e a -> TreePos e a
modifyTree :: (Tree e a -> Tree e a) -> TreePos e a -> TreePos e a
modifyTree Tree e a -> Tree e a
f TreePos e a
pos = TreePos e a
pos {current :: Tree e a
current = Tree e a -> Tree e a
f Tree e a
t}
  where
    t :: Tree e a
t = TreePos e a -> Tree e a
forall e a. TreePos e a -> Tree e a
current TreePos e a
pos

-- | Insert a new branch label into the current focus of the zipper.
insertBranch :: e -> TreePos e a -> TreePos e a
insertBranch :: e -> TreePos e a -> TreePos e a
insertBranch e
br TreePos e a
pos = case TreePos e a -> Tree e a
forall e a. TreePos e a -> Tree e a
current TreePos e a
pos of
  Node e
_ a
lb Forest e a
ts -> TreePos e a
pos {current :: Tree e a
current = e -> a -> Forest e a -> Tree e a
forall e a. e -> a -> Forest e a -> Tree e a
Node e
br a
lb Forest e a
ts}

-- | Insert a new node label into the current focus of the zipper.
insertLabel :: a -> TreePos e a -> TreePos e a
insertLabel :: a -> TreePos e a -> TreePos e a
insertLabel a
lb TreePos e a
pos = case TreePos e a -> Tree e a
forall e a. TreePos e a -> Tree e a
current TreePos e a
pos of
  Node e
br a
_ Forest e a
ts -> TreePos e a
pos {current :: Tree e a
current = e -> a -> Forest e a -> Tree e a
forall e a. e -> a -> Forest e a -> Tree e a
Node e
br a
lb Forest e a
ts}