Merge pull request #261 from Nadrieril/sequence-as-vec

charon-ml/src/CharonVersion.ml

@@ -1,3 +1,3 @@
 (* This is an automatically generated file, generated from `charon/Cargo.toml`. *)
 (* To re-generate this file, rune `make` in the root directory *)
-let supported_charon_version = "0.1.10"
+let supported_charon_version = "0.1.11"

charon-ml/src/LlbcOfJson.ml

@@ -60,10 +60,21 @@ and raw_statement_of_json (id_to_file : id_to_file_map) (js : json) :
         let* i = int_of_json i in
         Ok (Continue i)
     | `String "Nop" -> Ok Nop
-    | `Assoc [ ("Sequence", `List [ st1; st2 ]) ] ->
-        let* st1 = statement_of_json id_to_file st1 in
-        let* st2 = statement_of_json id_to_file st2 in
-        Ok (Sequence (st1, st2))
+    (* We get a list from the rust side, which we fold into our recursive `Sequence` representation. *)
+    | `Assoc [ ("Sequence", `List seq) ] -> (
+        let seq = List.map (statement_of_json id_to_file) seq in
+        match List.rev seq with
+        | [] -> Ok Nop
+        | last :: rest ->
+            let* seq =
+              List.fold_left
+                (fun acc st ->
+                  let* st = st in
+                  let* acc = acc in
+                  Ok { span = st.span; content = Sequence (st, acc) })
+                last rest
+            in
+            Ok seq.content)
     | `Assoc [ ("Switch", tgt) ] ->
         let* switch = switch_of_json id_to_file tgt in
         Ok (Switch switch)

charon/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "charon"
-version = "0.1.10"
+version = "0.1.11"
 authors = ["Son Ho <hosonmarc@gmail.com>"]
 edition = "2021"
 

charon/src/ast/gast.rs

@@ -43,7 +43,7 @@ pub struct Var {
 }
 
 /// Marker to indicate that a declaration is opaque (i.e. we don't inspect its body).
-#[derive(Debug, Clone, Serialize, Deserialize, Drive, DriveMut)]
+#[derive(Debug, Copy, Clone, Serialize, Deserialize, Drive, DriveMut)]
 pub struct Opaque;
 
 /// An expression body.

charon/src/ast/llbc_ast.rs

@@ -55,11 +55,9 @@ pub enum RawStatement {
     Continue(usize),
     /// No-op.
     Nop,
-    /// The left statement must NOT be a sequence.
-    /// For instance, `(s0; s1); s2` is forbidden and should be rewritten
-    /// to the semantically equivalent statement `s0; (s1; s2)`
-    /// To ensure that, use [crate::llbc_ast_utils::new_sequence] to build sequences.
-    Sequence(Box<Statement>, Box<Statement>),
+    /// The contained statements must NOT be sequences. To ensure that, use [Sequence::then] to
+    /// build sequences.
+    Sequence(Vec<Statement>),
     Switch(Switch),
     Loop(Box<Statement>),
     Error(String),

charon/src/ast/llbc_ast_utils.rs

@@ -3,55 +3,19 @@
 use crate::llbc_ast::{RawStatement, Statement, Switch};
 use crate::meta;
 use crate::meta::Span;
-use derive_visitor::{DriveMut, VisitorMut};
-use std::ops::DerefMut;
+use derive_visitor::{visitor_fn_mut, DriveMut, Event};
 use take_mut::take;
 
-/// Goes from e.g. `(A; B; C) ; D` to `(A; (B; (C; D)))`.
+/// Goes from e.g. `(A; B; C) ; D` to `(A; (B; (C; D)))`. Statements in `firsts` must not be
+/// sequences.
 pub fn chain_statements(firsts: Vec<Statement>, last: Statement) -> Statement {
-    firsts.into_iter().rev().fold(last, |cont, bind| {
-        assert!(!bind.content.is_sequence());
-        let span = meta::combine_span(&bind.span, &cont.span);
-        Statement::new(span, RawStatement::Sequence(Box::new(bind), Box::new(cont)))
-    })
-}
-
-/// Utility function for [new_sequence].
-/// Efficiently appends a new statement at the rightmost place of a well-formed sequence.
-fn append_rightmost(seq: &mut Statement, r: Box<Statement>) {
-    let (_l1, l2) = match &mut seq.content {
-        RawStatement::Sequence(l1, l2) => (l1, l2),
-        _ => unreachable!(),
-    };
-    if l2.content.is_sequence() {
-        append_rightmost(l2, r);
+    if let Some(firsts) = Statement::from_seq(firsts) {
+        firsts.then(last)
     } else {
-        take(l2.deref_mut(), move |l2| {
-            let span = meta::combine_span(&l2.span, &r.span);
-            Statement::new(span, RawStatement::Sequence(Box::new(l2), r))
-        });
+        last
     }
 }
 
-/// Builds a sequence from well-formed statements.
-/// Ensures that the left statement will not be a sequence in the new sequence:
-/// Must be used instead of the raw [RawStatement::Sequence] constructor,
-/// unless you're sure that the left statement is not a sequence.
-pub fn new_sequence(mut l: Statement, r: Statement) -> Statement {
-    let span = meta::combine_span(&l.span, &r.span);
-
-    let r = Box::new(r);
-    let nst = match l.content {
-        RawStatement::Sequence(_, _) => {
-            append_rightmost(&mut l, r);
-            l.content
-        }
-        lc => RawStatement::Sequence(Box::new(Statement::new(l.span, lc)), r),
-    };
-
-    Statement::new(span, nst)
-}
-
 /// Combine the span information from a [Switch]
 pub fn combine_switch_targets_span(targets: &Switch) -> Span {
     match targets {
@@ -106,78 +70,134 @@ impl Statement {
         Statement { span, content }
     }
 
-    /// Collect the statements this is made up of into a vec. If `self` is not a sequence, the vec
-    /// will be `vec![self]`.
-    pub fn sequence_to_vec(self) -> Vec<Self> {
-        fn collect_to_vec(x: Statement, vec: &mut Vec<Statement>) {
-            match x.content {
-                RawStatement::Sequence(a, b) => {
-                    assert!(!a.content.is_sequence());
-                    vec.push(*a);
-                    collect_to_vec(*b, vec);
-                }
-                _ => vec.push(x),
-            }
+    pub fn from_seq(seq: Vec<Statement>) -> Option<Statement> {
+        if seq.is_empty() {
+            None
+        } else {
+            let span = seq
+                .iter()
+                .map(|st| st.span)
+                .reduce(|a, b| meta::combine_span(&a, &b))
+                .unwrap();
+            Some(Statement {
+                span,
+                content: RawStatement::Sequence(seq),
+            })
         }
-        let mut vec = Vec::new();
-        collect_to_vec(self, &mut vec);
-        vec
     }
 
-    /// If `self` is a sequence whose first component is a sequence, refold the statements to
-    /// repair the invariant on sequences. All otehr sequences are expected to respect the
-    /// invariant. Use this when mutating a statement in a visitor if needed.
-    pub fn reparenthesize(&mut self) {
-        if let RawStatement::Sequence(st1, _) = &mut self.content {
-            if st1.content.is_sequence() {
-                // We did mess up the parenthesization. Fix it.
-                take(self, |st| {
-                    let (st1, st2) = st.content.to_sequence();
-                    let st1 = st1.sequence_to_vec();
-                    chain_statements(st1, *st2)
-                })
-            }
+    pub fn into_box(self) -> Box<Self> {
+        Box::new(self)
+    }
+
+    /// Builds a sequence from well-formed statements. Ensures that we don't incorrectly nest
+    /// sequences. Must be used instead of the raw [RawStatement::Sequence] constructor.
+    pub fn then(self, r: Statement) -> Statement {
+        let span = meta::combine_span(&self.span, &r.span);
+
+        let vec = if !self.content.is_sequence() && !r.content.is_sequence() {
+            vec![self, r]
+        } else {
+            let mut l = self.into_sequence();
+            let mut r = r.into_sequence();
+            l.append(&mut r);
+            l
+        };
+
+        Statement::new(span, RawStatement::Sequence(vec))
+    }
+
+    pub fn then_opt(self, other: Option<Statement>) -> Statement {
+        if let Some(other) = other {
+            self.then(other)
+        } else {
+            self
         }
     }
-}
 
-/// Helper for [transform_statements]
-#[derive(VisitorMut)]
-#[visitor(Statement(exit))]
-struct TransformStatements<'a, F: FnMut(&mut Statement) -> Option<Vec<Statement>>> {
-    tr: &'a mut F,
-}
+    /// If `self` is a sequence, return that sequence. Otherwise return `vec![self]`.
+    pub fn into_sequence(self) -> Vec<Self> {
+        match self.content {
+            RawStatement::Sequence(vec) => vec,
+            _ => vec![self],
+        }
+    }
 
-impl<'a, F> TransformStatements<'a, F>
-where
-    F: FnMut(&mut Statement) -> Option<Vec<Statement>>,
-{
-    fn exit_statement(&mut self, st: &mut Statement) {
-        // Reparenthesize sequences we messed up while traversing.
-        st.reparenthesize();
-
-        // Transform the current statement
-        let st_seq = (self.tr)(st);
-        if let Some(seq) = st_seq && !seq.is_empty() {
-            take(st, |st| chain_statements(seq, st))
+    pub fn as_sequence_slice_mut(&mut self) -> &mut [Self] {
+        match self.content {
+            RawStatement::Sequence(ref mut vec) => vec.as_mut_slice(),
+            _ => std::slice::from_mut(self),
+        }
+    }
+
+    /// If `self` is a sequence that contains sequences, flatten the nesting. Use this when
+    /// mutating a statement in a visitor if needed.
+    pub fn flatten(&mut self) {
+        if let RawStatement::Sequence(vec) = &self.content {
+            if vec.iter().any(|st| st.content.is_sequence()) {
+                take(&mut self.content, |rawst| {
+                    RawStatement::Sequence(
+                        rawst
+                            .to_sequence()
+                            .into_iter()
+                            .flat_map(|st| st.into_sequence())
+                            .collect(),
+                    )
+                })
+            }
         }
     }
-}
 
-impl Statement {
     /// Apply a transformer to all the statements, in a bottom-up manner.
     ///
     /// The transformer should:
     /// - mutate the current statement in place
     /// - return the sequence of statements to introduce before the current statement
-    ///
-    /// We do the transformation in such a way that the
-    /// sequences of statements are properly chained. In particuliar,
-    /// if in `s1; s2` we transform `s1` to the sequence `s1_1; s1_2`,
-    /// then the resulting statement is `s1_1; s1_2; s2` and **not**
-    /// `{ s1_1; s1_2 }; s2`.
     pub fn transform<F: FnMut(&mut Statement) -> Option<Vec<Statement>>>(&mut self, f: &mut F) {
-        let mut visitor = TransformStatements { tr: f };
-        self.drive_mut(&mut visitor);
+        self.drive_mut(&mut visitor_fn_mut(|st: &mut Statement, e: Event| {
+            if matches!(e, Event::Exit) {
+                // Flatten any nested sequences created while traversing this statement..
+                st.flatten();
+
+                // Transform the current statement
+                let prefix_seq = f(st);
+                if let Some(prefix_seq) = prefix_seq && !prefix_seq.is_empty() {
+                    take_mut::take(st, |st| chain_statements(prefix_seq, st))
+                }
+            }
+        }));
+    }
+
+    /// Apply a transformer to all the statements, in a bottom-up manner. Compared to `transform`,
+    /// this also gives access to the following statements if any. Statements that are not part of
+    /// a sequence will be traversed as `[st]`. Statements that are will be traversed twice: once
+    /// as `[st]`, and then as `[st, ..]` with the following statements if any.
+    ///
+    /// The transformer should:
+    /// - mutate the current statements in place
+    /// - return the sequence of statements to introduce before the current statements
+    pub fn transform_sequences<F: FnMut(&mut [Statement]) -> Vec<Statement>>(&mut self, f: &mut F) {
+        self.drive_mut(&mut visitor_fn_mut(|st: &mut Statement, e: Event| {
+            if matches!(e, Event::Exit) {
+                // Flatten any nested sequences created while traversing this statement..
+                st.flatten();
+
+                if let RawStatement::Sequence(seq) = &mut st.content {
+                    for i in (0..seq.len()).rev() {
+                        let prefix_to_insert = f(&mut seq[i..]);
+                        if !prefix_to_insert.is_empty() {
+                            // Insert the new elements at index `i`. This only modifies `vec[i..]`
+                            // so we can keep iterating `i` down as if nothing happened.
+                            seq.splice(i..i, prefix_to_insert);
+                        }
+                    }
+                } else {
+                    let prefix_to_insert = f(std::slice::from_mut(st));
+                    if !prefix_to_insert.is_empty() {
+                        take_mut::take(st, |st| chain_statements(prefix_to_insert, st))
+                    }
+                }
+            }
+        }));
     }
 }

charon/src/pretty/fmt_with_ctx.rs

@@ -6,6 +6,7 @@ use crate::{
     gast,
     ids::Vector,
     llbc_ast::{self as llbc, *},
+    meta::*,
     names::*,
     reorder_decls::*,
     translate_predicates::NonLocalTraitClause,
@@ -14,6 +15,7 @@ use crate::{
     values::*,
 };
 use hax_frontend_exporter as hax;
+use itertools::Itertools;
 
 /// Format the AST type as a string.
 pub trait FmtWithCtx<C> {
@@ -861,11 +863,10 @@ impl<C: AstFormatter> FmtWithCtx<C> for llbc::Statement {
             RawStatement::Break(index) => format!("{tab}break {index}"),
             RawStatement::Continue(index) => format!("{tab}continue {index}"),
             RawStatement::Nop => format!("{tab}nop"),
-            RawStatement::Sequence(st1, st2) => format!(
-                "{}\n{}",
-                st1.fmt_with_ctx_and_indent(tab, ctx),
-                st2.fmt_with_ctx_and_indent(tab, ctx)
-            ),
+            RawStatement::Sequence(vec) => vec
+                .iter()
+                .map(|st| st.fmt_with_ctx_and_indent(tab, ctx))
+                .join("\n"),
             RawStatement::Switch(switch) => match switch {
                 Switch::If(discr, true_st, false_st) => {
                     let inner_tab = format!("{tab}{TAB_INCR}");
@@ -1456,6 +1457,12 @@ impl std::fmt::Display for LiteralTy {
     }
 }
 
+impl std::fmt::Display for Loc {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::result::Result<(), std::fmt::Error> {
+        write!(f, "{}:{}", self.line, self.col)
+    }
+}
+
 impl std::fmt::Display for Operand {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::result::Result<(), std::fmt::Error> {
         write!(f, "{}", self.fmt_with_ctx(&FmtCtx::new()))

charon/src/transform/index_to_function_calls.rs

@@ -207,7 +207,7 @@ impl<'a> Visitor<'a> {
         self.span = None;
 
         // Reparenthesize sequences we messed up while traversing.
-        st.reparenthesize();
+        st.flatten();
 
         // We potentially collected statements to prepend to this one.
         let seq = std::mem::take(&mut self.statements);

charon/src/transform/insert_assign_return_unit.rs

@@ -24,7 +24,7 @@ fn transform_st(st: &mut Statement) -> Option<Vec<Statement>> {
         );
         let assign_st = Statement::new(st.span, RawStatement::Assign(ret_place, unit_value));
         let ret_st = Statement::new(st.span, RawStatement::Return);
-        st.content = RawStatement::Sequence(Box::new(assign_st), Box::new(ret_st));
+        st.content = assign_st.then(ret_st).content;
     };
     None
 }

charon/src/transform/reconstruct_asserts.rs

@@ -25,8 +25,7 @@ fn transform_st(st: &mut Statement) -> Option<Vec<Statement>> {
                         expected: false,
                     }),
                 );
-                let st1 = Box::new(st1);
-                RawStatement::Sequence(st1, st2)
+                st1.then(*st2).content
             });
         }
     }