1use anyhow::{anyhow, bail, Context, Result};
2use regex::Regex;
3use std::collections::HashSet;
4use std::path::Path;
5
6#[derive(Clone, Debug)]
7pub struct Function {
8 pub name: String,
9 pub args_count: i8,
10 pub locals_count: i8,
11 pub body: Vec<Stmt>,
12 pub raw: Vec<String>,
14}
15
16#[derive(Clone, Debug, PartialEq, Eq)]
17pub enum GlobalKind {
18 NonVolatile,
19 Volatile,
20}
21
22#[derive(Clone, Debug)]
23pub struct GlobalDecl {
24 pub name: String,
25 pub kind: GlobalKind,
26}
27
28#[derive(Clone, Debug)]
29pub struct Program {
30 pub globals: Vec<GlobalDecl>,
31 pub functions: Vec<Function>,
32}
33
34#[derive(Clone, Debug)]
35pub enum Stmt {
36 Simple(String),
37 Return(Option<String>),
38 Break,
39 If {
40 arms: Vec<(String, Vec<Stmt>)>,
41 else_arm: Option<Vec<Stmt>>,
42 },
43 While {
44 cond: String,
45 body: Vec<Stmt>,
46 },
47}
48
49fn is_comment_or_empty(s: &str) -> bool {
50 let t = s.trim();
51 t.is_empty() || t.starts_with("--")
52}
53
54fn strip_local(stmt: &str) -> &str {
55 let t = stmt.trim_start();
56 if let Some(rest) = t.strip_prefix("local ") {
57 return rest.trim_start();
58 }
59 t
60}
61
62fn is_if_start(s: &str) -> bool {
63 let t = s.trim();
64 t.starts_with("if ") && t.ends_with(" then")
65}
66
67fn is_elseif(s: &str) -> bool {
68 let t = s.trim();
69 t.starts_with("elseif ") && t.ends_with(" then")
70}
71
72fn is_else(s: &str) -> bool {
73 s.trim() == "else"
74}
75
76fn is_while_start(s: &str) -> bool {
77 let t = s.trim();
78 t.starts_with("while ") && t.ends_with(" do")
79}
80
81fn is_end(s: &str) -> bool {
82 s.trim() == "end"
83}
84
85fn extract_if_cond(line: &str) -> Result<String> {
86 let t = line.trim();
87 if !t.ends_with(" then") {
88 bail!("invalid if header: {t}");
89 }
90 let inner = t
91 .strip_prefix("if ")
92 .or_else(|| t.strip_prefix("elseif "))
93 .ok_or_else(|| anyhow!("invalid if header: {t}"))?;
94 let inner = inner.trim_end_matches(" then");
95 Ok(inner.trim().to_string())
96}
97
98fn extract_while_cond(line: &str) -> Result<String> {
99 let t = line.trim();
100 if !t.ends_with(" do") {
101 bail!("invalid while header: {t}");
102 }
103 let inner = t
104 .strip_prefix("while ")
105 .ok_or_else(|| anyhow!("invalid while header: {t}"))?;
106 let inner = inner.trim_end_matches(" do");
107 Ok(inner.trim().to_string())
108}
109
110fn parse_block(lines: &[String], i: &mut usize, stop_on: &[&str]) -> Result<Vec<Stmt>> {
111 let mut out: Vec<Stmt> = Vec::new();
112
113 while *i < lines.len() {
114 let line0 = lines[*i].clone();
115 let line = line0.trim();
116
117 if is_comment_or_empty(line) {
118 *i += 1;
119 continue;
120 }
121
122 if stop_on.iter().any(|tok| match *tok {
123 "elseif" => is_elseif(line),
124 "else" => is_else(line),
125 "end" => is_end(line),
126 _ => false,
127 }) {
128 break;
129 }
130
131 if is_if_start(line) {
132 let cond = extract_if_cond(line)?;
133 *i += 1;
134 let then_block = parse_block(lines, i, &["elseif", "else", "end"])?;
135
136 let mut arms: Vec<(String, Vec<Stmt>)> = vec![(cond, then_block)];
137 while *i < lines.len() && is_elseif(lines[*i].trim()) {
138 let c = extract_if_cond(lines[*i].trim())?;
139 *i += 1;
140 let b = parse_block(lines, i, &["elseif", "else", "end"])?;
141 arms.push((c, b));
142 }
143
144 let mut else_arm: Option<Vec<Stmt>> = None;
145 if *i < lines.len() && is_else(lines[*i].trim()) {
146 *i += 1;
147 let b = parse_block(lines, i, &["end"])?;
148 else_arm = Some(b);
149 }
150
151 if *i >= lines.len() || !is_end(lines[*i].trim()) {
152 bail!("if without closing end");
153 }
154 *i += 1;
155
156 out.push(Stmt::If { arms, else_arm });
157 continue;
158 }
159
160 if is_while_start(line) {
161 let cond = extract_while_cond(line)?;
162 *i += 1;
163 let body = parse_block(lines, i, &["end"])?;
164 if *i >= lines.len() || !is_end(lines[*i].trim()) {
165 bail!("while without closing end");
166 }
167 *i += 1;
168 out.push(Stmt::While { cond, body });
169 continue;
170 }
171
172 if line == "break" {
173 *i += 1;
174 out.push(Stmt::Break);
175 continue;
176 }
177
178 if line == "return" {
179 *i += 1;
180 out.push(Stmt::Return(None));
181 continue;
182 }
183
184 if let Some(rest) = line.strip_prefix("return ") {
185 *i += 1;
186 out.push(Stmt::Return(Some(rest.trim().to_string())));
187 continue;
188 }
189
190 if line.starts_with("local ") && !line.contains('=') {
192 *i += 1;
193 continue;
194 }
195
196 let simple = strip_local(line).to_string();
197 *i += 1;
198 out.push(Stmt::Simple(simple));
199 }
200
201 Ok(out)
202}
203
204fn split_functions(lines: &[String], start_idx: usize) -> Result<Vec<Vec<String>>> {
205 let head_re = Regex::new(r"^(?:local\s+)?function\s+").unwrap();
206
207 let mut out: Vec<Vec<String>> = Vec::new();
208 let mut i = start_idx;
209 while i < lines.len() {
210 if head_re.is_match(lines[i].trim()) {
211 let start = i;
212 let mut nest = 1i32;
213 i += 1;
214 while i < lines.len() && nest > 0 {
215 let t = lines[i].trim();
216 if head_re.is_match(t) {
217 nest += 1;
218 } else if is_if_start(t) {
219 nest += 1;
220 } else if is_while_start(t) {
221 nest += 1;
222 } else if is_end(t) {
223 nest -= 1;
224 }
225 i += 1;
226 }
227 out.push(lines[start..i].to_vec());
228 } else if is_comment_or_empty(lines[i].trim()) {
229 i += 1;
230 } else {
231 bail!("unsupported top-level statement: {}", lines[i].trim());
232 }
233 }
234
235 if out.is_empty() {
236 bail!("no functions found in Lua");
237 }
238
239 Ok(out)
240}
241
242fn parse_global_line(
243 line: &str,
244 seen: &mut HashSet<String>,
245 out: &mut Vec<GlobalDecl>,
246) -> Result<()> {
247 let t = line.trim();
248 let (kind, rest) = if let Some(rest) = t.strip_prefix("global ") {
249 (GlobalKind::NonVolatile, rest.trim())
250 } else if let Some(rest) = t.strip_prefix("volatile global ") {
251 (GlobalKind::Volatile, rest.trim())
252 } else {
253 bail!("unsupported top-level statement: {t}");
254 };
255
256 if rest.is_empty() {
257 bail!("empty global declaration: {t}");
258 }
259 if rest.contains('=') {
260 bail!("global initializers are not supported: {t}");
261 }
262
263 let re_g = Regex::new(r"^g\d+$").unwrap();
264 let re_vg = Regex::new(r"^vg\d+$").unwrap();
265
266 for raw_name in rest.split(',') {
267 let name = raw_name.trim();
268 if name.is_empty() {
269 bail!("empty global name in declaration: {t}");
270 }
271 match kind {
272 GlobalKind::NonVolatile => {
273 if !re_g.is_match(name) {
274 bail!("non-volatile globals must be named gN: {name}");
275 }
276 }
277 GlobalKind::Volatile => {
278 if !re_vg.is_match(name) {
279 bail!("volatile globals must be named vgN: {name}");
280 }
281 }
282 }
283 if !seen.insert(name.to_string()) {
284 bail!("duplicate global declaration: {name}");
285 }
286 out.push(GlobalDecl {
287 name: name.to_string(),
288 kind: kind.clone(),
289 });
290 }
291
292 Ok(())
293}
294
295pub fn parse_lua(path: &Path) -> Result<Program> {
296 let txt =
297 std::fs::read_to_string(path).with_context(|| format!("read lua: {}", path.display()))?;
298 let lines: Vec<String> = txt.lines().map(|s| s.to_string()).collect();
299 let head_re = Regex::new(r"^(?:local\s+)?function\s+").unwrap();
300
301 let mut globals: Vec<GlobalDecl> = Vec::new();
302 let mut seen_globals: HashSet<String> = HashSet::new();
303 let mut first_fn_idx = None;
304
305 for (i, line) in lines.iter().enumerate() {
306 let t = line.trim();
307 if is_comment_or_empty(t) {
308 continue;
309 }
310 if head_re.is_match(t) {
311 first_fn_idx = Some(i);
312 break;
313 }
314 parse_global_line(t, &mut seen_globals, &mut globals)?;
315 }
316
317 let start_idx = first_fn_idx.ok_or_else(|| anyhow!("no functions found in Lua"))?;
318 let funcs_lines = split_functions(&lines, start_idx)?;
319
320 let head_re = Regex::new(r"^(?:local\s+)?function\s+([A-Za-z_]\w*)\s*\(([^)]*)\)\s*$").unwrap();
321 let re_a = Regex::new(r"\ba(\d+)\b").unwrap();
322 let re_l = Regex::new(r"\bl(\d+)\b").unwrap();
323
324 let mut funs: Vec<Function> = Vec::new();
325 for fl in funcs_lines {
326 if fl.is_empty() {
327 continue;
328 }
329 let head = fl[0].trim();
330 let caps = head_re
331 .captures(head)
332 .ok_or_else(|| anyhow!("unexpected function header: {head}"))?;
333 let name = caps.get(1).unwrap().as_str().to_string();
334 let args_s = caps.get(2).unwrap().as_str().trim();
335 let args: Vec<&str> = if args_s.is_empty() {
336 vec![]
337 } else {
338 args_s
339 .split(',')
340 .map(|x| x.trim())
341 .filter(|x| !x.is_empty())
342 .collect()
343 };
344
345 let mut max_a: Option<u32> = None;
346 for a in &args {
347 if let Some(mm) = re_a.captures(a) {
348 let v: u32 = mm.get(1).unwrap().as_str().parse().unwrap_or(0);
349 max_a = Some(max_a.map(|x| x.max(v)).unwrap_or(v));
350 }
351 }
352 for ln in &fl {
353 for mm in re_a.captures_iter(ln) {
354 let v: u32 = mm.get(1).unwrap().as_str().parse().unwrap_or(0);
355 max_a = Some(max_a.map(|x| x.max(v)).unwrap_or(v));
356 }
357 }
358 let args_count = i8::try_from(max_a.map(|x| x + 1).unwrap_or(0))
359 .map_err(|_| anyhow!("args_count does not fit i8"))?;
360
361 let mut max_l: Option<u32> = None;
362 for ln in &fl {
363 for mm in re_l.captures_iter(ln) {
364 let v: u32 = mm.get(1).unwrap().as_str().parse().unwrap_or(0);
365 max_l = Some(max_l.map(|x| x.max(v)).unwrap_or(v));
366 }
367 }
368 let locals_count = i8::try_from(max_l.map(|x| x + 1).unwrap_or(0))
369 .map_err(|_| anyhow!("locals_count does not fit i8"))?;
370
371 if fl.len() < 2 {
372 bail!("function {name}: too short");
373 }
374 let body_lines: Vec<String> = fl[1..fl.len() - 1].to_vec();
375 let mut idx = 0usize;
376 let body = parse_block(&body_lines, &mut idx, &[])?;
377
378 funs.push(Function {
379 name,
380 args_count,
381 locals_count,
382 body,
383 raw: fl,
384 });
385 }
386
387 Ok(Program {
388 globals,
389 functions: funs,
390 })
391}