Recover comments in ullbc too

charon/src/transform/mod.rs

@@ -130,10 +130,6 @@ pub static LLBC_PASSES: &[Pass] = &[
     StructuredBody(&remove_read_discriminant::Transform),
     // Cleanup the cfg.
     StructuredBody(&prettify_cfg::Transform),
-    // # Micro-pass: take all the comments found in the original body and assign them to
-    // statements. This must be last after all the statement-affecting passes to avoid losing
-    // comments.
-    StructuredBody(&recover_body_comments::Transform),
 ];
 
 /// Cleanup passes useful for both llbc and ullbc.
@@ -142,6 +138,10 @@ pub static SHARED_FINALIZING_PASSES: &[Pass] = &[
     NonBody(&remove_unused_locals::Transform),
     // # Micro-pass: remove the useless `StatementKind::Nop`s.
     NonBody(&remove_nops::Transform),
+    // # Micro-pass: take all the comments found in the original body and assign them to
+    // statements. This must be last after all the statement-affecting passes to avoid losing
+    // comments.
+    NonBody(&recover_body_comments::Transform),
     // # Reorder the graph of dependencies and compute the strictly connex components to:
     // - compute the order in which to extract the definitions
     // - find the recursive definitions

charon/src/transform/recover_body_comments.rs

@@ -1,34 +1,92 @@
 //! Take all the comments found in the original body and assign them to statements.
 use std::mem;
 
-use crate::llbc_ast::*;
+use crate::ast::*;
 use crate::transform::TransformCtx;
 
-use super::ctx::LlbcPass;
+use super::ctx::TransformPass;
+
+trait IsStatement {
+    fn get_span(&self) -> Span;
+    fn get_comments_before(&mut self) -> &mut Vec<String>;
+}
+
+impl IsStatement for llbc_ast::Statement {
+    fn get_span(&self) -> Span {
+        self.span
+    }
+    fn get_comments_before(&mut self) -> &mut Vec<String> {
+        &mut self.comments_before
+    }
+}
+impl IsStatement for ullbc_ast::Statement {
+    fn get_span(&self) -> Span {
+        self.span
+    }
+    fn get_comments_before(&mut self) -> &mut Vec<String> {
+        &mut self.comments_before
+    }
+}
+impl IsStatement for ullbc_ast::Terminator {
+    fn get_span(&self) -> Span {
+        self.span
+    }
+    fn get_comments_before(&mut self) -> &mut Vec<String> {
+        &mut self.comments_before
+    }
+}
+
+struct CommentsCtx {
+    comments: Vec<(usize, Vec<String>)>,
+}
+impl CommentsCtx {
+    fn visit<St: IsStatement>(&mut self, st: &mut St) {
+        let st_line = st.get_span().span.beg.line;
+        self.comments = mem::take(&mut self.comments)
+            .into_iter()
+            .filter_map(|(line, comments)| {
+                if line <= st_line {
+                    st.get_comments_before().extend(comments);
+                    None
+                } else {
+                    Some((line, comments))
+                }
+            })
+            .collect();
+    }
+}
 
 pub struct Transform;
-impl LlbcPass for Transform {
-    fn transform_body(&self, _ctx: &mut TransformCtx, b: &mut ExprBody) {
-        // Constraints in the ideal case:
-        // - each comment should be assigned to exactly one statement;
-        // - the order of comments in the source should refine the partial order of control flow;
-        // - a comment should come before the statement it was applied to.
+impl TransformPass for Transform {
+    fn transform_ctx(&self, ctx: &mut TransformCtx) {
+        ctx.for_each_fun_decl(|_ctx, fun| {
+            if let Ok(body) = &mut fun.body {
+                // Constraints in the ideal case:
+                // - each comment should be assigned to exactly one statement;
+                // - the order of comments in the source should refine the partial order of control flow;
+                // - a comment should come before the statement it was applied to.
 
-        // This is a pretty simple heuristic which is good enough for now.
-        let mut comments: Vec<(usize, Vec<String>)> = b.comments.clone();
-        b.body.visit_statements(|st: &mut Statement| {
-            let st_line = st.span.span.beg.line;
-            comments = mem::take(&mut comments)
-                .into_iter()
-                .filter_map(|(line, comments)| {
-                    if line <= st_line {
-                        st.comments_before.extend(comments);
-                        None
-                    } else {
-                        Some((line, comments))
+                // This is a pretty simple heuristic which is good enough for now.
+                let mut ctx = CommentsCtx {
+                    comments: match body {
+                        Body::Unstructured(b) => b.comments.clone(),
+                        Body::Structured(b) => b.comments.clone(),
+                    },
+                };
+                match body {
+                    Body::Unstructured(b) => {
+                        for block in &mut b.body {
+                            for st in &mut block.statements {
+                                ctx.visit(st);
+                            }
+                            ctx.visit(&mut block.terminator);
+                        }
                     }
-                })
-                .collect();
+                    Body::Structured(b) => b.body.visit_statements(|st| {
+                        ctx.visit(st);
+                    }),
+                }
+            }
         });
     }
 }

charon/tests/ui/ullbc-control-flow.out

@@ -812,6 +812,7 @@ fn test_crate::nested_loops_enum(@1: usize, @2: usize) -> usize
     bb6: {
         @39 := ();
         @14 := move (@39);
+        // Test comment
         @storage_dead(@22);
         @storage_dead(@20);
         @storage_dead(@19);