Apply HIR local events in clean state

Author: olwang

src/checks/body.rs

@@ -1,8 +1,7 @@
 use crate::analyzer::Analyzer;
 use crate::diagnostic::{Diagnostic, code};
 use crate::hir::{
-    CallResolution, HirBindingKind, HirEffectEventKind, HirReturnProof, HirTypeKind,
-    ResolvedCalleeKind,
+    CallResolution, HirEffectEventKind, HirReturnProof, HirTypeKind, ResolvedCalleeKind,
 };
 use crate::syntax::ast::{Block, Callee, DataEffect, Expr, FunctionDecl, Item, LetKind, Stmt};
 
@@ -30,14 +29,7 @@ fn seed_function_bindings(analyzer: &Analyzer<'_>, function: &FunctionDecl, stat
     let Some(body) = analyzer.hir.function_body(&function.name) else {
         return;
     };
-    for binding in &body.bindings {
-        if binding.kind != HirBindingKind::Param {
-            continue;
-        }
-        if let Some(type_name) = &binding.type_name {
-            state.record_type(binding.name.clone(), type_name.clone());
-        }
-    }
+    state.seed_params(&body.bindings);
 }
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
@@ -242,33 +234,15 @@ fn apply_move_events(
     span: &crate::diagnostic::Span,
     state: &mut BodyState,
 ) {
-    for event in analyzer.hir.effect_events(span) {
-        if !matches!(
-            event.kind,
-            HirEffectEventKind::Manage | HirEffectEventKind::Take
-        ) {
-            continue;
-        }
-        if state.locals.contains(&event.binding_name) {
-            state.mark_moved(&event.binding_name, event.span.clone());
-        }
-    }
+    state.apply_move_events(analyzer.hir.effect_events(span));
 }
 
 fn apply_retention_events(
     analyzer: &Analyzer<'_>,
     span: &crate::diagnostic::Span,
     state: &mut BodyState,
 ) {
-    for event in analyzer.hir.effect_events(span) {
-        if !matches!(event.kind, HirEffectEventKind::Retain { .. }) {
-            continue;
-        }
-        if !state.locals.contains(&event.binding_name) {
-            continue;
-        }
-        state.mark_retained(&event.binding_name);
-    }
+    state.apply_retention_events(analyzer.hir.effect_events(span));
 }
 
 fn check_moved_uses_in_stmt(analyzer: &mut Analyzer<'_>, statement: &Stmt, state: &BodyState) {

src/checks/local.rs

@@ -1,6 +1,7 @@
 use std::collections::{HashMap, HashSet};
 
 use crate::diagnostic::Span;
+use crate::hir::{HirBinding, HirBindingKind, HirEffectEvent, HirEffectEventKind};
 
 use super::body::Flow;
 
@@ -14,6 +15,17 @@ pub(crate) struct BodyState {
 }
 
 impl BodyState {
+    pub(crate) fn seed_params(&mut self, bindings: &[HirBinding]) {
+        for binding in bindings {
+            if binding.kind != HirBindingKind::Param {
+                continue;
+            }
+            if let Some(type_name) = &binding.type_name {
+                self.record_type(binding.name.clone(), type_name.clone());
+            }
+        }
+    }
+
     pub(crate) fn bind_managed(&mut self, name: impl Into<String>) {
         self.managed.insert(name.into());
     }
@@ -36,6 +48,31 @@ impl BodyState {
     pub(crate) fn mark_retained(&mut self, name: &str) {
         self.clean_locals.remove(name);
     }
+
+    pub(crate) fn apply_move_events(&mut self, events: &[HirEffectEvent]) {
+        for event in events {
+            if !matches!(
+                event.kind,
+                HirEffectEventKind::Manage | HirEffectEventKind::Take
+            ) {
+                continue;
+            }
+            if self.locals.contains(&event.binding_name) {
+                self.mark_moved(&event.binding_name, event.span.clone());
+            }
+        }
+    }
+
+    pub(crate) fn apply_retention_events(&mut self, events: &[HirEffectEvent]) {
+        for event in events {
+            if !matches!(event.kind, HirEffectEventKind::Retain { .. }) {
+                continue;
+            }
+            if self.locals.contains(&event.binding_name) {
+                self.mark_retained(&event.binding_name);
+            }
+        }
+    }
 }
 
 pub(crate) fn merge_if_state(
@@ -158,3 +195,46 @@ fn merge_fallthrough_states(base: &BodyState, left: &BodyState, right: &BodyStat
             .collect(),
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn span(line: usize) -> Span {
+        Span {
+            file: "test.rss".to_string(),
+            line,
+            column: 1,
+            length: 1,
+        }
+    }
+
+    #[test]
+    fn applies_move_and_retain_events_to_clean_local_state() {
+        let mut state = BodyState::default();
+        state.bind_local("image");
+        state.bind_local("cached");
+
+        state.apply_retention_events(&[HirEffectEvent {
+            function_name: "run".to_string(),
+            kind: HirEffectEventKind::Retain {
+                callee: "Cache.store".to_string(),
+                param: "value".to_string(),
+            },
+            binding_name: "cached".to_string(),
+            span: span(10),
+            value_span: span(10),
+        }]);
+        state.apply_move_events(&[HirEffectEvent {
+            function_name: "run".to_string(),
+            kind: HirEffectEventKind::Manage,
+            binding_name: "image".to_string(),
+            span: span(11),
+            value_span: span(11),
+        }]);
+
+        assert!(!state.clean_locals.contains("cached"));
+        assert!(!state.clean_locals.contains("image"));
+        assert_eq!(state.moved["image"].line, 11);
+    }
+}