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, ClosureNode, Func};
9use typst::syntax::ast;
10use typst::syntax::ast::AstNode;
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    let closure = &closure.node;
278    let closure_ast = match closure {
279        ClosureNode::Closure(node) => node.cast::<ast::Closure>().unwrap(),
280        _ => return,
281    };
282
283    for param in closure_ast.params().children() {
284        match param {
285            ast::Param::Pos(pos) => {
286                let name = format!("{}", PatternDisplay(&pos));
287                add_param(Interned::new(ParamTy {
288                    name: name.as_str().into(),
289                    docs: None,
290                    default: None,
291                    ty: Ty::Any,
292                    attrs: ParamAttrs::positional(),
293                }));
294            }
295            // todo: pattern
296            ast::Param::Named(named) => {
297                let default = unwrap_parens(named.expr()).to_untyped().clone().into_text();
298                add_param(Interned::new(ParamTy {
299                    name: named.name().get().into(),
300                    docs: Some(eco_format!("Default value: {default}")),
301                    default: Some(default),
302                    ty: Ty::Any,
303                    attrs: ParamAttrs::named(),
304                }));
305            }
306            ast::Param::Spread(spread) => {
307                let sink = spread.sink_ident().map(|sink| sink.as_str());
308                add_param(Interned::new(ParamTy {
309                    name: sink.unwrap_or_default().into(),
310                    docs: None,
311                    default: None,
312                    ty: Ty::Any,
313                    attrs: ParamAttrs::variadic(),
314                }));
315            }
316        }
317    }
318}
319
320struct PatternDisplay<'a>(&'a ast::Pattern<'a>);
321
322impl fmt::Display for PatternDisplay<'_> {
323    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
324        match self.0 {
325            ast::Pattern::Normal(ast::Expr::Ident(ident)) => f.write_str(ident.as_str()),
326            ast::Pattern::Normal(_) => f.write_str("?"), // unreachable?
327            ast::Pattern::Placeholder(_) => f.write_str("_"),
328            ast::Pattern::Parenthesized(paren_expr) => {
329                write!(f, "{}", PatternDisplay(&paren_expr.pattern()))
330            }
331            ast::Pattern::Destructuring(destructing) => {
332                write!(f, "(")?;
333                let mut first = true;
334                for item in destructing.items() {
335                    if first {
336                        first = false;
337                    } else {
338                        write!(f, ", ")?;
339                    }
340                    match item {
341                        ast::DestructuringItem::Pattern(pos) => {
342                            write!(f, "{}", PatternDisplay(&pos))?
343                        }
344                        ast::DestructuringItem::Named(named) => write!(
345                            f,
346                            "{}: {}",
347                            named.name().as_str(),
348                            unwrap_parens(named.expr()).to_untyped().text()
349                        )?,
350                        ast::DestructuringItem::Spread(spread) => write!(
351                            f,
352                            "..{}",
353                            spread
354                                .sink_ident()
355                                .map(|sink| sink.as_str())
356                                .unwrap_or_default()
357                        )?,
358                    }
359                }
360                write!(f, ")")?;
361                Ok(())
362            }
363        }
364    }
365}
366
367fn unwrap_parens(mut expr: ast::Expr) -> ast::Expr {
368    while let ast::Expr::Parenthesized(paren_expr) = expr {
369        expr = paren_expr.expr();
370    }
371
372    expr
373}