{-|
Module      : Control.Monad.Tree
Description : Implementation of a non-deterministic tree monad.
Copyright   : (c) Nathan Bedell, 2021
License     : MIT
Maintainer  : nbedell@tulane.edu

This module contains the definition of a simple rose-tree-like datatype, parameterized by
 both the type of the labels for the branches, the type of elements contained
 in the leaves of the datatype, and the type @f@ of containers used for the branching: @Tree b f a@. 
 
Fixing the type of labels (usually @ b ~ () @ for simplicity), we have instances for
 Functor, Applicative, Monad, and Alternative. This is similar to the list monad, but because
 the underlying datastructure is a tree, and we can represent branchind non-determinism,
 there is more flexibility -- as when trying to extract data from a @Tree@, one can choose
 a search strategy for flattening the tree.
 
This module provides a simple breadth-first-search strategy @bfs@, and a depth-first
 strategy @dfs@, however, theoretically other strategies could be used, taking advantage
 of the information provided by the type of labels `b` for the tree. For instance, with
 @b ~ Double@, a best-first or probabilistic search could be used.
-}
module Control.Monad.Tree (
  -- ** Tree data type
  Tree(..),
  -- ** Search algorithms
  dfs,
  dfs',
  bfs,
  bfs'
) where

import Control.Applicative
import Data.Functor.Classes

-- | Simple rose-tree data structure, with labels of 
--   type @n@ for nodel, labels of types @b@ for branches, containers of type @f@
--   used to branch on tree nodes, and values of type @a@ at the leaves.
--
--   @f@ will usually be @[]@, but this was kept generic to allow for
--   the use of more efficent containers if relevant.
--
data Tree n f a = 
    Leaf a 
  | Node n (f (Tree n f a)) 
  
instance (Eq a, Eq n, Eq1 f) => Eq (Tree n f a) where
  (Leaf a
x) == :: Tree n f a -> Tree n f a -> Bool
== (Leaf a
y) = a
x a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
y
  (Node n
x f (Tree n f a)
xs) == (Node n
y f (Tree n f a)
ys) = n
x n -> n -> Bool
forall a. Eq a => a -> a -> Bool
== n
y Bool -> Bool -> Bool
&& (Tree n f a -> Tree n f a -> Bool)
-> f (Tree n f a) -> f (Tree n f a) -> Bool
forall (f :: * -> *) a b.
Eq1 f =>
(a -> b -> Bool) -> f a -> f b -> Bool
liftEq Tree n f a -> Tree n f a -> Bool
forall a. Eq a => a -> a -> Bool
(==) f (Tree n f a)
xs f (Tree n f a)
ys

instance Functor f => Functor (Tree n f) where
   fmap :: (a -> b) -> Tree n f a -> Tree n f b
fmap a -> b
f (Leaf a
x) = b -> Tree n f b
forall n (f :: * -> *) a. a -> Tree n f a
Leaf (a -> b
f a
x)
   fmap a -> b
f (Node n
x f (Tree n f a)
ys) = n -> f (Tree n f b) -> Tree n f b
forall n (f :: * -> *) a. n -> f (Tree n f a) -> Tree n f a
Node n
x (f (Tree n f b) -> Tree n f b) -> f (Tree n f b) -> Tree n f b
forall a b. (a -> b) -> a -> b
$ (a -> b
f (a -> b) -> Tree n f a -> Tree n f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>) (Tree n f a -> Tree n f b) -> f (Tree n f a) -> f (Tree n f b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> f (Tree n f a)
ys
 
-- Note: This does not take weights into account.
bind :: Functor f => Tree n f a -> (a -> Tree n f b) -> Tree n f b
bind :: Tree n f a -> (a -> Tree n f b) -> Tree n f b
bind (Leaf a
x) a -> Tree n f b
f = a -> Tree n f b
f a
x
bind (Node n
x f (Tree n f a)
ys) a -> Tree n f b
f = n -> f (Tree n f b) -> Tree n f b
forall n (f :: * -> *) a. n -> f (Tree n f a) -> Tree n f a
Node n
x (f (Tree n f b) -> Tree n f b) -> f (Tree n f b) -> Tree n f b
forall a b. (a -> b) -> a -> b
$ (\Tree n f a
x -> Tree n f a -> (a -> Tree n f b) -> Tree n f b
forall (f :: * -> *) n a b.
Functor f =>
Tree n f a -> (a -> Tree n f b) -> Tree n f b
bind Tree n f a
x a -> Tree n f b
f) (Tree n f a -> Tree n f b) -> f (Tree n f a) -> f (Tree n f b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> f (Tree n f a)
ys

instance Functor f => Applicative (Tree n f) where
  pure :: a -> Tree n f a
pure = a -> Tree n f a
forall n (f :: * -> *) a. a -> Tree n f a
Leaf
  <*> :: Tree n f (a -> b) -> Tree n f a -> Tree n f b
(<*>) Tree n f (a -> b)
fs Tree n f a
xs = 
      Tree n f a -> (a -> Tree n f b) -> Tree n f b
forall (f :: * -> *) n a b.
Functor f =>
Tree n f a -> (a -> Tree n f b) -> Tree n f b
bind Tree n f a
xs (\a
x -> 
      Tree n f (a -> b) -> ((a -> b) -> Tree n f b) -> Tree n f b
forall (f :: * -> *) n a b.
Functor f =>
Tree n f a -> (a -> Tree n f b) -> Tree n f b
bind Tree n f (a -> b)
fs (\a -> b
f ->
        b -> Tree n f b
forall (f :: * -> *) a. Applicative f => a -> f a
pure (b -> Tree n f b) -> b -> Tree n f b
forall a b. (a -> b) -> a -> b
$ a -> b
f a
x)
      )

instance Functor f => Monad (Tree n f) where
   return :: a -> Tree n f a
return = a -> Tree n f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
   >>= :: Tree n f a -> (a -> Tree n f b) -> Tree n f b
(>>=)  = Tree n f a -> (a -> Tree n f b) -> Tree n f b
forall (f :: * -> *) n a b.
Functor f =>
Tree n f a -> (a -> Tree n f b) -> Tree n f b
bind

instance Alternative f => Alternative (Tree () f) where
  empty :: Tree () f a
empty   = () -> f (Tree () f a) -> Tree () f a
forall n (f :: * -> *) a. n -> f (Tree n f a) -> Tree n f a
Node () f (Tree () f a)
forall (f :: * -> *) a. Alternative f => f a
empty
  Tree () f a
t <|> :: Tree () f a -> Tree () f a -> Tree () f a
<|> Tree () f a
s = () -> f (Tree () f a) -> Tree () f a
forall n (f :: * -> *) a. n -> f (Tree n f a) -> Tree n f a
Node () (f (Tree () f a) -> Tree () f a) -> f (Tree () f a) -> Tree () f a
forall a b. (a -> b) -> a -> b
$ (Tree () f a -> f (Tree () f a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Tree () f a
t) f (Tree () f a) -> f (Tree () f a) -> f (Tree () f a)
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> (Tree () f a -> f (Tree () f a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Tree () f a
s)
  
-- | dfs with a predicate on the labels
dfs' :: (n -> Bool) ->  Tree n [] a -> [a]
dfs' :: (n -> Bool) -> Tree n [] a -> [a]
dfs' n -> Bool
p (Leaf a
x) = [a
x]
dfs' n -> Bool
p (Node n
label [Tree n [] a]
st) = [a
s | Tree n [] a
t <- [Tree n [] a]
st, a
s <- (n -> Bool) -> Tree n [] a -> [a]
forall n a. (n -> Bool) -> Tree n [] a -> [a]
dfs' n -> Bool
p Tree n [] a
t, Tree n [] a -> Bool
forall (f :: * -> *) a. Tree n f a -> Bool
pTree Tree n [] a
t]
  where 
    pTree :: Tree n f a -> Bool
pTree (Leaf a
x) = Bool
True
    pTree (Node n
l f (Tree n f a)
_) = n -> Bool
p n
l
  
-- | Depth first search of a tree
dfs :: Tree n [] a -> [a]
dfs :: Tree n [] a -> [a]
dfs (Leaf a
x) = [a
x]
dfs (Node n
_ [Tree n [] a]
st) = [ a
s | Tree n [] a
t <- [Tree n [] a]
st, a
s <- Tree n [] a -> [a]
forall n a. Tree n [] a -> [a]
dfs Tree n [] a
t]

-- | Breadth first search with a predicate on the labels.
bfs' :: (n -> Bool) -> Tree n [] a -> [a]
bfs' :: (n -> Bool) -> Tree n [] a -> [a]
bfs' n -> Bool
p Tree n [] a
t = (n -> Bool) -> [Tree n [] a] -> [a]
forall t a. (t -> Bool) -> [Tree t [] a] -> [a]
trav n -> Bool
p [Tree n [] a
t]
  where 
    trav :: (t -> Bool) -> [Tree t [] a] -> [a]
trav t -> Bool
p [] = []
    trav t -> Bool
p ((Leaf a
x) : [Tree t [] a]
q) = a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
: (t -> Bool) -> [Tree t [] a] -> [a]
trav t -> Bool
p [Tree t [] a]
q
    trav t -> Bool
p ((Node t
label [Tree t [] a]
st) : [Tree t [] a]
q) 
      | t -> Bool
p t
label   = (t -> Bool) -> [Tree t [] a] -> [a]
trav t -> Bool
p ([Tree t [] a]
q [Tree t [] a] -> [Tree t [] a] -> [Tree t [] a]
forall a. [a] -> [a] -> [a]
++ [Tree t [] a]
st)
      | Bool
otherwise = (t -> Bool) -> [Tree t [] a] -> [a]
trav t -> Bool
p [Tree t [] a]
q
             
-- | Breadth first search algorithm for trees
bfs :: Tree n [] a -> [a]
bfs :: Tree n [] a -> [a]
bfs = (n -> Bool) -> Tree n [] a -> [a]
forall n a. (n -> Bool) -> Tree n [] a -> [a]
bfs' (\n
x -> Bool
True)