From 6621a57cb7850bd49ff28d2f2387c9fc0df522bf Mon Sep 17 00:00:00 2001
From: Amanieu d'Antras <amanieu@gmail.com>
Date: Wed, 1 Dec 2021 01:43:02 +0000
Subject: [PATCH] Fix liveranges for branch parameters

---
 src/ion/liveranges.rs | 59 +++++++++++++++++++++++--------------------
 1 file changed, 32 insertions(+), 27 deletions(-)

diff --git a/src/ion/liveranges.rs b/src/ion/liveranges.rs
index 1140c1e..4637318 100644
--- a/src/ion/liveranges.rs
+++ b/src/ion/liveranges.rs
@@ -315,6 +315,7 @@ impl<'a, F: Function> Env<'a, F> {
         while !workqueue.is_empty() {
             let block = workqueue.pop_front().unwrap();
             workqueue_set.remove(&block);
+            let insns = self.func.block_insns(block);
 
             log::trace!("computing liveins for block{}", block.index());
 
@@ -323,7 +324,16 @@ impl<'a, F: Function> Env<'a, F> {
             let mut live = self.liveouts[block.index()].clone();
             log::trace!(" -> initial liveout set: {:?}", live);
 
-            for inst in self.func.block_insns(block).rev().iter() {
+            // Include outgoing blockparams in the initial live set.
+            if self.func.is_branch(insns.last()) {
+                for i in 0..self.func.block_succs(block).len() {
+                    for param in self.func.branch_blockparams(block, insns.last(), i) {
+                        live.set(param.vreg(), true);
+                    }
+                }
+            }
+
+            for inst in insns.rev().iter() {
                 if let Some((src, dst)) = self.func.is_move(inst) {
                     live.set(dst.vreg().vreg(), false);
                     live.set(src.vreg().vreg(), true);
@@ -399,12 +409,33 @@ impl<'a, F: Function> Env<'a, F> {
 
         for i in (0..self.func.num_blocks()).rev() {
             let block = Block::new(i);
+            let insns = self.func.block_insns(block);
 
             self.stats.livein_blocks += 1;
 
             // Init our local live-in set.
             let mut live = self.liveouts[block.index()].clone();
 
+            // If the last instruction is a branch (rather than
+            // return), create blockparam_out entries.
+            if self.func.is_branch(insns.last()) {
+                for (i, &succ) in self.func.block_succs(block).iter().enumerate() {
+                    let blockparams_in = self.func.block_params(succ);
+                    let blockparams_out = self.func.branch_blockparams(block, insns.last(), i);
+                    for (&blockparam_in, &blockparam_out) in
+                        blockparams_in.iter().zip(blockparams_out)
+                    {
+                        let blockparam_out = VRegIndex::new(blockparam_out.vreg());
+                        let blockparam_in = VRegIndex::new(blockparam_in.vreg());
+                        self.blockparam_outs
+                            .push((blockparam_out, block, succ, blockparam_in));
+
+                        // Include outgoing blockparams in the initial live set.
+                        live.set(blockparam_out.index(), true);
+                    }
+                }
+            }
+
             // Initially, registers are assumed live for the whole block.
             for vreg in live.iter() {
                 let range = CodeRange {
@@ -426,25 +457,6 @@ impl<'a, F: Function> Env<'a, F> {
                 self.vregs[param.vreg()].blockparam = block;
             }
 
-            let insns = self.func.block_insns(block);
-
-            // If the last instruction is a branch (rather than
-            // return), create blockparam_out entries.
-            if self.func.is_branch(insns.last()) {
-                for (i, &succ) in self.func.block_succs(block).iter().enumerate() {
-                    let blockparams_in = self.func.block_params(succ);
-                    let blockparams_out = self.func.branch_blockparams(block, insns.last(), i);
-                    for (&blockparam_in, &blockparam_out) in
-                        blockparams_in.iter().zip(blockparams_out)
-                    {
-                        let blockparam_out = VRegIndex::new(blockparam_out.vreg());
-                        let blockparam_in = VRegIndex::new(blockparam_in.vreg());
-                        self.blockparam_outs
-                            .push((blockparam_out, block, succ, blockparam_in));
-                    }
-                }
-            }
-
             // For each instruction, in reverse order, process
             // operands and clobbers.
             for inst in insns.rev().iter() {
@@ -893,13 +905,6 @@ impl<'a, F: Function> Env<'a, F> {
                             (OperandKind::Def, OperandPos::Early) => ProgPoint::before(inst),
                             (OperandKind::Def, OperandPos::Late) => ProgPoint::after(inst),
                             (OperandKind::Use, OperandPos::Late) => ProgPoint::after(inst),
-                            // If this is a branch, extend `pos` to
-                            // the end of the block. (Branch uses are
-                            // blockparams and need to be live at the
-                            // end of the block.)
-                            (OperandKind::Use, _) if self.func.is_branch(inst) => {
-                                self.cfginfo.block_exit[block.index()]
-                            }
                             // If there are any reused inputs in this
                             // instruction, and this is *not* the
                             // reused input, force `pos` to