tinymist_query/syntax/
docs.rs

1use std::{collections::BTreeMap, ops::Deref, sync::LazyLock};
2
3use ecow::eco_format;
4use tinymist_analysis::stats::GLOBAL_STATS;
5use typst::foundations::{IntoValue, Module, Str, Type};
6
7use crate::{StrRef, adt::interner::Interned};
8use crate::{adt::snapshot_map::SnapshotMap, analysis::SharedContext};
9use crate::{
10    docs::{DocString, VarDoc, convert_docs, identify_pat_docs, identify_tidy_module_docs},
11    prelude::*,
12    syntax::{Decl, DefKind},
13    ty::{BuiltinTy, DynTypeBounds, InsTy, PackageId, SigTy, Ty, TypeVar, TypeVarBounds},
14};
15
16use super::DeclExpr;
17
18pub(crate) fn do_compute_docstring(
19    ctx: &Arc<SharedContext>,
20    fid: TypstFileId,
21    docs: String,
22    kind: DefKind,
23) -> Option<DocString> {
24    let _guard = GLOBAL_STATS.stat(fid, "compute_docstring");
25
26    let checker = DocsChecker {
27        fid,
28        ctx,
29        var_bounds: HashMap::new(),
30        globals: HashMap::default(),
31        locals: SnapshotMap::default(),
32        next_id: 0,
33    };
34    use DefKind::*;
35    match kind {
36        Function | Variable => checker.check_pat_docs(docs),
37        Module => checker.check_module_docs(docs),
38        Constant | Struct | Reference => None,
39    }
40}
41
42struct DocsChecker<'a> {
43    fid: TypstFileId,
44    ctx: &'a Arc<SharedContext>,
45    /// The bounds of type variables
46    var_bounds: HashMap<DeclExpr, TypeVarBounds>,
47    /// Global name bindings
48    globals: HashMap<EcoString, Option<Ty>>,
49    /// Local name bindings
50    locals: SnapshotMap<EcoString, Ty>,
51    /// Next generated variable id
52    next_id: u32,
53}
54
55static EMPTY_MODULE: LazyLock<Module> =
56    LazyLock::new(|| Module::new("stub", typst::foundations::Scope::new()));
57
58impl DocsChecker<'_> {
59    pub fn check_pat_docs(mut self, docs: String) -> Option<DocString> {
60        let converted = convert_docs(self.ctx, &docs, Some(self.fid))
61            .and_then(|converted| identify_pat_docs(&converted));
62
63        let converted = match Self::fallback_docs(converted, &docs) {
64            Ok(docs) => docs,
65            Err(err) => return Some(err),
66        };
67
68        let module = self.ctx.module_by_str(docs);
69        let module = module.as_ref().unwrap_or(EMPTY_MODULE.deref());
70
71        let mut params = BTreeMap::new();
72        for param in converted.params.into_iter() {
73            params.insert(
74                param.name.into(),
75                VarDoc {
76                    docs: self.ctx.remove_html(param.docs),
77                    ty: self.check_type_strings(module, &param.types),
78                },
79            );
80        }
81
82        let res_ty = converted
83            .return_ty
84            .and_then(|ty| self.check_type_strings(module, &ty));
85
86        Some(DocString {
87            docs: Some(self.ctx.remove_html(converted.docs)),
88            var_bounds: self.var_bounds,
89            vars: params,
90            res_ty,
91        })
92    }
93
94    pub fn check_module_docs(self, docs: String) -> Option<DocString> {
95        let converted =
96            convert_docs(self.ctx, &docs, Some(self.fid)).and_then(identify_tidy_module_docs);
97
98        let converted = match Self::fallback_docs(converted, &docs) {
99            Ok(docs) => docs,
100            Err(err) => return Some(err),
101        };
102
103        Some(DocString {
104            docs: Some(self.ctx.remove_html(converted.docs)),
105            var_bounds: self.var_bounds,
106            vars: BTreeMap::new(),
107            res_ty: None,
108        })
109    }
110
111    fn fallback_docs<T>(converted: Result<T, EcoString>, docs: &str) -> Result<T, DocString> {
112        match converted {
113            Ok(converted) => Ok(converted),
114            Err(err) => {
115                let err = err.replace("`", "\\`");
116                let max_consecutive_backticks = docs
117                    .chars()
118                    .fold((0, 0), |(max, count), ch| {
119                        if ch == '`' {
120                            (max.max(count + 1), count + 1)
121                        } else {
122                            (max, 0)
123                        }
124                    })
125                    .0;
126                let backticks = "`".repeat((max_consecutive_backticks + 1).max(3));
127                let fallback_docs = eco_format!(
128                    "```\nfailed to parse docs: {err}\n```\n\n{backticks}typ\n{docs}\n{backticks}\n"
129                );
130                Err(DocString {
131                    docs: Some(fallback_docs),
132                    var_bounds: HashMap::new(),
133                    vars: BTreeMap::new(),
134                    res_ty: None,
135                })
136            }
137        }
138    }
139
140    fn generate_var(&mut self, name: StrRef) -> Ty {
141        self.next_id += 1;
142        let encoded = Interned::new(Decl::generated(DefId(self.next_id as u64)));
143        crate::log_debug_ct!("generate var {name:?} {encoded:?}");
144        let var = TypeVar {
145            name,
146            def: encoded.clone(),
147        };
148        let bounds = TypeVarBounds::new(var, DynTypeBounds::default());
149        let var = bounds.as_type();
150        self.var_bounds.insert(encoded, bounds);
151        var
152    }
153
154    fn check_type_strings(&mut self, m: &Module, inputs: &str) -> Option<Ty> {
155        let mut terms = vec![];
156        for name in inputs.split(",").map(|ty| ty.trim()) {
157            let Some(ty) = self.check_type_ident(m, name) else {
158                continue;
159            };
160            terms.push(ty);
161        }
162
163        Some(Ty::from_types(terms.into_iter()))
164    }
165
166    fn check_type_ident(&mut self, m: &Module, name: &str) -> Option<Ty> {
167        static TYPE_REPRS: LazyLock<HashMap<&'static str, Ty>> = LazyLock::new(|| {
168            let values = Vec::from_iter(
169                [
170                    Value::None,
171                    Value::Auto,
172                    // Value::Bool(Default::default()),
173                    Value::Int(Default::default()),
174                    Value::Float(Default::default()),
175                    Value::Length(Default::default()),
176                    Value::Angle(Default::default()),
177                    Value::Ratio(Default::default()),
178                    Value::Relative(Default::default()),
179                    Value::Fraction(Default::default()),
180                    Value::Str(Default::default()),
181                ]
182                .map(|v| v.ty())
183                .into_iter()
184                .chain([
185                    Type::of::<typst::visualize::Color>(),
186                    Type::of::<typst::visualize::Gradient>(),
187                    Type::of::<typst::visualize::Tiling>(),
188                    Type::of::<typst::foundations::Symbol>(),
189                    Type::of::<typst::foundations::Version>(),
190                    Type::of::<typst::foundations::Bytes>(),
191                    Type::of::<typst::foundations::Label>(),
192                    Type::of::<typst::foundations::Datetime>(),
193                    Type::of::<typst::foundations::Duration>(),
194                    Type::of::<typst::foundations::Content>(),
195                    Type::of::<typst::foundations::Styles>(),
196                    Type::of::<typst::foundations::Array>(),
197                    Type::of::<typst::foundations::Dict>(),
198                    Type::of::<typst::foundations::Func>(),
199                    Type::of::<typst::foundations::Args>(),
200                    Type::of::<typst::foundations::Type>(),
201                    Type::of::<typst::foundations::Module>(),
202                ]),
203            );
204
205            let shorts = values
206                .clone()
207                .into_iter()
208                .map(|ty| (ty.short_name(), Ty::Builtin(BuiltinTy::Type(ty))));
209            let longs = values
210                .into_iter()
211                .map(|ty| (ty.long_name(), Ty::Builtin(BuiltinTy::Type(ty))));
212            let builtins = [
213                ("any", Ty::Any),
214                ("bool", Ty::Boolean(None)),
215                ("boolean", Ty::Boolean(None)),
216                ("false", Ty::Boolean(Some(false))),
217                ("true", Ty::Boolean(Some(true))),
218            ];
219            HashMap::from_iter(shorts.chain(longs).chain(builtins))
220        });
221
222        let builtin_ty = TYPE_REPRS.get(name).cloned();
223        builtin_ty
224            .or_else(|| self.locals.get(name).cloned())
225            .or_else(|| self.check_type_annotation(m, name))
226    }
227
228    fn check_type_annotation(&mut self, module: &Module, name: &str) -> Option<Ty> {
229        if let Some(term) = self.globals.get(name) {
230            return term.clone();
231        }
232
233        let val = module.scope().get(name)?;
234        crate::log_debug_ct!("check doc type annotation: {name:?}");
235        if let Value::Content(raw) = val.read() {
236            let annotated = raw.clone().unpack::<typst::text::RawElem>().ok()?;
237            let annotated = annotated.text.clone().into_value().cast::<Str>().ok()?;
238            let code = typst::syntax::parse_code(&annotated.as_str().replace('\'', "θ"));
239            let mut exprs = code.cast::<ast::Code>()?.exprs();
240            let term = self.check_type_expr(module, exprs.next()?);
241            self.globals.insert(name.into(), term.clone());
242            term
243        } else {
244            None
245        }
246    }
247
248    fn check_type_expr(&mut self, module: &Module, expr: ast::Expr) -> Option<Ty> {
249        crate::log_debug_ct!("check doc type expr: {expr:?}");
250        match expr {
251            ast::Expr::Ident(ident) => self.check_type_ident(module, ident.get().as_str()),
252            ast::Expr::None(_)
253            | ast::Expr::Auto(_)
254            | ast::Expr::Bool(..)
255            | ast::Expr::Int(..)
256            | ast::Expr::Float(..)
257            | ast::Expr::Numeric(..)
258            | ast::Expr::Str(..) => {
259                SharedContext::const_eval(expr).map(|v| Ty::Value(InsTy::new(v)))
260            }
261            ast::Expr::Binary(binary) => {
262                let mut components = Vec::with_capacity(2);
263                components.push(self.check_type_expr(module, binary.lhs())?);
264
265                let mut rhs = binary.rhs();
266                while let ast::Expr::Binary(binary) = rhs {
267                    if binary.op() != ast::BinOp::Or {
268                        break;
269                    }
270
271                    components.push(self.check_type_expr(module, binary.lhs())?);
272                    rhs = binary.rhs();
273                }
274
275                components.push(self.check_type_expr(module, rhs)?);
276                Some(Ty::from_types(components.into_iter()))
277            }
278            ast::Expr::FuncCall(call) => match call.callee() {
279                ast::Expr::Ident(callee) => {
280                    let name = callee.get().as_str();
281                    match name {
282                        "array" => Some({
283                            let ast::Arg::Pos(pos) = call.args().items().next()? else {
284                                return None;
285                            };
286
287                            Ty::Array(self.check_type_expr(module, pos)?.into())
288                        }),
289                        "tag" => Some({
290                            let ast::Arg::Pos(ast::Expr::Str(s)) = call.args().items().next()?
291                            else {
292                                return None;
293                            };
294                            let pkg_id = PackageId::try_from(self.fid).ok();
295                            Ty::Builtin(BuiltinTy::Tag(Box::new((
296                                s.get().into(),
297                                pkg_id.map(From::from),
298                            ))))
299                        }),
300                        _ => None,
301                    }
302                }
303                _ => None,
304            },
305            ast::Expr::Closure(closure) => {
306                crate::log_debug_ct!("check doc closure annotation: {closure:?}");
307                let mut pos_all = vec![];
308                let mut named_all = BTreeMap::new();
309                let mut spread_right = None;
310                let snap = self.locals.snapshot();
311
312                let sig = None.or_else(|| {
313                    for param in closure.params().children() {
314                        match param {
315                            ast::Param::Pos(ast::Pattern::Normal(ast::Expr::Ident(pos))) => {
316                                let name = pos.get().clone();
317                                let term = self.generate_var(name.as_str().into());
318                                self.locals.insert(name, term.clone());
319                                pos_all.push(term);
320                            }
321                            ast::Param::Pos(_pos) => {
322                                pos_all.push(Ty::Any);
323                            }
324                            ast::Param::Named(named) => {
325                                let term = self
326                                    .check_type_expr(module, named.expr())
327                                    .unwrap_or(Ty::Any);
328                                named_all.insert(named.name().into(), term);
329                            }
330                            // todo: spread left/right
331                            ast::Param::Spread(spread) => {
332                                let Some(sink) = spread.sink_ident() else {
333                                    continue;
334                                };
335                                let sink_name = sink.get().clone();
336                                let rest_term = self.generate_var(sink_name.as_str().into());
337                                self.locals.insert(sink_name, rest_term.clone());
338                                spread_right = Some(rest_term);
339                            }
340                        }
341                    }
342
343                    let body = self.check_type_expr(module, closure.body())?;
344                    let sig = SigTy::new(
345                        pos_all.into_iter(),
346                        named_all,
347                        None,
348                        spread_right,
349                        Some(body),
350                    )
351                    .into();
352
353                    Some(Ty::Func(sig))
354                });
355
356                self.locals.rollback_to(snap);
357                sig
358            }
359            ast::Expr::Dict(decl) => {
360                crate::log_debug_ct!("check doc dict annotation: {decl:?}");
361                None
362            }
363            _ => None,
364        }
365    }
366}