feat: cache docstring building (Part. 2) (#680)

* feat: cache docstring building (Part. 2)

* dev: fix import

* dev: change some variable names

crates/tinymist-query/src/analysis/global.rs

@@ -19,9 +19,10 @@ use typst_shim::syntax::LinkedNodeExt;
 use crate::adt::interner::Interned;
 use crate::analysis::{
     analyze_bib, analyze_dyn_signature, analyze_expr_, analyze_import_, post_type_check, BibInfo,
-    DefUseInfo, DefinitionLink, IdentRef, ImportInfo, PathPreference, SigTy, Signature,
+    DefUseInfo, DefinitionLink, DocString, IdentRef, ImportInfo, PathPreference, SigTy, Signature,
     SignatureTarget, Ty, TypeScheme,
 };
+use crate::docs::DocStringKind;
 use crate::prelude::*;
 use crate::syntax::{
     construct_module_dependencies, find_expr_in_import, get_deref_target, resolve_id_by_path,
@@ -85,6 +86,7 @@ pub struct AnalysisGlobalCaches {
     def_use: FxDashMap<u128, (u64, Option<Arc<DefUseInfo>>)>,
     type_check: FxDashMap<u128, (u64, Option<Arc<TypeScheme>>)>,
     static_signatures: FxDashMap<u128, (u64, Source, usize, Signature)>,
+    docstrings: FxDashMap<u128, Option<Arc<DocString>>>,
     signatures: FxDashMap<u128, (u64, foundations::Func, Signature)>,
 }
 
@@ -349,13 +351,13 @@ impl<'w> AnalysisContext<'w> {
     }
 
     /// Get a module by string.
-    pub fn module_by_str(&mut self, rr: String) -> Option<Module> {
+    pub fn module_by_str(&self, rr: String) -> Option<Module> {
         let src = Source::new(*DETACHED_ENTRY, rr);
         self.module_by_src(src).ok()
     }
 
     /// Get (Create) a module by source.
-    pub fn module_by_src(&mut self, source: Source) -> SourceResult<Module> {
+    pub fn module_by_src(&self, source: Source) -> SourceResult<Module> {
         let route = Route::default();
         let mut tracer = Tracer::default();
 
@@ -479,6 +481,23 @@ impl<'w> AnalysisContext<'w> {
         }
     }
 
+    pub(crate) fn compute_docstring(
+        &self,
+        docs: String,
+        kind: DocStringKind,
+    ) -> Option<Arc<DocString>> {
+        let h = hash128(&(&docs, &kind));
+        let res = if let Some(res) = self.analysis.caches.docstrings.get(&h) {
+            res.clone()
+        } else {
+            let res = crate::analysis::tyck::compute_docstring(self, docs, kind).map(Arc::new);
+            self.analysis.caches.docstrings.insert(h, res.clone());
+            res
+        };
+
+        res
+    }
+
     /// Compute the signature of a function.
     pub fn compute_signature(
         &self,

crates/tinymist-query/src/analysis/tyck.rs

@@ -25,10 +25,12 @@ use super::{
 };
 
 mod apply;
+mod docs;
 mod select;
 mod syntax;
 
 pub(crate) use apply::*;
+pub(crate) use docs::*;
 pub(crate) use select::*;
 
 /// Type checking at the source unit level.
@@ -44,9 +46,6 @@ pub(crate) fn type_check(ctx: &mut AnalysisContext, source: Source) -> Option<Ar
         def_use_info,
         info: &mut info,
         externals: HashMap::new(),
-        docs_scope: HashMap::new(),
-        documenting_id: None,
-        generated: HashMap::new(),
         mode: InterpretMode::Markup,
     };
     let lnk = LinkedNode::new(source.root());
@@ -74,9 +73,6 @@ struct TypeChecker<'a, 'w> {
     def_use_info: Arc<DefUseInfo>,
 
     info: &'a mut TypeScheme,
-    docs_scope: HashMap<EcoString, Option<Ty>>,
-    documenting_id: Option<DefId>,
-    generated: HashMap<DefId, u32>,
     externals: HashMap<DefId, Option<Ty>>,
     mode: InterpretMode,
 }
@@ -119,12 +115,16 @@ impl<'a, 'w> TypeChecker<'a, 'w> {
             .or_else(|| Some(self.def_use_info.get_def(s.id()?, r)?.0))
     }
 
-    fn generate_var(&mut self, name: StrRef, id: DefId) -> Ty {
-        let next_id = self.generated.entry(id).or_insert(0);
-        *next_id += 1;
-        let encoded = DefId((id.0 + 1) * 0x100_0000_0000 + (*next_id as u64));
-        log::debug!("generate var {name:?} {encoded:?}");
-        let bounds = TypeVarBounds::new(TypeVar { name, def: encoded }, TypeBounds::default());
+    fn copy_based_on(&mut self, fr: &TypeVarBounds, offset: u64, id: DefId) -> Ty {
+        let encoded = DefId((id.0 + 1) * 0x100_0000_0000 + offset + fr.id().0);
+        log::debug!("copy var {fr:?} as {encoded:?}");
+        let bounds = TypeVarBounds::new(
+            TypeVar {
+                name: fr.name().clone(),
+                def: encoded,
+            },
+            fr.bounds.bounds().read().clone(),
+        );
         let var = bounds.as_type();
         self.info.vars.insert(encoded, bounds);
         var
@@ -155,13 +155,6 @@ impl<'a, 'w> TypeChecker<'a, 'w> {
         Some(var.as_type())
     }
 
-    fn with_docs_scope<T>(&mut self, f: impl FnOnce(&mut Self) -> T) -> T {
-        let res = f(self);
-        self.docs_scope.clear();
-        self.documenting_id = None;
-        res
-    }
-
     fn import_ty(&mut self, def_id: DefId) -> Option<Ty> {
         if let Some(ty) = self.externals.get(&def_id) {
             return ty.clone();

crates/tinymist-query/src/analysis/tyck/docs.rs

@@ -0,0 +1,324 @@
+use std::{collections::BTreeMap, sync::LazyLock};
+
+use reflexo::TakeAs;
+use typst::foundations::{IntoValue, Module, Str, Type};
+
+use crate::{
+    docs::{convert_docs, identify_func_docs, DocStringKind},
+    syntax::{find_docs_of, get_non_strict_def_target},
+};
+
+use super::*;
+
+const DOC_VARS: u64 = 0;
+
+impl<'a, 'w> TypeChecker<'a, 'w> {
+    pub fn check_closure_docs(&mut self, root: &LinkedNode) -> Option<DocString> {
+        let closure = root.cast::<ast::Closure>()?;
+
+        // todo: cache docs capture
+        // use parent of params, todo: reliable way to get the def target
+        let def = get_non_strict_def_target(root.clone())?;
+        let docs = find_docs_of(&self.source, def)?;
+
+        let documenting_id = closure
+            .name()
+            .and_then(|n| self.get_def_id(n.span(), &to_ident_ref(root, n)?))?;
+        let docstring = self.ctx.compute_docstring(docs, DocStringKind::Function)?;
+        Some(docstring.take().rename_based_on(documenting_id, self))
+    }
+}
+
+#[derive(Debug, Clone, Default)]
+pub(crate) struct DocString {
+    /// The documentation of the item
+    pub docs: Option<String>,
+    /// The typing on definitions
+    pub var_bounds: HashMap<DefId, TypeVarBounds>,
+    /// The variable doc associated with the item
+    pub vars: HashMap<EcoString, VarDoc>,
+    /// The type of the resultant type
+    pub res_ty: Option<Ty>,
+}
+
+impl DocString {
+    pub fn get_var(&self, name: &str) -> Option<&VarDoc> {
+        self.vars.get(name)
+    }
+
+    pub fn var_ty(&self, name: &str) -> Option<&Ty> {
+        self.get_var(name).and_then(|v| v.ty.as_ref())
+    }
+
+    fn rename_based_on(self, documenting_id: DefId, base: &mut TypeChecker) -> DocString {
+        let DocString {
+            docs,
+            var_bounds,
+            vars,
+            mut res_ty,
+        } = self;
+        let mut renamer = IdRenamer {
+            base,
+            var_bounds: &var_bounds,
+            base_id: documenting_id,
+            offset: DOC_VARS,
+        };
+        let mut vars = vars;
+        for (_name, doc) in vars.iter_mut() {
+            if let Some(ty) = &mut doc.ty {
+                if let Some(mutated) = ty.mutate(true, &mut renamer) {
+                    *ty = mutated;
+                }
+            }
+        }
+        if let Some(ty) = res_ty.as_mut() {
+            if let Some(mutated) = ty.mutate(true, &mut renamer) {
+                *ty = mutated;
+            }
+        }
+        DocString {
+            docs,
+            var_bounds,
+            vars,
+            res_ty,
+        }
+    }
+}
+
+#[derive(Debug, Clone, Default)]
+pub(crate) struct VarDoc {
+    pub _docs: Option<EcoString>,
+    pub ty: Option<Ty>,
+    pub _default: Option<EcoString>,
+}
+
+pub(crate) fn compute_docstring(
+    ctx: &AnalysisContext,
+    docs: String,
+    kind: DocStringKind,
+) -> Option<DocString> {
+    let checker = DocsChecker {
+        ctx,
+        vars: HashMap::new(),
+        docs_scope: HashMap::new(),
+        next_id: 0,
+    };
+    match kind {
+        DocStringKind::Function => checker.check_closure_docs(docs),
+    }
+}
+
+struct DocsChecker<'a, 'w> {
+    ctx: &'a AnalysisContext<'w>,
+    /// The typing on definitions
+    vars: HashMap<DefId, TypeVarBounds>,
+    docs_scope: HashMap<EcoString, Option<Ty>>,
+    next_id: u32,
+}
+
+impl<'a, 'w> DocsChecker<'a, 'w> {
+    pub fn check_closure_docs(mut self, docs: String) -> Option<DocString> {
+        let converted = convert_docs(self.ctx.world(), &docs).ok()?;
+        let converted = identify_func_docs(&converted).ok()?;
+        let module = self.ctx.module_by_str(docs)?;
+
+        let mut params = HashMap::new();
+        for param in converted.params.into_iter() {
+            params.insert(
+                param.name,
+                VarDoc {
+                    _docs: Some(param.docs),
+                    ty: self.check_doc_types(&module, &param.types),
+                    _default: param.default,
+                },
+            );
+        }
+
+        let res_ty = converted
+            .return_ty
+            .and_then(|ty| self.check_doc_types(&module, &ty));
+
+        Some(DocString {
+            docs: Some(converted.docs),
+            var_bounds: self.vars,
+            vars: params,
+            res_ty,
+        })
+    }
+
+    fn check_doc_types(&mut self, m: &Module, strs: &str) -> Option<Ty> {
+        let mut types = vec![];
+        for name in strs.split(",").map(|e| e.trim()) {
+            let Some(ty) = self.check_doc_type_ident(m, name) else {
+                continue;
+            };
+            types.push(ty);
+        }
+
+        Some(Ty::from_types(types.into_iter()))
+    }
+
+    fn check_doc_type_ident(&mut self, m: &Module, name: &str) -> Option<Ty> {
+        static TYPE_REPRS: LazyLock<HashMap<&'static str, Ty>> = LazyLock::new(|| {
+            let values = Vec::from_iter(
+                [
+                    Value::None,
+                    Value::Auto,
+                    // Value::Bool(Default::default()),
+                    Value::Int(Default::default()),
+                    Value::Float(Default::default()),
+                    Value::Length(Default::default()),
+                    Value::Angle(Default::default()),
+                    Value::Ratio(Default::default()),
+                    Value::Relative(Default::default()),
+                    Value::Fraction(Default::default()),
+                    Value::Str(Default::default()),
+                ]
+                .map(|v| v.ty())
+                .into_iter()
+                .chain([
+                    Type::of::<typst::visualize::Color>(),
+                    Type::of::<typst::visualize::Gradient>(),
+                    Type::of::<typst::visualize::Pattern>(),
+                    Type::of::<typst::symbols::Symbol>(),
+                    Type::of::<typst::foundations::Version>(),
+                    Type::of::<typst::foundations::Bytes>(),
+                    Type::of::<typst::foundations::Label>(),
+                    Type::of::<typst::foundations::Datetime>(),
+                    Type::of::<typst::foundations::Duration>(),
+                    Type::of::<typst::foundations::Content>(),
+                    Type::of::<typst::foundations::Styles>(),
+                    Type::of::<typst::foundations::Array>(),
+                    Type::of::<typst::foundations::Dict>(),
+                    Type::of::<typst::foundations::Func>(),
+                    Type::of::<typst::foundations::Args>(),
+                    Type::of::<typst::foundations::Type>(),
+                    Type::of::<typst::foundations::Module>(),
+                ]),
+            );
+
+            let shorts = values
+                .clone()
+                .into_iter()
+                .map(|ty| (ty.short_name(), Ty::Builtin(BuiltinTy::Type(ty))));
+            let longs = values
+                .into_iter()
+                .map(|ty| (ty.long_name(), Ty::Builtin(BuiltinTy::Type(ty))));
+            let builtins = [
+                ("any", Ty::Any),
+                ("bool", Ty::Boolean(None)),
+                ("boolean", Ty::Boolean(None)),
+                ("false", Ty::Boolean(Some(false))),
+                ("true", Ty::Boolean(Some(true))),
+            ];
+            HashMap::from_iter(shorts.chain(longs).chain(builtins))
+        });
+
+        let builtin_ty = TYPE_REPRS.get(name).cloned();
+        builtin_ty.or_else(|| self.check_doc_type_anno(m, name))
+    }
+
+    fn check_doc_type_anno(&mut self, m: &Module, name: &str) -> Option<Ty> {
+        if let Some(v) = self.docs_scope.get(name) {
+            return v.clone();
+        }
+
+        let v = m.scope().get(name)?;
+        log::debug!("check doc type annotation: {name:?}");
+        if let Value::Content(c) = v {
+            let annotated = c.clone().unpack::<typst::text::RawElem>().ok()?;
+            let text = annotated.text().clone().into_value().cast::<Str>().ok()?;
+            let code = typst::syntax::parse_code(&text.as_str().replace('\'', "θ"));
+            let mut exprs = code.cast::<ast::Code>()?.exprs();
+            let ret = self.check_doc_type_expr(m, exprs.next()?);
+            self.docs_scope.insert(name.into(), ret.clone());
+            ret
+        } else {
+            None
+        }
+    }
+
+    fn generate_var(&mut self, name: StrRef) -> Ty {
+        self.next_id += 1;
+        let encoded = DefId(self.next_id as u64);
+        log::debug!("generate var {name:?} {encoded:?}");
+        let bounds = TypeVarBounds::new(TypeVar { name, def: encoded }, TypeBounds::default());
+        let var = bounds.as_type();
+        self.vars.insert(encoded, bounds);
+        var
+    }
+
+    fn check_doc_type_expr(&mut self, m: &Module, s: ast::Expr) -> Option<Ty> {
+        log::debug!("check doc type expr: {s:?}");
+        match s {
+            ast::Expr::Ident(i) => self.check_doc_type_ident(m, i.get().as_str()),
+            ast::Expr::Closure(c) => {
+                log::debug!("check doc closure annotation: {c:?}");
+                let mut pos = vec![];
+                let mut named = BTreeMap::new();
+                let mut rest = None;
+
+                for param in c.params().children() {
+                    match param {
+                        ast::Param::Pos(ast::Pattern::Normal(ast::Expr::Ident(i))) => {
+                            let base_ty = self.docs_scope.get(i.get().as_str()).cloned();
+                            pos.push(base_ty.flatten().unwrap_or(Ty::Any));
+                        }
+                        ast::Param::Pos(_) => {
+                            pos.push(Ty::Any);
+                        }
+                        ast::Param::Named(e) => {
+                            let exp = self.check_doc_type_expr(m, e.expr()).unwrap_or(Ty::Any);
+                            named.insert(e.name().into(), exp);
+                        }
+                        // todo: spread left/right
+                        ast::Param::Spread(s) => {
+                            let Some(i) = s.sink_ident() else {
+                                continue;
+                            };
+                            let name = i.get().clone();
+                            let rest_ty = self
+                                .docs_scope
+                                .get(i.get().as_str())
+                                .cloned()
+                                .flatten()
+                                .unwrap_or_else(|| self.generate_var(name.as_str().into()));
+                            self.docs_scope.insert(name, Some(rest_ty.clone()));
+                            rest = Some(rest_ty);
+                        }
+                    }
+                }
+
+                let body = self.check_doc_type_expr(m, c.body())?;
+                let sig = SigTy::new(pos, named, rest, Some(body)).into();
+
+                Some(Ty::Func(sig))
+            }
+            ast::Expr::Dict(d) => {
+                log::debug!("check doc dict annotation: {d:?}");
+                None
+            }
+            _ => None,
+        }
+    }
+}
+
+struct IdRenamer<'a, 'b, 'w> {
+    base: &'a mut TypeChecker<'b, 'w>,
+    var_bounds: &'a HashMap<DefId, TypeVarBounds>,
+    base_id: DefId,
+    offset: u64,
+}
+
+impl<'a, 'b, 'w> TyMutator for IdRenamer<'a, 'b, 'w> {
+    fn mutate(&mut self, ty: &Ty, pol: bool) -> Option<Ty> {
+        match ty {
+            Ty::Var(v) => Some(self.base.copy_based_on(
+                self.var_bounds.get(&v.def).unwrap(),
+                self.offset,
+                self.base_id,
+            )),
+            ty => self.mutate_rec(ty, pol),
+        }
+    }
+}

crates/tinymist-query/src/analysis/tyck/syntax.rs

@@ -3,7 +3,7 @@
 use std::{collections::BTreeMap, sync::LazyLock};
 
 use typst::{
-    foundations::{IntoValue, Module, Str, Type, Value},
+    foundations::Value,
     syntax::{
         ast::{self, AstNode},
         LinkedNode, SyntaxKind,
@@ -11,21 +11,9 @@ use typst::{
 };
 
 use super::*;
-use crate::{
-    adt::interner::Interned,
-    docs::{convert_docs, identify_func_docs},
-    syntax::{find_docs_of, get_non_strict_def_target},
-    ty::*,
-};
+use crate::{adt::interner::Interned, ty::*};
 
-#[derive(Debug)]
-struct ParamDoc {
-    // docs: String,
-    ty: Option<Ty>,
-    // default: Option<String>,
-}
-
-type ParamDocs = HashMap<String, ParamDoc>;
+static EMPTY_DOCSTRING: LazyLock<DocString> = LazyLock::new(DocString::default);
 
 impl<'a, 'w> TypeChecker<'a, 'w> {
     pub(crate) fn check_syntax(&mut self, root: LinkedNode) -> Option<Ty> {
@@ -383,50 +371,11 @@ impl<'a, 'w> TypeChecker<'a, 'w> {
     }
 
     fn check_closure(&mut self, root: LinkedNode<'_>) -> Option<Ty> {
+        let dostring = self.check_closure_docs(&root);
+        let dostring = dostring.as_ref().unwrap_or(&EMPTY_DOCSTRING);
         let closure: ast::Closure = root.cast()?;
 
-        let docs = None.or_else(|| {
-            // use parent of params, todo: reliable way to get the def target
-            let def = get_non_strict_def_target(root.clone())?;
-            find_docs_of(&self.source, def)
-        });
-
-        let parsed = docs.and_then(|docs| {
-            let documenting_id = closure
-                .name()
-                .and_then(|n| self.get_def_id(n.span(), &to_ident_ref(&root, n)?))?;
-
-            let converted = convert_docs(self.ctx.world(), &docs).ok()?;
-            let converted = identify_func_docs(&converted).ok()?;
-            let module = self.ctx.module_by_str(docs)?;
-
-            // Wrap a docs scope
-            self.with_docs_scope(|this| {
-                this.documenting_id = Some(documenting_id);
-                let mut params = ParamDocs::new();
-                for param in converted.params.into_iter() {
-                    params.insert(
-                        param.name,
-                        ParamDoc {
-                            // docs: param.docs,
-                            ty: this.check_doc_types(&module, param.types),
-                            // default: param.default,
-                        },
-                    );
-                }
-
-                Some((
-                    Some(converted.docs),
-                    params,
-                    converted
-                        .return_ty
-                        .map(|ty| this.check_doc_types(&module, ty)),
-                ))
-            })
-        });
-        let (_docs, param_docs, _ret) = parsed.unwrap_or_default();
-
-        log::debug!("check closure: {:?} -> {param_docs:#?}", closure.name());
+        log::debug!("check closure: {:?} -> {dostring:#?}", closure.name());
 
         let mut pos = vec![];
         let mut named = BTreeMap::new();
@@ -436,16 +385,14 @@ impl<'a, 'w> TypeChecker<'a, 'w> {
         for param in closure.params().children() {
             match param {
                 ast::Param::Pos(pattern) => {
-                    pos.push(self.check_pattern(pattern, Ty::Any, &param_docs, root.clone()));
+                    pos.push(self.check_pattern(pattern, Ty::Any, dostring, root.clone()));
                 }
                 ast::Param::Named(e) => {
                     let name = e.name().get();
                     let exp = self.check_expr_in(e.expr().span(), root.clone());
                     let v = self.get_var(e.name().span(), to_ident_ref(&root, e.name())?)?;
-                    let anno = param_docs.get(name.as_str()).and_then(|p| p.ty.clone());
-                    log::debug!("check closure param: {name} with {exp:?} and annotation {anno:?}");
-                    if let Some(anno) = anno {
-                        self.constrain(&v, &anno);
+                    if let Some(annotated) = dostring.var_ty(name.as_str()) {
+                        self.constrain(&v, annotated);
                     }
                     // todo: this is less efficient than v.lbs.push(exp), we may have some idea to
                     // optimize it, so I put a todo here.
@@ -458,9 +405,8 @@ impl<'a, 'w> TypeChecker<'a, 'w> {
                     if let Some(e) = a.sink_ident() {
                         let exp = Ty::Builtin(BuiltinTy::Args);
                         let v = self.get_var(e.span(), to_ident_ref(&root, e)?)?;
-                        let anno = param_docs.get(e.get().as_str()).and_then(|p| p.ty.clone());
-                        if let Some(anno) = anno {
-                            self.constrain(&v, &anno);
+                        if let Some(annotated) = dostring.var_ty(e.get().as_str()) {
+                            self.constrain(&v, annotated);
                         }
                         self.constrain(&exp, &v);
                         rest = Some(v);
@@ -471,6 +417,7 @@ impl<'a, 'w> TypeChecker<'a, 'w> {
         }
 
         let body = self.check_expr_in(closure.body().span(), root);
+        let _ = dostring.res_ty;
 
         let named: Vec<(Interned<str>, Ty)> = named.into_iter().collect();
 
@@ -521,7 +468,7 @@ impl<'a, 'w> TypeChecker<'a, 'w> {
                     .map(|init| self.check_expr_in(init.span(), root.clone()))
                     .unwrap_or_else(|| Ty::Builtin(BuiltinTy::Infer));
 
-                self.check_pattern(pattern, value, &ParamDocs::default(), root.clone());
+                self.check_pattern(pattern, value, &EMPTY_DOCSTRING, root.clone());
             }
         }
 
@@ -627,6 +574,7 @@ impl<'a, 'w> TypeChecker<'a, 'w> {
     fn check_destruct_assign(&mut self, _root: LinkedNode<'_>) -> Option<Ty> {
         Some(Ty::Builtin(BuiltinTy::None))
     }
+
     fn check_expr_in(&mut self, span: Span, root: LinkedNode<'_>) -> Ty {
         root.find(span)
             .map(|node| self.check(node))
@@ -637,10 +585,10 @@ impl<'a, 'w> TypeChecker<'a, 'w> {
         &mut self,
         pattern: ast::Pattern<'_>,
         value: Ty,
-        param_docs: &ParamDocs,
+        docs: &DocString,
         root: LinkedNode<'_>,
     ) -> Ty {
-        self.check_pattern_(pattern, value, param_docs, root)
+        self.check_pattern_(pattern, value, docs, root)
             .unwrap_or(Ty::Builtin(BuiltinTy::Undef))
     }
 
@@ -648,18 +596,16 @@ impl<'a, 'w> TypeChecker<'a, 'w> {
         &mut self,
         pattern: ast::Pattern<'_>,
         value: Ty,
-        param_docs: &ParamDocs,
+        docs: &DocString,
         root: LinkedNode<'_>,
     ) -> Option<Ty> {
         Some(match pattern {
             ast::Pattern::Normal(ast::Expr::Ident(ident)) => {
                 let v = self.get_var(ident.span(), to_ident_ref(&root, ident)?)?;
-                let anno = param_docs
-                    .get(ident.get().as_str())
-                    .and_then(|p| p.ty.clone());
-                log::debug!("check pattern: {ident:?} with {value:?} and annotation {anno:?}");
-                if let Some(anno) = anno {
-                    self.constrain(&v, &anno);
+                let annotated = docs.var_ty(ident.get().as_str());
+                log::debug!("check pattern: {ident:?} with {value:?} and annotation {annotated:?}");
+                if let Some(annotated) = annotated {
+                    self.constrain(&v, annotated);
                 }
                 self.constrain(&value, &v);
                 v
@@ -667,161 +613,10 @@ impl<'a, 'w> TypeChecker<'a, 'w> {
             ast::Pattern::Normal(_) => Ty::Any,
             ast::Pattern::Placeholder(_) => Ty::Any,
             ast::Pattern::Parenthesized(exp) => {
-                self.check_pattern(exp.pattern(), value, param_docs, root)
+                self.check_pattern(exp.pattern(), value, docs, root)
             }
             // todo: pattern
             ast::Pattern::Destructuring(_destruct) => Ty::Any,
         })
     }
-
-    fn check_doc_types(&mut self, m: &Module, strs: String) -> Option<Ty> {
-        let mut types = vec![];
-        for name in strs.split(",").map(|e| e.trim()) {
-            let Some(ty) = self.check_doc_type_ident(m, name) else {
-                continue;
-            };
-            types.push(ty);
-        }
-
-        Some(Ty::from_types(types.into_iter()))
-    }
-
-    fn check_doc_type_ident(&mut self, m: &Module, name: &str) -> Option<Ty> {
-        static TYPE_REPRS: LazyLock<HashMap<&'static str, Ty>> = LazyLock::new(|| {
-            let values = Vec::from_iter(
-                [
-                    Value::None,
-                    Value::Auto,
-                    // Value::Bool(Default::default()),
-                    Value::Int(Default::default()),
-                    Value::Float(Default::default()),
-                    Value::Length(Default::default()),
-                    Value::Angle(Default::default()),
-                    Value::Ratio(Default::default()),
-                    Value::Relative(Default::default()),
-                    Value::Fraction(Default::default()),
-                    Value::Str(Default::default()),
-                ]
-                .map(|v| v.ty())
-                .into_iter()
-                .chain([
-                    Type::of::<typst::visualize::Color>(),
-                    Type::of::<typst::visualize::Gradient>(),
-                    Type::of::<typst::visualize::Pattern>(),
-                    Type::of::<typst::symbols::Symbol>(),
-                    Type::of::<typst::foundations::Version>(),
-                    Type::of::<typst::foundations::Bytes>(),
-                    Type::of::<typst::foundations::Label>(),
-                    Type::of::<typst::foundations::Datetime>(),
-                    Type::of::<typst::foundations::Duration>(),
-                    Type::of::<typst::foundations::Content>(),
-                    Type::of::<typst::foundations::Styles>(),
-                    Type::of::<typst::foundations::Array>(),
-                    Type::of::<typst::foundations::Dict>(),
-                    Type::of::<typst::foundations::Func>(),
-                    Type::of::<typst::foundations::Args>(),
-                    Type::of::<typst::foundations::Type>(),
-                    Type::of::<typst::foundations::Module>(),
-                ]),
-            );
-
-            let shorts = values
-                .clone()
-                .into_iter()
-                .map(|ty| (ty.short_name(), Ty::Builtin(BuiltinTy::Type(ty))));
-            let longs = values
-                .into_iter()
-                .map(|ty| (ty.long_name(), Ty::Builtin(BuiltinTy::Type(ty))));
-            let builtins = [
-                ("any", Ty::Any),
-                ("bool", Ty::Boolean(None)),
-                ("boolean", Ty::Boolean(None)),
-                ("false", Ty::Boolean(Some(false))),
-                ("true", Ty::Boolean(Some(true))),
-            ];
-            HashMap::from_iter(shorts.chain(longs).chain(builtins))
-        });
-
-        let builtin_ty = TYPE_REPRS.get(name).cloned();
-        builtin_ty.or_else(|| self.check_doc_type_anno(m, name))
-    }
-
-    fn check_doc_type_anno(&mut self, m: &Module, name: &str) -> Option<Ty> {
-        if let Some(v) = self.docs_scope.get(name) {
-            return v.clone();
-        }
-
-        let v = m.scope().get(name)?;
-        log::debug!("check doc type annotation: {name:?}");
-        if let Value::Content(c) = v {
-            let anno = c.clone().unpack::<typst::text::RawElem>().ok()?;
-            let text = anno.text().clone().into_value().cast::<Str>().ok()?;
-            let code = typst::syntax::parse_code(&text.as_str().replace('\'', "θ"));
-            let mut exprs = code.cast::<ast::Code>()?.exprs();
-            let ret = self.check_doc_type_expr(m, exprs.next()?);
-            self.docs_scope.insert(name.into(), ret.clone());
-            ret
-        } else {
-            None
-        }
-    }
-
-    fn check_doc_type_expr(&mut self, m: &Module, s: ast::Expr) -> Option<Ty> {
-        log::debug!("check doc type expr: {s:?}");
-        match s {
-            ast::Expr::Ident(i) => self.check_doc_type_ident(m, i.get().as_str()),
-            ast::Expr::Closure(c) => {
-                log::debug!("check doc closure annotation: {c:?}");
-                let mut pos = vec![];
-                let mut named = BTreeMap::new();
-                let mut rest = None;
-
-                for param in c.params().children() {
-                    match param {
-                        ast::Param::Pos(ast::Pattern::Normal(ast::Expr::Ident(i))) => {
-                            let base_ty = self.docs_scope.get(i.get().as_str()).cloned();
-                            pos.push(base_ty.flatten().unwrap_or(Ty::Any));
-                        }
-                        ast::Param::Pos(_) => {
-                            pos.push(Ty::Any);
-                        }
-                        ast::Param::Named(e) => {
-                            let exp = self.check_doc_type_expr(m, e.expr()).unwrap_or(Ty::Any);
-                            named.insert(e.name().into(), exp);
-                        }
-                        // todo: spread left/right
-                        ast::Param::Spread(s) => {
-                            let Some(i) = s.sink_ident() else {
-                                continue;
-                            };
-                            let name = i.get().clone();
-                            let rest_ty = self
-                                .docs_scope
-                                .get(i.get().as_str())
-                                .cloned()
-                                .flatten()
-                                .unwrap_or_else(|| {
-                                    self.generate_var(
-                                        name.as_str().into(),
-                                        self.documenting_id.unwrap(),
-                                    )
-                                });
-                            self.docs_scope.insert(name, Some(rest_ty.clone()));
-                            rest = Some(rest_ty);
-                        }
-                    }
-                }
-
-                let body = self.check_doc_type_expr(m, c.body())?;
-                let sig = SigTy::new(pos, named, rest, Some(body)).into();
-
-                Some(Ty::Func(sig))
-            }
-            ast::Expr::Dict(d) => {
-                log::debug!("check doc dict annotation: {d:?}");
-                None
-            }
-            _ => None,
-        }
-    }
 }

crates/tinymist-query/src/docs/mod.rs

@@ -23,12 +23,12 @@ use typst::syntax::package::{PackageManifest, PackageSpec};
 use typst::syntax::{FileId, Span, VirtualPath};
 use typst::World;
 
-pub(crate) use self::tidy::*;
 use crate::analysis::analyze_dyn_signature;
 use crate::syntax::{find_docs_of, get_non_strict_def_target, IdentRef};
 use crate::ty::Ty;
 use crate::upstream::truncated_doc_repr;
 use crate::AnalysisContext;
+pub(crate) use tidy::*;
 
 /// Information about a package.
 #[derive(Debug, Serialize, Deserialize)]
@@ -54,6 +54,13 @@ impl From<(PathBuf, PackageSpec)> for PackageInfo {
     }
 }
 
+/// Kind of a docstring.
+#[derive(Debug, Clone, Hash, Serialize, Deserialize)]
+pub enum DocStringKind {
+    /// A docstring for a function.
+    Function,
+}
+
 /// Docs about a symbol.
 #[derive(Debug, Clone, Serialize, Deserialize)]
 #[serde(tag = "kind")]

crates/tinymist-query/src/docs/tidy.rs

@@ -1,19 +1,20 @@
+use ecow::EcoString;
 use itertools::Itertools;
 use serde::{Deserialize, Serialize};
 use typst::diag::StrResult;
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct TidyParamDocs {
-    pub name: String,
-    pub docs: String,
-    pub types: String,
-    pub default: Option<String>,
+    pub name: EcoString,
+    pub docs: EcoString,
+    pub types: EcoString,
+    pub default: Option<EcoString>,
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct TidyFuncDocs {
     pub docs: String,
-    pub return_ty: Option<String>,
+    pub return_ty: Option<EcoString>,
     pub params: Vec<TidyParamDocs>,
 }
 
@@ -57,7 +58,7 @@ pub fn identify_func_docs(converted: &str) -> StrResult<TidyFuncDocs> {
                     continue;
                 };
 
-                return_ty = Some(w.trim().to_string());
+                return_ty = Some(w.trim().into());
                 break;
             }
 
@@ -120,7 +121,7 @@ pub fn identify_func_docs(converted: &str) -> StrResult<TidyFuncDocs> {
                 name: param_line.0,
                 types: param_line.1,
                 default: None,
-                docs: buf.into_iter().join("\n"),
+                docs: buf.into_iter().join("\n").into(),
             });
 
             break;

crates/tinymist-query/src/ty/def.rs

@@ -1057,8 +1057,8 @@ impl TypeVarBounds {
     }
 
     /// Get the name of the type variable
-    pub fn name(&self) -> StrRef {
-        self.var.name.clone()
+    pub fn name(&self) -> &StrRef {
+        &self.var.name
     }
     /// Get the definition id of the type variable
     pub fn id(&self) -> DefId {