From 415e1d63e314b1049ba2868733a8b1a6cce18712 Mon Sep 17 00:00:00 2001
From: Chris Fallin <chris@cfallin.org>
Date: Fri, 24 Dec 2021 14:07:09 -0800
Subject: [PATCH] Fix liveranges for loops

---
 src/backend/locations.rs | 40 ++++++++++++++++++++++++++++++++++++++++
 1 file changed, 40 insertions(+)

diff --git a/src/backend/locations.rs b/src/backend/locations.rs
index 78e3a72..a50358e 100644
--- a/src/backend/locations.rs
+++ b/src/backend/locations.rs
@@ -12,9 +12,17 @@ pub struct Locations {
     pub new_locals: Vec<wasmparser::Type>,
 }
 
+#[derive(Debug)]
+struct Frame {
+    is_loop: bool,
+    start_loc: usize,
+    use_at_end: Vec<(Value, usize)>,
+}
+
 pub struct Allocator<'a> {
     locations: &'a mut Locations,
     f: &'a FunctionBody,
+    active_frames: Vec<Frame>,
     spans: FxHashMap<(Value, usize), ValueSpan>,
     starts: Vec<ValueSpan>,
     ends: Vec<ValueSpan>,
@@ -46,6 +54,7 @@ impl Locations {
         let mut allocator = Allocator {
             locations: &mut locations,
             f,
+            active_frames: vec![],
             freelist: FxHashMap::default(),
             spans: FxHashMap::default(),
             starts: vec![],
@@ -62,6 +71,31 @@ impl<'a> Allocator<'a> {
     fn handle_op(&mut self, location: usize, op: &SerializedOperator) {
         let mut reads = vec![];
         let mut writes = vec![];
+
+        match op {
+            &SerializedOperator::StartBlock { .. } => {
+                self.active_frames.push(Frame {
+                    is_loop: false,
+                    start_loc: location,
+                    use_at_end: vec![],
+                });
+            }
+            &SerializedOperator::StartLoop { .. } => {
+                self.active_frames.push(Frame {
+                    is_loop: true,
+                    start_loc: location,
+                    use_at_end: vec![],
+                });
+            }
+            &SerializedOperator::End { .. } => {
+                let frame = self.active_frames.pop().unwrap();
+                if frame.is_loop {
+                    reads.extend(frame.use_at_end);
+                }
+            }
+            _ => {}
+        }
+
         op.visit_value_locals(
             &mut |value, index| {
                 reads.push((value, index));
@@ -88,6 +122,12 @@ impl<'a> Allocator<'a> {
             };
             span.end = location + 1;
             log::trace!(" -> span for {}: {:?}", value, span);
+
+            for frame in &mut self.active_frames {
+                if frame.is_loop && span.start < frame.start_loc {
+                    frame.use_at_end.push((value, index));
+                }
+            }
         }
 
         for (value, index) in writes {