tinymist_analysis/ty/
mutate.rs

1use crate::syntax::UnaryOp;
2use crate::ty::def::*;
3
4/// A trait to mutate a type.
5pub trait TyMutator {
6    /// Mutates the given type.
7    fn mutate(&mut self, ty: &Ty, pol: bool) -> Option<Ty> {
8        self.mutate_rec(ty, pol)
9    }
10
11    /// Mutates the given type recursively.
12    fn mutate_rec(&mut self, ty: &Ty, pol: bool) -> Option<Ty> {
13        use Ty::*;
14        match ty {
15            Value(..) | Any | Boolean(..) | Builtin(..) => None,
16            Union(v) => Some(Union(self.mutate_vec(v, pol)?)),
17            Var(..) | Let(..) => None,
18            Array(arr) => Some(Array(self.mutate(arr, pol)?.into())),
19            Dict(dict) => Some(Dict(self.mutate_record(dict, pol)?.into())),
20            Tuple(tup) => self.mutate_tuple(tup, pol),
21            Func(func) => Some(Func(self.mutate_func(func, pol)?.into())),
22            Args(args) => Some(Args(self.mutate_func(args, pol)?.into())),
23            Pattern(pat) => Some(Pattern(self.mutate_func(pat, pol)?.into())),
24            Param(param) => Some(Param(self.mutate_param(param, pol)?.into())),
25            Select(sel) => Some(Select(self.mutate_select(sel, pol)?.into())),
26            With(sig) => Some(With(self.mutate_with_sig(sig, pol)?.into())),
27            Unary(unary) => self.mutate_unary_ty(unary, pol),
28            Binary(binary) => Some(Binary(self.mutate_binary(binary, pol)?.into())),
29            If(if_expr) => Some(If(self.mutate_if(if_expr, pol)?.into())),
30        }
31    }
32
33    /// Mutates the given vector of types.
34    fn mutate_vec(&mut self, ty: &[Ty], pol: bool) -> Option<Interned<Vec<Ty>>> {
35        let mut mutated = false;
36
37        let mut types = Vec::with_capacity(ty.len());
38        for ty in ty.iter() {
39            match self.mutate(ty, pol) {
40                Some(ty) => {
41                    types.push(ty);
42                    mutated = true;
43                }
44                None => types.push(ty.clone()),
45            }
46        }
47
48        if mutated { Some(types.into()) } else { None }
49    }
50
51    /// Mutates and normalizes a tuple type.
52    fn mutate_tuple(&mut self, ty: &[Ty], pol: bool) -> Option<Ty> {
53        let mut mutated = false;
54        let mut types = Vec::with_capacity(ty.len());
55
56        for ty in ty.iter() {
57            let ty = match self.mutate(ty, pol) {
58                Some(ty) => {
59                    mutated = true;
60                    ty
61                }
62                None => ty.clone(),
63            };
64
65            if Self::push_spread_tuple_elements(&mut types, &ty) {
66                mutated = true;
67            } else {
68                types.push(ty);
69            }
70        }
71
72        mutated.then(|| Ty::Tuple(types.into()))
73    }
74
75    /// Pushes known tuple elements from an internal spread marker.
76    fn push_spread_tuple_elements(types: &mut Vec<Ty>, ty: &Ty) -> bool {
77        let Ty::Unary(unary) = ty else {
78            return false;
79        };
80        if unary.op != UnaryOp::Spread {
81            return false;
82        }
83
84        match &unary.lhs {
85            Ty::Tuple(elems) => {
86                types.extend(elems.iter().cloned());
87                true
88            }
89            Ty::Args(args) => {
90                types.extend(args.positional_params().cloned());
91                if let Some(rest) = args.rest_param()
92                    && !Self::push_spread_tuple_elements(
93                        types,
94                        &Ty::Unary(TypeUnary::new(UnaryOp::Spread, rest.clone())),
95                    )
96                {
97                    types.push(Ty::Unary(TypeUnary::new(UnaryOp::Spread, rest.clone())));
98                }
99                true
100            }
101            _ => false,
102        }
103    }
104
105    /// Mutates the given option of type.
106    fn mutate_option(&mut self, ty: Option<&Ty>, pol: bool) -> Option<Option<Ty>> {
107        match ty {
108            Some(ty) => self.mutate(ty, pol).map(Some),
109            None => None,
110        }
111    }
112
113    /// Mutates the given function signature.
114    fn mutate_func(&mut self, ty: &Interned<SigTy>, pol: bool) -> Option<SigTy> {
115        let types = self.mutate_vec(&ty.inputs, pol);
116        let ret = self.mutate_option(ty.body.as_ref(), pol);
117
118        if types.is_none() && ret.is_none() {
119            return None;
120        }
121
122        let sig = ty.as_ref().clone();
123        let types = types.unwrap_or_else(|| ty.inputs.clone());
124        let ret = ret.unwrap_or_else(|| ty.body.clone());
125        Some(SigTy {
126            inputs: types,
127            body: ret,
128            ..sig
129        })
130    }
131
132    /// Mutates the given parameter type.
133    fn mutate_param(&mut self, param: &Interned<ParamTy>, pol: bool) -> Option<ParamTy> {
134        let ty = self.mutate(&param.ty, pol)?;
135        let mut param = param.as_ref().clone();
136        param.ty = ty;
137        Some(param)
138    }
139
140    /// Mutates the given record type.
141    fn mutate_record(&mut self, record: &Interned<RecordTy>, pol: bool) -> Option<RecordTy> {
142        let types = self.mutate_vec(&record.types, pol)?;
143
144        let rec = record.as_ref().clone();
145        Some(RecordTy { types, ..rec })
146    }
147
148    /// Mutates the given function signature with type.
149    fn mutate_with_sig(&mut self, ty: &Interned<SigWithTy>, pol: bool) -> Option<SigWithTy> {
150        let sig = self.mutate(ty.sig.as_ref(), pol);
151        let with = self.mutate_func(&ty.with, pol);
152
153        if sig.is_none() && with.is_none() {
154            return None;
155        }
156
157        let sig = sig.map(Interned::new).unwrap_or_else(|| ty.sig.clone());
158        let with = with.map(Interned::new).unwrap_or_else(|| ty.with.clone());
159
160        Some(SigWithTy { sig, with })
161    }
162
163    /// Mutates the given unary type.
164    fn mutate_unary_ty(&mut self, ty: &Interned<TypeUnary>, pol: bool) -> Option<Ty> {
165        let lhs = self.mutate(&ty.lhs, pol)?;
166        if ty.op == UnaryOp::ElementOf
167            && let Some(elem) = Self::known_element_type(&lhs)
168        {
169            return Some(elem);
170        }
171
172        Some(Ty::Unary(TypeUnary { lhs, op: ty.op }.into()))
173    }
174
175    /// Gets the known element type of an iterable-like type.
176    fn known_element_type(ty: &Ty) -> Option<Ty> {
177        match ty {
178            Ty::Array(elem) => Some(elem.as_ref().clone()),
179            Ty::Tuple(elems) => Self::known_tuple_element_type(elems),
180            Ty::Args(args) => Self::known_args_element_type(args),
181            Ty::Let(bounds) => Self::known_element_types(bounds.lbs.iter()),
182            Ty::Union(types) => Self::known_element_types(types.iter()),
183            _ => None,
184        }
185    }
186
187    /// Gets known element types from multiple iterable-like types.
188    fn known_element_types<'a>(types: impl Iterator<Item = &'a Ty>) -> Option<Ty> {
189        let types = types
190            .filter_map(Self::known_element_type)
191            .collect::<Vec<_>>();
192        (!types.is_empty()).then(|| Ty::from_types(types.into_iter()))
193    }
194
195    /// Gets the known element type of a tuple.
196    fn known_tuple_element_type(elems: &[Ty]) -> Option<Ty> {
197        let mut types = vec![];
198        for elem in elems {
199            if let Ty::Unary(unary) = elem
200                && unary.op == UnaryOp::Spread
201            {
202                if let Some(elem) = Self::known_element_type(&unary.lhs) {
203                    types.push(elem);
204                }
205                continue;
206            }
207
208            types.push(elem.clone());
209        }
210
211        (!types.is_empty()).then(|| Ty::from_types(types.into_iter()))
212    }
213
214    /// Gets the known positional element type of arguments.
215    fn known_args_element_type(args: &ArgsTy) -> Option<Ty> {
216        let mut types = args.positional_params().cloned().collect::<Vec<_>>();
217        if let Some(rest) = args.rest_param()
218            && let Some(elem) = Self::known_element_type(rest)
219        {
220            types.push(elem);
221        }
222
223        (!types.is_empty()).then(|| Ty::from_types(types.into_iter()))
224    }
225
226    /// Mutates the given binary type.
227    fn mutate_binary(&mut self, ty: &Interned<TypeBinary>, pol: bool) -> Option<TypeBinary> {
228        let (lhs, rhs) = &ty.operands;
229
230        let x = self.mutate(lhs, pol);
231        let y = self.mutate(rhs, pol);
232
233        if x.is_none() && y.is_none() {
234            return None;
235        }
236
237        let lhs = x.unwrap_or_else(|| lhs.clone());
238        let rhs = y.unwrap_or_else(|| rhs.clone());
239
240        Some(TypeBinary {
241            operands: (lhs, rhs),
242            op: ty.op,
243        })
244    }
245
246    /// Mutates the given if type.
247    fn mutate_if(&mut self, ty: &Interned<IfTy>, pol: bool) -> Option<IfTy> {
248        let cond = self.mutate(ty.cond.as_ref(), pol);
249        let then = self.mutate(ty.then.as_ref(), pol);
250        let else_ = self.mutate(ty.else_.as_ref(), pol);
251
252        if cond.is_none() && then.is_none() && else_.is_none() {
253            return None;
254        }
255
256        let cond = cond.map(Interned::new).unwrap_or_else(|| ty.cond.clone());
257        let then = then.map(Interned::new).unwrap_or_else(|| ty.then.clone());
258        let else_ = else_.map(Interned::new).unwrap_or_else(|| ty.else_.clone());
259
260        Some(IfTy { cond, then, else_ })
261    }
262
263    /// Mutates the given select type.
264    fn mutate_select(&mut self, ty: &Interned<SelectTy>, pol: bool) -> Option<SelectTy> {
265        let target = self.mutate(ty.ty.as_ref(), pol)?.into();
266
267        Some(SelectTy {
268            ty: target,
269            select: ty.select.clone(),
270        })
271    }
272}
273
274impl<T> TyMutator for T
275where
276    T: FnMut(&Ty, bool) -> Option<Ty>,
277{
278    fn mutate(&mut self, ty: &Ty, pol: bool) -> Option<Ty> {
279        self(ty, pol)
280    }
281}
282
283impl Ty {
284    /// Mutates the given type.
285    pub fn mutate(&self, pol: bool, checker: &mut impl TyMutator) -> Option<Ty> {
286        checker.mutate(self, pol)
287    }
288}