tinymist_analysis/
sig.rs

1//! Analysis of function signatures.
2
3use core::fmt;
4use std::collections::BTreeMap;
5use std::sync::Arc;
6
7use ecow::{EcoString, EcoVec, eco_format, eco_vec};
8use typst::foundations::{Closure, Func};
9use typst::syntax::ast::AstNode;
10use typst::syntax::{SyntaxKind, ast};
11use typst::utils::LazyHash;
12
13// use super::{BoundChecker, Definition};
14use crate::ty::{InsTy, ParamTy, SigTy, StrRef, Ty};
15use crate::ty::{Interned, ParamAttrs};
16use crate::upstream::truncated_repr;
17// use crate::upstream::truncated_repr;
18
19/// Describes a function signature.
20#[derive(Debug, Clone)]
21pub enum Signature {
22    /// A primary function signature.
23    Primary(Arc<PrimarySignature>),
24    /// A partially applied function signature.
25    Partial(Arc<PartialSignature>),
26}
27
28impl Signature {
29    /// Returns the primary signature if it is one.
30    pub fn primary(&self) -> &Arc<PrimarySignature> {
31        match self {
32            Signature::Primary(sig) => sig,
33            Signature::Partial(sig) => &sig.signature,
34        }
35    }
36
37    /// Returns the with bindings of the signature.
38    pub fn bindings(&self) -> &[ArgsInfo] {
39        match self {
40            Signature::Primary(_) => &[],
41            Signature::Partial(sig) => &sig.with_stack,
42        }
43    }
44
45    /// Returns the all parameters of the signature.
46    pub fn params(&self) -> impl Iterator<Item = (&Interned<ParamTy>, Option<&Ty>)> {
47        // todo: with stack
48        self.primary().params()
49    }
50
51    /// Returns the type of the signature.
52    pub fn type_sig(&self) -> Interned<SigTy> {
53        // todo: with stack
54        self.primary().sig_ty.clone()
55    }
56
57    /// Returns the shift applied to the signature.
58    pub fn param_shift(&self) -> usize {
59        match self {
60            Signature::Primary(_) => 0,
61            Signature::Partial(sig) => sig
62                .with_stack
63                .iter()
64                .map(|ws| ws.items.len())
65                .sum::<usize>(),
66        }
67    }
68}
69
70/// Describes a primary function signature.
71#[derive(Debug, Clone)]
72pub struct PrimarySignature {
73    /// The documentation of the function
74    pub docs: Option<EcoString>,
75    /// The documentation of the parameter.
76    pub param_specs: Vec<Interned<ParamTy>>,
77    /// Whether the function has fill, stroke, or size parameters.
78    pub has_fill_or_size_or_stroke: bool,
79    /// The associated signature type.
80    pub sig_ty: Interned<SigTy>,
81    /// Whether the signature is broken.
82    pub _broken: bool,
83}
84
85impl PrimarySignature {
86    /// Returns the number of positional parameters of the function.
87    pub fn pos_size(&self) -> usize {
88        self.sig_ty.name_started as usize
89    }
90
91    /// Returns the positional parameters of the function.
92    pub fn pos(&self) -> &[Interned<ParamTy>] {
93        &self.param_specs[..self.pos_size()]
94    }
95
96    /// Returns the positional parameters of the function.
97    pub fn get_pos(&self, offset: usize) -> Option<&Interned<ParamTy>> {
98        self.pos().get(offset)
99    }
100
101    /// Returns the named parameters of the function.
102    pub fn named(&self) -> &[Interned<ParamTy>] {
103        &self.param_specs[self.pos_size()..self.pos_size() + self.sig_ty.names.names.len()]
104    }
105
106    /// Returns the named parameters of the function.
107    pub fn get_named(&self, name: &StrRef) -> Option<&Interned<ParamTy>> {
108        self.named().get(self.sig_ty.names.find(name)?)
109    }
110
111    /// Returns the name of the rest parameter of the function.
112    pub fn has_spread_right(&self) -> bool {
113        self.sig_ty.spread_right
114    }
115
116    /// Returns the rest parameter of the function.
117    pub fn rest(&self) -> Option<&Interned<ParamTy>> {
118        self.has_spread_right()
119            .then(|| &self.param_specs[self.pos_size() + self.sig_ty.names.names.len()])
120    }
121
122    /// Returns the all parameters of the function.
123    pub fn params(&self) -> impl Iterator<Item = (&Interned<ParamTy>, Option<&Ty>)> {
124        let pos = self.pos();
125        let named = self.named();
126        let rest = self.rest();
127        let type_sig = &self.sig_ty;
128        let pos = pos
129            .iter()
130            .enumerate()
131            .map(|(idx, pos)| (pos, type_sig.pos(idx)));
132        let named = named.iter().map(|x| (x, type_sig.named(&x.name)));
133        let rest = rest.into_iter().map(|x| (x, type_sig.rest_param()));
134
135        pos.chain(named).chain(rest)
136    }
137}
138
139/// Describes a function argument instance
140#[derive(Debug, Clone)]
141pub struct ArgInfo {
142    /// The argument's name.
143    pub name: Option<StrRef>,
144    /// The argument's term.
145    pub term: Option<Ty>,
146}
147
148/// Describes a function argument list.
149#[derive(Debug, Clone)]
150pub struct ArgsInfo {
151    /// The arguments.
152    pub items: EcoVec<ArgInfo>,
153}
154
155/// Describes a function signature that is already partially applied.
156#[derive(Debug, Clone)]
157pub struct PartialSignature {
158    /// The positional parameters.
159    pub signature: Arc<PrimarySignature>,
160    /// The stack of `fn.with(..)` calls.
161    pub with_stack: EcoVec<ArgsInfo>,
162}
163
164/// Gets the signature of a function.
165#[comemo::memoize]
166pub fn func_signature(func: Func) -> Signature {
167    use typst::foundations::func::Repr;
168    let mut with_stack = eco_vec![];
169    let mut func = func;
170    while let Repr::With(with) = func.inner() {
171        let (inner, args) = with.as_ref();
172        with_stack.push(ArgsInfo {
173            items: args
174                .items
175                .iter()
176                .map(|arg| ArgInfo {
177                    name: arg.name.clone().map(From::from),
178                    term: Some(Ty::Value(InsTy::new(arg.value.v.clone()))),
179                })
180                .collect(),
181        });
182        func = inner.clone();
183    }
184
185    let mut pos_tys = vec![];
186    let mut named_tys = Vec::new();
187    let mut rest_ty = None;
188
189    let mut named_specs = BTreeMap::new();
190    let mut param_specs = Vec::new();
191    let mut rest_spec = None;
192
193    let mut broken = false;
194    let mut has_fill_or_size_or_stroke = false;
195
196    let mut add_param = |param: Interned<ParamTy>| {
197        let name = param.name.clone();
198        if param.attrs.named {
199            if matches!(name.as_ref(), "fill" | "stroke" | "size") {
200                has_fill_or_size_or_stroke = true;
201            }
202            named_tys.push((name.clone(), param.ty.clone()));
203            named_specs.insert(name.clone(), param.clone());
204        }
205
206        if param.attrs.variadic {
207            if rest_ty.is_some() {
208                broken = true;
209            } else {
210                rest_ty = Some(param.ty.clone());
211                rest_spec = Some(param);
212            }
213        } else if param.attrs.positional {
214            // todo: we have some params that are both positional and named
215            pos_tys.push(param.ty.clone());
216            param_specs.push(param);
217        }
218    };
219
220    let ret_ty = match func.inner() {
221        Repr::With(..) => unreachable!(),
222        Repr::Closure(closure) => {
223            analyze_closure_signature(closure.clone(), &mut add_param);
224            None
225        }
226        Repr::Element(..) | Repr::Native(..) | Repr::Plugin(..) => {
227            for param in func.params().unwrap_or_default() {
228                add_param(Interned::new(ParamTy {
229                    name: param.name.into(),
230                    docs: Some(param.docs.into()),
231                    default: param.default.map(|default| truncated_repr(&default())),
232                    ty: Ty::from_param_site(&func, param),
233                    attrs: param.into(),
234                }));
235            }
236
237            func.returns().map(|r| Ty::from_return_site(&func, r))
238        }
239    };
240
241    let sig_ty = SigTy::new(pos_tys.into_iter(), named_tys, None, rest_ty, ret_ty);
242
243    for name in &sig_ty.names.names {
244        let Some(param) = named_specs.get(name) else {
245            continue;
246        };
247        param_specs.push(param.clone());
248    }
249    if let Some(doc) = rest_spec {
250        param_specs.push(doc);
251    }
252
253    let signature = Arc::new(PrimarySignature {
254        docs: func.docs().map(From::from),
255        param_specs,
256        has_fill_or_size_or_stroke,
257        sig_ty: sig_ty.into(),
258        _broken: broken,
259    });
260
261    log::trace!("got signature {signature:?}");
262
263    if with_stack.is_empty() {
264        return Signature::Primary(signature);
265    }
266
267    Signature::Partial(Arc::new(PartialSignature {
268        signature,
269        with_stack,
270    }))
271}
272
273fn analyze_closure_signature(
274    closure: Arc<LazyHash<Closure>>,
275    add_param: &mut impl FnMut(Interned<ParamTy>),
276) {
277    log::trace!("closure signature for: {:?}", closure.node.kind());
278
279    let closure = &closure.node;
280    let closure_ast = match closure.kind() {
281        SyntaxKind::Closure => closure.cast::<ast::Closure>().unwrap(),
282        _ => return,
283    };
284
285    for param in closure_ast.params().children() {
286        match param {
287            ast::Param::Pos(pos) => {
288                let name = format!("{}", PatternDisplay(&pos));
289                add_param(Interned::new(ParamTy {
290                    name: name.as_str().into(),
291                    docs: None,
292                    default: None,
293                    ty: Ty::Any,
294                    attrs: ParamAttrs::positional(),
295                }));
296            }
297            // todo: pattern
298            ast::Param::Named(named) => {
299                let default = unwrap_parens(named.expr()).to_untyped().clone().into_text();
300                add_param(Interned::new(ParamTy {
301                    name: named.name().get().into(),
302                    docs: Some(eco_format!("Default value: {default}")),
303                    default: Some(default),
304                    ty: Ty::Any,
305                    attrs: ParamAttrs::named(),
306                }));
307            }
308            ast::Param::Spread(spread) => {
309                let sink = spread.sink_ident().map(|sink| sink.as_str());
310                add_param(Interned::new(ParamTy {
311                    name: sink.unwrap_or_default().into(),
312                    docs: None,
313                    default: None,
314                    ty: Ty::Any,
315                    attrs: ParamAttrs::variadic(),
316                }));
317            }
318        }
319    }
320}
321
322struct PatternDisplay<'a>(&'a ast::Pattern<'a>);
323
324impl fmt::Display for PatternDisplay<'_> {
325    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
326        match self.0 {
327            ast::Pattern::Normal(ast::Expr::Ident(ident)) => f.write_str(ident.as_str()),
328            ast::Pattern::Normal(_) => f.write_str("?"), // unreachable?
329            ast::Pattern::Placeholder(_) => f.write_str("_"),
330            ast::Pattern::Parenthesized(paren_expr) => {
331                write!(f, "{}", PatternDisplay(&paren_expr.pattern()))
332            }
333            ast::Pattern::Destructuring(destructing) => {
334                write!(f, "(")?;
335                let mut first = true;
336                for item in destructing.items() {
337                    if first {
338                        first = false;
339                    } else {
340                        write!(f, ", ")?;
341                    }
342                    match item {
343                        ast::DestructuringItem::Pattern(pos) => {
344                            write!(f, "{}", PatternDisplay(&pos))?
345                        }
346                        ast::DestructuringItem::Named(named) => write!(
347                            f,
348                            "{}: {}",
349                            named.name().as_str(),
350                            unwrap_parens(named.expr()).to_untyped().text()
351                        )?,
352                        ast::DestructuringItem::Spread(spread) => write!(
353                            f,
354                            "..{}",
355                            spread
356                                .sink_ident()
357                                .map(|sink| sink.as_str())
358                                .unwrap_or_default()
359                        )?,
360                    }
361                }
362                write!(f, ")")?;
363                Ok(())
364            }
365        }
366    }
367}
368
369fn unwrap_parens(mut expr: ast::Expr) -> ast::Expr {
370    while let ast::Expr::Parenthesized(paren_expr) = expr {
371        expr = paren_expr.expr();
372    }
373
374    expr
375}