From 52013a7f464be06dedb3e38a45676fb6d30e4bdd Mon Sep 17 00:00:00 2001
From: Chris Fallin <chris@cfallin.org>
Date: Tue, 29 Nov 2022 21:04:55 -0800
Subject: [PATCH] Fix subtle interaction of treeifier-induced code motion and
 local (register) allocation.

---
 src/backend/localify.rs | 67 ++++++++++++++++++++++++++++++++---------
 src/passes/basic_opt.rs |  1 +
 2 files changed, 53 insertions(+), 15 deletions(-)

diff --git a/src/backend/localify.rs b/src/backend/localify.rs
index 44e526f..3dea14e 100644
--- a/src/backend/localify.rs
+++ b/src/backend/localify.rs
@@ -241,6 +241,47 @@ impl<'a> Context<'a> {
             )
         });
 
+        fn visit_inst_uses(
+            body: &FunctionBody,
+            cfg: &CFGInfo,
+            trees: &Trees,
+            block: Block,
+            inst: Value,
+            live_values: &mut HashSet<Value>,
+            live_locals: &mut HashSet<Local>,
+            results: &mut Localifier,
+            affinities: &HashMap<Value, SmallVec<[Value; 4]>>,
+        ) {
+            body.values[inst].visit_uses(|u| {
+                // If treeified, then don't process use. However, do
+                // process uses of the treeified value.
+                if trees.owner.contains_key(&u) {
+                    visit_inst_uses(
+                        &body,
+                        &cfg,
+                        trees,
+                        block,
+                        u,
+                        live_values,
+                        live_locals,
+                        results,
+                        affinities,
+                    );
+                } else {
+                    handle_use(
+                        body,
+                        cfg,
+                        block,
+                        u,
+                        live_values,
+                        live_locals,
+                        results,
+                        affinities,
+                    )
+                }
+            });
+        }
+
         for &inst in self.body.blocks[block].insts.iter().rev() {
             handle_def(
                 self.body,
@@ -249,21 +290,17 @@ impl<'a> Context<'a> {
                 &mut live_locals,
                 &self.results,
             );
-            self.body.values[inst].visit_uses(|u| {
-                // If treeified, then don't process use.
-                if !self.trees.owner.contains_key(&u) {
-                    handle_use(
-                        self.body,
-                        self.cfg,
-                        block,
-                        u,
-                        &mut live_values,
-                        &mut live_locals,
-                        &mut self.results,
-                        &self.affinities,
-                    )
-                }
-            });
+            visit_inst_uses(
+                &self.body,
+                &self.cfg,
+                &self.trees,
+                block,
+                inst,
+                &mut live_values,
+                &mut live_locals,
+                &mut self.results,
+                &self.affinities,
+            );
         }
 
         for &(_, param) in &self.body.blocks[block].params {
diff --git a/src/passes/basic_opt.rs b/src/passes/basic_opt.rs
index cced07b..7b8b474 100644
--- a/src/passes/basic_opt.rs
+++ b/src/passes/basic_opt.rs
@@ -180,6 +180,7 @@ fn const_eval(op: &Operator, vals: &[ConstVal]) -> Option<ConstVal> {
         (Operator::I64Extend32S, [ConstVal::I64(a)]) => {
             Some(ConstVal::I64(*a as i32 as i64 as u64))
         }
+        (Operator::Select, [x, y, ConstVal::I32(k)]) => Some(if *k != 0 { *x } else { *y }),
 
         _ => None,
     }