Merge pull request #358 from Nadrieril/keep-body-comments

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.41"
+let supported_charon_version = "0.1.42"

charon-ml/src/GAstOfJson.ml

@@ -352,6 +352,7 @@ and gexpr_body_of_json :
           ("span", span);
           ("arg_count", arg_count);
           ("locals", locals);
+          ("comments", _);
           ("body", body);
         ] ->
         let* span = span_of_json id_to_file span in

charon-ml/src/LlbcAst.ml

@@ -30,7 +30,12 @@ type raw_statement =
   | Loop of statement
   | Error of string
 
-and statement = { span : span; content : raw_statement }
+and statement = {
+  span : span;
+  content : raw_statement;
+  comments_before : string list;  (** Comments that precede this statement. *)
+}
+
 and block = statement
 
 and switch =

charon-ml/src/LlbcAstUtils.ml

@@ -42,7 +42,7 @@ let fun_decl_has_loops (fd : fun_decl) : bool =
 let mk_sequence (st1 : statement) (st2 : statement) : statement =
   let span = MetaUtils.combine_span st1.span st2.span in
   let content = Sequence (st1, st2) in
-  { span; content }
+  { span; content; comments_before = [] }
 
 (** Chain two statements into a sequence, by pushing the second statement
     at the end of the first one (diving into sequences, switches, etc.).
@@ -60,7 +60,7 @@ let rec chain_statements (st1 : statement) (st2 : statement) : statement =
       (* Insert inside the switch *)
       let span = MetaUtils.combine_span st1.span st2.span in
       let content = Switch (chain_statements_in_switch switch st2) in
-      { span; content }
+      { span; content; comments_before = [] }
   | Sequence (st3, st4) ->
       (* Insert at the end of the statement *)
       mk_sequence st3 (chain_statements st4 st2)

charon-ml/src/LlbcOfJson.ml

@@ -58,10 +58,16 @@ and statement_of_json (id_to_file : id_to_file_map) (js : json) :
     (statement, string) result =
   combine_error_msgs js __FUNCTION__
     (match js with
-    | `Assoc [ ("span", span); ("content", content) ] ->
+    | `Assoc
+        [
+          ("span", span);
+          ("content", content);
+          ("comments_before", comments_before);
+        ] ->
         let* span = span_of_json id_to_file span in
         let* content = raw_statement_of_json id_to_file content in
-        Ok ({ span; content } : statement)
+        let* comments_before = list_of_json string_of_json comments_before in
+        Ok ({ span; content; comments_before } : statement)
     | _ -> Error "")
 
 and block_of_json (id_to_file : id_to_file_map) (js : json) :
@@ -74,11 +80,16 @@ and block_of_json (id_to_file : id_to_file_map) (js : json) :
           list_of_json (statement_of_json id_to_file) statements
         in
         match List.rev statements with
-        | [] -> Ok { span; content = Nop }
+        | [] -> Ok { span; content = Nop; comments_before = [] }
         | last :: rest ->
             let seq =
               List.fold_left
-                (fun acc st -> { span = st.span; content = Sequence (st, acc) })
+                (fun acc st ->
+                  {
+                    span = st.span;
+                    content = Sequence (st, acc);
+                    comments_before = [];
+                  })
                 last rest
             in
             Ok seq

charon/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "charon"
-version = "0.1.41"
+version = "0.1.42"
 authors = ["Son Ho <hosonmarc@gmail.com>"]
 edition = "2021"
 license = "Apache-2.0"

charon/src/ast/gast.rs

@@ -59,6 +59,10 @@ pub struct GExprBody<T> {
     /// - the input arguments
     /// - the remaining locals, used for the intermediate computations
     pub locals: Vector<VarId, Var>,
+    /// For each line inside the body, we record any whole-line `//` comments found before it. They
+    /// are added to statements in the late `recover_body_comments` pass.
+    #[charon::opaque]
+    pub comments: Vec<(usize, Vec<String>)>,
     pub body: T,
 }
 

charon/src/ast/llbc_ast.rs

@@ -54,6 +54,9 @@ pub enum RawStatement {
 pub struct Statement {
     pub span: Span,
     pub content: RawStatement,
+    /// Comments that precede this statement.
+    // This is filled in a late pass after all the control-flow manipulation.
+    pub comments_before: Vec<String>,
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize, Drive, DriveMut)]

charon/src/ast/llbc_ast_utils.rs

@@ -59,7 +59,11 @@ impl Switch {
 
 impl Statement {
     pub fn new(span: Span, content: RawStatement) -> Self {
-        Statement { span, content }
+        Statement {
+            span,
+            content,
+            comments_before: vec![],
+        }
     }
 
     pub fn into_box(self) -> Box<Self> {

charon/src/ast/meta.rs

@@ -9,7 +9,7 @@ use std::path::PathBuf;
 
 generate_index_type!(FileId);
 
-#[derive(Debug, Copy, Clone, Serialize, Deserialize, Drive, DriveMut)]
+#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash, Serialize, Deserialize, Drive, DriveMut)]
 pub struct Loc {
     /// The (1-based) line number.
     pub line: usize,
@@ -23,7 +23,7 @@ fn dummy_span_data() -> rustc_span::SpanData {
 }
 
 /// Span information
-#[derive(Debug, Copy, Clone, Serialize, Deserialize, Drive, DriveMut)]
+#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash, Serialize, Deserialize, Drive, DriveMut)]
 pub struct RawSpan {
     pub file_id: FileId,
     pub beg: Loc,
@@ -47,7 +47,7 @@ impl From<RawSpan> for rustc_error_messages::MultiSpan {
 }
 
 /// Meta information about a piece of code (block, statement, etc.)
-#[derive(Debug, Copy, Clone, Serialize, Deserialize, Drive, DriveMut)]
+#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash, Serialize, Deserialize, Drive, DriveMut)]
 pub struct Span {
     /// The source code span.
     ///

charon/src/bin/charon-driver/translate/translate_functions_to_ullbc.rs

@@ -18,6 +18,7 @@ use charon_lib::pretty::FmtWithCtx;
 use charon_lib::ullbc_ast::*;
 use hax_frontend_exporter as hax;
 use hax_frontend_exporter::{HasMirSetter, HasOwnerIdSetter, SInto};
+use itertools::Itertools;
 use rustc_hir::def_id::DefId;
 use rustc_middle::mir::START_BLOCK;
 use translate_types::translate_bound_region_kind_name;
@@ -1256,6 +1257,48 @@ impl<'tcx, 'ctx, 'ctx1> BodyTransCtx<'tcx, 'ctx, 'ctx1> {
         Ok(t_args)
     }
 
+    /// Gather all the lines that start with `//` inside the given span.
+    fn translate_body_comments(&mut self, charon_span: Span) -> Vec<(usize, Vec<String>)> {
+        let rust_span = charon_span.rust_span();
+        let source_map = self.t_ctx.tcx.sess.source_map();
+        if rust_span.ctxt().is_root()
+            && let Ok(body_text) = source_map.span_to_snippet(rust_span.into())
+        {
+            let mut comments = body_text
+                .lines()
+                // Iter through the lines of this body in reverse order.
+                .rev()
+                .enumerate()
+                // Compute the absolute line number
+                .map(|(i, line)| (charon_span.span.end.line - i, line))
+                // Extract the comment if this line starts with `//`
+                .map(|(line_nbr, line)| (line_nbr, line.trim_start().strip_prefix("//")))
+                .peekable()
+                .batching(|iter| {
+                    // Get the next line. This is not a comment: it's either the last line of the
+                    // body or a line that wasn't consumed by `peeking_take_while`.
+                    let (line_nbr, _first) = iter.next()?;
+                    // Collect all the comments before this line.
+                    let mut comments = iter
+                        // `peeking_take_while` ensures we don't consume a line that returns
+                        // `false`. It will be consumed by the next round of `batching`.
+                        .peeking_take_while(|(_, opt_comment)| opt_comment.is_some())
+                        .map(|(_, opt_comment)| opt_comment.unwrap())
+                        .map(str::trim_start)
+                        .map(str::to_owned)
+                        .collect_vec();
+                    comments.reverse();
+                    Some((line_nbr, comments))
+                })
+                .filter(|(_, comments)| !comments.is_empty())
+                .collect_vec();
+            comments.reverse();
+            comments
+        } else {
+            Vec::new()
+        }
+    }
+
     /// Translate a function body if we can (it has MIR) and we want to (we don't translate bodies
     /// declared opaque, and only translate non-local bodies if `extract_opaque_bodies` is set).
     fn translate_body(
@@ -1332,6 +1375,7 @@ impl<'tcx, 'ctx, 'ctx1> BodyTransCtx<'tcx, 'ctx, 'ctx1> {
             span,
             arg_count,
             locals: mem::take(&mut self.vars),
+            comments: self.translate_body_comments(span),
             body: blocks,
         })))
     }
@@ -1416,11 +1460,11 @@ impl<'tcx, 'ctx, 'ctx1> BodyTransCtx<'tcx, 'ctx, 'ctx1> {
                     hax::ClosureKind::FnMut => ClosureKind::FnMut,
                     hax::ClosureKind::FnOnce => ClosureKind::FnOnce,
                 };
-                let state = args
+                let state: Vector<TypeVarId, Ty> = args
                     .upvar_tys
                     .iter()
                     .map(|ty| self.translate_ty(span, erase_regions, &ty))
-                    .try_collect::<Vector<TypeVarId, Ty>>()?;
+                    .try_collect()?;
                 Some(ClosureInfo { kind, state })
             }
             hax::FullDefKind::Fn { .. } => None,

charon/src/bin/generate-ml/main.rs

@@ -1164,11 +1164,11 @@ fn generate_ml(crate_data: TranslatedCrate, output_dir: PathBuf) -> anyhow::Resu
                       list_of_json (statement_of_json id_to_file) statements
                     in
                     match List.rev statements with
-                    | [] -> Ok { span; content = Nop }
+                    | [] -> Ok { span; content = Nop; comments_before = [] }
                     | last :: rest ->
                         let seq =
                           List.fold_left
-                            (fun acc st -> { span = st.span; content = Sequence (st, acc) })
+                            (fun acc st -> { span = st.span; content = Sequence (st, acc); comments_before = [] })
                             last rest
                         in
                         Ok seq

charon/src/lib.rs

@@ -15,6 +15,7 @@
 // For rustdoc: prevents overflows
 #![recursion_limit = "256"]
 #![feature(box_patterns)]
+#![feature(extract_if)]
 #![feature(if_let_guard)]
 #![feature(impl_trait_in_assoc_type)]
 #![feature(iterator_try_collect)]

charon/src/pretty/fmt_with_ctx.rs

@@ -9,6 +9,7 @@ use crate::{
     ullbc_ast::{self as ullbc, *},
 };
 use itertools::Itertools;
+use std::fmt::Write;
 
 /// Format the AST type as a string.
 pub trait FmtWithCtx<C> {
@@ -994,39 +995,50 @@ impl<C: AstFormatter> FmtWithCtx<C> for llbc::Statement {
 
     fn fmt_with_ctx_and_indent(&self, tab: &str, ctx: &C) -> String {
         use llbc::RawStatement;
-        match &self.content {
-            RawStatement::Assign(place, rvalue) => format!(
+        let mut out = String::new();
+        for line in &self.comments_before {
+            let _ = writeln!(&mut out, "{tab}// {line}");
+        }
+        let _ = match &self.content {
+            RawStatement::Assign(place, rvalue) => write!(
+                &mut out,
                 "{}{} := {}",
                 tab,
                 place.fmt_with_ctx(ctx),
                 rvalue.fmt_with_ctx(ctx),
             ),
             RawStatement::FakeRead(place) => {
-                format!("{}@fake_read({})", tab, place.fmt_with_ctx(ctx))
+                write!(&mut out, "{}@fake_read({})", tab, place.fmt_with_ctx(ctx))
             }
-            RawStatement::SetDiscriminant(place, variant_id) => format!(
+            RawStatement::SetDiscriminant(place, variant_id) => write!(
+                &mut out,
                 "{}@discriminant({}) := {}",
                 tab,
                 place.fmt_with_ctx(ctx),
                 variant_id
             ),
             RawStatement::Drop(place) => {
-                format!("{}drop {}", tab, place.fmt_with_ctx(ctx))
+                write!(&mut out, "{}drop {}", tab, place.fmt_with_ctx(ctx))
+            }
+            RawStatement::Assert(assert) => {
+                write!(&mut out, "{}{}", tab, assert.fmt_with_ctx(ctx),)
             }
-            RawStatement::Assert(assert) => format!("{}{}", tab, assert.fmt_with_ctx(ctx),),
             RawStatement::Call(call) => {
                 let (call_s, _) = fmt_call(ctx, call);
-                format!("{tab}{} := {call_s}", call.dest.fmt_with_ctx(ctx),)
+                write!(&mut out, "{tab}{} := {call_s}", call.dest.fmt_with_ctx(ctx),)
+            }
+            RawStatement::Abort(kind) => {
+                write!(&mut out, "{}", kind.fmt_with_ctx_and_indent(tab, ctx))
             }
-            RawStatement::Abort(kind) => kind.fmt_with_ctx_and_indent(tab, ctx),
-            RawStatement::Return => format!("{tab}return"),
-            RawStatement::Break(index) => format!("{tab}break {index}"),
-            RawStatement::Continue(index) => format!("{tab}continue {index}"),
-            RawStatement::Nop => format!("{tab}nop"),
+            RawStatement::Return => write!(&mut out, "{tab}return"),
+            RawStatement::Break(index) => write!(&mut out, "{tab}break {index}"),
+            RawStatement::Continue(index) => write!(&mut out, "{tab}continue {index}"),
+            RawStatement::Nop => write!(&mut out, "{tab}nop"),
             RawStatement::Switch(switch) => match switch {
                 Switch::If(discr, true_st, false_st) => {
                     let inner_tab = format!("{tab}{TAB_INCR}");
-                    format!(
+                    write!(
+                        &mut out,
                         "{tab}if {} {{\n{}{tab}}}\n{tab}else {{\n{}{tab}}}",
                         discr.fmt_with_ctx(ctx),
                         true_st.fmt_with_ctx_and_indent(&inner_tab, ctx),
@@ -1053,7 +1065,8 @@ impl<C: AstFormatter> FmtWithCtx<C> for llbc::Statement {
                         otherwise.fmt_with_ctx_and_indent(&inner_tab2, ctx),
                     ));
 
-                    format!(
+                    write!(
+                        &mut out,
                         "{tab}switch {} {{\n{}{tab}}}",
                         discr.fmt_with_ctx(ctx),
                         maps.iter().format(""),
@@ -1081,7 +1094,8 @@ impl<C: AstFormatter> FmtWithCtx<C> for llbc::Statement {
                         ));
                     };
 
-                    format!(
+                    write!(
+                        &mut out,
                         "{tab}match {} {{\n{}{tab}}}",
                         discr.fmt_with_ctx(ctx),
                         maps.iter().format(""),
@@ -1090,13 +1104,15 @@ impl<C: AstFormatter> FmtWithCtx<C> for llbc::Statement {
             },
             RawStatement::Loop(body) => {
                 let inner_tab = format!("{tab}{TAB_INCR}");
-                format!(
+                write!(
+                    &mut out,
                     "{tab}loop {{\n{}{tab}}}",
                     body.fmt_with_ctx_and_indent(&inner_tab, ctx),
                 )
             }
-            RawStatement::Error(s) => format!("{tab}@ERROR({})", s),
-        }
+            RawStatement::Error(s) => write!(&mut out, "{tab}@ERROR({})", s),
+        };
+        out
     }
 }