rill_lang/types/
infer.rs

1//! Algorithm-W-style inference over scalar types, with bottom-up arity
2//! synthesis and combinator arity checking.
3
4use std::collections::HashMap;
5
6use super::ty::{Scalar, Scheme, Subst, Type, TypeVarId};
7use super::unify::unify_scalar;
8use crate::ast::{BinOp, Expr, Program};
9use crate::builtin::SignatureSource;
10use crate::error::{CompileError, Span};
11
12/// The typed result of inference: the program's definitions plus the resolved
13/// type of `process` and the final substitution.
14#[derive(Debug, Clone)]
15pub struct TypedProgram {
16    /// The original program (unchanged AST).
17    pub program: Program,
18    /// Resolved diagram type of `process` (arities + scalars, substitution applied).
19    pub process_ty: Type,
20}
21
22/// Inference context: fresh var supply, definition schemes, local bindings,
23/// and a signature source for built-in resolution.
24struct Ctx<'a> {
25    next: TypeVarId,
26    subst: Subst,
27    defs: HashMap<String, Scheme>,
28    locals: HashMap<String, Type>,
29    sigs: &'a dyn SignatureSource,
30}
31
32impl Ctx<'_> {
33    fn fresh(&mut self) -> Scalar {
34        let v = self.next;
35        self.next += 1;
36        Scalar::Var(v)
37    }
38
39    fn instantiate(&mut self, scheme: &Scheme) -> Type {
40        let mut remap: HashMap<TypeVarId, Scalar> = HashMap::new();
41        for v in &scheme.vars {
42            let f = self.fresh();
43            remap.insert(*v, f);
44        }
45        let rw = |s: &Scalar| match s {
46            Scalar::Var(v) => remap.get(v).cloned().unwrap_or_else(|| s.clone()),
47            _ => s.clone(),
48        };
49        Type {
50            ins: scheme.ty.ins.iter().map(&rw).collect(),
51            outs: scheme.ty.outs.iter().map(&rw).collect(),
52        }
53    }
54
55    fn free_vars(&self, t: &Type) -> Vec<TypeVarId> {
56        let mut acc = Vec::new();
57        for s in t.ins.iter().chain(t.outs.iter()) {
58            if let Scalar::Var(v) = self.subst.resolve_scalar(s) {
59                if !acc.contains(&v) {
60                    acc.push(v);
61                }
62            }
63        }
64        acc
65    }
66}
67
68/// Back-compat: infer with no built-ins.
69pub fn infer_program(program: &Program) -> Result<TypedProgram, CompileError> {
70    infer_program_with(program, &crate::builtin::NoSigs)
71}
72
73/// Infer with a signature source for built-in resolution.
74pub fn infer_program_with(
75    program: &Program,
76    sigs: &dyn SignatureSource,
77) -> Result<TypedProgram, CompileError> {
78    let mut ctx = Ctx {
79        next: 0,
80        subst: Subst::default(),
81        defs: HashMap::new(),
82        locals: HashMap::new(),
83        sigs,
84    };
85
86    let mut process_ty: Option<Type> = None;
87    for def in &program.defs {
88        ctx.locals.clear();
89        for p in &def.params {
90            let s = ctx.fresh();
91            ctx.locals.insert(p.clone(), Type::uniform(1, 1, s));
92        }
93        let ty = infer_expr(&mut ctx, &def.body)?;
94        let resolved = ctx.subst.apply(&ty);
95        let vars = ctx.free_vars(&resolved);
96        ctx.defs.insert(
97            def.name.clone(),
98            Scheme {
99                vars,
100                ty: resolved.clone(),
101            },
102        );
103        if def.name == "process" {
104            process_ty = Some(resolved);
105        }
106    }
107
108    let pty = process_ty.ok_or_else(|| CompileError::Type {
109        msg: "program has no `process` definition".into(),
110        span: Span::new(0, 0),
111    })?;
112    let vars: Vec<TypeVarId> = (0..ctx.next).collect();
113    for v in vars {
114        ctx.subst.map.entry(v).or_insert(Scalar::Float);
115    }
116    let pty = ctx.subst.apply(&pty);
117
118    if pty.arity_out() != 1 || pty.arity_in() > 1 {
119        return Err(CompileError::Type {
120            msg: format!(
121                "`process` must have arity (0|1)->1, found ({}->{})",
122                pty.arity_in(),
123                pty.arity_out()
124            ),
125            span: process_span(program),
126        });
127    }
128    Ok(TypedProgram {
129        program: program.clone(),
130        process_ty: pty,
131    })
132}
133
134fn process_span(program: &Program) -> Span {
135    program
136        .defs
137        .iter()
138        .find(|d| d.name == "process")
139        .map(|d| d.span)
140        .unwrap_or(Span::new(0, 0))
141}
142
143/// Infer the diagram type of an expression, synthesizing concrete arities.
144fn infer_expr(ctx: &mut Ctx<'_>, e: &Expr) -> Result<Type, CompileError> {
145    match e {
146        Expr::Int(_, _) => Ok(Type {
147            ins: vec![],
148            outs: vec![Scalar::Int],
149        }),
150        Expr::Float(_, _) => Ok(Type {
151            ins: vec![],
152            outs: vec![Scalar::Float],
153        }),
154        Expr::Wire(_) => {
155            let s = ctx.fresh();
156            Ok(Type::uniform(1, 1, s))
157        }
158        Expr::Cut(_) => {
159            let s = ctx.fresh();
160            Ok(Type {
161                ins: vec![s],
162                outs: vec![],
163            })
164        }
165        Expr::Ref(name, span) => infer_ref(ctx, name, *span),
166        Expr::Neg(inner, span) => {
167            let t = infer_expr(ctx, inner)?;
168            check_all_numeric(ctx, &t, *span)?;
169            Ok(t)
170        }
171        Expr::Apply { name, args, span } => infer_apply(ctx, name, args, *span),
172        Expr::Str(_, span) => Err(CompileError::Type {
173            msg: "string literal is only valid as a `param` name".into(),
174            span: *span,
175        }),
176        Expr::Bin { op, lhs, rhs, span } => {
177            let a = infer_expr(ctx, lhs)?;
178            let b = infer_expr(ctx, rhs)?;
179            infer_bin(ctx, *op, &a, &b, *span)
180        }
181    }
182}
183
184fn infer_ref(ctx: &mut Ctx<'_>, name: &str, span: Span) -> Result<Type, CompileError> {
185    if matches!(name, "+" | "-" | "*" | "/" | "%") {
186        let s = ctx.fresh();
187        return Ok(Type::uniform(2, 1, s));
188    }
189    if matches!(
190        name,
191        "sin" | "cos" | "tan" | "sqrt" | "exp" | "ln" | "tanh" | "abs"
192    ) {
193        return Ok(Type::uniform(1, 1, Scalar::Float));
194    }
195    if matches!(name, "min" | "max") {
196        let s = ctx.fresh();
197        return Ok(Type::uniform(2, 1, s));
198    }
199    if let Some(sig) = ctx.sigs.builtin_sig(name) {
200        if sig.num_params == 0 {
201            return Ok(Type::uniform(
202                sig.signal_ins,
203                sig.signal_outs,
204                Scalar::Float,
205            ));
206        }
207    }
208    if let Some(t) = ctx.locals.get(name) {
209        return Ok(t.clone());
210    }
211    if let Some(scheme) = ctx.defs.get(name).cloned() {
212        return Ok(ctx.instantiate(&scheme));
213    }
214    Err(CompileError::Type {
215        msg: format!("unknown identifier `{name}`"),
216        span,
217    })
218}
219
220fn infer_apply(
221    ctx: &mut Ctx<'_>,
222    name: &str,
223    args: &[Expr],
224    span: Span,
225) -> Result<Type, CompileError> {
226    if name == "param" {
227        if args.len() != 2 && args.len() != 4 {
228            return Err(CompileError::Type {
229                msg: "param expects (name, default[, min, max])".into(),
230                span,
231            });
232        }
233        if !matches!(args[0], Expr::Str(_, _)) {
234            return Err(CompileError::Type {
235                msg: "param name must be a string literal".into(),
236                span: args[0].span(),
237            });
238        }
239        for a in &args[1..] {
240            let at = infer_expr(ctx, a)?;
241            if at.arity_in() != 0 || at.arity_out() != 1 {
242                return Err(CompileError::Type {
243                    msg: "param default/min/max must be constants".into(),
244                    span: a.span(),
245                });
246            }
247        }
248        return Ok(Type {
249            ins: vec![],
250            outs: vec![Scalar::Float],
251        });
252    }
253    if name == "smooth" {
254        if args.len() != 2 {
255            return Err(CompileError::Type {
256                msg: "smooth expects exactly 2 arguments: smooth(signal, ms)".into(),
257                span,
258            });
259        }
260        let sig_ty = infer_expr(ctx, &args[0])?;
261        if sig_ty.arity_out() != 1 {
262            return Err(CompileError::Type {
263                msg: "smooth first argument must produce exactly one wire".into(),
264                span: args[0].span(),
265            });
266        }
267        let ms_ty = infer_expr(ctx, &args[1])?;
268        if ms_ty.arity_in() != 0 || ms_ty.arity_out() != 1 {
269            return Err(CompileError::Type {
270                msg: "smooth second argument (ms) must be a constant expression".into(),
271                span: args[1].span(),
272            });
273        }
274        unify_scalar(&ms_ty.outs[0], &Scalar::Float, &mut ctx.subst, span)?;
275        return Ok(Type {
276            ins: sig_ty.ins.clone(),
277            outs: vec![Scalar::Float],
278        });
279    }
280    if let Some(sig) = ctx.sigs.builtin_sig(name) {
281        let sig = sig.clone();
282        if args.len() != sig.num_params {
283            return Err(CompileError::Type {
284                msg: format!(
285                    "built-in `{name}` expects {} param(s), got {}",
286                    sig.num_params,
287                    args.len()
288                ),
289                span,
290            });
291        }
292        for a in args {
293            let at = infer_expr(ctx, a)?;
294            if at.arity_in() != 0 || at.arity_out() != 1 {
295                return Err(CompileError::Type {
296                    msg: format!("param to `{name}` must be a constant expression"),
297                    span: a.span(),
298                });
299            }
300        }
301        return Ok(Type::uniform(
302            sig.signal_ins,
303            sig.signal_outs,
304            Scalar::Float,
305        ));
306    }
307    let mut combined: Option<Type> = None;
308    for arg in args {
309        let at = infer_expr(ctx, arg)?;
310        combined = Some(match combined {
311            None => at,
312            Some(acc) => par(&acc, &at),
313        });
314    }
315    let callee = infer_ref(ctx, name, span)?;
316    match combined {
317        Some(args_ty) => seq(ctx, &args_ty, &callee, span),
318        None => Ok(callee),
319    }
320}
321
322fn infer_bin(
323    ctx: &mut Ctx<'_>,
324    op: BinOp,
325    a: &Type,
326    b: &Type,
327    span: Span,
328) -> Result<Type, CompileError> {
329    match op {
330        BinOp::Seq => seq(ctx, a, b, span),
331        BinOp::Par => Ok(par(a, b)),
332        BinOp::Split => split(ctx, a, b, span),
333        BinOp::Merge => merge(ctx, a, b, span),
334        BinOp::Feedback => feedback(ctx, a, b, span),
335        BinOp::Delay => delay(ctx, a, b, span),
336        BinOp::Add | BinOp::Sub | BinOp::Mul | BinOp::Div | BinOp::Rem => arith(ctx, a, b, span),
337    }
338}
339
340fn par(a: &Type, b: &Type) -> Type {
341    let mut ins = a.ins.clone();
342    ins.extend(b.ins.clone());
343    let mut outs = a.outs.clone();
344    outs.extend(b.outs.clone());
345    Type { ins, outs }
346}
347
348fn seq(ctx: &mut Ctx<'_>, a: &Type, b: &Type, span: Span) -> Result<Type, CompileError> {
349    if a.arity_out() != b.arity_in() {
350        return Err(CompileError::Type {
351            msg: format!(
352                "sequential `:` arity mismatch: lhs outputs {}, rhs inputs {}",
353                a.arity_out(),
354                b.arity_in()
355            ),
356            span,
357        });
358    }
359    for (x, y) in a.outs.iter().zip(b.ins.iter()) {
360        unify_scalar(x, y, &mut ctx.subst, span)?;
361    }
362    Ok(Type {
363        ins: a.ins.clone(),
364        outs: b.outs.clone(),
365    })
366}
367
368fn split(ctx: &mut Ctx<'_>, a: &Type, b: &Type, span: Span) -> Result<Type, CompileError> {
369    let (ao, bi) = (a.arity_out(), b.arity_in());
370    if ao == 0 || bi % ao != 0 {
371        return Err(CompileError::Type {
372            msg: format!(
373                "split `<:` requires rhs inputs ({bi}) be a multiple of lhs outputs ({ao})"
374            ),
375            span,
376        });
377    }
378    let reps = bi / ao;
379    for r in 0..reps {
380        for k in 0..ao {
381            unify_scalar(&a.outs[k], &b.ins[r * ao + k], &mut ctx.subst, span)?;
382        }
383    }
384    Ok(Type {
385        ins: a.ins.clone(),
386        outs: b.outs.clone(),
387    })
388}
389
390fn merge(ctx: &mut Ctx<'_>, a: &Type, b: &Type, span: Span) -> Result<Type, CompileError> {
391    let (ao, bi) = (a.arity_out(), b.arity_in());
392    if bi == 0 || ao % bi != 0 {
393        return Err(CompileError::Type {
394            msg: format!(
395                "merge `:>` requires lhs outputs ({ao}) be a multiple of rhs inputs ({bi})"
396            ),
397            span,
398        });
399    }
400    let groups = ao / bi;
401    for g in 0..groups {
402        for k in 0..bi {
403            unify_scalar(&a.outs[g * bi + k], &b.ins[k], &mut ctx.subst, span)?;
404        }
405    }
406    Ok(Type {
407        ins: a.ins.clone(),
408        outs: b.outs.clone(),
409    })
410}
411
412fn feedback(ctx: &mut Ctx<'_>, a: &Type, b: &Type, span: Span) -> Result<Type, CompileError> {
413    let (ai, ao, bi, bo) = (a.arity_in(), a.arity_out(), b.arity_in(), b.arity_out());
414    if bi > ao || bo > ai {
415        return Err(CompileError::Type {
416            msg: format!(
417                "feedback `~` arity mismatch: need B.in({bi})<=A.out({ao}) and B.out({bo})<=A.in({ai})"
418            ),
419            span,
420        });
421    }
422    for k in 0..bi {
423        unify_scalar(&b.ins[k], &a.outs[k], &mut ctx.subst, span)?;
424    }
425    for k in 0..bo {
426        unify_scalar(&b.outs[k], &a.ins[k], &mut ctx.subst, span)?;
427    }
428    Ok(Type {
429        ins: a.ins[bo..].to_vec(),
430        outs: a.outs.clone(),
431    })
432}
433
434fn delay(ctx: &mut Ctx<'_>, a: &Type, b: &Type, span: Span) -> Result<Type, CompileError> {
435    if a.arity_out() != 1 {
436        return Err(CompileError::Type {
437            msg: format!(
438                "`@` left side must have output arity 1, found {}",
439                a.arity_out()
440            ),
441            span,
442        });
443    }
444    if b.arity_in() != 0 || b.arity_out() != 1 {
445        return Err(CompileError::Type {
446            msg: "`@` delay length must be a constant expression".into(),
447            span,
448        });
449    }
450    unify_scalar(&b.outs[0], &Scalar::Int, &mut ctx.subst, span)?;
451    Ok(Type {
452        ins: a.ins.clone(),
453        outs: a.outs.clone(),
454    })
455}
456
457fn arith(ctx: &mut Ctx<'_>, a: &Type, b: &Type, span: Span) -> Result<Type, CompileError> {
458    if a.arity_out() != 1 || b.arity_out() != 1 {
459        return Err(CompileError::Type {
460            msg: "arithmetic operands must each produce exactly one wire".into(),
461            span,
462        });
463    }
464    unify_scalar(&a.outs[0], &b.outs[0], &mut ctx.subst, span)?;
465    let mut ins = a.ins.clone();
466    ins.extend(b.ins.clone());
467    Ok(Type {
468        ins,
469        outs: vec![a.outs[0].clone()],
470    })
471}
472
473fn check_all_numeric(_ctx: &mut Ctx<'_>, _t: &Type, _span: Span) -> Result<(), CompileError> {
474    Ok(())
475}
476
477#[cfg(test)]
478mod tests {
479    use super::*;
480    use crate::lexer::tokenize;
481    use crate::parser::parse;
482
483    fn ty_of(src: &str) -> Result<TypedProgram, CompileError> {
484        infer_program(&parse(&tokenize(src).unwrap()).unwrap())
485    }
486
487    #[test]
488    fn wire_is_1_to_1() {
489        let t = ty_of("process = _;").unwrap();
490        assert_eq!((t.process_ty.arity_in(), t.process_ty.arity_out()), (1, 1));
491    }
492
493    #[test]
494    fn gain_is_1_to_1() {
495        let t = ty_of("process = _ * 0.5;").unwrap();
496        assert_eq!((t.process_ty.arity_in(), t.process_ty.arity_out()), (1, 1));
497    }
498
499    #[test]
500    fn integrator_is_1_to_1() {
501        let t = ty_of("process = + ~ _;").unwrap();
502        assert_eq!((t.process_ty.arity_in(), t.process_ty.arity_out()), (1, 1));
503    }
504
505    #[test]
506    fn rejects_seq_arity_mismatch() {
507        assert!(ty_of("process = (_ , _) : _;").is_err());
508    }
509
510    #[test]
511    fn rejects_bad_process_arity() {
512        assert!(ty_of("process = _ , _;").is_err());
513    }
514
515    #[test]
516    fn split_then_merge_ok() {
517        let t = ty_of("process = _ <: (_ , _) :> + ;").unwrap();
518        assert_eq!((t.process_ty.arity_in(), t.process_ty.arity_out()), (1, 1));
519    }
520
521    #[test]
522    fn delay_constant_ok_variable_errors() {
523        assert!(ty_of("process = _ @ 1;").is_ok());
524        assert!(ty_of("process = _ @ _;").is_err());
525    }
526
527    #[test]
528    fn user_def_alias_resolves() {
529        let t = ty_of("gain = _ * 0.5; process = gain;").unwrap();
530        assert_eq!((t.process_ty.arity_in(), t.process_ty.arity_out()), (1, 1));
531    }
532
533    #[test]
534    fn function_application_ok() {
535        let t = ty_of("g(x) = x * 0.5; process = g(_);").unwrap();
536        assert_eq!((t.process_ty.arity_in(), t.process_ty.arity_out()), (1, 1));
537    }
538
539    struct TestSigs;
540    impl crate::builtin::SignatureSource for TestSigs {
541        fn builtin_sig(&self, name: &str) -> Option<&crate::builtin::BuiltinSig> {
542            use crate::builtin::{BuiltinKind, BuiltinSig};
543            match name {
544                "lowpass" => Some(Box::leak(Box::new(BuiltinSig {
545                    name: "lowpass",
546                    signal_ins: 1,
547                    signal_outs: 1,
548                    num_params: 2,
549                    kind: BuiltinKind::Block,
550                }))),
551                "onepole" => Some(Box::leak(Box::new(BuiltinSig {
552                    name: "onepole",
553                    signal_ins: 1,
554                    signal_outs: 1,
555                    num_params: 2,
556                    kind: BuiltinKind::Sample,
557                }))),
558                _ => None,
559            }
560        }
561    }
562
563    fn ty_with(src: &str) -> Result<TypedProgram, CompileError> {
564        infer_program_with(&parse(&tokenize(src).unwrap()).unwrap(), &TestSigs)
565    }
566
567    #[test]
568    fn builtin_call_is_1_to_1() {
569        let t = ty_with("process = _ : lowpass(1000.0, 0.7);").unwrap();
570        assert_eq!((t.process_ty.arity_in(), t.process_ty.arity_out()), (1, 1));
571    }
572
573    #[test]
574    fn builtin_wrong_param_count_errors() {
575        assert!(ty_with("process = _ : lowpass(1000.0);").is_err());
576    }
577
578    #[test]
579    fn builtin_non_const_param_errors() {
580        assert!(ty_with("process = _ : lowpass(_, 0.7);").is_err());
581    }
582
583    #[test]
584    fn sample_builtin_in_feedback_typechecks() {
585        assert!(ty_with("process = + ~ onepole(200.0, 0.5);").is_ok());
586    }
587}
rill_lang/types/infer.rs

rill_lang/types/
infer.rs
