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
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
//! Take all the comments found in the original body and assign them to statements.
use std::mem;

use crate::ast::*;
use crate::transform::TransformCtx;

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 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 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);
                        }
                    }
                    Body::Structured(b) => b.body.visit_statements(|st| {
                        ctx.visit(st);
                    }),
                }
            }
        });
    }
}