diatom/interpreter/scanner/
capture_scanner.rs

1use ahash::AHashMap;
2
3use crate::interpreter::Capture;
4
5use super::*;
6
7/// Scan and declare all (nested or not) captured variable before compile closure
8pub struct CaptureScanner<'a, Buffer: IoWrite> {
9    pub register_table: &'a mut RegisterTable,
10    pub gc: &'a mut Gc<Buffer>,
11    pub insts: &'a mut Vec<VmInst>,
12    pub overridden: AHashMap<String, usize>,
13}
14
15impl<'a, Buffer: IoWrite> CaptureScanner<'a, Buffer> {
16    pub fn scan_expr(&mut self, expr: &Expr) {
17        match expr {
18            Expr::Block { body, .. } => body.iter().for_each(|stmt| self.scan_stmt(stmt)),
19            Expr::If {
20                conditional,
21                default,
22                ..
23            } => conditional.iter().for_each(|(expr, stmts)| {
24                self.scan_expr(expr);
25                stmts.iter().for_each(|stmt| self.scan_stmt(stmt));
26                if let Some(stmts) = default {
27                    stmts.iter().for_each(|stmt| self.scan_stmt(stmt))
28                }
29            }),
30            Expr::Prefix { rhs, .. } => self.scan_expr(rhs),
31            Expr::Call {
32                lhs, parameters, ..
33            } => {
34                self.scan_expr(lhs);
35                parameters.iter().for_each(|expr| self.scan_expr(expr));
36            }
37            Expr::Index { lhs, rhs, .. } => {
38                self.scan_expr(lhs);
39                self.scan_expr(rhs);
40            }
41            Expr::OpenRange { lhs, .. } => {
42                self.scan_expr(&Expr::Id {
43                    loc: Loc {
44                        start: 0,
45                        end: 0,
46                        fid: usize::MAX,
47                    },
48                    name: "Range".to_string(),
49                });
50                self.scan_expr(lhs);
51            }
52            // Member rhs can not have legal constant values
53            Expr::Infix {
54                lhs,
55                op: OpInfix::Member | OpInfix::DoubleColon,
56                ..
57            } => {
58                self.scan_expr(lhs);
59            }
60            Expr::Infix { lhs, rhs, op, .. } => {
61                // x..y implicitly use `Range`
62                if matches!(op, OpInfix::Range) {
63                    self.scan_expr(&Expr::Id {
64                        loc: Loc {
65                            start: 0,
66                            end: 0,
67                            fid: usize::MAX,
68                        },
69                        name: "Range".to_string(),
70                    });
71                }
72                self.scan_expr(lhs);
73                self.scan_expr(rhs);
74            }
75            Expr::Fn {
76                parameters, body, ..
77            } => {
78                // Parameters will override upper scope variables
79                parameters.iter().for_each(|(name, _)| {
80                    if let Some(count) = self.overridden.get_mut(name) {
81                        *count += 1
82                    } else {
83                        self.overridden.insert(name.clone(), 1);
84                    }
85                });
86
87                self.scan_expr(body);
88
89                parameters.iter().for_each(|(name, _)| {
90                    let count = self.overridden.get_mut(name).unwrap();
91                    *count -= 1;
92                    if *count == 0 {
93                        self.overridden.remove(name);
94                    }
95                })
96            }
97            Expr::Id { name, .. } => {
98                if self.overridden.get(name).is_some() {
99                    return;
100                }
101                if let Some((id, depth, loc)) = self.register_table.lookup_variable(name) {
102                    // variable is captured
103                    // make a local copy and register capture info
104                    if depth > 0 {
105                        assert_eq!(depth, 1);
106                        let local_id = self.register_table.declare_captured_variable(name, loc);
107                        self.register_table.capture.push(Capture {
108                            rd: local_id,
109                            rs: id,
110                        });
111                    }
112                }
113            }
114            Expr::Parentheses { content, .. } => self.scan_expr(content),
115            Expr::Const { value, .. } => self.scan_const(value),
116            Expr::_Module { .. } => todo!(),
117            Expr::Error => unreachable!(),
118        }
119    }
120
121    pub fn scan_stmt(&mut self, stmt: &Stmt) {
122        match stmt {
123            Stmt::Expr { expr, .. } => self.scan_expr(expr),
124            Stmt::Continue { .. } => (),
125            Stmt::Break { .. } => (),
126            Stmt::Return { value, .. } => {
127                if let Some(expr) = value {
128                    self.scan_expr(expr)
129                }
130            }
131            Stmt::Loop {
132                condition, body, ..
133            } => {
134                if let Some(expr) = condition {
135                    self.scan_expr(expr)
136                }
137                body.iter().for_each(|stmt| self.scan_stmt(stmt));
138            }
139            Stmt::For { iterator, body, .. } => {
140                // For macro implicitly use `Option`
141                self.scan_expr(&Expr::Id {
142                    loc: Loc {
143                        start: 0,
144                        end: 0,
145                        fid: usize::MAX,
146                    },
147                    name: "Option".to_string(),
148                });
149                self.scan_expr(iterator);
150                body.iter().for_each(|stmt| self.scan_stmt(stmt));
151            }
152            Stmt::Def {
153                variable,
154                parameters,
155                body,
156                ..
157            } => {
158                self.scan_expr(variable);
159
160                // Parameters will override upper scope variables
161                parameters.iter().for_each(|(name, _)| {
162                    if let Some(count) = self.overridden.get_mut(name) {
163                        *count += 1
164                    } else {
165                        self.overridden.insert(name.clone(), 1);
166                    }
167                });
168
169                body.iter().for_each(|stmt| self.scan_stmt(stmt));
170
171                parameters.iter().for_each(|(name, _)| {
172                    let count = self.overridden.get_mut(name).unwrap();
173                    *count -= 1;
174                    if *count == 0 {
175                        self.overridden.remove(name);
176                    }
177                })
178            }
179            Stmt::Error => unreachable!(),
180        }
181    }
182
183    fn scan_const(&mut self, constant: &Const) {
184        match constant {
185            Const::List(list) => {
186                list.iter().for_each(|expr| self.scan_expr(expr));
187            }
188            Const::Table(entries) => {
189                entries.iter().for_each(|(_, expr, _)| self.scan_expr(expr));
190            }
191            _ => (),
192        }
193    }
194}