type_inference/
type_inference.rs

1use std::collections::HashMap;
2
3use m2_syn::visit::Visit;
4use m2_syn::{
5    Array, BinaryExpression, BinaryOperator, Component, Expr, ForLoop, IfStatement, IntegerLiteral,
6    List, LoopBody, ParenthesizedExpression, ParseError, Punctuated, ReturnStatement, Sequence,
7    SourceId, StringLiteral, Symbol, WhileLoop, parse_file,
8};
9
10#[derive(Debug, Clone, PartialEq, Eq)]
11enum TypeRange<T, Q> {
12    Known(Vec<T>),
13    Deferred(Q),
14}
15
16#[derive(Debug, Clone, PartialEq, Eq)]
17struct Fact<T, V, Q> {
18    types: TypeRange<T, Q>,
19    value: Option<V>,
20}
21
22#[derive(Debug, Clone, Copy)]
23enum CollectionKind {
24    Array,
25    List,
26    Sequence,
27}
28
29#[derive(Debug, Clone, Copy)]
30enum ControlKind {
31    If,
32    For,
33    While,
34    Return,
35}
36
37trait InferenceContext {
38    type Type: Clone;
39    type Value: Clone;
40    type Query: Clone;
41
42    fn integer(&self, literal: &IntegerLiteral) -> FactOf<Self>;
43    fn string(&self, literal: &StringLiteral) -> FactOf<Self>;
44    fn symbol(&self, symbol: &Symbol) -> FactOf<Self>;
45    fn binary(
46        &self,
47        node: &BinaryExpression,
48        left: FactOf<Self>,
49        right: FactOf<Self>,
50    ) -> FactOf<Self>;
51    fn collection(&self, kind: CollectionKind, elements: Vec<FactOf<Self>>) -> FactOf<Self>;
52    fn control(&self, kind: ControlKind, neighbours: Vec<FactOf<Self>>) -> FactOf<Self>;
53}
54
55type FactOf<C> = Fact<
56    <C as InferenceContext>::Type,
57    <C as InferenceContext>::Value,
58    <C as InferenceContext>::Query,
59>;
60
61struct TypeChecker<'context, C: InferenceContext> {
62    context: &'context C,
63    results: Vec<FactOf<C>>,
64}
65
66impl<C: InferenceContext> TypeChecker<'_, C> {
67    fn infer(&mut self, expression: &Expr) -> FactOf<C> {
68        let old_len = self.results.len();
69        self.visit_expr(expression);
70        assert_eq!(self.results.len(), old_len + 1);
71        self.results.pop().expect("expression produced one fact")
72    }
73
74    fn components(&mut self, components: &Punctuated<Component>) -> Vec<FactOf<C>> {
75        components
76            .iter()
77            .filter_map(|component| match component {
78                Component::Expr(expression) => Some(self.infer(expression)),
79                Component::EmptyComponent(_) => None,
80            })
81            .collect()
82    }
83
84    fn loop_body(&mut self, body: &LoopBody) -> Vec<FactOf<C>> {
85        body.listed_value
86            .iter()
87            .chain(&body.ignored_value)
88            .map(|expression| self.infer(expression))
89            .collect()
90    }
91}
92
93impl<'ast, C: InferenceContext> Visit<'ast> for TypeChecker<'_, C> {
94    fn visit_integer_literal(&mut self, literal: &'ast IntegerLiteral) {
95        self.results.push(self.context.integer(literal));
96    }
97
98    fn visit_string_literal(&mut self, literal: &'ast StringLiteral) {
99        self.results.push(self.context.string(literal));
100    }
101
102    fn visit_symbol(&mut self, symbol: &'ast Symbol) {
103        self.results.push(self.context.symbol(symbol));
104    }
105
106    fn visit_binary_expression(&mut self, expression: &'ast BinaryExpression) {
107        let left = self.infer(&expression.left);
108        let right = self.infer(&expression.right);
109        self.results
110            .push(self.context.binary(expression, left, right));
111    }
112
113    fn visit_parenthesized_expression(&mut self, expression: &'ast ParenthesizedExpression) {
114        let fact = expression.value.as_deref().map(|value| self.infer(value));
115        self.results.push(fact.unwrap_or_else(|| {
116            self.context
117                .collection(CollectionKind::Sequence, Vec::new())
118        }));
119    }
120
121    fn visit_array(&mut self, array: &'ast Array) {
122        let elements = self.components(&array.elements);
123        self.results
124            .push(self.context.collection(CollectionKind::Array, elements));
125    }
126
127    fn visit_list(&mut self, list: &'ast List) {
128        let elements = self.components(&list.elements);
129        self.results
130            .push(self.context.collection(CollectionKind::List, elements));
131    }
132
133    fn visit_sequence(&mut self, sequence: &'ast Sequence) {
134        let elements = self.components(&sequence.elements);
135        self.results
136            .push(self.context.collection(CollectionKind::Sequence, elements));
137    }
138
139    fn visit_if_statement(&mut self, statement: &'ast IfStatement) {
140        let mut neighbours = vec![
141            self.infer(&statement.condition),
142            self.infer(&statement.then_clause.value),
143        ];
144        if let Some(clause) = &statement.else_clause {
145            neighbours.push(self.infer(&clause.value));
146        }
147        self.results
148            .push(self.context.control(ControlKind::If, neighbours));
149    }
150
151    fn visit_for_loop(&mut self, statement: &'ast ForLoop) {
152        let mut neighbours = Vec::new();
153        if let Some(range) = &statement.range {
154            neighbours.extend(
155                range
156                    .iterated_collection
157                    .iter()
158                    .chain(&range.range_start)
159                    .chain(&range.range_end)
160                    .map(|expression| self.infer(expression)),
161            );
162        }
163        neighbours.extend(statement.filter.iter().map(|value| self.infer(value)));
164        neighbours.extend(self.loop_body(&statement.body));
165        self.results
166            .push(self.context.control(ControlKind::For, neighbours));
167    }
168
169    fn visit_while_loop(&mut self, statement: &'ast WhileLoop) {
170        let mut neighbours = vec![self.infer(&statement.condition)];
171        neighbours.extend(self.loop_body(&statement.body));
172        self.results
173            .push(self.context.control(ControlKind::While, neighbours));
174    }
175
176    fn visit_return_statement(&mut self, statement: &'ast ReturnStatement) {
177        let neighbours = statement
178            .value
179            .iter()
180            .map(|value| self.infer(value))
181            .collect();
182        self.results
183            .push(self.context.control(ControlKind::Return, neighbours));
184    }
185}
186
187#[derive(Debug, Clone, PartialEq, Eq)]
188enum DemoType {
189    Integer,
190    Number,
191    String,
192    Array,
193    List,
194    Sequence,
195    Thing,
196}
197
198#[derive(Debug, Clone, PartialEq, Eq)]
199enum DemoValue {
200    Integer(i64),
201}
202
203type DemoFact = Fact<DemoType, DemoValue, String>;
204
205struct DemoRegistry {
206    symbols: HashMap<String, DemoFact>,
207}
208
209impl DemoRegistry {
210    fn known(types: Vec<DemoType>, value: Option<DemoValue>) -> DemoFact {
211        Fact {
212            types: TypeRange::Known(types),
213            value,
214        }
215    }
216
217    fn deferred(query: impl Into<String>) -> DemoFact {
218        Fact {
219            types: TypeRange::Deferred(query.into()),
220            value: None,
221        }
222    }
223}
224
225impl InferenceContext for DemoRegistry {
226    type Type = DemoType;
227    type Value = DemoValue;
228    type Query = String;
229
230    fn integer(&self, literal: &IntegerLiteral) -> DemoFact {
231        Self::known(
232            vec![DemoType::Integer, DemoType::Number, DemoType::Thing],
233            literal.text.parse().ok().map(DemoValue::Integer),
234        )
235    }
236
237    fn string(&self, _literal: &StringLiteral) -> DemoFact {
238        Self::known(vec![DemoType::String, DemoType::Thing], None)
239    }
240
241    fn symbol(&self, symbol: &Symbol) -> DemoFact {
242        self.symbols
243            .get(&symbol.text)
244            .cloned()
245            .unwrap_or_else(|| Self::deferred(format!("symbol `{}`", symbol.text)))
246    }
247
248    fn binary(&self, node: &BinaryExpression, left: DemoFact, right: DemoFact) -> DemoFact {
249        match (&node.operator, left.value, right.value) {
250            (
251                BinaryOperator::Add(_),
252                Some(DemoValue::Integer(left)),
253                Some(DemoValue::Integer(right)),
254            ) => Self::known(
255                vec![DemoType::Integer, DemoType::Number, DemoType::Thing],
256                Some(DemoValue::Integer(left + right)),
257            ),
258            _ => Self::deferred("installed binary signature"),
259        }
260    }
261
262    fn collection(&self, kind: CollectionKind, _elements: Vec<DemoFact>) -> DemoFact {
263        let ty = match kind {
264            CollectionKind::Array => DemoType::Array,
265            CollectionKind::List => DemoType::List,
266            CollectionKind::Sequence => DemoType::Sequence,
267        };
268        Self::known(vec![ty, DemoType::Thing], None)
269    }
270
271    fn control(&self, kind: ControlKind, _neighbours: Vec<DemoFact>) -> DemoFact {
272        Self::deferred(format!("{kind:?} typing rule"))
273    }
274}
275
276fn infer(source: &str, context: &DemoRegistry) -> Result<DemoFact, ParseError> {
277    let file = parse_file(source, SourceId(1))?;
278    let expression = match &file.elements[0] {
279        m2_syn::AnyCell::ExpressionCell(cell) => cell.value.as_ref(),
280        _ => panic!("expected an expression cell"),
281    };
282    Ok(TypeChecker {
283        context,
284        results: Vec::new(),
285    }
286    .infer(expression))
287}
288
289fn main() -> Result<(), ParseError> {
290    let context = DemoRegistry {
291        symbols: HashMap::from([(
292            "answer".into(),
293            DemoRegistry::known(
294                vec![DemoType::Integer, DemoType::Number, DemoType::Thing],
295                Some(DemoValue::Integer(40)),
296            ),
297        )]),
298    };
299
300    let known = infer("(answer + 2)", &context)?;
301    let deferred = infer("missing + 2", &context)?;
302    let collection = infer("{answer, missing}", &context)?;
303
304    println!("known: {known:?}");
305    println!("deferred: {deferred:?}");
306    println!("collection: {collection:?}");
307    Ok(())
308}
type_inference/type_inference.rs

type_inference/
type_inference.rs
