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

tinymist_query/syntax/
docs.rs
