1use tatara_lisp::{Atom, Span, Spanned, SpannedForm};
26
27#[derive(Clone, Copy, Debug, PartialEq, Eq)]
29pub enum HeadKind {
30 SpecialForm,
31 Macro,
32 Value,
33}
34
35#[derive(Clone, Copy, Debug, PartialEq, Eq)]
37pub enum BinderKind {
38 Parameter,
39 Assignment,
41 Function,
43 Let,
44 Catch,
45}
46
47#[derive(Clone, Debug, PartialEq)]
54pub struct Import {
55 pub package: String,
56 pub names: Vec<(String, Span)>,
59 pub listed: bool,
61 pub span: Span,
63 pub package_span: Span,
65}
66
67
68#[must_use]
77pub fn use_target(form: &Spanned) -> Option<Import> {
78 let items = form.as_list()?;
79 let (head, arg, list) = match items {
80 [head, arg] => (head, arg, None),
81 [head, arg, list] => (head, arg, Some(list)),
82 _ => return None,
83 };
84 let (SpannedForm::Atom(Atom::Symbol(s)), SpannedForm::Atom(Atom::Str(name))) =
85 (&head.form, &arg.form)
86 else {
87 return None;
88 };
89 if s != "use" {
90 return None;
91 }
92 let mut names = Vec::new();
93 if let Some(list) = list {
94 let parts = list.as_list()?;
95 if parts.first().and_then(Spanned::as_symbol) != Some("list") {
96 return None;
97 }
98 for p in &parts[1..] {
99 match &p.form {
100 SpannedForm::Atom(Atom::Keyword(k)) => names.push((k.clone(), p.span)),
101 _ => return None,
102 }
103 }
104 }
105 Some(Import {
106 package: name.clone(),
107 names,
108 listed: list.is_some(),
109 span: form.span,
110 package_span: arg.span,
111 })
112}
113
114pub const SHAPED_FORMS: &[&str] = &[
119 "lambda",
120 "let",
121 "let*",
122 "letrec",
123 "try",
124 "cond",
125 "quote",
126 "quasiquote",
127 "provide",
128 "require",
129];
130
131pub trait Scopes {
133 fn head_kind(&self, name: &str) -> Option<HeadKind>;
136 fn open(&mut self);
138 fn bind(&mut self, name: &str, span: Span, kind: BinderKind);
140 fn close(&mut self);
142 fn reference(&mut self, node: &Spanned, name: &str, opaque: bool);
145 fn head(&mut self, _node: &Spanned, _name: &str, _kind: HeadKind) {}
148 fn template_symbol(&mut self, _node: &Spanned, _name: &str) {}
151}
152
153pub fn walk_program(forms: &[Spanned], pass: &mut dyn Scopes) {
157 for form in forms {
158 walk_top(form, pass);
159 }
160}
161
162pub fn walk_top(form: &Spanned, pass: &mut dyn Scopes) {
164 Walker { pass, depth: 0 }.walk(form, false);
165}
166
167#[must_use]
177pub fn definitions_of(form: &Spanned) -> Vec<(String, Span, HeadKind)> {
178 let mut out = Vec::new();
179 definition_nodes(form, &mut |node, kind| {
180 if let Some(n) = node.as_symbol() {
181 out.push((n.to_string(), node.span, kind));
182 }
183 });
184 out
185}
186
187pub fn definition_nodes<'a>(form: &'a Spanned, out: &mut dyn FnMut(&'a Spanned, HeadKind)) {
190 let Some(items) = form.as_list() else { return };
191 let head = items.first().and_then(Spanned::as_symbol);
192 match head {
193 Some("defmacro") => {
194 if let Some(n) = items.get(1).filter(|n| n.as_symbol().is_some()) {
195 out(n, HeadKind::Macro);
196 }
197 }
198 Some("define" | "define-typed") => {
199 if let Some(n) = defined_name_node(items) {
200 out(n, HeadKind::Value);
201 }
202 if items.get(1).is_some_and(|t| t.as_symbol().is_some()) {
205 for item in &items[2..] {
206 definition_nodes(item, out);
207 }
208 }
209 }
210 Some("lambda" | "let" | "let*" | "letrec" | "deftest" | "quote" | "quasiquote") => {}
211 Some("try") => {
212 for item in &items[1..] {
214 let is_clause = item
215 .as_list()
216 .and_then(|c| c.first())
217 .and_then(Spanned::as_symbol)
218 .is_some_and(|h| h == "catch" || h == "finally");
219 if !is_clause {
220 definition_nodes(item, out);
221 }
222 }
223 }
224 _ => {
225 for item in items {
226 definition_nodes(item, out);
227 }
228 }
229 }
230}
231
232#[must_use]
234pub fn defined_name_node(items: &[Spanned]) -> Option<&Spanned> {
235 let head = items.first()?.as_symbol()?;
236 if head != "define" && head != "define-typed" {
237 return None;
238 }
239 let target = items.get(1)?;
240 if target.as_symbol().is_some() {
241 return Some(target);
242 }
243 let name = target.as_list()?.first()?;
244 name.as_symbol().map(|_| name)
245}
246
247#[must_use]
250pub fn params_of(list: &Spanned) -> Vec<(String, Span)> {
251 if let Some(n) = list.as_symbol() {
252 return vec![(n.to_string(), list.span)];
253 }
254 list.as_list()
255 .map(|ps| ps.iter().filter_map(param).collect())
256 .unwrap_or_default()
257}
258
259fn param(p: &Spanned) -> Option<(String, Span)> {
260 let n = p.as_symbol()?;
261 if n.starts_with('&') || n == "." {
262 return None;
263 }
264 Some((n.to_string(), p.span))
265}
266
267struct Walker<'p> {
268 pass: &'p mut dyn Scopes,
269 depth: usize,
271}
272
273impl Walker<'_> {
274 fn open(&mut self) {
275 self.depth += 1;
276 self.pass.open();
277 }
278
279 fn close(&mut self) {
280 self.depth -= 1;
281 self.pass.close();
282 }
283
284 fn hoist(&mut self, body: &[Spanned]) {
288 for f in body {
289 let function = f.as_list().is_some_and(|items| {
290 items.first().and_then(Spanned::as_symbol) == Some("define")
291 && items.get(1).is_some_and(|t| t.as_list().is_some())
292 });
293 let pass = &mut *self.pass;
294 definition_nodes(f, &mut |node, kind| {
295 let Some(name) = node.as_symbol() else { return };
296 let bk = match kind {
297 HeadKind::Value if !function => BinderKind::Assignment,
298 _ => BinderKind::Function,
299 };
300 pass.bind(name, node.span, bk);
301 });
302 }
303 }
304
305 fn body(&mut self, params: &[(String, Span)], forms: &[Spanned], opaque: bool) {
308 self.open();
309 for (p, s) in params {
310 self.pass.bind(p, *s, BinderKind::Parameter);
311 }
312 self.hoist(forms);
313 for f in forms {
314 self.walk(f, opaque);
315 }
316 self.close();
317 }
318
319 fn walk(&mut self, form: &Spanned, opaque: bool) {
320 match &form.form {
321 SpannedForm::Atom(Atom::Symbol(n)) => self.pass.reference(form, n, opaque),
322 SpannedForm::Quote(_) => {}
323 SpannedForm::Quasiquote(inner) => self.walk_template(inner, opaque),
324 SpannedForm::Unquote(inner) | SpannedForm::UnquoteSplice(inner) => {
325 self.walk(inner, opaque);
326 }
327 SpannedForm::List(items) if !items.is_empty() => self.walk_list(items, opaque),
328 _ => {}
329 }
330 }
331
332 fn walk_template(&mut self, form: &Spanned, opaque: bool) {
334 match &form.form {
335 SpannedForm::Atom(Atom::Symbol(n)) => self.pass.template_symbol(form, n),
336 SpannedForm::Unquote(inner) | SpannedForm::UnquoteSplice(inner) => {
337 self.walk(inner, opaque);
338 }
339 SpannedForm::List(items) => {
340 for i in items {
341 self.walk_template(i, opaque);
342 }
343 }
344 SpannedForm::Quasiquote(inner) => self.walk_template(inner, opaque),
345 _ => {}
346 }
347 }
348
349 fn walk_all(&mut self, items: &[Spanned], opaque: bool) {
350 for i in items {
351 self.walk(i, opaque);
352 }
353 }
354
355 fn walk_list(&mut self, items: &[Spanned], opaque: bool) {
356 let Some(head) = items[0].as_symbol() else {
357 self.walk_all(items, opaque);
358 return;
359 };
360 if head == "deftest" && self.depth == 0 {
363 if let Some(body) = items.get(2) {
364 self.body(&[], std::slice::from_ref(body), opaque);
365 }
366 return;
367 }
368 if head == "use" && self.depth == 0 {
371 self.pass.head(&items[0], head, HeadKind::SpecialForm);
372 return;
373 }
374 if head == "define" || head == "define-typed" {
377 self.pass.head(&items[0], head, HeadKind::SpecialForm);
378 self.walk_define(head, items, opaque);
379 return;
380 }
381 if head == "defmacro" && self.depth == 0 {
385 self.pass.head(&items[0], head, HeadKind::SpecialForm);
386 let params = items.get(2).map(params_of).unwrap_or_default();
387 self.body(¶ms, items.get(3..).unwrap_or(&[]), opaque);
388 return;
389 }
390 match self.pass.head_kind(head) {
391 Some(HeadKind::SpecialForm) => {
392 self.pass.head(&items[0], head, HeadKind::SpecialForm);
393 self.walk_special(head, items, opaque);
394 }
395 Some(HeadKind::Macro) => {
396 self.pass.head(&items[0], head, HeadKind::Macro);
398 for i in &items[1..] {
399 self.walk(i, true);
400 }
401 }
402 Some(HeadKind::Value) | None => {
403 self.walk(&items[0], opaque);
404 self.walk_all(&items[1..], opaque);
405 }
406 }
407 }
408
409 fn walk_define(&mut self, head: &str, items: &[Spanned], opaque: bool) {
410 let Some(target) = items.get(1) else { return };
411 if target.as_symbol().is_some() {
412 self.walk_all(&items[2..], opaque);
414 return;
415 }
416 let Some(sig) = target.as_list() else {
417 self.walk_all(&items[1..], opaque);
418 return;
419 };
420 let params: Vec<(String, Span)> = sig
421 .iter()
422 .skip(1)
423 .filter_map(|p| {
424 if head == "define-typed" {
425 let pair = p.as_list()?;
426 let n = pair.first()?;
427 Some((n.as_symbol()?.to_string(), n.span))
428 } else {
429 param(p)
430 }
431 })
432 .collect();
433 let body_from = if head == "define-typed" { 3 } else { 2 };
435 self.body(¶ms, items.get(body_from..).unwrap_or(&[]), opaque);
436 }
437
438 fn walk_special(&mut self, head: &str, items: &[Spanned], opaque: bool) {
439 let args = &items[1..];
440 match head {
441 "quote" | "provide" | "require" => {}
442 "quasiquote" => {
443 for a in args {
444 self.walk_template(a, opaque);
445 }
446 }
447 "lambda" => {
448 let params = args.first().map(params_of).unwrap_or_default();
449 self.body(¶ms, args.get(1..).unwrap_or(&[]), opaque);
450 }
451 "cond" => {
452 for clause in args {
453 match clause.as_list() {
454 Some(parts) => {
455 let rest = match parts.first().and_then(Spanned::as_symbol) {
456 Some("else") => &parts[1..],
457 _ => parts,
458 };
459 for p in rest {
460 if p.as_symbol() != Some("=>") {
461 self.walk(p, opaque);
462 }
463 }
464 }
465 None => self.walk(clause, opaque),
466 }
467 }
468 }
469 "let" | "let*" | "letrec" => self.walk_let(head, args, opaque),
470 "try" => {
471 for a in args {
472 let clause = a.as_list();
473 match clause.and_then(|c| c.first()).and_then(Spanned::as_symbol) {
474 Some("catch") => {
475 let c = clause.expect("matched");
476 let var = c.get(1).and_then(|v| {
477 v.as_list()
478 .and_then(|l| l.first())
479 .and_then(|s| s.as_symbol().map(|n| (n.to_string(), s.span)))
480 .or_else(|| v.as_symbol().map(|n| (n.to_string(), v.span)))
481 });
482 self.open();
483 if let Some((n, s)) = var {
484 self.pass.bind(&n, s, BinderKind::Catch);
485 }
486 self.walk_all(c.get(2..).unwrap_or(&[]), opaque);
487 self.close();
488 }
489 Some("finally") => {
490 let c = clause.expect("matched");
491 self.walk_all(&c[1..], opaque);
492 }
493 _ => self.walk(a, opaque),
494 }
495 }
496 }
497 _ => self.walk_all(args, opaque),
500 }
501 }
502
503 fn walk_let(&mut self, head: &str, args: &[Spanned], opaque: bool) {
504 let (name, rest) = match args.first() {
506 Some(n) if n.as_symbol().is_some() && head == "let" => (Some(n), &args[1..]),
507 _ => (None, args),
508 };
509 let Some(bindings) = rest.first().and_then(Spanned::as_list) else {
510 self.walk_all(rest, opaque);
511 return;
512 };
513 let pairs: Vec<(String, Span, Option<&Spanned>)> = bindings
514 .iter()
515 .filter_map(|b| match b.as_list() {
516 Some(pair) => {
517 let n = pair.first()?;
518 Some((n.as_symbol()?.to_string(), n.span, pair.get(1)))
519 }
520 None => b.as_symbol().map(|n| (n.to_string(), b.span, None)),
521 })
522 .collect();
523 let body = rest.get(1..).unwrap_or(&[]);
524 if head == "let" {
525 for (_, _, init) in &pairs {
527 if let Some(e) = init {
528 self.walk(e, opaque);
529 }
530 }
531 }
532 self.open();
533 match head {
534 "let" => {
535 if let Some(n) = name {
536 self.pass
537 .bind(n.as_symbol().expect("matched"), n.span, BinderKind::Function);
538 }
539 for (n, s, _) in &pairs {
540 self.pass.bind(n, *s, BinderKind::Let);
541 }
542 }
543 "let*" => {
544 for (n, s, init) in &pairs {
545 if let Some(e) = init {
546 self.walk(e, opaque);
547 }
548 self.pass.bind(n, *s, BinderKind::Let);
549 }
550 }
551 _ => {
552 for (n, s, _) in &pairs {
553 self.pass.bind(n, *s, BinderKind::Let);
554 }
555 for (_, _, init) in &pairs {
556 if let Some(e) = init {
557 self.walk(e, opaque);
558 }
559 }
560 }
561 }
562 self.hoist(body);
563 self.walk_all(body, opaque);
564 self.close();
565 }
566}