tinymist_query/analysis/
signature.rs

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