From c960d728bf85bb1d09f0ec27a885dffd16f966ed Mon Sep 17 00:00:00 2001
From: Chris Fallin <chris@cfallin.org>
Date: Tue, 7 May 2024 15:22:18 -0700
Subject: [PATCH] Support treeification of adjacent side-effecting ops in some
 cases.

---
 src/backend/treeify.rs | 87 +++++++++++++++++++++++-------------------
 1 file changed, 47 insertions(+), 40 deletions(-)

diff --git a/src/backend/treeify.rs b/src/backend/treeify.rs
index ab87295..bcc9b7d 100644
--- a/src/backend/treeify.rs
+++ b/src/backend/treeify.rs
@@ -41,50 +41,57 @@ impl Trees {
         let mut remat = HashSet::default();
         let mut multi_use = HashSet::default();
 
-        for (value, def) in body.values.entries() {
-            match def {
-                &ValueDef::Operator(op, args, _) => {
-                    // Ignore operators with invalid args: these must
-                    // always be unreachable.
-                    if body.arg_pool[args].iter().any(|arg| arg.is_invalid()) {
-                        continue;
-                    }
-                    // If this is an always-rematerialized operator,
-                    // mark it as such and continue.
-                    if is_remat(&op) {
-                        remat.insert(value);
-                        continue;
-                    }
-
-                    // For each of the args, if the value is produced
-                    // by a single-output op and is movable, and is
-                    // not already recorded in `multi_use`, place it
-                    // in the arg slot. Otherwise if owned already
-                    // somewhere else, undo that and put in
-                    // `multi_use`.
-                    for (i, &arg) in body.arg_pool[args].iter().enumerate() {
-                        let arg = body.resolve_alias(arg);
-                        if multi_use.contains(&arg) {
+        for block_def in body.blocks.values() {
+            let mut last_non_pure = None;
+            for &value in &block_def.insts {
+                match &body.values[value] {
+                    &ValueDef::Operator(op, args, _) => {
+                        // Ignore operators with invalid args: these must
+                        // always be unreachable.
+                        if body.arg_pool[args].iter().any(|arg| arg.is_invalid()) {
                             continue;
-                        } else if let Some(old_owner) = owner.remove(&arg) {
-                            owned.remove(&old_owner);
-                            multi_use.insert(arg);
-                        } else if Self::is_movable(body, arg) {
-                            let pos = u16::try_from(i).unwrap();
-                            let value_arg = ValueArg(value, pos);
-                            owner.insert(arg, value_arg);
-                            owned.insert(value_arg, arg);
+                        }
+                        // If this is an always-rematerialized operator,
+                        // mark it as such and continue.
+                        if is_remat(&op) {
+                            remat.insert(value);
+                            continue;
+                        }
+
+                        // For each of the args, if the value is produced
+                        // by a single-output op and is movable, and is
+                        // not already recorded in `multi_use`, place it
+                        // in the arg slot. Otherwise if owned already
+                        // somewhere else, undo that and put in
+                        // `multi_use`.
+                        for (i, &arg) in body.arg_pool[args].iter().enumerate() {
+                            let arg = body.resolve_alias(arg);
+                            if multi_use.contains(&arg) {
+                                continue;
+                            } else if let Some(old_owner) = owner.remove(&arg) {
+                                owned.remove(&old_owner);
+                                multi_use.insert(arg);
+                            } else if Self::is_movable(body, arg) || Some(arg) == last_non_pure {
+                                let pos = u16::try_from(i).unwrap();
+                                let value_arg = ValueArg(value, pos);
+                                owner.insert(arg, value_arg);
+                                owned.insert(value_arg, arg);
+                            }
+                        }
+
+                        if !op.is_pure() {
+                            last_non_pure = Some(value);
                         }
                     }
+                    &ValueDef::PickOutput(..) => {
+                        // Can ignore use: multi-arity values are never treeified.
+                    }
+                    &ValueDef::BlockParam(..)
+                    | &ValueDef::Alias(..)
+                    | &ValueDef::Placeholder(..)
+                    | &ValueDef::Trace(..)
+                    | &ValueDef::None => {}
                 }
-                &ValueDef::PickOutput(..) => {
-                    // Can ignore use: multi-arity values are never treeified.
-                }
-                &ValueDef::BlockParam(..)
-                | &ValueDef::Alias(..)
-                | &ValueDef::Placeholder(..)
-                | &ValueDef::Trace(..)
-                | &ValueDef::None => {}
             }
         }
         for block in body.blocks.values() {