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 var_bounds: HashMap<DeclExpr, TypeVarBounds>,
47 globals: HashMap<EcoString, Option<Ty>>,
49 locals: SnapshotMap<EcoString, Ty>,
51 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, ¶m.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::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 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}