{-# LANGUAGE DeriveTraversable #-}

{- |
Description: Implementation of a rose tree we use for 'AST's.
License: GNU LGPLv3
Maintainer: paul.bittner@uni-ulm.de

Implementation of a rose tree (i.e., a tree whose nodes can have an arbitrary amount of children, including 0) we use for 'AST's .
-}
module Tree where

import Data.List
import Data.Maybe
import Data.Set

import ListUtil
import Util

-- | Rose tree of elements of type a where children of nodes are given as lists.
data Tree a = Tree a [Tree a] deriving (Tree a -> Tree a -> Bool
(Tree a -> Tree a -> Bool)
-> (Tree a -> Tree a -> Bool) -> Eq (Tree a)
forall a. Eq a => Tree a -> Tree a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Tree a -> Tree a -> Bool
$c/= :: forall a. Eq a => Tree a -> Tree a -> Bool
== :: Tree a -> Tree a -> Bool
$c== :: forall a. Eq a => Tree a -> Tree a -> Bool
Eq, Functor Tree
Foldable Tree
(Functor Tree, Foldable Tree) =>
(forall (f :: * -> *) a b.
 Applicative f =>
 (a -> f b) -> Tree a -> f (Tree b))
-> (forall (f :: * -> *) a.
    Applicative f =>
    Tree (f a) -> f (Tree a))
-> (forall (m :: * -> *) a b.
    Monad m =>
    (a -> m b) -> Tree a -> m (Tree b))
-> (forall (m :: * -> *) a. Monad m => Tree (m a) -> m (Tree a))
-> Traversable Tree
(a -> f b) -> Tree a -> f (Tree b)
forall (t :: * -> *).
(Functor t, Foldable t) =>
(forall (f :: * -> *) a b.
 Applicative f =>
 (a -> f b) -> t a -> f (t b))
-> (forall (f :: * -> *) a. Applicative f => t (f a) -> f (t a))
-> (forall (m :: * -> *) a b.
    Monad m =>
    (a -> m b) -> t a -> m (t b))
-> (forall (m :: * -> *) a. Monad m => t (m a) -> m (t a))
-> Traversable t
forall (m :: * -> *) a. Monad m => Tree (m a) -> m (Tree a)
forall (f :: * -> *) a. Applicative f => Tree (f a) -> f (Tree a)
forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Tree a -> m (Tree b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Tree a -> f (Tree b)
sequence :: Tree (m a) -> m (Tree a)
$csequence :: forall (m :: * -> *) a. Monad m => Tree (m a) -> m (Tree a)
mapM :: (a -> m b) -> Tree a -> m (Tree b)
$cmapM :: forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Tree a -> m (Tree b)
sequenceA :: Tree (f a) -> f (Tree a)
$csequenceA :: forall (f :: * -> *) a. Applicative f => Tree (f a) -> f (Tree a)
traverse :: (a -> f b) -> Tree a -> f (Tree b)
$ctraverse :: forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Tree a -> f (Tree b)
$cp2Traversable :: Foldable Tree
$cp1Traversable :: Functor Tree
Traversable)

instance Show a => Show (Tree a) where
  show :: Tree a -> String
show = Int -> ShowS -> (a -> String) -> Tree a -> String
forall a b.
(Show a, Monoid b) =>
Int -> (String -> b) -> (a -> b) -> Tree a -> b
prettyPrint 0 ShowS
forall a. a -> a
id a -> String
forall a. Show a => a -> String
show

instance Functor Tree where
  fmap :: (a -> b) -> Tree a -> Tree b
fmap f :: a -> b
f (Tree n :: a
n c :: [Tree a]
c) = b -> [Tree b] -> Tree b
forall a. a -> [Tree a] -> Tree a
Tree (a -> b
f a
n) ((Tree a -> Tree b) -> [Tree a] -> [Tree b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((a -> b) -> Tree a -> Tree b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f) [Tree a]
c)

instance Foldable Tree where
  foldMap :: (a -> m) -> Tree a -> m
foldMap f :: a -> m
f {- a to Monoid -} (Tree x :: a
x c :: [Tree a]
c) = m -> m -> m
forall a. Monoid a => a -> a -> a
mappend (a -> m
f a
x) ([m] -> m
forall a. Monoid a => [a] -> a
mconcat ([m] -> m) -> [m] -> m
forall a b. (a -> b) -> a -> b
$ (Tree a -> m) -> [Tree a] -> [m]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((a -> m) -> Tree a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap a -> m
f) [Tree a]
c)

instance Applicative Tree where
   pure :: a -> Tree a
pure a :: a
a = a -> [Tree a] -> Tree a
forall a. a -> [Tree a] -> Tree a
Tree a
a []
   (Tree f :: a -> b
f cf :: [Tree (a -> b)]
cf) <*> :: Tree (a -> b) -> Tree a -> Tree b
<*> (Tree x :: a
x cx :: [Tree a]
cx) = b -> [Tree b] -> Tree b
forall a. a -> [Tree a] -> Tree a
Tree (a -> b
f a
x) (((Tree a -> Tree b) -> [Tree a] -> [Tree b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((a -> b) -> Tree a -> Tree b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f) [Tree a]
cx)[Tree b] -> [Tree b] -> [Tree b]
forall a. [a] -> [a] -> [a]
++((Tree (a -> b) -> [Tree b]) -> [Tree (a -> b)] -> [Tree b]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (\g :: Tree (a -> b)
g -> (Tree a -> Tree b) -> [Tree a] -> [Tree b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\c :: Tree a
c -> Tree (a -> b)
g Tree (a -> b) -> Tree a -> Tree b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Tree a
c) [Tree a]
cx) [Tree (a -> b)]
cf))

-- | Returns true iff the given tree is a leaf (i.e., it has no children).
isleaf :: Tree a -> Bool
isleaf :: Tree a -> Bool
isleaf (Tree _ children :: [Tree a]
children) = [Tree a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
Data.List.null [Tree a]
children

-- | Returns the value stored in the root of the given tree.
element :: Tree a -> a
element :: Tree a -> a
element (Tree n :: a
n _) = a
n

-- | Finds a subtree in the given tree (first argument) whose root has the given value (second argument).
-- Throws an error iff no such tree exists.
tree :: (Eq a, Show a) => Tree a -> a -> Tree a
tree :: Tree a -> a -> Tree a
tree t :: Tree a
t x :: a
x = case Tree a -> a -> Maybe (Tree a)
forall a. Eq a => Tree a -> a -> Maybe (Tree a)
safetree Tree a
t a
x of
  Nothing -> String -> Tree a
forall a. HasCallStack => String -> a
error (String -> Tree a) -> String -> Tree a
forall a b. (a -> b) -> a -> b
$ "The element "String -> ShowS
forall a. [a] -> [a] -> [a]
++(a -> String
forall a. Show a => a -> String
show a
x)String -> ShowS
forall a. [a] -> [a] -> [a]
++" is not part of tree "String -> ShowS
forall a. [a] -> [a] -> [a]
++(Tree a -> String
forall a. Show a => a -> String
show Tree a
t)
  Just t' :: Tree a
t' -> Tree a
t'

-- | Same as 'tree' but returns @Nothing@ in the error case (i.e., when no subtree contains the given value).
safetree :: Eq a => Tree a -> a -> Maybe(Tree a)
safetree :: Tree a -> a -> Maybe (Tree a)
safetree t :: Tree a
t x :: a
x = (Tree a -> Bool) -> Tree a -> Maybe (Tree a)
forall a. (Tree a -> Bool) -> Tree a -> Maybe (Tree a)
Tree.find (\(Tree y :: a
y _) -> a
x a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
y) Tree a
t

-- | Transforms a tree into a set.
-- The returned set contains exactly the values previously held in the tree.
toset :: Ord a => Tree a -> Set a
toset :: Tree a -> Set a
toset = [a] -> Set a
forall a. Ord a => [a] -> Set a
fromList([a] -> Set a) -> (Tree a -> [a]) -> Tree a -> Set a
forall b c a. (b -> c) -> (a -> b) -> a -> c
.((a -> [a]) -> Tree a -> [a]
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap a -> [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure)

-- | Find a subtree in the given tree (first argument) whose root matches the predicate (second argument).
find :: (Tree a -> Bool) -> Tree a -> Maybe(Tree a)
find :: (Tree a -> Bool) -> Tree a -> Maybe (Tree a)
find predicate :: Tree a -> Bool
predicate x :: Tree a
x@(Tree _ children :: [Tree a]
children) = case Tree a -> Bool
predicate Tree a
x of
  True -> Tree a -> Maybe (Tree a)
forall a. a -> Maybe a
Just Tree a
x
  False -> [Tree a] -> Maybe (Tree a)
forall a. [a] -> Maybe a
safehead ([Tree a] -> Maybe (Tree a)) -> [Tree a] -> Maybe (Tree a)
forall a b. (a -> b) -> a -> b
$ [Maybe (Tree a)] -> [Tree a]
forall a. [Maybe a] -> [a]
catMaybes ([Maybe (Tree a)] -> [Tree a]) -> [Maybe (Tree a)] -> [Tree a]
forall a b. (a -> b) -> a -> b
$ (Tree a -> Maybe (Tree a)) -> [Tree a] -> [Maybe (Tree a)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\t :: Tree a
t -> (Tree a -> Bool) -> Tree a -> Maybe (Tree a)
forall a. (Tree a -> Bool) -> Tree a -> Maybe (Tree a)
Tree.find Tree a -> Bool
predicate Tree a
t) [Tree a]
children

-- | Same as 'find' but takes a predicate over elements instead of a predicate over trees to identify the subtree of interest.
findByNode :: (a -> Bool) -> Tree a -> Maybe(Tree a)
findByNode :: (a -> Bool) -> Tree a -> Maybe (Tree a)
findByNode p :: a -> Bool
p = (Tree a -> Bool) -> Tree a -> Maybe (Tree a)
forall a. (Tree a -> Bool) -> Tree a -> Maybe (Tree a)
Tree.find (\(Tree n :: a
n _) -> a -> Bool
p a
n)

-- | Returns the parent of the given node (second argument) in the given tree (first argument).
-- The returned parent is a subtree of the first given tree and has the second given tree as child.
-- Throws an error iff no parent exists.
parent :: Eq a => Tree a -> Tree a -> Maybe(Tree a)
parent :: Tree a -> Tree a -> Maybe (Tree a)
parent root :: Tree a
root t :: Tree a
t = (Tree a -> Bool) -> Tree a -> Maybe (Tree a)
forall a. (Tree a -> Bool) -> Tree a -> Maybe (Tree a)
Tree.find (\(Tree _ children :: [Tree a]
children) -> Tree a -> [Tree a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
elem Tree a
t [Tree a]
children) Tree a
root

{- |
Retrieves all nodes that are above the given node (second argument) in the given tree (first argument)
(This is the transitive closure of 'parent'.)
-}
ancestors :: Eq a => Tree a -> Tree a -> [Tree a]
ancestors :: Tree a -> Tree a -> [Tree a]
ancestors root :: Tree a
root t :: Tree a
t = case Tree a -> Tree a -> Maybe (Tree a)
forall a. Eq a => Tree a -> Tree a -> Maybe (Tree a)
parent Tree a
root Tree a
t of
  Nothing -> []
  Just p :: Tree a
p -> (Tree a -> Tree a -> [Tree a]
forall a. Eq a => Tree a -> Tree a -> [Tree a]
ancestors Tree a
root Tree a
p)[Tree a] -> [Tree a] -> [Tree a]
forall a. [a] -> [a] -> [a]
++[Tree a
p]

-- | Applies the given function to each node from bottom to top (i.e., leaves are converted first, root last).
manipulate :: (Tree a -> Tree a) -> Tree a -> Tree a
manipulate :: (Tree a -> Tree a) -> Tree a -> Tree a
manipulate f :: Tree a -> Tree a
f (Tree x :: a
x children :: [Tree a]
children) = Tree a -> Tree a
f (a -> [Tree a] -> Tree a
forall a. a -> [Tree a] -> Tree a
Tree a
x ((Tree a -> Tree a) -> [Tree a] -> [Tree a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Tree a -> Tree a) -> Tree a -> Tree a
forall a. (Tree a -> Tree a) -> Tree a -> Tree a
manipulate Tree a -> Tree a
f) [Tree a]
children))

{- |
Removes all subtrees not meeting the imposed condition.
The root remains untouched.
-}
filterTrees :: (Tree a -> Bool) -> Tree a -> Tree a
filterTrees :: (Tree a -> Bool) -> Tree a -> Tree a
filterTrees p :: Tree a -> Bool
p = (Tree a -> Tree a) -> Tree a -> Tree a
forall a. (Tree a -> Tree a) -> Tree a -> Tree a
manipulate (\(Tree n :: a
n c :: [Tree a]
c) -> a -> [Tree a] -> Tree a
forall a. a -> [Tree a] -> Tree a
Tree a
n ((Tree a -> Bool) -> [Tree a] -> [Tree a]
forall a. (a -> Bool) -> [a] -> [a]
Data.List.filter Tree a -> Bool
p [Tree a]
c))

{- |
Removes all nodes not meeting the imposed condition.
Children of removed nodes are moved up and become children of the parent of the removed node.
The root remains untouched.
-}
filterNodes :: (Tree a -> Bool) -> Tree a -> Tree a
filterNodes :: (Tree a -> Bool) -> Tree a -> Tree a
filterNodes p :: Tree a -> Bool
p = (Tree a -> Tree a) -> Tree a -> Tree a
forall a. (Tree a -> Tree a) -> Tree a -> Tree a
manipulate (\tree :: Tree a
tree@(Tree node :: a
node children :: [Tree a]
children) ->
    a -> [Tree a] -> Tree a
forall a. a -> [Tree a] -> Tree a
Tree a
node ([[Tree a]] -> [Tree a]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[Tree a]] -> [Tree a]) -> [[Tree a]] -> [Tree a]
forall a b. (a -> b) -> a -> b
$ (Tree a -> [Tree a]) -> [Tree a] -> [[Tree a]]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\c :: Tree a
c@(Tree _ cc :: [Tree a]
cc) -> if Tree a -> Bool
p Tree a
c then [Tree a
c] else [Tree a]
cc) [Tree a]
children))


-- | Pretty print function that folds the tree into a monoidial type @b@.
-- @b@ is the type that is the output of the print (e.g., String or Text).
-- First argument is the length of the indent in spaces (e.g., 2 will produce an indent of @"  "@ per level).
-- Second argument is a function that lifts strings to the output type b (e.g., a constructor of @b@ accepting a string).
-- Third argument is a function to convert tree elements to the output type (e.g., @show@ if @b@ is @String@).
prettyPrint :: (Show a, Monoid b) => Int -> (String -> b) -> (a -> b) -> Tree a -> b
prettyPrint :: Int -> (String -> b) -> (a -> b) -> Tree a -> b
prettyPrint i :: Int
i strToB :: String -> b
strToB nodePrinter :: a -> b
nodePrinter (Tree n :: a
n []) = (String -> b
strToB (String -> b) -> String -> b
forall a b. (a -> b) -> a -> b
$ Int -> String
genIndent Int
i) b -> b -> b
forall a. Semigroup a => a -> a -> a
<> (a -> b
nodePrinter a
n) b -> b -> b
forall a. Semigroup a => a -> a -> a
<>  (String -> b
strToB " []\n")
prettyPrint i :: Int
i strToB :: String -> b
strToB nodePrinter :: a -> b
nodePrinter (Tree n :: a
n children :: [Tree a]
children) = (String -> b
strToB (String -> b) -> String -> b
forall a b. (a -> b) -> a -> b
$ Int -> String
genIndent Int
i)
  b -> b -> b
forall a. Semigroup a => a -> a -> a
<> (a -> b
nodePrinter a
n)
  b -> b -> b
forall a. Semigroup a => a -> a -> a
<> (String -> b
strToB " [\n")
  b -> b -> b
forall a. Semigroup a => a -> a -> a
<> ([b] -> b
forall a. Monoid a => [a] -> a
mconcat ([b] -> b) -> [b] -> b
forall a b. (a -> b) -> a -> b
$ (Tree a -> b) -> [Tree a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Int -> (String -> b) -> (a -> b) -> Tree a -> b
forall a b.
(Show a, Monoid b) =>
Int -> (String -> b) -> (a -> b) -> Tree a -> b
prettyPrint (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+2) String -> b
strToB a -> b
nodePrinter) [Tree a]
children)
  b -> b -> b
forall a. Semigroup a => a -> a -> a
<> (String -> b
strToB (String -> b) -> String -> b
forall a b. (a -> b) -> a -> b
$ Int -> String
genIndent Int
i)
  b -> b -> b
forall a. Semigroup a => a -> a -> a
<> (String -> b
strToB "]\n")