charon_lib/transform/add_missing_info/
recover_body_comments.rs

1//! Take all the comments found in the original body and assign them to statements.
2use std::mem;
3
4use crate::ast::*;
5use crate::transform::TransformCtx;
6
7use crate::transform::ctx::TransformPass;
8
9trait IsStatement {
10    fn get_span(&self) -> Span;
11    fn get_comments_before(&mut self) -> &mut Vec<String>;
12}
13
14impl IsStatement for llbc_ast::Statement {
15    fn get_span(&self) -> Span {
16        self.span
17    }
18    fn get_comments_before(&mut self) -> &mut Vec<String> {
19        &mut self.comments_before
20    }
21}
22impl IsStatement for ullbc_ast::Statement {
23    fn get_span(&self) -> Span {
24        self.span
25    }
26    fn get_comments_before(&mut self) -> &mut Vec<String> {
27        &mut self.comments_before
28    }
29}
30impl IsStatement for ullbc_ast::Terminator {
31    fn get_span(&self) -> Span {
32        self.span
33    }
34    fn get_comments_before(&mut self) -> &mut Vec<String> {
35        &mut self.comments_before
36    }
37}
38
39struct CommentsCtx {
40    comments: Vec<(usize, Vec<String>)>,
41}
42impl CommentsCtx {
43    fn visit<St: IsStatement>(&mut self, st: &mut St) {
44        let st_line = st.get_span().data.beg.line;
45        self.comments = mem::take(&mut self.comments)
46            .into_iter()
47            .filter_map(|(line, comments)| {
48                if line <= st_line {
49                    st.get_comments_before().extend(comments);
50                    None
51                } else {
52                    Some((line, comments))
53                }
54            })
55            .collect();
56    }
57}
58
59pub struct Transform;
60impl TransformPass for Transform {
61    fn transform_ctx(&self, ctx: &mut TransformCtx) {
62        ctx.for_each_fun_decl(|_ctx, fun| {
63            let body = &mut fun.body;
64            if !body.has_contents() {
65                return;
66            }
67            // Constraints in the ideal case:
68            // - each comment should be assigned to exactly one statement;
69            // - the order of comments in the source should refine the partial order of control flow;
70            // - a comment should come before the statement it was applied to.
71
72            // This is a pretty simple heuristic which is good enough for now.
73            let mut ctx = CommentsCtx {
74                comments: match body {
75                    Body::Unstructured(b) => b.comments.clone(),
76                    Body::Structured(b) => b.comments.clone(),
77                    _ => unreachable!(),
78                },
79            };
80            match body {
81                Body::Unstructured(b) => {
82                    for block in &mut b.body {
83                        for st in &mut block.statements {
84                            // Many assignments have a `storage_live` before them; we don't
85                            // want to put the comment there.
86                            if !st.kind.is_storage_live() {
87                                ctx.visit(st);
88                            }
89                        }
90                        ctx.visit(&mut block.terminator);
91                    }
92                }
93                Body::Structured(b) => b.body.visit_statements(|st| {
94                    if !st.kind.is_storage_live() {
95                        ctx.visit(st);
96                    }
97                }),
98                _ => unreachable!(),
99            }
100        });
101    }
102}
charon_lib/transform/add_missing_info/recover_body_comments.rs

charon_lib/transform/add_missing_info/
recover_body_comments.rs
