Merge pull request #504 from AeneasVerif/son/cfg

Improve the CFG reconstruction

charon/src/transform/duplicate_return.rs

@@ -0,0 +1,79 @@
+//! The MIR uses a unique `return` node, which can be an issue when reconstructing
+//! the control-flow.
+//!
+//! For instance, it often leads to code of the following shape:
+//! ```text
+//! if b {
+//!   ...
+//!   x = 0;
+//! }
+//! else {
+//!   ...
+//!   x = 1;
+//! }
+//! return x;
+//! ```
+//!
+//! while a more natural reconstruction would be:
+//! ```text
+//! if b {
+//!   ...
+//!   return 0;
+//! }
+//! else {
+//!   ...
+//!   return 1;
+//! }
+//! ```
+
+use crate::ids::Generator;
+use crate::transform::TransformCtx;
+use crate::ullbc_ast::*;
+use derive_visitor::{visitor_enter_fn_mut, DriveMut};
+use std::collections::HashMap;
+
+use super::ctx::UllbcPass;
+
+pub struct Transform;
+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
+            .body
+            .iter_indexed()
+            .filter_map(|(bid, block)| {
+                if block.statements.is_empty() && block.terminator.content.is_return() {
+                    Some((bid, block.terminator.span))
+                } else {
+                    None
+                }
+            })
+            .collect();
+
+        // Whenever we find a goto the return block, introduce an auxiliary block
+        // for this (remark: in the end, it makes the return block dangling).
+        // We do this in two steps.
+        // First, introduce fresh ids.
+        assert!(usize::from(b.body.next_id()) == b.body.len());
+        let mut generator = Generator::new_with_init_value(b.body.len());
+        let mut new_spans = Vec::new();
+        b.body
+            .drive_mut(&mut visitor_enter_fn_mut(|bid: &mut BlockId| {
+                if let Some(span) = returns.get(bid) {
+                    *bid = generator.fresh_id();
+                    new_spans.push(*span);
+                }
+            }));
+
+        // Then introduce the new blocks
+        for span in new_spans {
+            let _ = b.body.push(BlockData {
+                statements: Vec::new(),
+                terminator: Terminator {
+                    span,
+                    content: RawTerminator::Return,
+                },
+            });
+        }
+    }
+}

charon/src/transform/filter_useless_blocks.rs

@@ -0,0 +1,42 @@
+//! Some passes like [`reconstruct_assert`] lead to the apparition of "dangling" blocks,
+//! which are referenced nowhere and thus become useless. This pass filters those out.
+
+use std::collections::{HashMap, HashSet};
+
+use crate::transform::TransformCtx;
+use crate::ullbc_ast::*;
+use derive_visitor::{visitor_enter_fn_mut, DriveMut};
+
+use super::ctx::UllbcPass;
+
+pub struct Transform;
+impl UllbcPass for Transform {
+    fn transform_body(&self, _ctx: &mut TransformCtx, b: &mut ExprBody) {
+        // Perform a depth-first search to identify all the blocks reachable
+        // from the first block.
+        let mut explored: HashSet<BlockId> = HashSet::new();
+        let mut to_explore: Vec<BlockId> = vec![BlockId::from_usize(0)];
+        while let Some(bid) = to_explore.pop() {
+            if explored.contains(&bid) {
+                continue;
+            }
+            explored.insert(bid);
+            to_explore.append(&mut b.body[bid].targets())
+        }
+
+        // Renumerotate
+        let mut bid_map: HashMap<BlockId, BlockId> = HashMap::new();
+        for (bid, block) in std::mem::take(&mut b.body).into_iter_indexed() {
+            if explored.contains(&bid) {
+                let nbid = b.body.push(block);
+                bid_map.insert(bid, nbid);
+            }
+        }
+
+        // Update all block ids
+        b.body
+            .drive_mut(&mut visitor_enter_fn_mut(|bid: &mut BlockId| {
+                *bid = *bid_map.get(bid).unwrap();
+            }));
+    }
+}

charon/src/transform/mod.rs

@@ -1,6 +1,8 @@
 pub mod check_generics;
 pub mod ctx;
+pub mod duplicate_return;
 pub mod filter_invisible_trait_impls;
+pub mod filter_useless_blocks;
 pub mod graphs;
 pub mod hide_marker_traits;
 pub mod index_intermediate_assigns;
@@ -77,11 +79,16 @@ pub static ULLBC_PASSES: &[Pass] = &[
     UnstructuredBody(&ops_to_function_calls::Transform),
     // # Micro-pass: make sure the block ids used in the ULLBC are consecutive
     UnstructuredBody(&update_block_indices::Transform),
+    // # Micro-pass: reconstruct the asserts
+    UnstructuredBody(&reconstruct_asserts::Transform),
+    // # Micro-pass: duplicate the return blocks
+    UnstructuredBody(&duplicate_return::Transform),
+    // # Micro-pass: filter the "dangling" blocks. Those might have been introduced by,
+    // for instance, [`reconstruct_asserts`].
+    UnstructuredBody(&filter_useless_blocks::Transform),
 ];
 
 pub static LLBC_PASSES: &[Pass] = &[
-    // # Micro-pass: reconstruct the asserts
-    StructuredBody(&reconstruct_asserts::Transform),
     // # Micro-pass: `panic!()` expands to a new function definition each time. This pass cleans
     // those up.
     StructuredBody(&inline_local_panic_functions::Transform),

charon/src/transform/reconstruct_asserts.rs

@@ -3,60 +3,60 @@
 //! introduce `if ... then { panic!(...) } else { ...}`.
 //! This pass introduces `assert` instead in order to make the code shorter.
 
-use std::mem;
+use std::collections::HashSet;
 
-use crate::llbc_ast::*;
 use crate::transform::TransformCtx;
+use crate::ullbc_ast::*;
 
-use super::ctx::LlbcPass;
-
-fn transform_st(st: &mut Statement) -> Vec<Statement> {
-    if let RawStatement::Switch(Switch::If(_, then_block, else_block)) = &mut st.content {
-        if let Some(first) = then_block.statements.first()
-            && first.content.is_abort()
-        {
-            // Replace the `if` with an `assert`.
-            let (op, then_block, else_block) = mem::replace(&mut st.content, RawStatement::Nop)
-                .to_switch()
-                .unwrap()
-                .to_if()
-                .unwrap();
-            let assert = Statement::new(
-                then_block.span,
-                RawStatement::Assert(Assert {
-                    cond: op,
-                    expected: false,
-                }),
-            );
-            [assert].into_iter().chain(else_block.statements).collect()
-        } else if let Some(first) = else_block.statements.first()
-            && first.content.is_abort()
-        {
-            // Replace the `if` with an `assert`.
-            let (op, then_block, else_block) = mem::replace(&mut st.content, RawStatement::Nop)
-                .to_switch()
-                .unwrap()
-                .to_if()
-                .unwrap();
-            let assert = Statement::new(
-                else_block.span,
-                RawStatement::Assert(Assert {
-                    cond: op,
-                    expected: true,
-                }),
-            );
-            [assert].into_iter().chain(then_block.statements).collect()
-        } else {
-            Vec::new()
-        }
-    } else {
-        Vec::new()
-    }
-}
+use super::ctx::UllbcPass;
 
 pub struct Transform;
-impl LlbcPass for Transform {
+impl UllbcPass for Transform {
     fn transform_body(&self, _ctx: &mut TransformCtx, b: &mut ExprBody) {
-        b.body.transform(&mut transform_st);
+        // Start by computing the set of blocks which are actually panics.
+        // Remark: doing this in two steps because reading the blocks at random
+        // while doing in-place updates is not natural to do in Rust.
+        let panics: HashSet<BlockId> = b
+            .body
+            .iter_indexed()
+            .filter_map(|(bid, block)| {
+                if block.statements.is_empty() && block.terminator.content.is_abort() {
+                    Some(bid)
+                } else {
+                    None
+                }
+            })
+            .collect();
+
+        for block in b.body.iter_mut() {
+            match &block.terminator.content {
+                RawTerminator::Switch {
+                    discr: _,
+                    targets: SwitchTargets::If(bid0, bid1),
+                } => {
+                    let (nbid, expected) = if panics.contains(bid0) {
+                        (*bid1, false)
+                    } else if panics.contains(bid1) {
+                        (*bid0, true)
+                    } else {
+                        continue;
+                    };
+
+                    let content = std::mem::replace(
+                        &mut block.terminator.content,
+                        RawTerminator::Goto { target: nbid },
+                    );
+                    let (discr, _) = content.as_switch().unwrap();
+                    block.statements.push(Statement {
+                        span: block.terminator.span,
+                        content: RawStatement::Assert(Assert {
+                            cond: discr.clone(),
+                            expected,
+                        }),
+                    });
+                }
+                _ => (),
+            }
+        }
     }
 }

charon/tests/cargo/toml.out

@@ -19,12 +19,13 @@ where
 
     match *(self@1) {
         1 => {
-            @0 := const (true)
         },
         0 => {
             @0 := const (false)
+            return
         },
     }
+    @0 := const (true)
     return
 }
 

charon/tests/ui/arrays.out

@@ -1859,15 +1859,17 @@ fn test_crate::slice_pattern_4<'_0>(@1: &'_0 (Slice<()>))
     @3 := const (1 : usize)
     @4 := move (@2) == move (@3)
     if move (@4) {
-        @6 := &*(x@1)
-        @7 := @SliceIndexShared<'_, ()>(move (@6), const (0 : usize))
-        _named@5 := &*(@7)
-        @0 := ()
-        drop _named@5
     }
     else {
         @0 := ()
+        @0 := ()
+        return
     }
+    @6 := &*(x@1)
+    @7 := @SliceIndexShared<'_, ()>(move (@6), const (0 : usize))
+    _named@5 := &*(@7)
+    @0 := ()
+    drop _named@5
     @0 := ()
     return
 }

charon/tests/ui/closures.out

@@ -110,7 +110,6 @@ where
     @fake_read(x@1)
     match *(x@1) {
         0 => {
-            @0 := core::option::Option::None {  }
         },
         1 => {
             x@3 := &(*(x@1) as variant @1).0
@@ -123,8 +122,10 @@ where
             drop @4
             drop @4
             drop x@3
+            return
         },
     }
+    @0 := core::option::Option::None {  }
     return
 }
 

charon/tests/ui/issue-114-opaque-bodies.out

@@ -74,14 +74,15 @@ where
     let @3: T; // anonymous local
 
     if copy (self@1) {
-        @3 := move (t@2)
-        @0 := core::option::Option::Some { 0: move (@3) }
-        drop @3
     }
     else {
         @0 := core::option::Option::None {  }
         drop t@2
+        return
     }
+    @3 := move (t@2)
+    @0 := core::option::Option::Some { 0: move (@3) }
+    drop @3
     return
 }
 

charon/tests/ui/matches.out

@@ -14,12 +14,13 @@ fn test_crate::test1(@1: test_crate::E1) -> bool
     @fake_read(x@1)
     match x@1 {
         0 | 1 => {
-            @0 := const (true)
         },
         2 => {
             @0 := const (false)
+            return
         },
     }
+    @0 := const (true)
     return
 }
 
@@ -53,18 +54,17 @@ fn test_crate::test2(@1: test_crate::E2) -> u32
     match x@1 {
         0 => {
             n@2 := copy ((x@1 as variant @0).0)
-            @0 := copy (n@2)
-            drop n@2
         },
         1 => {
             n@2 := copy ((x@1 as variant @1).0)
-            @0 := copy (n@2)
-            drop n@2
         },
         2 => {
             @0 := const (0 : u32)
+            return
         },
     }
+    @0 := copy (n@2)
+    drop n@2
     return
 }