{- |
Module      : ConClusion.BinaryTree
Description : Custom binary tree type with some special functions
Copyright   : Phillip Seeber, 2023
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