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, } +#[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, spans: FxHashMap<(Value, usize), ValueSpan>, starts: Vec, ends: Vec, @@ -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 {