tinymist_query/syntax/
docs.rs

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

tinymist_query/syntax/
docs.rs
