third_party/bazel/rules_haskell/examples/vector/benchmarks/Algo/Rootfix.hs


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15

module Algo.Rootfix where

import Data.Vector.Unboxed as V

rootfix :: (V.Vector Int, V.Vector Int) -> V.Vector Int
{-# NOINLINE rootfix #-}
rootfix (ls, rs) = rootfix (V.replicate (V.length ls) 1) ls rs
    where
      rootfix xs ls rs
        = let zs   = V.replicate (V.length ls * 2) 0
              vs   = V.update_ (V.update_ zs ls xs) rs (V.map negate xs)
              sums = V.prescanl' (+) 0 vs
          in
          V.backpermute sums ls