Extract clean local state machine

Author: olwang

src/checks/body.rs

@@ -1,5 +1,3 @@
-use std::collections::{HashMap, HashSet};
-
 use crate::analyzer::Analyzer;
 use crate::diagnostic::{Diagnostic, code};
 use crate::hir::{
@@ -8,6 +6,8 @@ use crate::hir::{
 };
 use crate::syntax::ast::{Block, Callee, DataEffect, Expr, FunctionDecl, Item, LetKind, Stmt};
 
+use super::local::{BodyState, merge_if_state, merge_loop_state};
+
 pub(crate) fn check(analyzer: &mut Analyzer<'_>) {
     let functions: Vec<FunctionDecl> = analyzer
         .syntax_program
@@ -26,15 +26,6 @@ pub(crate) fn check(analyzer: &mut Analyzer<'_>) {
     }
 }
 
-#[derive(Debug, Clone, Default)]
-struct BodyState {
-    locals: HashSet<String>,
-    clean_locals: HashSet<String>,
-    managed: HashSet<String>,
-    moved: HashMap<String, crate::diagnostic::Span>,
-    value_types: HashMap<String, String>,
-}
-
 fn seed_function_bindings(analyzer: &Analyzer<'_>, function: &FunctionDecl, state: &mut BodyState) {
     let Some(body) = analyzer.hir.function_body(&function.name) else {
         return;
@@ -44,15 +35,13 @@ fn seed_function_bindings(analyzer: &Analyzer<'_>, function: &FunctionDecl, stat
             continue;
         }
         if let Some(type_name) = &binding.type_name {
-            state
-                .value_types
-                .insert(binding.name.clone(), type_name.clone());
+            state.record_type(binding.name.clone(), type_name.clone());
         }
     }
 }
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
-enum Flow {
+pub(crate) enum Flow {
     Fallthrough,
     Return,
     Break,
@@ -180,18 +169,17 @@ fn apply_stmt_effects(analyzer: &mut Analyzer<'_>, statement: &Stmt, state: &mut
         Stmt::Let(stmt) => {
             match stmt.kind {
                 LetKind::Managed => {
-                    state.managed.insert(stmt.name.clone());
+                    state.bind_managed(stmt.name.clone());
                 }
                 LetKind::Local => {
-                    state.locals.insert(stmt.name.clone());
-                    state.clean_locals.insert(stmt.name.clone());
+                    state.bind_local(stmt.name.clone());
                 }
             }
             if let Some(value) = &stmt.value {
                 if let Some(type_name) = hir_binding_type(analyzer, &stmt.span)
                     .or_else(|| infer_expr_type(analyzer, value, state))
                 {
-                    state.value_types.insert(stmt.name.clone(), type_name);
+                    state.record_type(stmt.name.clone(), type_name);
                 }
                 apply_expr_effects(analyzer, value, state);
             }
@@ -219,127 +207,6 @@ fn hir_binding_type(analyzer: &Analyzer<'_>, span: &crate::diagnostic::Span) ->
         .and_then(|binding| binding.type_name.clone())
 }
 
-fn merge_if_state(
-    state: &mut BodyState,
-    base: &BodyState,
-    then_state: BodyState,
-    then_flow: Flow,
-    else_branch: Option<(BodyState, Flow)>,
-) -> Flow {
-    let (else_state, else_flow) = else_branch.unwrap_or_else(|| (base.clone(), Flow::Fallthrough));
-    let mut fallthrough_states = Vec::new();
-    if then_flow == Flow::Fallthrough {
-        fallthrough_states.push(then_state.clone());
-    }
-    if else_flow == Flow::Fallthrough {
-        fallthrough_states.push(else_state.clone());
-    }
-
-    match fallthrough_states.as_slice() {
-        [] => {
-            *state = base.clone();
-            merge_non_fallthrough(then_flow, else_flow)
-        }
-        [only] => {
-            *state = fallthrough_projection(base, only);
-            Flow::Fallthrough
-        }
-        [left, right] => {
-            *state = merge_fallthrough_states(base, left, right);
-            Flow::Fallthrough
-        }
-        _ => unreachable!("if has at most two branches"),
-    }
-}
-
-fn merge_loop_state(
-    state: &mut BodyState,
-    base: &BodyState,
-    body_state: BodyState,
-    body_flow: Flow,
-    may_skip: bool,
-) -> Flow {
-    if !may_skip && body_flow == Flow::Return {
-        *state = base.clone();
-        return Flow::Return;
-    }
-    if !may_skip && body_flow == Flow::Break {
-        *state = fallthrough_projection(base, &body_state);
-        return Flow::Fallthrough;
-    }
-
-    let mut moved = base.moved.clone();
-    if body_flow != Flow::Return {
-        for (name, span) in &body_state.moved {
-            if base.locals.contains(name) || base.moved.contains_key(name) {
-                moved.entry(name.clone()).or_insert_with(|| span.clone());
-            }
-        }
-    }
-
-    state.locals = base.locals.clone();
-    state.managed = base.managed.clone();
-    state.value_types = base.value_types.clone();
-    state.moved = moved;
-    state.clean_locals = base
-        .clean_locals
-        .intersection(&body_state.clean_locals)
-        .filter(|name| base.locals.contains(*name))
-        .cloned()
-        .collect();
-    Flow::Fallthrough
-}
-
-fn merge_non_fallthrough(left: Flow, right: Flow) -> Flow {
-    if left == right { left } else { Flow::Return }
-}
-
-fn fallthrough_projection(base: &BodyState, branch: &BodyState) -> BodyState {
-    let mut moved = base.moved.clone();
-    for (name, span) in &branch.moved {
-        if base.locals.contains(name) || base.moved.contains_key(name) {
-            moved.entry(name.clone()).or_insert_with(|| span.clone());
-        }
-    }
-
-    BodyState {
-        locals: base.locals.clone(),
-        managed: base.managed.clone(),
-        value_types: base.value_types.clone(),
-        moved,
-        clean_locals: branch
-            .clean_locals
-            .intersection(&base.clean_locals)
-            .filter(|name| base.locals.contains(*name))
-            .cloned()
-            .collect(),
-    }
-}
-
-fn merge_fallthrough_states(base: &BodyState, left: &BodyState, right: &BodyState) -> BodyState {
-    let mut moved = base.moved.clone();
-    for branch in [left, right] {
-        for (name, span) in &branch.moved {
-            if base.locals.contains(name) || base.moved.contains_key(name) {
-                moved.entry(name.clone()).or_insert_with(|| span.clone());
-            }
-        }
-    }
-
-    BodyState {
-        locals: base.locals.clone(),
-        managed: base.managed.clone(),
-        value_types: base.value_types.clone(),
-        moved,
-        clean_locals: left
-            .clean_locals
-            .intersection(&right.clean_locals)
-            .filter(|name| base.locals.contains(*name))
-            .cloned()
-            .collect(),
-    }
-}
-
 fn apply_expr_effects(analyzer: &mut Analyzer<'_>, expr: &Expr, state: &mut BodyState) {
     match expr {
         Expr::Manage { value, span } => {
@@ -383,10 +250,7 @@ fn apply_move_events(
             continue;
         }
         if state.locals.contains(&event.binding_name) {
-            state
-                .moved
-                .insert(event.binding_name.clone(), event.span.clone());
-            state.clean_locals.remove(&event.binding_name);
+            state.mark_moved(&event.binding_name, event.span.clone());
         }
     }
 }
@@ -403,7 +267,7 @@ fn apply_retention_events(
         if !state.locals.contains(&event.binding_name) {
             continue;
         }
-        state.clean_locals.remove(&event.binding_name);
+        state.mark_retained(&event.binding_name);
     }
 }
 

src/checks/local.rs

@@ -0,0 +1,160 @@
+use std::collections::{HashMap, HashSet};
+
+use crate::diagnostic::Span;
+
+use super::body::Flow;
+
+#[derive(Debug, Clone, Default)]
+pub(crate) struct BodyState {
+    pub(crate) locals: HashSet<String>,
+    pub(crate) clean_locals: HashSet<String>,
+    pub(crate) managed: HashSet<String>,
+    pub(crate) moved: HashMap<String, Span>,
+    pub(crate) value_types: HashMap<String, String>,
+}
+
+impl BodyState {
+    pub(crate) fn bind_managed(&mut self, name: impl Into<String>) {
+        self.managed.insert(name.into());
+    }
+
+    pub(crate) fn bind_local(&mut self, name: impl Into<String>) {
+        let name = name.into();
+        self.locals.insert(name.clone());
+        self.clean_locals.insert(name);
+    }
+
+    pub(crate) fn record_type(&mut self, name: impl Into<String>, type_name: impl Into<String>) {
+        self.value_types.insert(name.into(), type_name.into());
+    }
+
+    pub(crate) fn mark_moved(&mut self, name: &str, span: Span) {
+        self.moved.insert(name.to_string(), span);
+        self.clean_locals.remove(name);
+    }
+
+    pub(crate) fn mark_retained(&mut self, name: &str) {
+        self.clean_locals.remove(name);
+    }
+}
+
+pub(crate) fn merge_if_state(
+    state: &mut BodyState,
+    base: &BodyState,
+    then_state: BodyState,
+    then_flow: Flow,
+    else_branch: Option<(BodyState, Flow)>,
+) -> Flow {
+    let (else_state, else_flow) = else_branch.unwrap_or_else(|| (base.clone(), Flow::Fallthrough));
+    let mut fallthrough_states = Vec::new();
+    if then_flow == Flow::Fallthrough {
+        fallthrough_states.push(then_state.clone());
+    }
+    if else_flow == Flow::Fallthrough {
+        fallthrough_states.push(else_state.clone());
+    }
+
+    match fallthrough_states.as_slice() {
+        [] => {
+            *state = base.clone();
+            merge_non_fallthrough(then_flow, else_flow)
+        }
+        [only] => {
+            *state = fallthrough_projection(base, only);
+            Flow::Fallthrough
+        }
+        [left, right] => {
+            *state = merge_fallthrough_states(base, left, right);
+            Flow::Fallthrough
+        }
+        _ => unreachable!("if has at most two branches"),
+    }
+}
+
+pub(crate) fn merge_loop_state(
+    state: &mut BodyState,
+    base: &BodyState,
+    body_state: BodyState,
+    body_flow: Flow,
+    may_skip: bool,
+) -> Flow {
+    if !may_skip && body_flow == Flow::Return {
+        *state = base.clone();
+        return Flow::Return;
+    }
+    if !may_skip && body_flow == Flow::Break {
+        *state = fallthrough_projection(base, &body_state);
+        return Flow::Fallthrough;
+    }
+
+    let mut moved = base.moved.clone();
+    if body_flow != Flow::Return {
+        for (name, span) in &body_state.moved {
+            if base.locals.contains(name) || base.moved.contains_key(name) {
+                moved.entry(name.clone()).or_insert_with(|| span.clone());
+            }
+        }
+    }
+
+    state.locals = base.locals.clone();
+    state.managed = base.managed.clone();
+    state.value_types = base.value_types.clone();
+    state.moved = moved;
+    state.clean_locals = base
+        .clean_locals
+        .intersection(&body_state.clean_locals)
+        .filter(|name| base.locals.contains(*name))
+        .cloned()
+        .collect();
+    Flow::Fallthrough
+}
+
+fn merge_non_fallthrough(left: Flow, right: Flow) -> Flow {
+    if left == right { left } else { Flow::Return }
+}
+
+fn fallthrough_projection(base: &BodyState, branch: &BodyState) -> BodyState {
+    let mut moved = base.moved.clone();
+    for (name, span) in &branch.moved {
+        if base.locals.contains(name) || base.moved.contains_key(name) {
+            moved.entry(name.clone()).or_insert_with(|| span.clone());
+        }
+    }
+
+    BodyState {
+        locals: base.locals.clone(),
+        managed: base.managed.clone(),
+        value_types: base.value_types.clone(),
+        moved,
+        clean_locals: branch
+            .clean_locals
+            .intersection(&base.clean_locals)
+            .filter(|name| base.locals.contains(*name))
+            .cloned()
+            .collect(),
+    }
+}
+
+fn merge_fallthrough_states(base: &BodyState, left: &BodyState, right: &BodyState) -> BodyState {
+    let mut moved = base.moved.clone();
+    for branch in [left, right] {
+        for (name, span) in &branch.moved {
+            if base.locals.contains(name) || base.moved.contains_key(name) {
+                moved.entry(name.clone()).or_insert_with(|| span.clone());
+            }
+        }
+    }
+
+    BodyState {
+        locals: base.locals.clone(),
+        managed: base.managed.clone(),
+        value_types: base.value_types.clone(),
+        moved,
+        clean_locals: left
+            .clean_locals
+            .intersection(&right.clean_locals)
+            .filter(|name| base.locals.contains(*name))
+            .cloned()
+            .collect(),
+    }
+}

src/checks/mod.rs

@@ -1,4 +1,5 @@
 pub(crate) mod body;
 pub(crate) mod calls;
 pub(crate) mod forbidden;
+pub(crate) mod local;
 pub(crate) mod mode;