tinymist_analysis/ty/
subst.rs

1use super::{Sig, SigShape, TyMutator};
2use crate::ty::prelude::*;
3
4impl Sig<'_> {
5    /// Calls the signature with the given arguments.
6    pub fn call(&self, args: &Interned<ArgsTy>, pol: bool, ctx: &mut impl TyCtxMut) -> Option<Ty> {
7        crate::log_debug_ct!("call {self:?} {args:?} {pol:?}");
8        ctx.with_scope(|ctx| {
9            let body = self.check_bind(args, ctx)?;
10
11            // Substitute the bound variables in the body or just body
12            let mut checker = SubstituteChecker::new(ctx);
13            Some(checker.ty(&body, pol).unwrap_or(body))
14        })
15    }
16
17    /// Checks the binding of the signature.
18    pub fn check_bind(&self, args: &Interned<ArgsTy>, ctx: &mut impl TyCtxMut) -> Option<Ty> {
19        let SigShape { sig, withs } = self.shape(ctx)?;
20
21        // todo: check if the signature has free variables
22        // let has_free_vars = sig.has_free_variables;
23
24        let rest_bind = Self::rest_bind(&sig, args, withs);
25
26        for (arg_recv, arg_ins) in sig.matches(args, withs) {
27            if let Ty::Var(arg_recv) = arg_recv {
28                crate::log_debug_ct!("bind {arg_recv:?} {arg_ins:?}");
29                ctx.bind_local(arg_recv, arg_ins.clone());
30            }
31        }
32
33        if let Some((rest_var, rest_ty)) = rest_bind {
34            crate::log_debug_ct!("bind rest {rest_var:?} {rest_ty:?}");
35            ctx.bind_local(&rest_var, rest_ty);
36        }
37
38        sig.body.clone()
39    }
40
41    fn rest_bind(
42        sig: &Interned<SigTy>,
43        args: &Interned<ArgsTy>,
44        withs: Option<&Vec<Interned<SigTy>>>,
45    ) -> Option<(Interned<TypeVar>, Ty)> {
46        let Ty::Var(rest_var) = sig.rest_param()? else {
47            return None;
48        };
49
50        let fixed_pos = sig.positional_params().len();
51        let rest_pos = withs
52            .into_iter()
53            .flat_map(|withs| withs.iter().rev())
54            .flat_map(|with| with.positional_params())
55            .chain(args.positional_params())
56            .skip(fixed_pos)
57            .cloned()
58            .collect::<Vec<_>>();
59
60        let rest_named = args
61            .named_params()
62            .filter(|(name, _)| sig.named(name).is_none())
63            .map(|(name, ty)| (name.clone(), ty.clone()))
64            .collect::<Vec<_>>();
65
66        let rest = args.rest_param().cloned();
67        let rest_args = ArgsTy::new(rest_pos.into_iter(), rest_named, None, rest, None);
68
69        Some((rest_var.clone(), Ty::Args(rest_args.into())))
70    }
71}
72
73/// A checker to substitute the bound variables.
74struct SubstituteChecker<'a, T: TyCtxMut> {
75    ctx: &'a mut T,
76    memo: FxHashMap<(Ty, bool), Option<Ty>>,
77}
78
79impl<T: TyCtxMut> SubstituteChecker<'_, T> {
80    fn new(ctx: &mut T) -> SubstituteChecker<'_, T> {
81        SubstituteChecker {
82            ctx,
83            memo: FxHashMap::default(),
84        }
85    }
86
87    /// Substitutes the bound variables in the given type.
88    fn ty(&mut self, body: &Ty, pol: bool) -> Option<Ty> {
89        body.mutate(pol, self)
90    }
91}
92
93impl<T: TyCtxMut> TyMutator for SubstituteChecker<'_, T> {
94    fn mutate(&mut self, ty: &Ty, pol: bool) -> Option<Ty> {
95        let key = (ty.clone(), pol);
96        if let Some(result) = self.memo.get(&key) {
97            return result.clone();
98        }
99
100        // todo: extrude the type into a polarized type
101        let result = match ty {
102            Ty::Var(var) => self.ctx.local_bind_of(var),
103            Ty::Let(bounds) => {
104                let mut lbs = bounds
105                    .lbs
106                    .iter()
107                    .map(|bound| self.mutate(bound, !pol).unwrap_or_else(|| bound.clone()))
108                    .collect::<Vec<_>>();
109                let mut ubs = bounds
110                    .ubs
111                    .iter()
112                    .map(|bound| self.mutate(bound, pol).unwrap_or_else(|| bound.clone()))
113                    .collect::<Vec<_>>();
114                if ubs.is_empty() && lbs.len() == 1 {
115                    lbs.pop()
116                } else if lbs.is_empty() && ubs.len() == 1 {
117                    ubs.pop()
118                } else {
119                    Some(Ty::Let(TypeBounds { lbs, ubs }.into()))
120                }
121            }
122            _ => self.mutate_rec(ty, pol),
123        };
124        self.memo.insert(key, result.clone());
125        result
126    }
127}
128
129#[cfg(test)]
130mod tests {
131    use insta::{assert_debug_snapshot, assert_snapshot};
132    use tinymist_derive::BindTyCtx;
133
134    use super::{DynTypeBounds, Interned, Ty, TyCtx, TypeInfo, TypeVar};
135    use crate::ty::ApplyChecker;
136    use crate::ty::tests::*;
137    #[test]
138    fn test_ty() {
139        use super::*;
140        let ty = Ty::Builtin(BuiltinTy::Clause);
141        let ty_ref = TyRef::new(ty.clone());
142        assert_debug_snapshot!(ty_ref, @"Clause");
143    }
144
145    #[derive(Default, BindTyCtx)]
146    #[bind(0)]
147    struct CallCollector(TypeInfo, Vec<Ty>);
148
149    impl ApplyChecker for CallCollector {
150        fn apply(
151            &mut self,
152            sig: super::Sig,
153            arguments: &crate::adt::interner::Interned<super::ArgsTy>,
154            pol: bool,
155        ) {
156            let ty = sig.call(arguments, pol, &mut self.0);
157            if let Some(ty) = ty {
158                self.1.push(ty);
159            }
160        }
161    }
162
163    #[test]
164    fn test_sig_call() {
165        use super::*;
166
167        fn call(sig: Interned<SigTy>, args: Interned<SigTy>) -> String {
168            let sig_ty = Ty::Func(sig);
169            let mut collector = CallCollector::default();
170            sig_ty.call(&args, false, &mut collector);
171
172            collector.1.iter().fold(String::new(), |mut acc, ty| {
173                if !acc.is_empty() {
174                    acc.push_str(", ");
175                }
176
177                acc.push_str(&format!("{ty:?}"));
178                acc
179            })
180        }
181
182        assert_snapshot!(call(literal_sig!(p1 -> p1), literal_args!(q1)), @"@q1");
183        assert_snapshot!(call(literal_sig!(!u1: w1 -> w1), literal_args!(!u1: w2)), @"@w2");
184    }
185
186    #[test]
187    fn test_substitute_checker_memoizes_shared_type_dag() {
188        use super::*;
189        use crate::syntax::Decl;
190
191        let var = TypeVar::new("input".into(), Decl::lit("input").into());
192        let var_ty = Ty::Var(var.clone());
193        let changed = Ty::If(IfTy::new(
194            var_ty.clone().into(),
195            var_ty.clone().into(),
196            var_ty.clone().into(),
197        ));
198        let unchanged = Ty::If(IfTy::new(
199            Ty::Boolean(Some(true)).into(),
200            Ty::Boolean(Some(true)).into(),
201            Ty::Boolean(Some(true)).into(),
202        ));
203        let body = Ty::Tuple(
204            vec![
205                changed.clone(),
206                changed.clone(),
207                unchanged.clone(),
208                unchanged.clone(),
209            ]
210            .into(),
211        );
212
213        let mut ctx = TypeInfo::default();
214        ctx.bind_local(&var, Ty::Boolean(Some(false)));
215        let mut checker = SubstituteChecker::new(&mut ctx);
216        let result = checker.ty(&body, false);
217
218        assert!(result.is_some());
219        assert!(matches!(
220            checker.memo.get(&(changed, false)),
221            Some(Some(Ty::If(_)))
222        ));
223        assert_eq!(checker.memo.get(&(unchanged, false)), Some(&None));
224    }
225}