Use local HIR return proofs for freshness checks

Author: olwang

README.md

@@ -25,7 +25,7 @@ It currently implements:
 - Per-function HIR body views that group bindings, calls, fields, effects, and returns
 - AST-driven mode and call checks for local-only features, named arguments, data effects, and retaining APIs
 - AST-driven body checks for local moves, early-exit-aware `fresh` returns, resource escape, resolved handle-field `take`, and managed closure captures
-- Local ownership, resource-retain escape, and handle-field `take` checks now index binding types, field accesses, and move/retain events from resolved HIR body trees
+- Local ownership, `fresh` return, resource-retain escape, and handle-field `take` checks now index binding types, return proofs, field accesses, and move/retain events from resolved HIR body trees
 - Focused check modules for mode, calls, body semantics, and forbidden operator behavior
 - Fixture-based pass/fail scenario tests under `tests/fixtures`
 

src/checks/body.rs

@@ -101,7 +101,7 @@ fn check_stmt_semantics(
         Stmt::Return(stmt) => {
             if let Some(value) = &stmt.value {
                 if function.returns_fresh {
-                    check_fresh_return(analyzer, function, value, state);
+                    check_fresh_return(analyzer, local_analysis, function, value, state);
                 }
                 check_take_of_handle_field(analyzer, local_analysis, value, state);
             }
@@ -534,6 +534,7 @@ fn resource_capture_diagnostic(
 
 fn check_fresh_return(
     analyzer: &mut Analyzer<'_>,
+    local_analysis: &LocalAnalysis,
     function: &FunctionDecl,
     value: &Expr,
     state: &BodyState,
@@ -568,7 +569,7 @@ fn check_fresh_return(
             );
         }
         Expr::Call { callee, span, .. } => {
-            if !hir_return_proves_fresh_call(analyzer, callee, span) {
+            if !hir_return_proves_fresh_call(analyzer, local_analysis, callee, span) {
                 analyzer.diagnostics.push(
                     Diagnostic::warning(
                         code::FRESHNESS_UNKNOWN,
@@ -584,7 +585,7 @@ fn check_fresh_return(
             }
         }
         Expr::Effect { value, .. } | Expr::Manage { value, .. } => {
-            check_fresh_return(analyzer, function, value, state);
+            check_fresh_return(analyzer, local_analysis, function, value, state);
         }
         Expr::Field { span, .. } | Expr::Closure { span, .. } | Expr::Unknown(span) => {
             analyzer.diagnostics.push(
@@ -607,12 +608,13 @@ fn check_fresh_return(
 
 fn hir_return_proves_fresh_call(
     analyzer: &Analyzer<'_>,
+    local_analysis: &LocalAnalysis,
     callee: &Callee,
     span: &crate::diagnostic::Span,
 ) -> bool {
-    if let Some(return_fact) = analyzer.hir.return_fact(span) {
+    if let Some(return_proof) = local_analysis.return_proof(span) {
         return matches!(
-            return_fact.proof,
+            return_proof,
             HirReturnProof::StructConstructor | HirReturnProof::FreshCall
         );
     }

src/checks/local.rs

@@ -3,7 +3,7 @@ use std::collections::{HashMap, HashSet};
 use crate::diagnostic::Span;
 use crate::hir::{
     HirBinding, HirBindingKind, HirBlock, HirEffectEvent, HirEffectEventKind, HirExpr,
-    HirFieldAccess, HirFunctionBody, HirStmt,
+    HirFieldAccess, HirFunctionBody, HirReturnProof, HirStmt,
 };
 
 use super::body::Flow;
@@ -22,6 +22,7 @@ pub(crate) struct LocalAnalysis {
     events_by_span: HashMap<Span, Vec<HirEffectEvent>>,
     binding_types_by_span: HashMap<Span, String>,
     field_accesses_by_span: HashMap<Span, HirFieldAccess>,
+    return_proofs_by_span: HashMap<Span, HirReturnProof>,
 }
 
 impl LocalAnalysis {
@@ -39,12 +40,17 @@ impl LocalAnalysis {
             .as_ref()
             .and_then(|body| body.block.as_ref())
             .map_or_else(HashMap::new, index_field_accesses_from_block);
+        let return_proofs_by_span = body
+            .as_ref()
+            .and_then(|body| body.block.as_ref())
+            .map_or_else(HashMap::new, index_return_proofs_from_block);
 
         Self {
             body,
             events_by_span,
             binding_types_by_span,
             field_accesses_by_span,
+            return_proofs_by_span,
         }
     }
 
@@ -87,6 +93,10 @@ impl LocalAnalysis {
         self.field_accesses_by_span.get(span)
     }
 
+    pub(crate) fn return_proof(&self, span: &Span) -> Option<&HirReturnProof> {
+        self.return_proofs_by_span.get(span)
+    }
+
     fn effect_events(&self, span: &Span) -> &[HirEffectEvent] {
         self.events_by_span.get(span).map_or(&[], Vec::as_slice)
     }
@@ -317,6 +327,95 @@ fn collect_expr_field_accesses(expr: &HirExpr, accesses: &mut HashMap<Span, HirF
     }
 }
 
+fn index_return_proofs_from_block(block: &HirBlock) -> HashMap<Span, HirReturnProof> {
+    let mut proofs = HashMap::new();
+    collect_block_return_proofs(block, &mut proofs);
+    proofs
+}
+
+fn collect_block_return_proofs(block: &HirBlock, proofs: &mut HashMap<Span, HirReturnProof>) {
+    for statement in &block.statements {
+        collect_stmt_return_proofs(statement, proofs);
+    }
+}
+
+fn collect_stmt_return_proofs(statement: &HirStmt, proofs: &mut HashMap<Span, HirReturnProof>) {
+    match statement {
+        HirStmt::Return { value, proof, span } => {
+            let proof_span = value.as_ref().map_or(span, hir_expr_span);
+            proofs.insert(proof_span.clone(), proof.clone());
+            if let Some(value) = value {
+                collect_expr_return_proofs(value, proofs);
+            }
+        }
+        HirStmt::Let {
+            value: Some(value), ..
+        }
+        | HirStmt::Expr(value) => collect_expr_return_proofs(value, proofs),
+        HirStmt::With { resource, body, .. } => {
+            collect_expr_return_proofs(resource, proofs);
+            collect_block_return_proofs(body, proofs);
+        }
+        HirStmt::If {
+            condition,
+            then_body,
+            else_body,
+            ..
+        } => {
+            collect_expr_return_proofs(condition, proofs);
+            collect_block_return_proofs(then_body, proofs);
+            if let Some(else_body) = else_body {
+                collect_block_return_proofs(else_body, proofs);
+            }
+        }
+        HirStmt::Loop {
+            condition, body, ..
+        } => {
+            if let Some(condition) = condition {
+                collect_expr_return_proofs(condition, proofs);
+            }
+            collect_block_return_proofs(body, proofs);
+        }
+        HirStmt::Let { value: None, .. }
+        | HirStmt::Break(_)
+        | HirStmt::Continue(_)
+        | HirStmt::Unknown(_) => {}
+    }
+}
+
+fn collect_expr_return_proofs(expr: &HirExpr, proofs: &mut HashMap<Span, HirReturnProof>) {
+    match expr {
+        HirExpr::Call { args, .. } => {
+            for arg in args {
+                collect_expr_return_proofs(&arg.value, proofs);
+            }
+        }
+        HirExpr::Effect { value, .. } | HirExpr::Manage { value, .. } => {
+            collect_expr_return_proofs(value, proofs);
+        }
+        HirExpr::Field { base, .. } => collect_expr_return_proofs(base, proofs),
+        HirExpr::Closure { body, .. } => collect_block_return_proofs(body, proofs),
+        HirExpr::Ident { .. }
+        | HirExpr::Number { .. }
+        | HirExpr::String { .. }
+        | HirExpr::Unknown(_) => {}
+    }
+}
+
+fn hir_expr_span(expr: &HirExpr) -> &Span {
+    match expr {
+        HirExpr::Ident { span, .. }
+        | HirExpr::Number { span, .. }
+        | HirExpr::String { span, .. }
+        | HirExpr::Field { span, .. }
+        | HirExpr::Call { span, .. }
+        | HirExpr::Effect { span, .. }
+        | HirExpr::Manage { span, .. }
+        | HirExpr::Closure { span, .. }
+        | HirExpr::Unknown(span) => span,
+    }
+}
+
 impl BodyState {
     pub(crate) fn seed_params(&mut self, bindings: &[HirBinding]) {
         for binding in bindings {
@@ -635,6 +734,18 @@ mod tests {
                         },
                         span: span(22),
                     }),
+                    HirStmt::Return {
+                        value: Some(HirExpr::Call {
+                            callee: Callee::Name("Image.load".to_string()),
+                            args: Vec::new(),
+                            resolution: CallResolution::Unknown,
+                            events: Vec::new(),
+                            type_name: Some("Image".to_string()),
+                            span: span(23),
+                        }),
+                        proof: HirReturnProof::FreshCall,
+                        span: span(23),
+                    },
                 ],
                 span: span(1),
             }),
@@ -656,6 +767,10 @@ mod tests {
                 .field_access(&span(22))
                 .is_some_and(|access| access.is_handle)
         );
+        assert!(matches!(
+            local_analysis.return_proof(&span(23)),
+            Some(HirReturnProof::FreshCall)
+        ));
 
         local_analysis.apply_retention_events(&span(20), &mut state);
         local_analysis.apply_move_events(&span(21), &mut state);

src/hir.rs

@@ -283,7 +283,6 @@ pub struct Hir {
     field_accesses: Vec<HirFieldAccess>,
     effect_events: Vec<HirEffectEvent>,
     returns: Vec<HirReturn>,
-    returns_by_span: HashMap<Span, HirReturn>,
     function_bodies: HashMap<String, HirFunctionBody>,
 }
 
@@ -369,10 +368,6 @@ impl Hir {
         self.function_bodies.get(function_name)
     }
 
-    pub fn return_fact(&self, span: &Span) -> Option<&HirReturn> {
-        self.returns_by_span.get(span)
-    }
-
     pub fn resolve_call(&self, callee: &Callee) -> CallResolution {
         let call_name = callee_name(callee);
         if is_enum_variant_call(call_name) {
@@ -440,11 +435,6 @@ impl Hir {
             .iter()
             .map(|site| (site.span.clone(), site.resolution.clone()))
             .collect();
-        self.returns_by_span = facts
-            .returns
-            .iter()
-            .map(|return_fact| (return_fact.span.clone(), return_fact.clone()))
-            .collect();
         self.function_bodies = build_function_bodies(&facts);
         self.call_sites = facts.call_sites;
         self.bindings = facts.bindings;
@@ -1791,12 +1781,6 @@ fn render(body: read String) -> Result<fresh Response, HttpError> {
             returns[0].proof,
             HirReturnProof::Ident { ref name } if name == "response"
         ));
-        assert!(matches!(
-            hir.return_fact(&returns[0].span)
-                .expect("return lookup by span works")
-                .proof,
-            HirReturnProof::Ident { .. }
-        ));
 
         let body = hir.function_body("render").expect("function body exists");
         assert_eq!(body.function_name, "render");