diff --git a/charon/Cargo.lock b/charon/Cargo.lock
index dc33f1f0a..6baa0ee87 100644
--- a/charon/Cargo.lock
+++ b/charon/Cargo.lock
@@ -165,6 +165,12 @@ dependencies = [
  "serde",
 ]
 
+[[package]]
+name = "by_address"
+version = "1.2.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "64fa3c856b712db6612c019f14756e64e4bcea13337a6b33b696333a9eaa2d06"
+
 [[package]]
 name = "cc"
 version = "1.0.101"
@@ -183,6 +189,7 @@ version = "0.1.44"
 dependencies = [
  "anyhow",
  "assert_cmd",
+ "by_address",
  "clap",
  "colored",
  "convert_case 0.6.0",
diff --git a/charon/Cargo.toml b/charon/Cargo.toml
index 6bd342df5..b15848aeb 100644
--- a/charon/Cargo.toml
+++ b/charon/Cargo.toml
@@ -77,6 +77,7 @@ which = "6.0.1"
 
 hax-frontend-exporter = { git = "https://github.com/hacspec/hax", branch = "main", optional = true }
 macros = { path = "./macros" }
+by_address = "1.2.1"
 
 [features]
 default = ["rustc"]
diff --git a/charon/src/transform/recover_body_comments.rs b/charon/src/transform/recover_body_comments.rs
index 742e21273..3fabc93ef 100644
--- a/charon/src/transform/recover_body_comments.rs
+++ b/charon/src/transform/recover_body_comments.rs
@@ -21,6 +21,16 @@ impl LlbcPass for Transform {
     // We 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).
     fn transform_body(&self, _ctx: &mut TransformCtx<'_>, b: &mut ExprBody) {
+        let mut statements_for_line: HashMap<usize, HashSet<ByAddress<&Statement>>> =
+            Default::default();
+        b.body.drive(&mut visitor_enter_fn(|st: &Statement| {
+            let span = st.span;
+            let end_line = statements_for_line
+                .entry(span.span.beg.line)
+                .or_insert(span.span.beg.line);
+            *end_line = max(*end_line, span.span.end.line);
+        }));
+
         // For each source line (that a comment may apply to), we try to compute the set of lines
         // that are spanned by the statement/expression that starts on that line. This assumes
         // standard one-statement-per-line rust formatting.
@@ -34,8 +44,6 @@ impl LlbcPass for Transform {
             *end_line = max(*end_line, span.span.end.line);
         }));
 
-        // TODO: for each syntactic line, find the span of the corresponding semantic line if
-        // possible.
         // TODO: order by statement kind: call, assign
         // 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.
diff --git a/charon/tests/ui/comments.rs b/charon/tests/ui/comments.rs
index d556abe7c..c00bcda75 100644
--- a/charon/tests/ui/comments.rs
+++ b/charon/tests/ui/comments.rs
@@ -23,7 +23,7 @@ pub fn sum(s: &[u32]) -> u32 {
         // sum + 100
         sum + 100
     } else {
-        // let sum untouched
+        // return sum untouched
         sum
     };
     // Function call