module Bench.Vector.Algo.Rootfix where

import Data.Vector.Unboxed as V

rootfix :: (V.Vector Int, V.Vector Int) -> V.Vector Int
{-# NOINLINE rootfix #-}
rootfix :: (Vector Int, Vector Int) -> Vector Int
rootfix (Vector Int
ls, Vector Int
rs) = Vector Int -> Vector Int -> Vector Int -> Vector Int
forall {a}.
(Unbox a, Num a) =>
Vector a -> Vector Int -> Vector Int -> Vector a
rootfix (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
      rootfix :: Vector a -> Vector Int -> Vector Int -> Vector a
rootfix Vector a
xs Vector Int
ls Vector Int
rs
        = let zs :: Vector a
zs   = Int -> a -> Vector a
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) a
0
              vs :: Vector a
vs   = Vector a -> Vector Int -> Vector a -> Vector a
forall a. Unbox a => Vector a -> Vector Int -> Vector a -> Vector a
V.update_ (Vector a -> Vector Int -> Vector a -> Vector a
forall a. Unbox a => Vector a -> Vector Int -> Vector a -> Vector a
V.update_ Vector a
zs Vector Int
ls Vector a
xs) Vector Int
rs ((a -> a) -> Vector a -> Vector a
forall a b. (Unbox a, Unbox b) => (a -> b) -> Vector a -> Vector b
V.map a -> a
forall a. Num a => a -> a
negate Vector a
xs)
              sums :: Vector a
sums = (a -> a -> a) -> a -> Vector a -> Vector a
forall a b.
(Unbox a, Unbox b) =>
(a -> b -> a) -> a -> Vector b -> Vector a
V.prescanl' a -> a -> a
forall a. Num a => a -> a -> a
(+) a
0 Vector a
vs
          in
          Vector a -> Vector Int -> Vector a
forall a. Unbox a => Vector a -> Vector Int -> Vector a
V.backpermute Vector a
sums Vector Int
ls