Improve comments heuristic

charon/src/ast/meta.rs

@@ -9,7 +9,20 @@ use std::path::PathBuf;
 
 generate_index_type!(FileId);
 
-#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash, Serialize, Deserialize, Drive, DriveMut)]
+#[derive(
+    Debug,
+    Copy,
+    Clone,
+    PartialEq,
+    Eq,
+    PartialOrd,
+    Ord,
+    Hash,
+    Serialize,
+    Deserialize,
+    Drive,
+    DriveMut,
+)]
 pub struct Loc {
     /// The (1-based) line number.
     pub line: usize,

charon/src/transform/recover_body_comments.rs

@@ -1,6 +1,8 @@
 //! Take all the comments found in the original body and assign them to statements.
 
-use derive_visitor::{visitor_enter_fn_mut, DriveMut};
+use std::collections::HashMap;
+
+use derive_visitor::{visitor_enter_fn, visitor_enter_fn_mut, Drive, DriveMut};
 
 use crate::llbc_ast::*;
 use crate::transform::TransformCtx;
@@ -9,20 +11,52 @@ use super::ctx::LlbcPass;
 
 pub struct Transform;
 impl LlbcPass for Transform {
+    // 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.
+    // We approximate this with a reasonable heuristic.
     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.
+        // Find for each line the statement span that starts the earliest as this is more likely to
+        // correspond to what the comment was intended to point to.
+        let mut best_span_for_line: HashMap<usize, Span> = Default::default();
+        b.body.drive(&mut visitor_enter_fn(|st: &Statement| {
+            let span = st.span;
+            best_span_for_line
+                .entry(span.span.beg.line)
+                .and_modify(|best_span| {
+                    // Find the span that starts the earliest, and among these the largest span.
+                    if span.span.beg.col < best_span.span.beg.col
+                        || (span.span.beg.col == best_span.span.beg.col
+                            && span.span.end > best_span.span.end)
+                    {
+                        *best_span = span
+                    }
+                })
+                .or_insert(span);
+        }));
 
-        // This is a pretty simple heuristic which is good enough for now.
-        let mut comments: Vec<(Loc, Vec<String>)> = b.comments.clone();
+        // The map of lines to comments that apply to it.
+        let mut comments_per_line: HashMap<usize, Vec<String>> = b
+            .comments
+            .iter()
+            .cloned()
+            .map(|(loc, comments)| (loc.line, comments))
+            .collect();
+        // Assign each comment to the first statement that has the best span for its starting line.
+        // This may drop some comments if no statement starts with the relevant line (this can
+        // happen if e.g. the statement was optimized out or the comment applied to an item
+        // instead).
         b.body
             .drive_mut(&mut visitor_enter_fn_mut(|st: &mut Statement| {
-                st.comments_before = comments
-                    .extract_if(|(loc, _)| *loc == st.span.span.beg)
-                    .flat_map(|(_, comments)| comments)
-                    .collect();
+                if best_span_for_line
+                    .get(&st.span.span.beg.line)
+                    .is_some_and(|best_span| *best_span == st.span)
+                {
+                    st.comments_before = comments_per_line
+                        .remove(&st.span.span.beg.line)
+                        .unwrap_or_default()
+                }
             }));
     }
 }

charon/tests/cargo/workspace.out

@@ -13,8 +13,8 @@ fn crate2::extra_random_number() -> u32
     let @0: u32; // return
     let @1: u32; // anonymous local
 
-    // Even more random.
     @1 := crate1::random_number()
+    // Even more random.
     @0 := move (@1) + const (1 : u32)
     drop @1
     return

charon/tests/ui/comments.out

@@ -45,8 +45,12 @@ fn test_crate::sum<'_0>(@1: &'_0 (Slice<u32>)) -> u32
     let @26: &'_ (u32); // anonymous local
 
     sum@2 := const (0 : u32)
+    // Comment1
     @fake_read(sum@2)
     i@3 := const (0 : usize)
+    // Comment2
+    //  Comment2.1
+    // Comment2.2
     @fake_read(i@3)
     @6 := copy (sum@2)
     @5 := move (@6) + const (2 : u32)
@@ -96,8 +100,8 @@ fn test_crate::sum<'_0>(@1: &'_0 (Slice<u32>)) -> u32
     @17 := move (@18) > const (10 : u32)
     if move (@17) {
         drop @18
-        // Comment9
         @19 := copy (sum@2)
+        // Comment9
         @16 := move (@19) + const (100 : u32)
         drop @19
     }
@@ -145,6 +149,7 @@ fn test_crate::thing()
     @2 := move (@5) >> const (3 : i32)
     x@1 := move (@2) + const (12 : u32)
     drop @2
+    // This comment belongs above the assignment to `x` and not above intermediate computations.
     @fake_read(x@1)
     @4 := copy (x@1)
     @3 := test_crate::function_call(move (@4))
@@ -163,6 +168,7 @@ fn test_crate::fake_read(@1: u32)
     let x@1: u32; // arg #1
     let @2: (); // anonymous local
 
+    // This statement is translated to a `fake_read`.
     @fake_read(x@1)
     @2 := ()
     @0 := move (@2)

charon/tests/ui/predicates-on-late-bound-vars.out

@@ -65,6 +65,7 @@ fn test_crate::foo()
     ref_b@1 := core::cell::{core::cell::RefCell<T>}#21::new<bool>(const (false))
     @fake_read(ref_b@1)
     @3 := &ref_b@1
+    // `try_borrow` has a type that includes predicates on late bound regions.
     @2 := core::cell::{core::cell::RefCell<T>}#22::try_borrow<bool>(move (@3))
     drop @3
     @fake_read(@2)

charon/tests/ui/rvalues.out

@@ -255,6 +255,7 @@ fn test_crate::transmute(@1: Array<u32, 2 : usize>) -> u64
     let @0: u64; // return
     let x@1: Array<u32, 2 : usize>; // arg #1
 
+    // When optimized, this becomes a built-in cast. Otherwise this is just a call to `transmute`.
     @0 := transmute<Array<u32, 2 : usize>, u64>(copy (x@1))
     return
 }