use tatara_lisp::{Atom, Sexp, Spanned, SpannedForm};
#[must_use]
pub fn erase_types(forms: &[Spanned]) -> Vec<Spanned> {
forms.iter().map(erase_form).collect()
}
fn erase_form(s: &Spanned) -> Spanned {
match &s.form {
SpannedForm::List(items) if is_define_typed(items) => {
let SpannedForm::List(sig) = &items[1].form else {
return s.clone();
};
let mut plain = Vec::with_capacity(sig.len());
plain.push(sig[0].clone());
for p in &sig[1..] {
match &p.form {
SpannedForm::List(pair) if !pair.is_empty() => plain.push(pair[0].clone()),
_ => plain.push(p.clone()),
}
}
Spanned::new(
s.span,
SpannedForm::List(vec![
Spanned::new(
items[0].span,
SpannedForm::Atom(Atom::Symbol("define".into())),
),
Spanned::new(items[1].span, SpannedForm::List(plain)),
erase_form(&items[3]),
]),
)
}
SpannedForm::List(items) => Spanned::new(
s.span,
SpannedForm::List(items.iter().map(erase_form).collect()),
),
_ => s.clone(),
}
}
fn is_define_typed(items: &[Spanned]) -> bool {
items.len() == 4
&& matches!(&items[0].form, SpannedForm::Atom(Atom::Symbol(n)) if &**n == "define-typed")
&& matches!(&items[1].form, SpannedForm::List(sig) if !sig.is_empty())
}
#[must_use]
pub fn to_sexps(forms: &[Spanned]) -> Vec<Sexp> {
forms.iter().map(Spanned::to_sexp).collect()
}
#[cfg(test)]
mod tests {
use super::*;
fn tree(src: &str) -> Vec<Spanned> {
blue_lang_syntax::parse_program_tree(src).expect("parse")
}
fn erased(src: &str) -> String {
to_sexps(&erase_types(&tree(src)))
.iter()
.map(ToString::to_string)
.collect::<Vec<_>>()
.join("\n")
}
#[test]
fn an_annotated_def_erases_to_a_plain_def() {
assert_eq!(
erased("def add(a: Int, b: Int) -> Int\n a + b\nend"),
"(define (add a b) (+ a b))"
);
}
#[test]
fn annotating_does_not_change_the_erased_program() {
assert_eq!(
erased("def add(a: Int, b: Int) -> Int\n a + b\nend"),
erased("def add(a, b)\n a + b\nend"),
"an annotation must not survive into the code that runs"
);
}
#[test]
fn erasure_preserves_arity() {
assert_eq!(
erased("def three(a: Int, b, c: Str) -> dyn\n a\nend"),
"(define (three a b c) a)"
);
}
#[test]
fn a_nested_annotated_def_is_erased() {
let out = erased("def outer(x: Int) -> Int\n def inner(y: Int) -> Int\n y\n end\nend");
assert!(
!out.contains("define-typed"),
"no annotation may survive anywhere in the tree: {out}"
);
}
#[test]
fn erasure_leaves_untyped_code_untouched() {
for src in ["1 + 2", "def f(x)\n x\nend", "[1, 2, 3]", "a.b(c)"] {
let forms = tree(src);
assert_eq!(
erase_types(&forms),
forms,
"erasure changed untyped source {src:?}"
);
}
}
#[test]
fn erasure_leaves_no_synthetic_span_behind() {
fn walk(node: &Spanned, src: &str, seen: &mut usize) {
assert!(
!node.span.is_synthetic(),
"a synthetic span survived erasure of {src:?}"
);
assert!(
src.get(node.span.start..node.span.end).is_some(),
"span {:?} is not a range in {src:?}",
node.span
);
*seen += 1;
match &node.form {
SpannedForm::List(xs) => {
for x in xs {
walk(x, src, seen);
}
}
SpannedForm::Quote(x)
| SpannedForm::Quasiquote(x)
| SpannedForm::Unquote(x)
| SpannedForm::UnquoteSplice(x) => walk(x, src, seen),
SpannedForm::Nil | SpannedForm::Atom(_) => {}
}
}
for src in [
"def add(a: Int, b: Int) -> Int\n a + b\nend",
"def three(a: Int, b, c: Str) -> dyn\n a\nend",
"def outer(x: Int) -> Int\n def inner(y: Int) -> Int\n y\n end\nend",
"def f(x)\n x\nend\nf(1)",
] {
let mut seen = 0;
for form in &erase_types(&tree(src)) {
walk(form, src, &mut seen);
}
assert!(seen >= 4, "{src:?} walked only {seen} nodes");
}
}
#[test]
fn the_invented_define_carries_the_annotations_own_position() {
let src = "x = 1\ndef add(a: Int) -> Int\n a\nend";
let erased = erase_types(&tree(src));
let head = erased[1].as_list().expect("a list")[0].clone();
assert_eq!(head.as_symbol(), Some("define"));
assert_eq!(
src.get(head.span.start..head.span.end),
Some("def"),
"the invented `define` must stand where the annotated `def` stood"
);
}
}