-- |
-- Module      :  ELynx.Tree.Parallel
-- Description :  Evaluation strategies for trees
-- Copyright   :  (c) Dominik Schrempf, 2021
-- License     :  GPL-3.0-or-later
--
-- Maintainer  :  dominik.schrempf@gmail.com
-- Stability   :  unstable
-- Portability :  portable
--
-- Creation date: Mon Sep  7 13:36:45 2020.
--
-- Parallel evaluation up to a given layer. By convention layer 0 only has one
-- element: the root node. The layer 1 includes the daughter nodes of the root
-- node, and so on.
--
-- The layer to which a node belongs should not be confused with the 'depth' of
-- a tree.
module ELynx.Tree.Parallel
  ( parTree,
    parBranchFoldMap,
    parLabelFoldMap,
  )
where

import Control.Parallel.Strategies
import Data.Foldable
import ELynx.Tree.Rooted

myParList :: Strategy a -> Strategy [a]
myParList :: Strategy a -> Strategy [a]
myParList Strategy a
_ [] = Strategy [a]
forall (m :: * -> *) a. Monad m => a -> m a
return []
myParList Strategy a
s [a]
xs = do
  [a]
ys <- Strategy a -> Strategy [a]
forall a. Strategy a -> Strategy [a]
parList Strategy a
s Strategy [a] -> Strategy [a]
forall a b. (a -> b) -> a -> b
$ [a] -> [a]
forall a. [a] -> [a]
tail [a]
xs
  a
y <- Strategy a
s Strategy a -> Strategy a
forall a b. (a -> b) -> a -> b
$ [a] -> a
forall a. [a] -> a
head [a]
xs
  Strategy [a]
forall (m :: * -> *) a. Monad m => a -> m a
return Strategy [a] -> Strategy [a]
forall a b. (a -> b) -> a -> b
$ a
y a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
ys

-- | Parallel evaluation strategy for a tree into normal form.
--
-- Evaluate the sub trees up to given layer in parallel.
parTree :: (NFData e, NFData a) => Int -> Strategy (Tree e a)
parTree :: Int -> Strategy (Tree e a)
parTree Int
n t :: Tree e a
t@(Node e
br a
lb Forest e a
ts)
  | Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1 = do
    Forest e a
ts' <- Strategy (Tree e a) -> Strategy (Forest e a)
forall a. Strategy a -> Strategy [a]
myParList Strategy (Tree e a)
forall a. NFData a => Strategy a
rdeepseq Forest e a
ts
    Strategy (Tree e a)
forall (m :: * -> *) a. Monad m => a -> m a
return Strategy (Tree e a) -> Strategy (Tree e a)
forall a b. (a -> b) -> a -> b
$ e -> a -> Forest e a -> Tree e a
forall e a. e -> a -> Forest e a -> Tree e a
Node e
br a
lb Forest e a
ts'
  | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
2 = do
    Forest e a
ts' <- Strategy (Tree e a) -> Strategy (Forest e a)
forall a. Strategy a -> Strategy [a]
myParList (Int -> Strategy (Tree e a)
forall e a. (NFData e, NFData a) => Int -> Strategy (Tree e a)
parTree (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)) Forest e a
ts
    Strategy (Tree e a)
forall (m :: * -> *) a. Monad m => a -> m a
return Strategy (Tree e a) -> Strategy (Tree e a)
forall a b. (a -> b) -> a -> b
$ e -> a -> Forest e a -> Tree e a
forall e a. e -> a -> Forest e a -> Tree e a
Node e
br a
lb Forest e a
ts'
  | Bool
otherwise = Strategy (Tree e a)
forall a. NFData a => Strategy a
rdeepseq Tree e a
t

branchFoldMap :: (e -> f) -> (f -> f -> f) -> Tree e a -> f
branchFoldMap :: (e -> f) -> (f -> f -> f) -> Tree e a -> f
branchFoldMap e -> f
f f -> f -> f
op (Node e
br a
_ Forest e a
ts) = (f -> f -> f) -> f -> [f] -> f
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' f -> f -> f
op (e -> f
f e
br) ([f] -> f) -> [f] -> f
forall a b. (a -> b) -> a -> b
$ (Tree e a -> f) -> Forest e a -> [f]
forall a b. (a -> b) -> [a] -> [b]
map ((e -> f) -> (f -> f -> f) -> Tree e a -> f
forall e f a. (e -> f) -> (f -> f -> f) -> Tree e a -> f
branchFoldMap e -> f
f f -> f -> f
op) Forest e a
ts

-- IDEA: Use and benchmark branch and node specific instances with parFoldMaps.
--
-- @
-- parFoldMap' = blabla
-- parBranchFoldMap' = parFoldMap' . ZipBranchTree
-- parNodeFoldMap' = parFoldMap' . ZipNodeTree
-- @

-- | Map and fold over branches. Evaluate the sub trees up to given layer in parallel.
parBranchFoldMap :: NFData f => Int -> (e -> f) -> (f -> f -> f) -> Tree e a -> f
parBranchFoldMap :: Int -> (e -> f) -> (f -> f -> f) -> Tree e a -> f
parBranchFoldMap Int
n e -> f
f f -> f -> f
op t :: Tree e a
t@(Node e
br a
_ Forest e a
ts)
  | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
1 = (f -> f -> f) -> f -> [f] -> f
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' f -> f -> f
op (e -> f
f e
br) ((Tree e a -> f) -> Forest e a -> [f]
forall a b. (a -> b) -> [a] -> [b]
map (Int -> (e -> f) -> (f -> f -> f) -> Tree e a -> f
forall f e a.
NFData f =>
Int -> (e -> f) -> (f -> f -> f) -> Tree e a -> f
parBranchFoldMap (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) e -> f
f f -> f -> f
op) Forest e a
ts [f] -> Strategy [f] -> [f]
forall a. a -> Strategy a -> a
`using` Strategy f -> Strategy [f]
forall a. Strategy a -> Strategy [a]
myParList Strategy f
forall a. NFData a => Strategy a
rdeepseq)
  | Bool
otherwise = (e -> f) -> (f -> f -> f) -> Tree e a -> f
forall e f a. (e -> f) -> (f -> f -> f) -> Tree e a -> f
branchFoldMap e -> f
f f -> f -> f
op Tree e a
t

nodeFoldMap :: (a -> b) -> (b -> b -> b) -> Tree e a -> b
nodeFoldMap :: (a -> b) -> (b -> b -> b) -> Tree e a -> b
nodeFoldMap a -> b
f b -> b -> b
op (Node e
_ a
lb Forest e a
ts) = (b -> b -> b) -> b -> [b] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' b -> b -> b
op (a -> b
f a
lb) ([b] -> b) -> [b] -> b
forall a b. (a -> b) -> a -> b
$ (Tree e a -> b) -> Forest e a -> [b]
forall a b. (a -> b) -> [a] -> [b]
map ((a -> b) -> (b -> b -> b) -> Tree e a -> b
forall a b e. (a -> b) -> (b -> b -> b) -> Tree e a -> b
nodeFoldMap a -> b
f b -> b -> b
op) Forest e a
ts

-- | Map and fold over labels. Evaluate the sub trees up to given layer in parallel.
parLabelFoldMap :: NFData b => Int -> (a -> b) -> (b -> b -> b) -> Tree e a -> b
parLabelFoldMap :: Int -> (a -> b) -> (b -> b -> b) -> Tree e a -> b
parLabelFoldMap Int
n a -> b
f b -> b -> b
op t :: Tree e a
t@(Node e
_ a
lb Forest e a
ts)
  | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
1 = (b -> b -> b) -> b -> [b] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' b -> b -> b
op (a -> b
f a
lb) ((Tree e a -> b) -> Forest e a -> [b]
forall a b. (a -> b) -> [a] -> [b]
map (Int -> (a -> b) -> (b -> b -> b) -> Tree e a -> b
forall b a e.
NFData b =>
Int -> (a -> b) -> (b -> b -> b) -> Tree e a -> b
parLabelFoldMap (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) a -> b
f b -> b -> b
op) Forest e a
ts [b] -> Strategy [b] -> [b]
forall a. a -> Strategy a -> a
`using` Strategy b -> Strategy [b]
forall a. Strategy a -> Strategy [a]
myParList Strategy b
forall a. NFData a => Strategy a
rdeepseq)
  | Bool
otherwise = (a -> b) -> (b -> b -> b) -> Tree e a -> b
forall a b e. (a -> b) -> (b -> b -> b) -> Tree e a -> b
nodeFoldMap a -> b
f b -> b -> b
op Tree e a
t