-- |
-- Module      : ConClusion.Numeric.Data
-- Description : Type castings, conversions and utilities
-- Copyright   : Phillip Seeber, 2021
-- License     : AGPL-3
-- Maintainer  : phillip.seeber@googlemail.com
-- Stability   : experimental
-- Portability : POSIX, Windows
module ConClusion.Numeric.Data
  ( -- * Conversion of array types.

    -- Conversion between HMatrix and Massiv's array types
    IndexException (..),
    vecH2M,
    vecM2H,
    matH2M,
    matM2H,

    -- * Array Processing
    magnitude,
    normalise,
    angle,
    minDistAt,
    minDistAtVec,
    iMinimumM,

    -- * Utilities
    printMat,

    -- * Binary Trees
    BinTree (..),
    root,
    takeBranchesWhile,
    takeLeafyBranchesWhile,
  )
where

import Data.Aeson hiding (Array)
import Data.Massiv.Array as Massiv hiding (IndexException)
import Data.Massiv.Array.Manifest.Vector as Massiv
import Formatting
import Numeric.LinearAlgebra as LA hiding (magnitude, (<>))
import RIO
import qualified RIO.Vector.Storable as VectorS
import System.IO.Unsafe (unsafePerformIO)

-- | Exception regarding indexing in some kind of aaray.
newtype IndexException = IndexException String deriving (Int -> IndexException -> ShowS
[IndexException] -> ShowS
IndexException -> String
(Int -> IndexException -> ShowS)
-> (IndexException -> String)
-> ([IndexException] -> ShowS)
-> Show IndexException
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [IndexException] -> ShowS
$cshowList :: [IndexException] -> ShowS
show :: IndexException -> String
$cshow :: IndexException -> String
showsPrec :: Int -> IndexException -> ShowS
$cshowsPrec :: Int -> IndexException -> ShowS
Show)

instance Exception IndexException

-- | Converts a vector from the HMatrix package to the Massiv representation.
{-# SCC vecH2M #-}
vecH2M :: (Element e, Mutable r Ix1 e) => VectorS.Vector e -> Massiv.Vector r e
vecH2M :: Vector e -> Vector r e
vecH2M Vector e
hVec = Comp -> Sz Int -> Vector e -> Vector r e
forall (v :: * -> *) a ix r.
(Typeable v, Vector v a, Mutable (ARepr v) ix a, Mutable r ix a) =>
Comp -> Sz ix -> v a -> Array r ix a
fromVector' Comp
Seq (Int -> Sz Int
forall ix. Index ix => ix -> Sz ix
Sz (Int -> Sz Int) -> Int -> Sz Int
forall a b. (a -> b) -> a -> b
$ Vector e -> Int
forall a. Storable a => Vector a -> Int
VectorS.length Vector e
hVec) Vector e
hVec

-- | Converts a vector from the Massiv representation to the HMatrix representation.
{-# SCC vecM2H #-}
vecM2H :: (Manifest r Ix1 e, Element e) => Massiv.Vector r e -> LA.Vector e
vecM2H :: Vector r e -> Vector e
vecM2H = Vector r e -> Vector e
forall r ix e (v :: * -> *).
(Manifest r ix e, Mutable (ARepr v) ix e, Vector v e,
 VRepr (ARepr v) ~ v) =>
Array r ix e -> v e
Massiv.toVector

-- | Converts a matrix from the HMatrix representation to the Massiv representation.
{-# SCC matH2M #-}
matH2M :: (Mutable r Ix1 e, Element e) => LA.Matrix e -> Massiv.Matrix r e
matH2M :: Matrix e -> Matrix r e
matH2M Matrix e
hMat = Sz Ix2 -> Array r Int e -> Matrix r e
forall ix' r ix e.
(Index ix', Load r ix e, Resize r ix) =>
Sz ix' -> Array r ix e -> Array r ix' e
Massiv.resize' (Ix2 -> Sz Ix2
forall ix. Index ix => ix -> Sz ix
Sz (Ix2 -> Sz Ix2) -> Ix2 -> Sz Ix2
forall a b. (a -> b) -> a -> b
$ Int
nRows Int -> Int -> Ix2
:. Int
nCols) (Array r Int e -> Matrix r e)
-> (Matrix e -> Array r Int e) -> Matrix e -> Matrix r e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector e -> Array r Int e
forall e r. (Element e, Mutable r Int e) => Vector e -> Vector r e
vecH2M (Vector e -> Array r Int e)
-> (Matrix e -> Vector e) -> Matrix e -> Array r Int e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Matrix e -> Vector e
forall t. Element t => Matrix t -> Vector t
LA.flatten (Matrix e -> Matrix r e) -> Matrix e -> Matrix r e
forall a b. (a -> b) -> a -> b
$ Matrix e
hMat
  where
    nRows :: Int
nRows = Matrix e -> Int
forall t. Matrix t -> Int
LA.rows Matrix e
hMat
    nCols :: Int
nCols = Matrix e -> Int
forall t. Matrix t -> Int
LA.cols Matrix e
hMat

-- | Converts a matrix from Massiv to HMatrix representation.
{-# SCC matM2H #-}
matM2H :: (Manifest r Ix1 e, Element e, Resize r Ix2, Load r Ix2 e) => Massiv.Matrix r e -> LA.Matrix e
matM2H :: Matrix r e -> Matrix e
matM2H Matrix r e
mMat = Int -> Vector e -> Matrix e
forall t. Storable t => Int -> Vector t -> Matrix t
LA.reshape Int
nCols (Vector e -> Matrix e)
-> (Matrix r e -> Vector e) -> Matrix r e -> Matrix e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector r e -> Vector e
forall r e. (Manifest r Int e, Element e) => Vector r e -> Vector e
vecM2H (Vector r e -> Vector e)
-> (Matrix r e -> Vector r e) -> Matrix r e -> Vector e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Matrix r e -> Vector r e
forall r ix e.
(Load r ix e, Resize r ix) =>
Array r ix e -> Array r Int e
Massiv.flatten (Matrix r e -> Matrix e) -> Matrix r e -> Matrix e
forall a b. (a -> b) -> a -> b
$ Matrix r e
mMat
  where
    Sz (Int
_nRows :. Int
nCols) = Matrix r e -> Sz Ix2
forall r ix e. Load r ix e => Array r ix e -> Sz ix
Massiv.size Matrix r e
mMat

-- | Magnitude of a vector (length).
magnitude :: (Massiv.Numeric r e, Source r Ix1 e, Floating e) => Massiv.Vector r e -> e
magnitude :: Vector r e -> e
magnitude Vector r e
v = e -> e
forall a. Floating a => a -> a
sqrt (e -> e) -> e -> e
forall a b. (a -> b) -> a -> b
$ Vector r e
v Vector r e -> Vector r e -> e
forall r e.
(Numeric r e, Source r Int e) =>
Vector r e -> Vector r e -> e
!.! Vector r e
v

-- | Normalise a vector.
normalise :: (Massiv.Numeric r e, Source r Ix1 e, Floating e) => Massiv.Vector r e -> Massiv.Vector r e
normalise :: Vector r e -> Vector r e
normalise Vector r e
v = Vector r e
v Vector r e -> e -> Vector r e
forall ix r e.
(Index ix, Numeric r e) =>
Array r ix e -> e -> Array r ix e
.* (e
1 e -> e -> e
forall a. Fractional a => a -> a -> a
/ Vector r e -> e
forall r e.
(Numeric r e, Source r Int e, Floating e) =>
Vector r e -> e
magnitude Vector r e
v)

-- | Angle between two vectors.
angle :: (Massiv.Numeric r e, Source r Ix1 e, Floating e) => Massiv.Vector r e -> Massiv.Vector r e -> e
angle :: Vector r e -> Vector r e -> e
angle Vector r e
a Vector r e
b = e -> e
forall a. Floating a => a -> a
acos (e -> e) -> e -> e
forall a b. (a -> b) -> a -> b
$ Vector r e
a Vector r e -> Vector r e -> e
forall r e.
(Numeric r e, Source r Int e) =>
Vector r e -> Vector r e -> e
!.! Vector r e
b e -> e -> e
forall a. Fractional a => a -> a -> a
/ (Vector r e -> e
forall r e.
(Numeric r e, Source r Int e, Floating e) =>
Vector r e -> e
magnitude Vector r e
a e -> e -> e
forall a. Num a => a -> a -> a
* Vector r e -> e
forall r e.
(Numeric r e, Source r Int e, Floating e) =>
Vector r e -> e
magnitude Vector r e
b)

-- | Find the minimal distance in a distance matrix, which is not the main diagonal.
{-# SCC minDistAt #-}
minDistAt ::
  ( Manifest r Ix2 e,
    MonadThrow m,
    Ord e
  ) =>
  Massiv.Matrix r e ->
  m (e, Ix2)
minDistAt :: Matrix r e -> m (e, Ix2)
minDistAt Matrix r e
arr
  | Matrix r e -> Bool
forall r ix e. Load r ix e => Array r ix e -> Bool
isEmpty Matrix r e
arr = SizeException -> m (e, Ix2)
forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM (SizeException -> m (e, Ix2)) -> SizeException -> m (e, Ix2)
forall a b. (a -> b) -> a -> b
$ Sz Ix2 -> SizeException
forall ix. Index ix => Sz ix -> SizeException
SizeEmptyException (Matrix r e -> Sz Ix2
forall r ix e. Load r ix e => Array r ix e -> Sz ix
Massiv.size Matrix r e
arr)
  | Bool
otherwise = (e, Ix2) -> m (e, Ix2)
forall (m :: * -> *) a. Monad m => a -> m a
return ((e, Ix2) -> m (e, Ix2))
-> (IO (e, Ix2) -> (e, Ix2)) -> IO (e, Ix2) -> m (e, Ix2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IO (e, Ix2) -> (e, Ix2)
forall a. IO a -> a
unsafePerformIO (IO (e, Ix2) -> m (e, Ix2)) -> IO (e, Ix2) -> m (e, Ix2)
forall a b. (a -> b) -> a -> b
$ ((e, Ix2) -> Ix2 -> e -> (e, Ix2))
-> (e, Ix2)
-> ((e, Ix2) -> (e, Ix2) -> (e, Ix2))
-> (e, Ix2)
-> Matrix r e
-> IO (e, Ix2)
forall (m :: * -> *) r ix e a b.
(MonadIO m, Source r ix e) =>
(a -> ix -> e -> a)
-> a -> (b -> a -> b) -> b -> Array r ix e -> m b
ifoldlP (e, Ix2) -> Ix2 -> e -> (e, Ix2)
forall a. Ord a => (a, Ix2) -> Ix2 -> a -> (a, Ix2)
minFold (e, Ix2)
start (e, Ix2) -> (e, Ix2) -> (e, Ix2)
forall a b. Ord a => (a, b) -> (a, b) -> (a, b)
chFold (e, Ix2)
start Matrix r e
arr
  where
    ix0 :: Ix2
ix0 = Int -> Ix2
forall ix. Index ix => Int -> ix
pureIndex Int
0
    e0 :: e
e0 = Matrix r e
arr Matrix r e -> Ix2 -> e
forall r ix e. Manifest r ix e => Array r ix e -> ix -> e
Massiv.! Ix2
ix0
    start :: (e, Ix2)
start = (e
e0, Ix2
ix0)
    minFold :: (a, Ix2) -> Ix2 -> a -> (a, Ix2)
minFold acc :: (a, Ix2)
acc@(a
eA, Ix2
_) ix :: Ix2
ix@(Int
m :. Int
n) a
e = if a
e a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
eA Bool -> Bool -> Bool
&& Int
m Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
n then (a
e, Ix2
ix) else (a, Ix2)
acc
    chFold :: (a, b) -> (a, b) -> (a, b)
chFold acc :: (a, b)
acc@(a
eA, b
_) ch :: (a, b)
ch@(a
e, b
_) = if a
e a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
eA then (a, b)
ch else (a, b)
acc

-- | Find the minimal element of a vector, which is at a larger than the supplied index.
minDistAtVec ::
  ( Manifest r Ix1 e,
    MonadThrow m,
    Ord e
  ) =>
  Ix1 ->
  Massiv.Vector r e ->
  m (e, Ix1)
minDistAtVec :: Int -> Vector r e -> m (e, Int)
minDistAtVec Int
ixStart Vector r e
vec
  | Vector r e -> Bool
forall r ix e. Load r ix e => Array r ix e -> Bool
isEmpty Vector r e
vec = SizeException -> m (e, Int)
forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM (SizeException -> m (e, Int)) -> SizeException -> m (e, Int)
forall a b. (a -> b) -> a -> b
$ Sz Int -> SizeException
forall ix. Index ix => Sz ix -> SizeException
SizeEmptyException (Vector r e -> Sz Int
forall r ix e. Load r ix e => Array r ix e -> Sz ix
Massiv.size Vector r e
vec)
  | Int
ixStart Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
nElems = IndexException -> m (e, Int)
forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM (IndexException -> m (e, Int)) -> IndexException -> m (e, Int)
forall a b. (a -> b) -> a -> b
$ Sz Int -> Int -> IndexException
forall ix. Index ix => Sz ix -> ix -> IndexException
IndexOutOfBoundsException (Int -> Sz Int
forall ix. Index ix => ix -> Sz ix
Sz Int
nElems) Int
ixStart
  | Bool
otherwise = do
    let (e
minE, Int
minIx) = IO (e, Int) -> (e, Int)
forall a. IO a -> a
unsafePerformIO (IO (e, Int) -> (e, Int)) -> IO (e, Int) -> (e, Int)
forall a b. (a -> b) -> a -> b
$ ((e, Int) -> Int -> e -> (e, Int))
-> (e, Int)
-> ((e, Int) -> (e, Int) -> (e, Int))
-> (e, Int)
-> Vector r e
-> IO (e, Int)
forall (m :: * -> *) r ix e a b.
(MonadIO m, Source r ix e) =>
(a -> ix -> e -> a)
-> a -> (b -> a -> b) -> b -> Array r ix e -> m b
ifoldlP (e, Int) -> Int -> e -> (e, Int)
forall a b. Ord a => (a, b) -> b -> a -> (a, b)
minFold (e, Int)
startAcc (e, Int) -> (e, Int) -> (e, Int)
forall a. Ord a => a -> a -> a
chFold (e, Int)
startAcc Vector r e
searchVec
    (e, Int) -> m (e, Int)
forall (m :: * -> *) a. Monad m => a -> m a
return (e
minE, Int
minIx Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
ixStart Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
  where
    Sz Int
nElems = Vector r e -> Sz Int
forall r ix e. Load r ix e => Array r ix e -> Sz ix
Massiv.size Vector r e
vec
    searchVec :: Vector r e
searchVec = Sz Int -> Vector r e -> Vector r e
forall r e. Source r Int e => Sz Int -> Vector r e -> Vector r e
Massiv.drop (Int -> Sz Int
forall ix. Index ix => ix -> Sz ix
Sz (Int -> Sz Int) -> Int -> Sz Int
forall a b. (a -> b) -> a -> b
$ Int
ixStart Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Vector r e
vec
    ix0 :: Int
ix0 = Int
0
    e0 :: e
e0 = Vector r e
searchVec Vector r e -> Int -> e
forall r ix e. Manifest r ix e => Array r ix e -> ix -> e
Massiv.! Int
ix0
    startAcc :: (e, Int)
startAcc = (e
e0, Int
ix0)
    minFold :: (a, b) -> b -> a -> (a, b)
minFold acc :: (a, b)
acc@(a
eA, b
_) b
ix a
e = if a
e a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
eA then (a
e, b
ix) else (a, b)
acc
    chFold :: a -> a -> a
chFold a
acc a
ch = a -> a -> a
forall a. Ord a => a -> a -> a
min a
acc a
ch

-- | Like 'Massiv.minimumM' but also returns the index of the minimal element.
iMinimumM ::
  ( Manifest r ix a,
    MonadThrow m,
    Ord a
  ) =>
  Array r ix a ->
  m (a, ix)
iMinimumM :: Array r ix a -> m (a, ix)
iMinimumM Array r ix a
arr
  | Array r ix a -> Bool
forall r ix e. Load r ix e => Array r ix e -> Bool
isEmpty Array r ix a
arr = SizeException -> m (a, ix)
forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM (SizeException -> m (a, ix)) -> SizeException -> m (a, ix)
forall a b. (a -> b) -> a -> b
$ Sz ix -> SizeException
forall ix. Index ix => Sz ix -> SizeException
SizeEmptyException (Array r ix a -> Sz ix
forall r ix e. Load r ix e => Array r ix e -> Sz ix
Massiv.size Array r ix a
arr)
  | Bool
otherwise = (a, ix) -> m (a, ix)
forall (m :: * -> *) a. Monad m => a -> m a
return ((a, ix) -> m (a, ix))
-> (IO (a, ix) -> (a, ix)) -> IO (a, ix) -> m (a, ix)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IO (a, ix) -> (a, ix)
forall a. IO a -> a
unsafePerformIO (IO (a, ix) -> m (a, ix)) -> IO (a, ix) -> m (a, ix)
forall a b. (a -> b) -> a -> b
$ ((a, ix) -> ix -> a -> (a, ix))
-> (a, ix)
-> ((a, ix) -> (a, ix) -> (a, ix))
-> (a, ix)
-> Array r ix a
-> IO (a, ix)
forall (m :: * -> *) r ix e a b.
(MonadIO m, Source r ix e) =>
(a -> ix -> e -> a)
-> a -> (b -> a -> b) -> b -> Array r ix e -> m b
ifoldlP (a, ix) -> ix -> a -> (a, ix)
forall a b. Ord a => (a, b) -> b -> a -> (a, b)
minFold (a, ix)
start (a, ix) -> (a, ix) -> (a, ix)
forall a b. Ord a => (a, b) -> (a, b) -> (a, b)
chFold (a, ix)
start Array r ix a
arr
  where
    ix0 :: ix
ix0 = Int -> ix
forall ix. Index ix => Int -> ix
pureIndex Int
0
    e0 :: a
e0 = Array r ix a
arr Array r ix a -> ix -> a
forall r ix e. Manifest r ix e => Array r ix e -> ix -> e
Massiv.! ix
ix0
    start :: (a, ix)
start = (a
e0, ix
ix0)

    minFold :: (a, b) -> b -> a -> (a, b)
minFold acc :: (a, b)
acc@(a
eA, b
_) b
ix a
e = if a
e a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
eA then (a
e, b
ix) else (a, b)
acc
    chFold :: (a, b) -> (a, b) -> (a, b)
chFold acc :: (a, b)
acc@(a
eA, b
_) ch :: (a, b)
ch@(a
e, b
_) = if a
e a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
eA then (a, b)
ch else (a, b)
acc

-- | Quickly print a matrix with numerical values
printMat :: (Source r Ix2 e, Real e) => Massiv.Matrix r e -> Massiv.Matrix D Text
printMat :: Matrix r e -> Matrix D Text
printMat Matrix r e
mat = (e -> Text) -> Matrix r e -> Matrix D Text
forall r ix e' e.
Source r ix e' =>
(e' -> e) -> Array r ix e' -> Array D ix e
Massiv.map (Format Text (e -> Text) -> e -> Text
forall a. Format Text a -> a
sformat (Int -> Char -> Format Text (Builder -> Text)
forall a r. Buildable a => Int -> Char -> Format r (a -> r)
left Int
4 Char
' ' Format Text (Builder -> Text)
-> Format Text (e -> Text) -> Format Text (e -> Text)
forall r r' a.
Format r (Builder -> r') -> Format r' a -> Format r a
%. Int -> Format Text (e -> Text)
forall a r. Real a => Int -> Format r (a -> r)
fixed Int
2)) Matrix r e
mat

----------------------------------------------------------------------------------------------------
-- Binary Trees.

-- | A binary tree.
data BinTree e = Leaf e | Node e (BinTree e) (BinTree e)
  deriving (BinTree e -> BinTree e -> Bool
(BinTree e -> BinTree e -> Bool)
-> (BinTree e -> BinTree e -> Bool) -> Eq (BinTree e)
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
[BinTree e] -> ShowS
BinTree e -> String
(Int -> BinTree e -> ShowS)
-> (BinTree e -> String)
-> ([BinTree e] -> ShowS)
-> Show (BinTree e)
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 x. BinTree e -> Rep (BinTree e) x)
-> (forall x. Rep (BinTree e) x -> BinTree e)
-> Generic (BinTree e)
forall x. Rep (BinTree e) x -> BinTree e
forall x. BinTree e -> Rep (BinTree e) x
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 :: (a -> b) -> BinTree a -> BinTree b
fmap a -> b
f (Leaf a
a) = b -> BinTree b
forall e. e -> BinTree e
Leaf (a -> b
f a
a)
  fmap a -> b
f (Node a
a BinTree a
l BinTree a
r) = b -> BinTree b -> BinTree b -> BinTree b
forall e. e -> BinTree e -> BinTree e -> BinTree e
Node (a -> b
f a
a) ((a -> b) -> BinTree a -> BinTree b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f BinTree a
l) ((a -> b) -> BinTree a -> BinTree b
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 :: 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 :: (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 ix e. Construct DL ix e => Array DL ix e
forall r ix e. Construct 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 Vector DL a -> a -> Vector DL a
forall r e. 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 Vector DL a -> a -> Vector DL a
forall r e. 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 :: (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 ix e. Construct DL ix e => Array DL ix e
forall r ix e. Construct 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 Vector DL a -> a -> Vector DL a
forall r e. 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 Vector DL a -> a -> Vector DL a
forall r e. 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