charon_lib/transform/
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 super::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().span.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            if let Ok(body) = &mut fun.body {
64                // Constraints in the ideal case:
65                // - each comment should be assigned to exactly one statement;
66                // - the order of comments in the source should refine the partial order of control flow;
67                // - a comment should come before the statement it was applied to.
68
69                // This is a pretty simple heuristic which is good enough for now.
70                let mut ctx = CommentsCtx {
71                    comments: match body {
72                        Body::Unstructured(b) => b.comments.clone(),
73                        Body::Structured(b) => b.comments.clone(),
74                    },
75                };
76                match body {
77                    Body::Unstructured(b) => {
78                        for block in &mut b.body {
79                            for st in &mut block.statements {
80                                // Many assignments have a `storage_live` before them; we don't
81                                // want to put the comment there.
82                                if !st.content.is_storage_live() {
83                                    ctx.visit(st);
84                                }
85                            }
86                            ctx.visit(&mut block.terminator);
87                        }
88                    }
89                    Body::Structured(b) => b.body.visit_statements(|st| {
90                        if !st.content.is_storage_live() {
91                            ctx.visit(st);
92                        }
93                    }),
94                }
95            }
96        });
97    }
98}
charon_lib/transform/recover_body_comments.rs

charon_lib/transform/
recover_body_comments.rs
