From 585993063959551cb282f0120110e4f1922f9686 Mon Sep 17 00:00:00 2001
From: Chris Fallin <chris@cfallin.org>
Date: Tue, 29 Nov 2022 00:27:38 -0800
Subject: [PATCH] WIP.

---
 src/backend/localify.rs | 72 +++++++++++++++++++++++++++++------------
 src/cfg/mod.rs          | 20 +++++++++++-
 2 files changed, 70 insertions(+), 22 deletions(-)

diff --git a/src/backend/localify.rs b/src/backend/localify.rs
index 2260961..34902c3 100644
--- a/src/backend/localify.rs
+++ b/src/backend/localify.rs
@@ -52,9 +52,9 @@ impl<'a> Context<'a> {
                     .zip(body.blocks[target.block].params.iter().map(|(_, val)| val))
                 {
                     affinities
-                        .entry(param)
+                        .entry(arg)
                         .or_insert_with(|| smallvec![])
-                        .push(arg);
+                        .push(param);
                 }
             });
         }
@@ -124,6 +124,8 @@ impl<'a> Context<'a> {
         // use (first observed), allocate a local.
         fn handle_use(
             body: &FunctionBody,
+            cfg: &CFGInfo,
+            block: Block,
             u: Value,
             live_values: &mut HashSet<Value>,
             live_locals: &mut HashSet<Local>,
@@ -145,32 +147,56 @@ impl<'a> Context<'a> {
                         unreachable!();
                     }
                     &ValueDef::PickOutput(value, idx, _) => {
-                        handle_use(body, value, live_values, live_locals, results, affinities);
+                        handle_use(
+                            body,
+                            cfg,
+                            block,
+                            value,
+                            live_values,
+                            live_locals,
+                            results,
+                            affinities,
+                        );
                         results.values[u] = smallvec![results.values[value][idx]];
                     }
                     &ValueDef::BlockParam(..) | &ValueDef::Operator(..) => {
+                        // Uses from dominating blocks may be live
+                        // over a loop backedge, so for simplicity,
+                        // let's not reuse any locals for uses of
+                        // non-local defs.
+                        let can_reuse = cfg.def_block[u] == block;
+
                         let locals = def
                             .tys()
                             .iter()
                             .map(|&ty| {
-                                // Try to find a local of the right type that is not live.
-                                let affinities = affinities.get(&u).map(|v| &v[..]).unwrap_or(&[]);
-                                let local = affinities
-                                    .iter()
-                                    .filter_map(|&aff_val| {
-                                        let local = *results.values[aff_val].get(0)?;
-                                        Some((local, results.locals[local]))
-                                    })
-                                    .chain(results.locals.entries().map(|(local, &ty)| (local, ty)))
-                                    .filter_map(|(local, local_ty)| {
-                                        if local_ty == ty && live_locals.insert(local) {
-                                            Some(local)
-                                        } else {
-                                            None
-                                        }
-                                    })
-                                    .next();
-                                local.unwrap_or_else(|| {
+                                let reused = if can_reuse {
+                                    // Try to find a local of the right type that is not live.
+                                    let affinities =
+                                        affinities.get(&u).map(|v| &v[..]).unwrap_or(&[]);
+                                    let mut try_list = affinities
+                                        .iter()
+                                        .filter_map(|&aff_val| {
+                                            let local = *results.values[aff_val].get(0)?;
+                                            Some((local, results.locals[local]))
+                                        })
+                                        .chain(
+                                            results
+                                                .locals
+                                                .entries()
+                                                .map(|(local, &ty)| (local, ty)),
+                                        );
+
+                                    try_list
+                                        .find(|&(local, local_ty)| {
+                                            local_ty == ty && live_locals.insert(local)
+                                        })
+                                        .map(|(local, _)| local)
+                                } else {
+                                    None
+                                };
+
+                                reused.unwrap_or_else(|| {
                                     let local = results.locals.push(ty);
                                     live_locals.insert(local);
                                     local
@@ -205,6 +231,8 @@ impl<'a> Context<'a> {
         self.body.blocks[block].terminator.visit_uses(|u| {
             handle_use(
                 self.body,
+                self.cfg,
+                block,
                 u,
                 &mut live_values,
                 &mut live_locals,
@@ -226,6 +254,8 @@ impl<'a> Context<'a> {
                 if !self.trees.owner.contains_key(&u) {
                     handle_use(
                         self.body,
+                        self.cfg,
+                        block,
                         u,
                         &mut live_values,
                         &mut live_locals,
diff --git a/src/cfg/mod.rs b/src/cfg/mod.rs
index 8214103..fdf624b 100644
--- a/src/cfg/mod.rs
+++ b/src/cfg/mod.rs
@@ -4,7 +4,7 @@
 // LLVM exception.
 
 use crate::entity::{EntityRef, PerEntity};
-use crate::ir::{Block, FunctionBody, Terminator};
+use crate::ir::{Block, FunctionBody, Terminator, Value, ValueDef};
 use smallvec::SmallVec;
 
 pub mod domtree;
@@ -28,6 +28,8 @@ pub struct CFGInfo {
     pub domtree: PerEntity<Block, Block>,
     /// Domtree children.
     pub domtree_children: PerEntity<Block, DomtreeChildren>,
+    /// Defining block for a given value.
+    pub def_block: PerEntity<Value, Block>,
 }
 
 #[derive(Clone, Debug, Default)]
@@ -91,6 +93,21 @@ impl CFGInfo {
             }
         }
 
+        let mut def_block: PerEntity<Value, Block> = PerEntity::default();
+        for (block, block_def) in f.blocks.entries() {
+            for &value in &block_def.insts {
+                def_block[value] = block;
+            }
+        }
+        for value in f.values.iter() {
+            let orig_value = f.resolve_alias(value);
+            let underlying_value = match &f.values[orig_value] {
+                &ValueDef::PickOutput(value, ..) => value,
+                _ => orig_value,
+            };
+            def_block[value] = def_block[underlying_value];
+        }
+
         CFGInfo {
             entry: f.entry,
             block_preds,
@@ -100,6 +117,7 @@ impl CFGInfo {
             postorder_pos,
             domtree,
             domtree_children,
+            def_block,
         }
     }