Skip to main content

type_inference/
type_inference.rs

1//! Demonstrates expression type inference with the generated visitor.
2
3use m2_syn::nodes::{
4    Cell, Expr, ExprAssign, ExprAssignSimple, ExprCollection, ExprLiteral, ExprOp, ExprOpBin,
5    PackedContent,
6};
7use m2_syn::visit::Visit;
8use m2_syn::{ParseError, parse_file};
9
10#[derive(Default)]
11struct TypeVisitor {
12    inferred: Option<String>,
13}
14
15impl TypeVisitor {
16    fn type_of(&mut self, expression: &Expr) -> String {
17        assert!(self.inferred.is_none());
18        self.visit_expr(expression);
19        self.inferred
20            .take()
21            .expect("visiting an expression must infer one type")
22    }
23
24    fn type_of_cell(&mut self, cell: &Cell) -> String {
25        if cell.semicolon.is_some() {
26            return "Nothing".into();
27        }
28
29        match &cell.content {
30            PackedContent::Expr(expression) => self.type_of(expression),
31            PackedContent::Pack(_) => "Sequence".into(),
32        }
33    }
34}
35
36impl<'ast> Visit<'ast> for TypeVisitor {
37    fn visit_expr(&mut self, expression: &'ast Expr) {
38        let inferred = match expression {
39            Expr::Literal(ExprLiteral::Integer(_)) => "Int".into(),
40            Expr::Literal(ExprLiteral::Float(_)) => "Float".into(),
41            Expr::Literal(ExprLiteral::String(_) | ExprLiteral::RawString(_)) => "String".into(),
42            Expr::Symbol(symbol) => match symbol.text() {
43                "null" => "Nothing".into(),
44                "true" | "false" => "Bool".into(),
45                _ => "Unknown".into(),
46            },
47            Expr::Assign(ExprAssign::Simple(assignment)) => match assignment {
48                ExprAssignSimple::Global(assignment) => self.type_of(&assignment.value),
49                ExprAssignSimple::Local(assignment) => self.type_of(&assignment.value),
50            },
51            Expr::Op(ExprOp::Bin(ExprOpBin::Regular(binary)))
52                if binary.operator.spelling() == "==" =>
53            {
54                "Bool".into()
55            }
56            Expr::If(expression) => {
57                let then_type = self.type_of(&expression.then_clause.expr);
58                let else_type = expression
59                    .else_clause
60                    .as_ref()
61                    .map_or_else(|| "Nothing".into(), |clause| self.type_of(&clause.expr));
62                union(then_type, else_type)
63            }
64            Expr::Holder(holder) => match &holder.value {
65                Some(expression) => self.type_of(expression),
66                None if !holder.muted.is_empty() => "Nothing".into(),
67                None => "Sequence".into(),
68            },
69            Expr::Collection(ExprCollection::Sequence(_)) => "Sequence".into(),
70            Expr::Collection(ExprCollection::List(_)) => "List".into(),
71            Expr::Collection(ExprCollection::Array(_)) => "Array".into(),
72            Expr::Collection(ExprCollection::AngleBar(_)) => "List".into(),
73            _ => "Unknown".into(),
74        };
75        self.inferred = Some(inferred);
76    }
77}
78
79fn union(left: String, right: String) -> String {
80    if left == right {
81        left
82    } else {
83        format!("{left}|{right}")
84    }
85}
86
87fn infer(source: &str) -> Result<String, ParseError> {
88    let file = parse_file(source)?;
89    let [cell] = file.cells.as_slice() else {
90        panic!("expected one expression");
91    };
92    Ok(TypeVisitor::default().type_of_cell(cell))
93}
94
95fn main() -> Result<(), ParseError> {
96    for source in [
97        "1",
98        "1.5",
99        "answer = 1",
100        "a == b",
101        "if condition then 1 else 1.5",
102        "if condition then 1 else 2",
103        "if condition then 1",
104        "null",
105        "(1)",
106        "1;",
107        "(1, 2)",
108        "1 + 2",
109    ] {
110        println!("{source} => {}", infer(source)?);
111    }
112    Ok(())
113}
114
115#[cfg(test)]
116mod tests {
117    use super::infer;
118
119    #[test]
120    fn infers_the_demonstrated_expression_types() {
121        for (source, expected) in [
122            ("1", "Int"),
123            ("1.5", "Float"),
124            ("answer = 1", "Int"),
125            ("a == b", "Bool"),
126            ("if condition then 1 else 1.5", "Int|Float"),
127            ("if condition then 1 else 2", "Int"),
128            ("if condition then 1", "Int|Nothing"),
129            ("null", "Nothing"),
130            ("(1)", "Int"),
131            ("1;", "Nothing"),
132            ("(1, 2)", "Sequence"),
133            ("1 + 2", "Unknown"),
134        ] {
135            assert_eq!(infer(source).unwrap(), expected, "{source}");
136        }
137    }
138}