charon_lib/transform/
recover_body_comments.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
//! Take all the comments found in the original body and assign them to statements.

use derive_visitor::{visitor_enter_fn_mut, DriveMut};

use crate::llbc_ast::*;
use crate::transform::TransformCtx;

use super::ctx::LlbcPass;

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.

        // This is a pretty simple heuristic which is good enough for now.
        let mut comments: Vec<(usize, Vec<String>)> = b.comments.clone();
        b.body
            .drive_mut(&mut visitor_enter_fn_mut(|st: &mut Statement| {
                let st_line = st.span.span.beg.line;
                st.comments_before = comments
                    .extract_if(|(i, _)| *i <= st_line)
                    .flat_map(|(_, comments)| comments)
                    .collect();
            }));
    }
}