about summary refs log tree commit diff
path: root/third_party/bazel/rules_haskell/examples/vector/benchmarks/Algo/Rootfix.hs
diff options
context:
space:
mode:
Diffstat (limited to 'third_party/bazel/rules_haskell/examples/vector/benchmarks/Algo/Rootfix.hs')
-rw-r--r--third_party/bazel/rules_haskell/examples/vector/benchmarks/Algo/Rootfix.hs15
1 files changed, 15 insertions, 0 deletions
diff --git a/third_party/bazel/rules_haskell/examples/vector/benchmarks/Algo/Rootfix.hs b/third_party/bazel/rules_haskell/examples/vector/benchmarks/Algo/Rootfix.hs
new file mode 100644
index 0000000000..1b112a801a
--- /dev/null
+++ b/third_party/bazel/rules_haskell/examples/vector/benchmarks/Algo/Rootfix.hs
@@ -0,0 +1,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
+