Skip to main content

type_inference/
type_inference.rs

1//! Demonstrates type inference.
2
3use std::collections::HashMap;
4
5use m2_syn::nodes::{
6    Array, BinaryOperator, Component, Expr, ExprControlReturn, ExprForBare, ExprForIn,
7    ExprForRange, ExprIf, ExprOpBinRegular, ExprWhile, Holder, IntegerLiteral, List, LoopBody,
8    PackedContent, Sequence, StringLiteral, Symbol,
9};
10use m2_syn::visit::Visit;
11use m2_syn::{ParseError, parse_file};
12
13#[derive(Debug, Clone, PartialEq, Eq)]
14/// Alternatives used by the TypeRange example.
15enum TypeRange<T, Q> {
16    /// The Known alternative.
17    Known(Vec<T>),
18    /// The Deferred alternative.
19    Deferred(Q),
20}
21
22#[derive(Debug, Clone, PartialEq, Eq)]
23/// State used by the Fact example.
24struct Fact<T, V, Q> {
25    /// The types.
26    types: TypeRange<T, Q>,
27    /// The inferred value, when evaluation produced one.
28    value: Option<V>,
29}
30
31#[derive(Debug, Clone, Copy)]
32/// Alternatives used by the CollectionKind example.
33enum CollectionKind {
34    /// The Array alternative.
35    Array,
36    /// The List alternative.
37    List,
38    /// The Sequence alternative.
39    Sequence,
40}
41
42#[derive(Debug, Clone, Copy)]
43/// Alternatives used by the ControlKind example.
44enum ControlKind {
45    /// The If alternative.
46    If,
47    /// The For alternative.
48    For,
49    /// The While alternative.
50    While,
51    /// The Return alternative.
52    Return,
53}
54
55/// Behavior demonstrated by InferenceContext.
56trait InferenceContext {
57    /// The representation of an inferred M2 type.
58    type Type: Clone;
59    /// The representation of a known compile-time value.
60    type Value: Clone;
61    /// The representation of deferred type information.
62    type Query: Clone;
63
64    /// Infers an integer literal.
65    fn integer(&self, literal: &IntegerLiteral) -> FactOf<Self>;
66    /// Infers a string literal.
67    fn string(&self, literal: &StringLiteral) -> FactOf<Self>;
68    /// Looks up a symbol's inferred fact.
69    fn symbol(&self, symbol: &Symbol) -> FactOf<Self>;
70    /// Combines the facts inferred for a binary expression's operands.
71    fn binary(
72        &self,
73        node: &ExprOpBinRegular,
74        left: FactOf<Self>,
75        right: FactOf<Self>,
76    ) -> FactOf<Self>;
77    /// Infers a collection from its element facts.
78    fn collection(&self, kind: CollectionKind, elements: Vec<FactOf<Self>>) -> FactOf<Self>;
79    /// Infers a control form from the facts of its neighboring expressions.
80    fn control(&self, kind: ControlKind, neighbours: Vec<FactOf<Self>>) -> FactOf<Self>;
81}
82
83/// Convenient alias for FactOf.
84type FactOf<C> = Fact<
85    <C as InferenceContext>::Type,
86    <C as InferenceContext>::Value,
87    <C as InferenceContext>::Query,
88>;
89
90/// State used by the TypeChecker example.
91struct TypeChecker<'context, C: InferenceContext> {
92    /// The context.
93    context: &'context C,
94    /// The results.
95    results: Vec<FactOf<C>>,
96}
97
98impl<C: InferenceContext> TypeChecker<'_, C> {
99    /// Infers one expression and removes its completed fact from the visit stack.
100    fn infer(&mut self, expression: &Expr) -> FactOf<C> {
101        let old_len = self.results.len();
102        self.visit_expr(expression);
103        assert_eq!(self.results.len(), old_len + 1);
104        self.results.pop().expect("expression produced one fact")
105    }
106
107    /// Infers every nonempty component of optional packed content.
108    fn components(&mut self, content: &Option<PackedContent>) -> Vec<FactOf<C>> {
109        match content {
110            None => Vec::new(),
111            Some(PackedContent::Expr(expression)) => vec![self.infer(expression)],
112            Some(PackedContent::Pack(pack)) => pack
113                .pack
114                .iter()
115                .filter_map(|component| match component {
116                    Component::Expr(expression) => Some(self.infer(expression)),
117                    Component::Empty(_) => None,
118                })
119                .collect(),
120        }
121    }
122
123    /// Infers the expressions contributed by a loop body.
124    fn loop_body(&mut self, body: &LoopBody) -> Vec<FactOf<C>> {
125        let mut neighbours = Vec::new();
126        if let Some(step) = &body.list_step {
127            neighbours.push(self.infer(&step.expr));
128        }
129        if let Some(step) = &body.do_step {
130            neighbours.push(self.infer(&step.statement));
131        }
132        neighbours
133    }
134}
135
136impl<'ast, C: InferenceContext> Visit<'ast> for TypeChecker<'_, C> {
137    fn visit_integer_literal(&mut self, literal: &'ast IntegerLiteral) {
138        self.results.push(self.context.integer(literal));
139    }
140
141    fn visit_string_literal(&mut self, literal: &'ast StringLiteral) {
142        self.results.push(self.context.string(literal));
143    }
144
145    fn visit_symbol(&mut self, symbol: &'ast Symbol) {
146        self.results.push(self.context.symbol(symbol));
147    }
148
149    fn visit_expr_op_bin_regular(&mut self, expression: &'ast ExprOpBinRegular) {
150        let left = self.infer(&expression.left_operand);
151        let right = self.infer(&expression.right_operand);
152        self.results
153            .push(self.context.binary(expression, left, right));
154    }
155
156    fn visit_holder(&mut self, expression: &'ast Holder) {
157        let fact = expression.value.as_deref().map(|value| self.infer(value));
158        self.results.push(fact.unwrap_or_else(|| {
159            self.context
160                .collection(CollectionKind::Sequence, Vec::new())
161        }));
162    }
163
164    fn visit_array(&mut self, array: &'ast Array) {
165        let elements = self.components(&array.elements);
166        self.results
167            .push(self.context.collection(CollectionKind::Array, elements));
168    }
169
170    fn visit_list(&mut self, list: &'ast List) {
171        let elements = self.components(&list.elements);
172        self.results
173            .push(self.context.collection(CollectionKind::List, elements));
174    }
175
176    fn visit_sequence(&mut self, sequence: &'ast Sequence) {
177        let elements = self.components(&sequence.elements);
178        self.results
179            .push(self.context.collection(CollectionKind::Sequence, elements));
180    }
181
182    fn visit_expr_if(&mut self, statement: &'ast ExprIf) {
183        let mut neighbours = vec![
184            self.infer(&statement.condition),
185            self.infer(&statement.then_clause.expr),
186        ];
187        if let Some(clause) = &statement.else_clause {
188            neighbours.push(self.infer(&clause.expr));
189        }
190        self.results
191            .push(self.context.control(ControlKind::If, neighbours));
192    }
193
194    fn visit_expr_for_bare(&mut self, statement: &'ast ExprForBare) {
195        let mut neighbours = Vec::new();
196        neighbours.extend(
197            statement
198                .condition
199                .iter()
200                .map(|condition| self.infer(&condition.condition)),
201        );
202        neighbours.extend(self.loop_body(&statement.body));
203        self.results
204            .push(self.context.control(ControlKind::For, neighbours));
205    }
206
207    fn visit_expr_for_range(&mut self, statement: &'ast ExprForRange) {
208        let mut neighbours = Vec::new();
209        if let Some(bound) = &statement.lower_bound {
210            neighbours.push(self.infer(&bound.value));
211        }
212        if let Some(bound) = &statement.upper_bound {
213            neighbours.push(self.infer(&bound.value));
214        }
215        neighbours.extend(
216            statement
217                .condition
218                .iter()
219                .map(|condition| self.infer(&condition.condition)),
220        );
221        neighbours.extend(self.loop_body(&statement.body));
222        self.results
223            .push(self.context.control(ControlKind::For, neighbours));
224    }
225
226    fn visit_expr_for_in(&mut self, statement: &'ast ExprForIn) {
227        let mut neighbours = vec![self.infer(&statement.collection)];
228        neighbours.extend(
229            statement
230                .condition
231                .iter()
232                .map(|condition| self.infer(&condition.condition)),
233        );
234        neighbours.extend(self.loop_body(&statement.body));
235        self.results
236            .push(self.context.control(ControlKind::For, neighbours));
237    }
238
239    fn visit_expr_while(&mut self, statement: &'ast ExprWhile) {
240        let mut neighbours = vec![self.infer(&statement.condition)];
241        neighbours.extend(self.loop_body(&statement.body));
242        self.results
243            .push(self.context.control(ControlKind::While, neighbours));
244    }
245
246    fn visit_expr_control_return(&mut self, statement: &'ast ExprControlReturn) {
247        let neighbours = statement
248            .value
249            .iter()
250            .map(|value| self.infer(value))
251            .collect();
252        self.results
253            .push(self.context.control(ControlKind::Return, neighbours));
254    }
255}
256
257#[derive(Debug, Clone, PartialEq, Eq)]
258/// Alternatives used by the DemoType example.
259enum DemoType {
260    /// The Integer alternative.
261    Integer,
262    /// The Number alternative.
263    Number,
264    /// The String alternative.
265    String,
266    /// The Array alternative.
267    Array,
268    /// The List alternative.
269    List,
270    /// The Sequence alternative.
271    Sequence,
272    /// The Thing alternative.
273    Thing,
274}
275
276#[derive(Debug, Clone, PartialEq, Eq)]
277/// Alternatives used by the DemoValue example.
278enum DemoValue {
279    /// The Integer alternative.
280    Integer(i64),
281}
282
283/// Convenient alias for DemoFact.
284type DemoFact = Fact<DemoType, DemoValue, String>;
285
286/// State used by the DemoRegistry example.
287struct DemoRegistry {
288    /// The symbols.
289    symbols: HashMap<String, DemoFact>,
290}
291
292impl DemoRegistry {
293    /// Handles known.
294    fn known(types: Vec<DemoType>, value: Option<DemoValue>) -> DemoFact {
295        Fact {
296            types: TypeRange::Known(types),
297            value,
298        }
299    }
300
301    /// Handles deferred.
302    fn deferred(query: impl Into<String>) -> DemoFact {
303        Fact {
304            types: TypeRange::Deferred(query.into()),
305            value: None,
306        }
307    }
308}
309
310impl InferenceContext for DemoRegistry {
311    type Type = DemoType;
312    type Value = DemoValue;
313    type Query = String;
314
315    fn integer(&self, literal: &IntegerLiteral) -> DemoFact {
316        Self::known(
317            vec![DemoType::Integer, DemoType::Number, DemoType::Thing],
318            literal.text().parse().ok().map(DemoValue::Integer),
319        )
320    }
321
322    fn string(&self, _literal: &StringLiteral) -> DemoFact {
323        Self::known(vec![DemoType::String, DemoType::Thing], None)
324    }
325
326    fn symbol(&self, symbol: &Symbol) -> DemoFact {
327        self.symbols
328            .get(symbol.text())
329            .cloned()
330            .unwrap_or_else(|| Self::deferred(format!("symbol `{}`", symbol.text())))
331    }
332
333    fn binary(&self, node: &ExprOpBinRegular, left: DemoFact, right: DemoFact) -> DemoFact {
334        match (&node.operator, left.value, right.value) {
335            (
336                BinaryOperator::Add(_),
337                Some(DemoValue::Integer(left)),
338                Some(DemoValue::Integer(right)),
339            ) => Self::known(
340                vec![DemoType::Integer, DemoType::Number, DemoType::Thing],
341                Some(DemoValue::Integer(left + right)),
342            ),
343            _ => Self::deferred("installed binary signature"),
344        }
345    }
346
347    fn collection(&self, kind: CollectionKind, _elements: Vec<DemoFact>) -> DemoFact {
348        let ty = match kind {
349            CollectionKind::Array => DemoType::Array,
350            CollectionKind::List => DemoType::List,
351            CollectionKind::Sequence => DemoType::Sequence,
352        };
353        Self::known(vec![ty, DemoType::Thing], None)
354    }
355
356    fn control(&self, kind: ControlKind, _neighbours: Vec<DemoFact>) -> DemoFact {
357        Self::deferred(format!("{kind:?} typing rule"))
358    }
359}
360
361/// Handles infer.
362fn infer(source: &str, context: &DemoRegistry) -> Result<DemoFact, ParseError> {
363    let file = parse_file(source)?;
364    let expression = match &file.cells[0].content {
365        PackedContent::Expr(expression) => expression.as_ref(),
366        PackedContent::Pack(_) => panic!("expected one expression"),
367    };
368    Ok(TypeChecker {
369        context,
370        results: Vec::new(),
371    }
372    .infer(expression))
373}
374
375fn main() -> Result<(), ParseError> {
376    let context = DemoRegistry {
377        symbols: HashMap::from([(
378            "answer".into(),
379            DemoRegistry::known(
380                vec![DemoType::Integer, DemoType::Number, DemoType::Thing],
381                Some(DemoValue::Integer(40)),
382            ),
383        )]),
384    };
385
386    let known = infer("(answer + 2)", &context)?;
387    let deferred = infer("missing + 2", &context)?;
388    let collection = infer("{answer, missing}", &context)?;
389
390    println!("known: {known:?}");
391    println!("deferred: {deferred:?}");
392    println!("collection: {collection:?}");
393    Ok(())
394}