1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75

use std::collections::HashMap;

use crate::expr::{format_value, Expr};

// #insight The optimizer does not err.

// #insight
// The optimizer converts general Expr::List expressions into execution-friendly
// expressions like Expr::Array, Expr::Dict, etc. It also strips unnecessary
// annotations.

// #todo what does optimize do? I think it just removes some annotations.

pub fn optimize_fn(expr: Expr) -> Expr {
    match expr.unpack() {
        Expr::List(ref terms) => {
            if !terms.is_empty() {
                if let Expr::Symbol(s) = &terms[0].unpack() {
                    if s == "Array" {
                        let items = terms[1..].iter().map(|ax| ax.unpack().clone()).collect();
                        return Expr::maybe_annotated(Expr::Array(items), expr.annotations());
                    } else if s == "Dict" {
                        let items: Vec<Expr> =
                            terms[1..].iter().map(|ax| ax.unpack().clone()).collect();
                        let mut dict = HashMap::new();
                        for pair in items.chunks(2) {
                            let k = pair[0].clone();
                            let v = pair[1].clone();
                            dict.insert(format_value(k), v);
                        }
                        return Expr::maybe_annotated(Expr::Dict(dict), expr.annotations());
                    }
                }
            }
            expr
        }
        _ => expr,
    }
}

pub fn optimize(expr: Expr) -> Expr {
    expr.transform(&optimize_fn)
}

#[cfg(test)]
mod tests {
    use crate::{api::parse_string, optimize::optimize};

    #[test]
    fn optimize_rewrites_array_expressions() {
        let input = r#"(do (let a [1 2 3 4]) (writeln (+ 2 3)))"#;

        let expr = parse_string(input).unwrap();

        let expr_optimized = optimize(expr);

        let s = format!("{expr_optimized:?}");

        assert!(s.contains("Array([Int(1), Int(2), Int(3), Int(4)])"));
    }

    // #todo the test is flaky for some reason, temporarily disabled, investigate.
    // #[test]
    // fn optimize_rewrites_dict_expressions() {
    //     let input = r#"(let a {:name "George" :age 25})"#;

    //     let expr = parse_string(input).unwrap();

    //     let expr_optimized = optimize(expr);

    //     let s = format!("{expr_optimized:?}");

    //     assert!(s.contains(r#"Dict({"name": String("George"), "age": Int(25)})"#));
    // }
}