-- |
-- Module      : ConClusion.BinaryTree
-- Description : Custom binary tree type with some special functions
-- Copyright   : Phillip Seeber, 2022
-- License     : AGPL-3
-- Maintainer  : phillip.seeber@googlemail.com
-- Stability   : experimental
-- Portability : POSIX, Windows
module ConClusion.BinaryTree
  ( BinTree (..),
    root,
    takeBranchesWhile,
    takeLeafyBranchesWhile,
  )
where

import Data.Aeson hiding (Array)
import Data.Massiv.Array as Massiv hiding (IndexException)
import RIO

-- | A binary tree.
data BinTree e = Leaf e | Node e (BinTree e) (BinTree e)
  deriving (BinTree e -> BinTree e -> Bool
forall e. Eq e => BinTree e -> BinTree e -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: BinTree e -> BinTree e -> Bool
$c/= :: forall e. Eq e => BinTree e -> BinTree e -> Bool
== :: BinTree e -> BinTree e -> Bool
$c== :: forall e. Eq e => BinTree e -> BinTree e -> Bool
Eq, Int -> BinTree e -> ShowS
forall e. Show e => Int -> BinTree e -> ShowS
forall e. Show e => [BinTree e] -> ShowS
forall e. Show e => BinTree e -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [BinTree e] -> ShowS
$cshowList :: forall e. Show e => [BinTree e] -> ShowS
show :: BinTree e -> String
$cshow :: forall e. Show e => BinTree e -> String
showsPrec :: Int -> BinTree e -> ShowS
$cshowsPrec :: forall e. Show e => Int -> BinTree e -> ShowS
Show, forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
forall e x. Rep (BinTree e) x -> BinTree e
forall e x. BinTree e -> Rep (BinTree e) x
$cto :: forall e x. Rep (BinTree e) x -> BinTree e
$cfrom :: forall e x. BinTree e -> Rep (BinTree e) x
Generic)

instance (FromJSON e) => FromJSON (BinTree e)

instance (ToJSON e) => ToJSON (BinTree e)

instance Functor BinTree where
  fmap :: forall a b. (a -> b) -> BinTree a -> BinTree b
fmap a -> b
f (Leaf a
a) = forall e. e -> BinTree e
Leaf (a -> b
f a
a)
  fmap a -> b
f (Node a
a BinTree a
l BinTree a
r) = forall e. e -> BinTree e -> BinTree e -> BinTree e
Node (a -> b
f a
a) (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f BinTree a
l) (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f BinTree a
r)

-- | Look at the root of a binary tree.
root :: BinTree e -> e
root :: forall e. BinTree e -> e
root (Leaf e
e) = e
e
root (Node e
e BinTree e
_ BinTree e
_) = e
e

-- | Steps down each branch of a tree until some criterion is satisfied or the 
-- end of the branch is reached. Each end of the branch is added to a result.
takeBranchesWhile :: (a -> Bool) -> BinTree a -> Massiv.Vector DL a
takeBranchesWhile :: forall a. (a -> Bool) -> BinTree a -> Vector DL a
takeBranchesWhile a -> Bool
chk BinTree a
tree = BinTree a -> Vector DL a -> Vector DL a
go BinTree a
tree (forall r ix e. Load r ix e => Array r ix e
Massiv.empty @DL)
  where
    go :: BinTree a -> Vector DL a -> Vector DL a
go (Leaf a
v) Vector DL a
acc = if a -> Bool
chk a
v then Vector DL a
acc forall r e.
(Size r, Load r Int e) =>
Vector r e -> e -> Vector DL e
`snoc` a
v else Vector DL a
acc
    go (Node a
v BinTree a
l BinTree a
r) Vector DL a
acc =
      let vAcc :: Vector DL a
vAcc = if a -> Bool
chk a
v then Vector DL a
acc forall r e.
(Size r, Load r Int e) =>
Vector r e -> e -> Vector DL e
`snoc` a
v else Vector DL a
acc
          lAcc :: Vector DL a
lAcc = BinTree a -> Vector DL a -> Vector DL a
go BinTree a
l Vector DL a
vAcc
          rAcc :: Vector DL a
rAcc = BinTree a -> Vector DL a -> Vector DL a
go BinTree a
r Vector DL a
lAcc
       in if a -> Bool
chk a
v then Vector DL a
rAcc else Vector DL a
vAcc

-- | Takes the first value in each branch, that does not fullfill the criterion 
-- anymore and adds it to the result. Terminal leafes of the branches are always
-- taken.
takeLeafyBranchesWhile :: (a -> Bool) -> BinTree a -> Massiv.Vector DL a
takeLeafyBranchesWhile :: forall a. (a -> Bool) -> BinTree a -> Vector DL a
takeLeafyBranchesWhile a -> Bool
chk BinTree a
tree = BinTree a -> Vector DL a -> Vector DL a
go BinTree a
tree (forall r ix e. Load r ix e => Array r ix e
Massiv.empty @DL)
  where
    go :: BinTree a -> Vector DL a -> Vector DL a
go (Leaf a
v) Vector DL a
acc = Vector DL a
acc forall r e.
(Size r, Load r Int e) =>
Vector r e -> e -> Vector DL e
`snoc` a
v
    go (Node a
v BinTree a
l BinTree a
r) Vector DL a
acc =
      let vAcc :: Vector DL a
vAcc = if a -> Bool
chk a
v then Vector DL a
acc else Vector DL a
acc forall r e.
(Size r, Load r Int e) =>
Vector r e -> e -> Vector DL e
`snoc` a
v
          lAcc :: Vector DL a
lAcc = BinTree a -> Vector DL a -> Vector DL a
go BinTree a
l Vector DL a
vAcc
          rAcc :: Vector DL a
rAcc = BinTree a -> Vector DL a -> Vector DL a
go BinTree a
r Vector DL a
lAcc
       in if a -> Bool
chk a
v then Vector DL a
rAcc else Vector DL a
vAcc