Insert StorageDeads for all locals

charon-ml/src/generated/Generated_LlbcAst.ml

@@ -33,15 +33,17 @@ and statement_kind =
           modelled as a function call as it cannot diverge *)
   | StorageLive of local_id
       (** Indicates that this local should be allocated; if it is already
-          allocated, this frees the local and re-allocates it. The return value
-          and arguments do not receive a [StorageLive]. We ensure in the
-          micro-pass [insert_storage_lives] that all other locals have a
-          [StorageLive] associated with them. *)
+          allocated, this frees the local and re-allocates it. The arguments do
+          not receive a [StorageLive]. We ensure in the micro-pass
+          [insert_storage_statements] that all other locals have a [StorageLive]
+          associated with them. *)
   | StorageDead of local_id
       (** Indicates that this local should be deallocated; if it is already
           deallocated, this is a no-op. A local may not have a [StorageDead] in
           the function's body, in which case it is implicitly deallocated at the
-          end of the function. *)
+          end of the function. The return local does not receive a
+          [StorageDead]. We ensure in the micro-pass [insert_storage_statements]
+          that all other locals have a [StorageDead] before function exits. *)
   | PlaceMention of place
       (** A place is mentioned, but not accessed. The place itself must still be
           valid though, so this statement is not a no-op: it can trigger UB if

charon-ml/src/generated/Generated_UllbcAst.ml

@@ -27,15 +27,17 @@ and statement_kind =
           modelled as a function call as it cannot diverge *)
   | StorageLive of local_id
       (** Indicates that this local should be allocated; if it is already
-          allocated, this frees the local and re-allocates it. The return value
-          and arguments do not receive a [StorageLive]. We ensure in the
-          micro-pass [insert_storage_lives] that all other locals have a
-          [StorageLive] associated with them. *)
+          allocated, this frees the local and re-allocates it. The arguments do
+          not receive a [StorageLive]. We ensure in the micro-pass
+          [insert_storage_statements] that all other locals have a [StorageLive]
+          associated with them. *)
   | StorageDead of local_id
       (** Indicates that this local should be deallocated; if it is already
           deallocated, this is a no-op. A local may not have a [StorageDead] in
           the function's body, in which case it is implicitly deallocated at the
-          end of the function. *)
+          end of the function. The return local does not receive a
+          [StorageDead]. We ensure in the micro-pass [insert_storage_statements]
+          that all other locals have a [StorageDead] before function exits. *)
   | PlaceMention of place
       (** A place is mentioned, but not accessed. The place itself must still be
           valid though, so this statement is not a no-op: it can trigger UB if

charon/src/ast/llbc_ast.rs

@@ -38,13 +38,15 @@ pub enum StatementKind {
     /// call as it cannot diverge
     CopyNonOverlapping(Box<CopyNonOverlapping>),
     /// Indicates that this local should be allocated; if it is already allocated, this frees
-    /// the local and re-allocates it. The return value and arguments do not receive a
-    /// `StorageLive`. We ensure in the micro-pass `insert_storage_lives` that all other locals
-    /// have a `StorageLive` associated with them.
+    /// the local and re-allocates it. The arguments do not receive a `StorageLive`. We ensure in
+    /// the micro-pass `insert_storage_statements` that all other locals have a `StorageLive`
+    /// associated with them.
     StorageLive(LocalId),
     /// Indicates that this local should be deallocated; if it is already deallocated, this is
     /// a no-op. A local may not have a `StorageDead` in the function's body, in which case it
-    /// is implicitly deallocated at the end of the function.
+    /// is implicitly deallocated at the end of the function. The return local does not receive a
+    /// `StorageDead`. We ensure in the micro-pass `insert_storage_statements` that all other locals
+    /// have a `StorageDead` before function exits.
     StorageDead(LocalId),
     /// A place is mentioned, but not accessed. The place itself must still be valid though, so
     /// this statement is not a no-op: it can trigger UB if the place's projections are not valid

charon/src/ast/ullbc_ast.rs

@@ -38,13 +38,15 @@ pub enum StatementKind {
     /// call as it cannot diverge
     CopyNonOverlapping(Box<CopyNonOverlapping>),
     /// Indicates that this local should be allocated; if it is already allocated, this frees
-    /// the local and re-allocates it. The return value and arguments do not receive a
-    /// `StorageLive`. We ensure in the micro-pass `insert_storage_lives` that all other locals
-    /// have a `StorageLive` associated with them.
+    /// the local and re-allocates it. The arguments do not receive a `StorageLive`. We ensure in
+    /// the micro-pass `insert_storage_statements` that all other locals have a `StorageLive`
+    /// associated with them.
     StorageLive(LocalId),
     /// Indicates that this local should be deallocated; if it is already deallocated, this is
     /// a no-op. A local may not have a `StorageDead` in the function's body, in which case it
-    /// is implicitly deallocated at the end of the function.
+    /// is implicitly deallocated at the end of the function. The return local does not receive a
+    /// `StorageDead`. We ensure in the micro-pass `insert_storage_statements` that all other locals
+    /// have a `StorageDead` before function exits.
     StorageDead(LocalId),
     /// A place is mentioned, but not accessed. The place itself must still be valid though, so
     /// this statement is not a no-op: it can trigger UB if the place's projections are not valid

charon/src/transform/control_flow/duplicate_return.rs

@@ -34,15 +34,23 @@ use std::collections::HashMap;
 use crate::transform::ctx::UllbcPass;
 
 pub struct Transform;
+fn is_return_block(block: &BlockData) -> bool {
+    block.terminator.kind.is_return()
+        && block
+            .statements
+            .iter()
+            .all(|st| matches!(st.kind, StatementKind::StorageDead(_)))
+}
+
 impl UllbcPass for Transform {
     fn transform_body(&self, _ctx: &mut TransformCtx, b: &mut ExprBody) {
         // Find the return block id (there should be one).
-        let returns: HashMap<BlockId, Span> = b
+        let returns: HashMap<BlockId, BlockData> = b
             .body
             .iter_enumerated()
             .filter_map(|(bid, block)| {
-                if block.statements.is_empty() && block.terminator.kind.is_return() {
-                    Some((bid, block.terminator.span))
+                if is_return_block(block) {
+                    Some((bid, block.clone()))
                 } else {
                     None
                 }
@@ -54,20 +62,17 @@ impl UllbcPass for Transform {
         // We do this in two steps.
         // First, introduce fresh ids.
         let mut generator = Generator::new_with_init_value(b.body.next_idx());
-        let mut new_spans = Vec::new();
+        let mut new_blocks = Vec::new();
         b.body.dyn_visit_in_body_mut(|bid: &mut BlockId| {
-            if let Some(span) = returns.get(bid) {
+            if let Some(block) = returns.get(bid) {
                 *bid = generator.fresh_id();
-                new_spans.push(*span);
+                new_blocks.push(block.clone());
             }
         });
 
         // Then introduce the new blocks
-        for span in new_spans {
-            let _ = b.body.push(BlockData {
-                statements: Vec::new(),
-                terminator: Terminator::new(span, TerminatorKind::Return),
-            });
+        for block in new_blocks {
+            let _ = b.body.push(block);
         }
     }
 }