module Bench.Vector.Algo.Leaffix where

import Data.Vector.Unboxed as V

leaffix :: (Vector Int, Vector Int) -> Vector Int
{-# NOINLINE leaffix #-}
leaffix :: (Vector Int, Vector Int) -> Vector Int
leaffix (Vector Int
ls,Vector Int
rs)
    = Vector Int -> Vector Int -> Vector Int -> Vector Int
forall {c}.
(Unbox c, Num c) =>
Vector c -> Vector Int -> Vector Int -> Vector c
leaffix (Int -> Int -> Vector Int
forall a. Unbox a => Int -> a -> Vector a
V.replicate (Vector Int -> Int
forall a. Unbox a => Vector a -> Int
V.length Vector Int
ls) Int
1) Vector Int
ls Vector Int
rs
    where
      leaffix :: Vector c -> Vector Int -> Vector Int -> Vector c
leaffix Vector c
xs Vector Int
ls Vector Int
rs
        = let zs :: Vector c
zs   = Int -> c -> Vector c
forall a. Unbox a => Int -> a -> Vector a
V.replicate (Vector Int -> Int
forall a. Unbox a => Vector a -> Int
V.length Vector Int
ls Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
2) c
0
              vs :: Vector c
vs   = Vector c -> Vector Int -> Vector c -> Vector c
forall a. Unbox a => Vector a -> Vector Int -> Vector a -> Vector a
V.update_ Vector c
zs Vector Int
ls Vector c
xs
              sums :: Vector c
sums = (c -> c -> c) -> c -> Vector c -> Vector c
forall a b.
(Unbox a, Unbox b) =>
(a -> b -> a) -> a -> Vector b -> Vector a
V.prescanl' c -> c -> c
forall a. Num a => a -> a -> a
(+) c
0 Vector c
vs
          in
          (c -> c -> c) -> Vector c -> Vector c -> Vector c
forall a b c.
(Unbox a, Unbox b, Unbox c) =>
(a -> b -> c) -> Vector a -> Vector b -> Vector c
V.zipWith (-) (Vector c -> Vector Int -> Vector c
forall a. Unbox a => Vector a -> Vector Int -> Vector a
V.backpermute Vector c
sums Vector Int
ls) (Vector c -> Vector Int -> Vector c
forall a. Unbox a => Vector a -> Vector Int -> Vector a
V.backpermute Vector c
sums Vector Int
rs)