From 730f7b0d527a7727d33d4e41b2c162fb55d5e8f8 Mon Sep 17 00:00:00 2001
From: Troels Henriksen <athas@sigkill.dk>
Date: Fri, 8 Dec 2023 20:48:49 +0100
Subject: [PATCH] Fix #2058.

---
 CHANGELOG.md            |  3 +++
 src/Futhark/IR/SegOp.hs | 14 +++++++-------
 tests/issue2058.fut     |  5 +++++
 3 files changed, 15 insertions(+), 7 deletions(-)
 create mode 100644 tests/issue2058.fut

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 48335a3a0d..a0a9bebac7 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -32,6 +32,9 @@ and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0.
 
 * Futhark's LSP server should work better with Eglot.
 
+* Incorrect copy removal inside histograms could cause compiler error
+  (#2058).
+
 ## [0.25.9]
 
 ### Added
diff --git a/src/Futhark/IR/SegOp.hs b/src/Futhark/IR/SegOp.hs
index e4ba18c084..a56b6da6b2 100644
--- a/src/Futhark/IR/SegOp.hs
+++ b/src/Futhark/IR/SegOp.hs
@@ -1157,8 +1157,8 @@ simplifySegOp (SegScan lvl space scans ts kbody) = do
 simplifySegOp (SegHist lvl space ops ts kbody) = do
   (lvl', space', ts') <- Engine.simplify (lvl, space, ts)
 
-  (ops', ops_hoisted) <- fmap unzip $
-    forM ops $
+  Engine.localVtable (flip (foldr ST.consume) $ concatMap histDest ops) $ do
+    (ops', ops_hoisted) <- fmap unzip . forM ops $
       \(HistOp w rf arrs nes dims lam) -> do
         w' <- Engine.simplify w
         rf' <- Engine.simplify rf
@@ -1174,12 +1174,12 @@ simplifySegOp (SegHist lvl space ops ts kbody) = do
             op_hoisted
           )
 
-  (kbody', body_hoisted) <- simplifyKernelBody space kbody
+    (kbody', body_hoisted) <- simplifyKernelBody space kbody
 
-  pure
-    ( SegHist lvl' space' ops' ts' kbody',
-      mconcat ops_hoisted <> body_hoisted
-    )
+    pure
+      ( SegHist lvl' space' ops' ts' kbody',
+        mconcat ops_hoisted <> body_hoisted
+      )
   where
     scope = scopeOfSegSpace space
     scope_vtable = ST.fromScope scope
diff --git a/tests/issue2058.fut b/tests/issue2058.fut
new file mode 100644
index 0000000000..40b691d6cf
--- /dev/null
+++ b/tests/issue2058.fut
@@ -0,0 +1,5 @@
+-- We neglected to mark the target arrays as consumed while
+-- simplifying the body.
+
+entry problem [n] (arr: *[n]i64) : [n]i64 =
+  reduce_by_index arr (+) 0 (iota n) (copy arr)