type_inference/
type_inference.rs1use 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}