tinymist_query/analysis/
signature.rs

1//! Analysis of function signatures.
2
3use std::cell::RefCell;
4
5use itertools::Either;
6use tinymist_analysis::{ArgInfo, ArgsInfo, PartialSignature, func_signature};
7use tinymist_derive::BindTyCtx;
8use tinymist_std::hash::FxHashSet;
9
10use super::{Definition, SharedContext, prelude::*};
11use crate::analysis::PostTypeChecker;
12use crate::docs::{DocText, UntypedDefDocs, UntypedSignatureDocs, UntypedVarDocs};
13use crate::syntax::{DeclExpr, classify_def_loosely};
14use crate::ty::{
15    BoundChecker, DocSource, DynTypeBounds, ParamAttrs, ParamTy, SigTy, SigWithTy, TyCtx, TyCtxMut,
16    TypeInfo, TypeVar,
17};
18
19pub use tinymist_analysis::{PrimarySignature, Signature};
20
21/// The language object that the signature is being analyzed for.
22#[derive(Debug, Clone)]
23pub enum SignatureTarget {
24    /// A static node without knowing the function at runtime.
25    Def(Option<Source>, Definition),
26    /// A static node without knowing the function at runtime.
27    SyntaxFast(Source, Span),
28    /// A static node without knowing the function at runtime.
29    Syntax(Source, Span),
30    /// A function that is known at runtime.
31    Runtime(Func),
32    /// A function that is known at runtime.
33    Convert(Func),
34}
35
36impl SignatureTarget {
37    /// Returns the span of the callee node.
38    pub fn span(&self) -> Span {
39        match self {
40            SignatureTarget::Def(_, def) => def.decl.span(),
41            SignatureTarget::SyntaxFast(_, span) | SignatureTarget::Syntax(_, span) => *span,
42            SignatureTarget::Runtime(func) | SignatureTarget::Convert(func) => func.span(),
43        }
44    }
45}
46
47#[typst_macros::time(span = callee_node.span())]
48pub(crate) fn analyze_signature(
49    ctx: &Arc<SharedContext>,
50    callee_node: SignatureTarget,
51) -> Option<Signature> {
52    ctx.compute_signature(callee_node.clone(), move |ctx| {
53        crate::log_debug_ct!("analyzing signature for {callee_node:?}");
54        analyze_type_signature(ctx, &callee_node)
55            .or_else(|| analyze_dyn_signature(ctx, &callee_node))
56    })
57}
58
59#[typst_macros::time(span = callee_node.span())]
60fn analyze_type_signature(
61    ctx: &Arc<SharedContext>,
62    callee_node: &SignatureTarget,
63) -> Option<Signature> {
64    let (type_info, ty) = match callee_node {
65        SignatureTarget::Convert(..) => return None,
66        SignatureTarget::SyntaxFast(source, span) | SignatureTarget::Syntax(source, span) => {
67            let type_info = ctx.type_check(source);
68            let ty = type_info.type_of_span(*span)?;
69            Some((type_info, ty))
70        }
71        SignatureTarget::Def(source, def) => {
72            let span = def.decl.span();
73            let type_info = ctx.type_check(source.as_ref()?);
74            let ty = type_info.type_of_span(span)?;
75            Some((type_info, ty))
76        }
77        SignatureTarget::Runtime(func) => {
78            let source = ctx.source_by_id(func.span().id()?).ok()?;
79            let node = source.find(func.span())?;
80            let def = classify_def_loosely(node.parent()?.clone())?;
81            let type_info = ctx.type_check(&source);
82            let ty = type_info.type_of_span(def.name()?.span())?;
83            Some((type_info, ty))
84        }
85    }?;
86
87    sig_of_type(ctx, &type_info, ty)
88}
89
90pub(crate) fn sig_of_type(
91    ctx: &Arc<SharedContext>,
92    type_info: &TypeInfo,
93    ty: Ty,
94) -> Option<Signature> {
95    // todo multiple sources
96    let mut srcs = ty.sources();
97    srcs.sort();
98    crate::log_debug_ct!("check type signature of ty: {ty:?} => {srcs:?}");
99    let type_var = srcs.into_iter().next()?;
100    match type_var {
101        DocSource::Var(v) => {
102            let mut ty_ctx =
103                SignatureTypeContext::new(PostTypeChecker::new(ctx.clone(), type_info));
104            let sig_ty = Ty::Func(ty.sig_repr(true, &mut ty_ctx)?);
105            let sig_ty = type_info.simplify(sig_ty, false);
106            let Ty::Func(sig_ty) = sig_ty else {
107                static WARN_ONCE: std::sync::Once = std::sync::Once::new();
108                WARN_ONCE.call_once(|| {
109                    // todo: seems like a bug
110                    log::warn!("expected function type, got {sig_ty:?}");
111                });
112                return None;
113            };
114
115            // todo: this will affect inlay hint: _var_with
116            ty_ctx.reset_cycle_guard();
117            let (var_with, docstring) = match type_info.var_docs.get(&v.def).map(|x| x.as_ref()) {
118                Some(UntypedDefDocs::Function(sig)) => (vec![], Either::Left(sig.as_ref())),
119                Some(UntypedDefDocs::Variable(docs)) => find_alias_stack(&mut ty_ctx, &v, docs)?,
120                _ => return None,
121            };
122
123            let docstring = match docstring {
124                Either::Left(docstring) => docstring,
125                Either::Right(func) => return Some(wind_stack(var_with, ctx.type_of_func(func))),
126            };
127
128            let mut param_specs = Vec::new();
129            let mut has_fill_or_size_or_stroke = false;
130            let mut _broken = false;
131
132            if docstring.pos.len() != sig_ty.positional_params().len() {
133                static WARN_ONCE: std::sync::Once = std::sync::Once::new();
134                WARN_ONCE.call_once(|| {
135                    // todo: seems like a bug
136                    log::warn!("positional params mismatch: {docstring:#?} != {sig_ty:#?}");
137                });
138                return None;
139            }
140
141            for (doc, ty) in docstring.pos.iter().zip(sig_ty.positional_params()) {
142                let default = doc.default.clone();
143                let ty = ty.clone();
144
145                let name = doc.name.clone();
146                if matches!(name.as_ref(), "fill" | "stroke" | "size") {
147                    has_fill_or_size_or_stroke = true;
148                }
149
150                param_specs.push(Interned::new(ParamTy {
151                    name,
152                    docs: Some(DocText::plain(doc.docs.clone())),
153                    default,
154                    ty,
155                    attrs: ParamAttrs::positional(),
156                }));
157            }
158
159            for (name, ty) in sig_ty.named_params() {
160                let docstring = docstring.named.get(name);
161                let default = Some(
162                    docstring
163                        .and_then(|doc| doc.default.clone())
164                        .unwrap_or_else(|| "unknown".into()),
165                );
166                let ty = ty.clone();
167
168                if matches!(name.as_ref(), "fill" | "stroke" | "size") {
169                    has_fill_or_size_or_stroke = true;
170                }
171
172                param_specs.push(Interned::new(ParamTy {
173                    name: name.clone(),
174                    docs: docstring.map(|doc| DocText::plain(doc.docs.clone())),
175                    default,
176                    ty,
177                    attrs: ParamAttrs::named(),
178                }));
179            }
180
181            if let Some(doc) = docstring.rest.as_ref() {
182                let default = doc.default.clone();
183
184                param_specs.push(Interned::new(ParamTy {
185                    name: doc.name.clone(),
186                    docs: Some(DocText::plain(doc.docs.clone())),
187                    default,
188                    ty: sig_ty.rest_param().cloned().unwrap_or(Ty::Any),
189                    attrs: ParamAttrs::variadic(),
190                }));
191            }
192
193            let sig = Signature::Primary(Arc::new(PrimarySignature {
194                docs: Some(DocText::plain(docstring.docs.clone())),
195                param_specs,
196                has_fill_or_size_or_stroke,
197                sig_ty,
198                _broken,
199            }));
200            Some(wind_stack(var_with, sig))
201        }
202        src @ (DocSource::Builtin(..) | DocSource::Ins(..)) => {
203            Some(ctx.type_of_func(src.as_func()?))
204        }
205    }
206}
207
208/// A type context for signature rendering that memoizes `(TypeVar, polarity)`
209/// visits within a single traversal so recursive bounds are only expanded once.
210struct SignatureTypeContext<'a> {
211    inner: PostTypeChecker<'a>,
212    visited: RefCell<FxHashSet<(DeclExpr, bool)>>,
213}
214
215impl<'a> SignatureTypeContext<'a> {
216    fn new(inner: PostTypeChecker<'a>) -> Self {
217        Self {
218            inner,
219            visited: RefCell::new(FxHashSet::default()),
220        }
221    }
222
223    fn reset_cycle_guard(&self) {
224        self.visited.borrow_mut().clear();
225    }
226}
227
228impl TyCtx for SignatureTypeContext<'_> {
229    fn global_bounds(&self, var: &Interned<TypeVar>, pol: bool) -> Option<DynTypeBounds> {
230        if !self.visited.borrow_mut().insert((var.def.clone(), pol)) {
231            return None;
232        }
233
234        self.inner.global_bounds(var, pol)
235    }
236
237    fn local_bind_of(&self, var: &Interned<TypeVar>) -> Option<Ty> {
238        self.inner.local_bind_of(var)
239    }
240}
241
242impl TyCtxMut for SignatureTypeContext<'_> {
243    type Snap = <TypeInfo as TyCtxMut>::Snap;
244
245    fn start_scope(&mut self) -> Self::Snap {
246        self.inner.start_scope()
247    }
248
249    fn end_scope(&mut self, snap: Self::Snap) {
250        self.inner.end_scope(snap)
251    }
252
253    fn bind_local(&mut self, var: &Interned<TypeVar>, ty: Ty) {
254        self.inner.bind_local(var, ty);
255    }
256
257    fn type_of_func(&mut self, func: &Func) -> Option<Interned<SigTy>> {
258        self.inner.type_of_func(func)
259    }
260
261    fn type_of_value(&mut self, val: &Value) -> Ty {
262        self.inner.type_of_value(val)
263    }
264
265    fn check_module_item(&mut self, module: TypstFileId, key: &StrRef) -> Option<Ty> {
266        self.inner.check_module_item(module, key)
267    }
268}
269
270fn wind_stack(var_with: Vec<WithElem>, sig: Signature) -> Signature {
271    if var_with.is_empty() {
272        return sig;
273    }
274
275    let (primary, mut base_args) = match sig {
276        Signature::Primary(primary) => (primary, eco_vec![]),
277        Signature::Partial(partial) => (partial.signature.clone(), partial.with_stack.clone()),
278    };
279
280    let mut accepting = primary.pos().iter().skip(base_args.len());
281
282    // Ignoring docs at the moment
283    for (_d, w) in var_with {
284        if let Some(w) = w {
285            let mut items = eco_vec![];
286            for pos in w.with.positional_params() {
287                let Some(arg) = accepting.next() else {
288                    break;
289                };
290                items.push(ArgInfo {
291                    name: Some(arg.name.clone()),
292                    term: Some(pos.clone()),
293                });
294            }
295            // todo: ignored spread arguments
296            if !items.is_empty() {
297                base_args.push(ArgsInfo { items });
298            }
299        }
300    }
301
302    Signature::Partial(Arc::new(PartialSignature {
303        signature: primary,
304        with_stack: base_args,
305    }))
306}
307
308type WithElem<'a> = (&'a UntypedVarDocs, Option<Interned<SigWithTy>>);
309
310fn find_alias_stack<'a>(
311    ctx: &'a mut SignatureTypeContext,
312    var: &Interned<TypeVar>,
313    docs: &'a UntypedVarDocs,
314) -> Option<(Vec<WithElem<'a>>, Either<&'a UntypedSignatureDocs, Func>)> {
315    let mut checker = AliasStackChecker {
316        ctx,
317        stack: vec![(docs, None)],
318        res: None,
319        checking_with: true,
320    };
321    Ty::Var(var.clone()).bounds(true, &mut checker);
322
323    checker.res.map(|res| (checker.stack, res))
324}
325
326#[derive(BindTyCtx)]
327#[bind(ctx)]
328struct AliasStackChecker<'a, 'b> {
329    ctx: &'a mut SignatureTypeContext<'b>,
330    stack: Vec<WithElem<'a>>,
331    res: Option<Either<&'a UntypedSignatureDocs, Func>>,
332    checking_with: bool,
333}
334
335impl BoundChecker for AliasStackChecker<'_, '_> {
336    fn check_var(&mut self, u: &Interned<TypeVar>, pol: bool) {
337        crate::log_debug_ct!("collecting var {u:?} {pol:?}");
338        if self.res.is_some() {
339            return;
340        }
341
342        if self.checking_with {
343            self.check_var_rec(u, pol);
344            return;
345        }
346
347        let docs = self.ctx.inner.info.var_docs.get(&u.def).map(|x| x.as_ref());
348
349        crate::log_debug_ct!("collecting var {u:?} {pol:?} => {docs:?}");
350        // todo: bind builtin functions
351        match docs {
352            Some(UntypedDefDocs::Function(sig)) => {
353                self.res = Some(Either::Left(sig));
354            }
355            Some(UntypedDefDocs::Variable(docs)) => {
356                self.checking_with = true;
357                self.stack.push((docs, None));
358                self.check_var_rec(u, pol);
359                self.stack.pop();
360                self.checking_with = false;
361            }
362            _ => {}
363        }
364    }
365
366    fn collect(&mut self, ty: &Ty, pol: bool) {
367        if self.res.is_some() {
368            return;
369        }
370
371        match (self.checking_with, ty) {
372            (true, Ty::With(w)) => {
373                crate::log_debug_ct!("collecting with {ty:?} {pol:?}");
374                self.stack.last_mut().unwrap().1 = Some(w.clone());
375                self.checking_with = false;
376                w.sig.bounds(pol, self);
377                self.checking_with = true;
378            }
379            (false, ty) => {
380                if let Some(src) = ty.as_source() {
381                    match src {
382                        DocSource::Var(u) => {
383                            self.check_var(&u, pol);
384                        }
385                        src @ (DocSource::Builtin(..) | DocSource::Ins(..)) => {
386                            if let Some(func) = src.as_func() {
387                                self.res = Some(Either::Right(func));
388                            }
389                        }
390                    }
391                }
392            }
393            _ => {}
394        }
395    }
396}
397
398#[typst_macros::time(span = callee_node.span())]
399fn analyze_dyn_signature(
400    ctx: &Arc<SharedContext>,
401    callee_node: &SignatureTarget,
402) -> Option<Signature> {
403    let func = match callee_node {
404        SignatureTarget::Def(_source, def) => def.value()?.to_func()?,
405        SignatureTarget::SyntaxFast(..) => return None,
406        SignatureTarget::Syntax(source, span) => {
407            let def = ctx.def_of_span(source, *span)?;
408            def.value()?.to_func()?
409        }
410        SignatureTarget::Convert(func) | SignatureTarget::Runtime(func) => func.clone(),
411    };
412
413    Some(func_signature(func))
414}
tinymist_query/analysis/signature.rs

tinymist_query/analysis/
signature.rs
