tinymist_analysis/
sig.rs

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