Skip to main content

lunar_lang/lua/
walker.rs

1use std::collections::HashMap;
2
3use purua::{parser::ast::*, TokenType};
4
5use super::lunarir::*;
6
7#[derive(Debug)]
8pub struct Walker {
9    pub msg_stack: Vec<LunarIR>,
10    pub idx_of_irep: usize,
11    pub current_irep: usize,
12    pub idx_of_ireps: HashMap<usize, IrepIndices>,
13}
14
15#[derive(Debug, Clone)]
16pub struct IrepIndices {
17    pub locals: usize,
18    pub syms: usize,
19    pub pool: usize,
20}
21
22impl Walker {
23    pub fn new() -> Self {
24        Walker {
25            msg_stack: Vec::new(),
26            idx_of_irep: 0,
27            current_irep: 0,
28            idx_of_ireps: HashMap::from([(
29                0,
30                IrepIndices {
31                    locals: 0,
32                    syms: 0,
33                    pool: 0,
34                },
35            )]),
36        }
37    }
38
39    pub fn push_msg(&mut self, msg: LunarIR) {
40        self.msg_stack.push(msg);
41    }
42
43    pub fn walk(&mut self, root: &Block) {
44        self.walk_block(root);
45    }
46
47    pub fn walk_block(&mut self, block: &Block) {
48        let chunk = &block.0;
49        self.walk_chunk(chunk);
50    }
51
52    pub fn walk_chunk(&mut self, chunk: &Chunk) {
53        self.push_msg(LunarIR::ChunkStart(self.idx_of_irep));
54        let before_irep = self.current_irep;
55        self.current_irep = self.idx_of_irep;
56        self.idx_of_irep += 1;
57        self.idx_of_ireps.insert(
58            self.idx_of_irep,
59            IrepIndices {
60                locals: 0,
61                syms: 0,
62                pool: 0,
63            },
64        );
65        if self.current_irep != 0 {
66            // TODO: replace dummy operand of OP_ENTER
67            self.push_msg(LunarIR::Enter(0x40000));
68            // dummy locals = 1
69            self.idx_of_ireps.get_mut(&self.current_irep).unwrap().locals += 1;
70        }
71
72        let statements = &chunk.0;
73        for statement in statements {
74            self.walk_stat(statement);
75        }
76
77        if let Some(last_stat) = &chunk.1 {
78            self.walk_laststat(last_stat);
79        } else {
80            self.push_msg(LunarIR::NoReturn);
81        }
82
83        if self.current_irep == 0 {
84            self.push_msg(LunarIR::Stop);
85        } else {
86            self.current_irep = before_irep;
87        }
88
89        self.push_msg(LunarIR::ChunkEnd);
90    }
91
92    pub fn walk_stat(&mut self, stat: &Stat) {
93        match stat {
94            Stat::FunctionCall(function_call) => {
95                let func_name = &function_call.0;
96                let args = &function_call.2;
97
98                // walk_var
99                self.walk_prefixexpr(func_name);
100
101                self.push_msg(LunarIR::FunctionCallStart(
102                    self.idx_of_ireps[&self.current_irep].syms - 1,
103                ));
104
105                self.walk_args(args);
106
107                self.push_msg(LunarIR::FunctionCallEnd);
108            },
109            Stat::For(_token, expr, expr1, expr2, block) => {
110                let begin = self.ensure_expr_as_number(expr) as usize;
111                let end = self.ensure_expr_as_number(expr1) as usize;
112                let step = match expr2 {
113                    Some(expr) => self.ensure_expr_as_number(expr),
114                    None => 1.0,
115                } as usize;
116
117                self.push_msg(LunarIR::StoreSym(
118                    self.idx_of_ireps[&self.current_irep].syms,
119                    "each".to_string(),
120                ));
121                self.push_msg(LunarIR::ForStart(self.idx_of_ireps[&self.current_irep].syms));
122                self.idx_of_ireps.get_mut(&self.current_irep).unwrap().syms += 1;
123
124                self.push_msg(LunarIR::ForParam(
125                    begin,
126                    end,
127                    step,
128                ));
129                // TODO: token is_a varname of each,
130                // which can be used inside of this block!
131                // TODO: real handle of children reps... :(
132                self.push_msg(LunarIR::Block(self.idx_of_irep - 1));
133                self.walk_block(block);
134                self.push_msg(LunarIR::ForEnd);
135            },
136            _ => {
137                // Handle other types of statements
138                panic!("Unsupported statement: {:?}", stat);
139            }
140        }
141    }
142
143    pub fn walk_args(&mut self, args: &Args) {
144        match args {
145            Args::ArgsString(string) => {
146                self.push_msg(LunarIR::PoolString(
147                    self.idx_of_ireps[&self.current_irep].pool,
148                    string.to_owned(),
149                ));
150                self.push_msg(LunarIR::FunctionCallArg(
151                    0,
152                    LunarValue::String(self.idx_of_ireps[&self.current_irep].pool),
153                ));
154                self.idx_of_ireps.get_mut(&self.current_irep).unwrap().pool += 1;
155            },
156            _ => {
157                // Handle other types of arguments
158                panic!("Unsupported argument: {:?}", args);
159            }
160        }
161    }
162
163    // TODO: replace TryInto?
164    pub fn ensure_expr_as_number(&mut self, expr: &Expr) -> f64 {
165        match expr {
166            Expr::Number(f) => *f,
167            _ => {
168                // Handle other types of expressions
169                panic!("Unsupported expression: {:?}", expr);
170            }
171        }
172    }
173
174    pub fn walk_prefixexpr(&mut self, prefix_expr: &PrefixExp) {
175        match prefix_expr {
176            PrefixExp::PrefixVar(var) => {
177                // Handle variable prefix expression
178                let var = var.as_ref();
179                self.walk_var(var);
180            },
181            _ => {
182                // Handle other types of prefix expressions
183                panic!("Unsupported prefix expression: {:?}", prefix_expr);
184            }
185        }
186    }
187
188    pub fn walk_var(&mut self, var: &Var) {
189        match var {
190            Var::VarName(name) => {
191                // Handle variable name
192                let name = match name.token_type {
193                    TokenType::Name => {
194                        name.lexeme.clone()
195                    },
196                    _ => {
197                        // Handle other types of variable names
198                        panic!("Unsupported variable name: {:?}", name);
199                    }
200                };
201                self.push_msg(LunarIR::StoreSym(
202                    self.idx_of_ireps[&self.current_irep].syms,
203                    name.clone().try_into().unwrap(),
204                ));
205                self.idx_of_ireps.get_mut(&self.current_irep).unwrap().syms += 1;
206            },
207            _ => {
208                // Handle other types of variables
209                panic!("Unsupported variable: {:?}", var);
210            }
211        }
212    }
213
214    pub fn walk_laststat(&mut self, last_stat: &LastStat) {
215        match last_stat {
216            LastStat::Return(_ret) => todo!(),
217            LastStat::Break => todo!(),
218        }
219    }
220}