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


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

module Algo.Leaffix where

import Data.Vector.Unboxed as V

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