1use std::collections::HashMap;
2
3use m2_syn::nodes::{
4 Array, BinaryExpr, BinaryOperator, Component, Expr, ExprPack, ForLoop, Holder, IfStatement,
5 IntegerLiteral, IterationDomain, List, LoopBody, ReturnExpr, Sequence, StringLiteral, Symbol,
6 WhileLoop,
7};
8use m2_syn::visit::Visit;
9use m2_syn::{ParseError, parse_file};
10
11#[derive(Debug, Clone, PartialEq, Eq)]
12enum TypeRange<T, Q> {
13 Known(Vec<T>),
14 Deferred(Q),
15}
16
17#[derive(Debug, Clone, PartialEq, Eq)]
18struct Fact<T, V, Q> {
19 types: TypeRange<T, Q>,
20 value: Option<V>,
21}
22
23#[derive(Debug, Clone, Copy)]
24enum CollectionKind {
25 Array,
26 List,
27 Sequence,
28}
29
30#[derive(Debug, Clone, Copy)]
31enum ControlKind {
32 If,
33 For,
34 While,
35 Return,
36}
37
38trait InferenceContext {
39 type Type: Clone;
40 type Value: Clone;
41 type Query: Clone;
42
43 fn integer(&self, literal: &IntegerLiteral) -> FactOf<Self>;
44 fn string(&self, literal: &StringLiteral) -> FactOf<Self>;
45 fn symbol(&self, symbol: &Symbol) -> FactOf<Self>;
46 fn binary(&self, node: &BinaryExpr, left: FactOf<Self>, right: FactOf<Self>) -> FactOf<Self>;
47 fn collection(&self, kind: CollectionKind, elements: Vec<FactOf<Self>>) -> FactOf<Self>;
48 fn control(&self, kind: ControlKind, neighbours: Vec<FactOf<Self>>) -> FactOf<Self>;
49}
50
51type FactOf<C> = Fact<
52 <C as InferenceContext>::Type,
53 <C as InferenceContext>::Value,
54 <C as InferenceContext>::Query,
55>;
56
57struct TypeChecker<'context, C: InferenceContext> {
58 context: &'context C,
59 results: Vec<FactOf<C>>,
60}
61
62impl<C: InferenceContext> TypeChecker<'_, C> {
63 fn infer(&mut self, expression: &Expr) -> FactOf<C> {
64 let old_len = self.results.len();
65 self.visit_expr(expression);
66 assert_eq!(self.results.len(), old_len + 1);
67 self.results.pop().expect("expression produced one fact")
68 }
69
70 fn components(&mut self, components: &Option<ExprPack>) -> Vec<FactOf<C>> {
71 components
72 .iter()
73 .flat_map(|pack| pack.pack.iter())
74 .filter_map(|component| match component {
75 Component::Expr(expression) => Some(self.infer(expression)),
76 Component::Empty(_) => None,
77 })
78 .collect()
79 }
80
81 fn loop_body(&mut self, body: &LoopBody) -> Vec<FactOf<C>> {
82 let mut neighbours = Vec::new();
83 if let Some(step) = &body.list_step {
84 neighbours.push(self.infer(&step.expr));
85 }
86 if let Some(step) = &body.do_step {
87 neighbours.push(self.infer(&step.statement));
88 }
89 neighbours
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_expr(&mut self, expression: &'ast BinaryExpr) {
107 let left = self.infer(&expression.left_operand);
108 let right = self.infer(&expression.right_operand);
109 self.results
110 .push(self.context.binary(expression, left, right));
111 }
112
113 fn visit_holder(&mut self, expression: &'ast Holder) {
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.expr),
143 ];
144 if let Some(clause) = &statement.else_clause {
145 neighbours.push(self.infer(&clause.expr));
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(domain) = &statement.iteration_domain {
154 match domain {
155 IterationDomain::IterRange(range) => {
156 if let Some(bound) = &range.lower_bound {
157 neighbours.push(self.infer(&bound.value));
158 }
159 if let Some(bound) = &range.upper_bound {
160 neighbours.push(self.infer(&bound.value));
161 }
162 }
163 IterationDomain::IterIn(domain) => {
164 neighbours.push(self.infer(&domain.collection));
165 }
166 }
167 }
168 neighbours.extend(
169 statement
170 .when_condition
171 .iter()
172 .map(|condition| self.infer(&condition.condition)),
173 );
174 neighbours.extend(self.loop_body(&statement.body));
175 self.results
176 .push(self.context.control(ControlKind::For, neighbours));
177 }
178
179 fn visit_while_loop(&mut self, statement: &'ast WhileLoop) {
180 let mut neighbours = vec![self.infer(&statement.condition)];
181 neighbours.extend(self.loop_body(&statement.body));
182 self.results
183 .push(self.context.control(ControlKind::While, neighbours));
184 }
185
186 fn visit_return_expr(&mut self, statement: &'ast ReturnExpr) {
187 let neighbours = statement
188 .value
189 .iter()
190 .map(|value| self.infer(value))
191 .collect();
192 self.results
193 .push(self.context.control(ControlKind::Return, neighbours));
194 }
195}
196
197#[derive(Debug, Clone, PartialEq, Eq)]
198enum DemoType {
199 Integer,
200 Number,
201 String,
202 Array,
203 List,
204 Sequence,
205 Thing,
206}
207
208#[derive(Debug, Clone, PartialEq, Eq)]
209enum DemoValue {
210 Integer(i64),
211}
212
213type DemoFact = Fact<DemoType, DemoValue, String>;
214
215struct DemoRegistry {
216 symbols: HashMap<String, DemoFact>,
217}
218
219impl DemoRegistry {
220 fn known(types: Vec<DemoType>, value: Option<DemoValue>) -> DemoFact {
221 Fact {
222 types: TypeRange::Known(types),
223 value,
224 }
225 }
226
227 fn deferred(query: impl Into<String>) -> DemoFact {
228 Fact {
229 types: TypeRange::Deferred(query.into()),
230 value: None,
231 }
232 }
233}
234
235impl InferenceContext for DemoRegistry {
236 type Type = DemoType;
237 type Value = DemoValue;
238 type Query = String;
239
240 fn integer(&self, literal: &IntegerLiteral) -> DemoFact {
241 Self::known(
242 vec![DemoType::Integer, DemoType::Number, DemoType::Thing],
243 literal.text().parse().ok().map(DemoValue::Integer),
244 )
245 }
246
247 fn string(&self, _literal: &StringLiteral) -> DemoFact {
248 Self::known(vec![DemoType::String, DemoType::Thing], None)
249 }
250
251 fn symbol(&self, symbol: &Symbol) -> DemoFact {
252 self.symbols
253 .get(symbol.text())
254 .cloned()
255 .unwrap_or_else(|| Self::deferred(format!("symbol `{}`", symbol.text())))
256 }
257
258 fn binary(&self, node: &BinaryExpr, left: DemoFact, right: DemoFact) -> DemoFact {
259 match (&node.operator, left.value, right.value) {
260 (
261 BinaryOperator::Add(_),
262 Some(DemoValue::Integer(left)),
263 Some(DemoValue::Integer(right)),
264 ) => Self::known(
265 vec![DemoType::Integer, DemoType::Number, DemoType::Thing],
266 Some(DemoValue::Integer(left + right)),
267 ),
268 _ => Self::deferred("installed binary signature"),
269 }
270 }
271
272 fn collection(&self, kind: CollectionKind, _elements: Vec<DemoFact>) -> DemoFact {
273 let ty = match kind {
274 CollectionKind::Array => DemoType::Array,
275 CollectionKind::List => DemoType::List,
276 CollectionKind::Sequence => DemoType::Sequence,
277 };
278 Self::known(vec![ty, DemoType::Thing], None)
279 }
280
281 fn control(&self, kind: ControlKind, _neighbours: Vec<DemoFact>) -> DemoFact {
282 Self::deferred(format!("{kind:?} typing rule"))
283 }
284}
285
286fn infer(source: &str, context: &DemoRegistry) -> Result<DemoFact, ParseError> {
287 let file = parse_file(source)?;
288 let expression = match file.cells[0].content.pack.first() {
289 Component::Expr(expression) => expression.as_ref(),
290 Component::Empty(_) => panic!("expected an expression component"),
291 };
292 Ok(TypeChecker {
293 context,
294 results: Vec::new(),
295 }
296 .infer(expression))
297}
298
299fn main() -> Result<(), ParseError> {
300 let context = DemoRegistry {
301 symbols: HashMap::from([(
302 "answer".into(),
303 DemoRegistry::known(
304 vec![DemoType::Integer, DemoType::Number, DemoType::Thing],
305 Some(DemoValue::Integer(40)),
306 ),
307 )]),
308 };
309
310 let known = infer("(answer + 2)", &context)?;
311 let deferred = infer("missing + 2", &context)?;
312 let collection = infer("{answer, missing}", &context)?;
313
314 println!("known: {known:?}");
315 println!("deferred: {deferred:?}");
316 println!("collection: {collection:?}");
317 Ok(())
318}