tinymist_query/
code_context.rs

1use std::ops::Deref;
2
3use comemo::Track;
4use serde::{Deserialize, Serialize};
5use tinymist_analysis::analyze_expr;
6use tinymist_project::{DiagnosticFormat, PathPattern};
7use tinymist_std::error::prelude::*;
8use tinymist_world::vfs::WorkspaceResolver;
9use tinymist_world::{EntryReader, EntryState, ShadowApi, diag::print_diagnostics_to_string};
10use typst::diag::{At, SourceResult};
11use typst::foundations::{Args, Dict, NativeFunc, eco_format};
12use typst::syntax::RangeMapper;
13use typst::utils::LazyHash;
14use typst::{
15    foundations::{Bytes, IntoValue, StyleChain},
16    text::TextElem,
17};
18use typst_shim::eval::{Eval, Vm};
19use typst_shim::syntax::LinkedNodeExt;
20
21use crate::{
22    prelude::*,
23    syntax::{InterpretMode, interpret_mode_at},
24};
25
26/// A query to get the mode at a specific position in a text document.
27#[derive(Debug, Clone, Deserialize)]
28#[serde(tag = "kind", rename_all = "camelCase")]
29pub enum InteractCodeContextQuery {
30    /// (Experimental) Evaluate a path expression at a specific position in a
31    /// text document.
32    PathAt {
33        /// Code to evaluate. If the code starts with `{` and ends with `}`, it
34        /// will be evaluated as a code expression, otherwise it will be
35        /// evaluated as a path pattern.
36        ///
37        /// ## Example
38        ///
39        /// evaluate a path pattern, which could use following definitions:
40        ///
41        /// ```plain
42        /// $root/x/$dir/../$name // is evaluated as
43        /// /path/to/root/x/dir/../main
44        /// ```
45        ///
46        /// ## Example
47        ///
48        /// evaluate a code expression, which could use following definitions:
49        /// - `root`: the root of the workspace
50        /// - `dir`: the directory of the current file
51        /// - `name`: the name of the current file
52        /// - `join(a, b, ...)`: join the arguments with the path separator
53        ///
54        /// ```plain
55        /// { join(root, "x", dir, "y", name) } // is evaluated as
56        /// /path/to/root/x/dir/y/main
57        /// ```
58        code: String,
59        /// The extra `sys.inputs` for the code expression.
60        inputs: Dict,
61    },
62    /// Get the mode at a specific position in a text document.
63    ModeAt {
64        /// The position inside the text document.
65        position: LspPosition,
66    },
67    /// Get the style at a specific position in a text document.
68    StyleAt {
69        /// The position inside the text document.
70        position: LspPosition,
71        /// Style to query
72        style: Vec<String>,
73    },
74}
75
76/// A response to a `InteractCodeContextQuery`.
77#[derive(Debug, Clone, Serialize, Deserialize)]
78#[serde(tag = "kind", rename_all = "camelCase")]
79pub enum InteractCodeContextResponse {
80    /// Evaluate a path expression at a specific position in a text document.
81    PathAt(QueryResult<serde_json::Value>),
82    /// Get the mode at a specific position in a text document.
83    ModeAt {
84        /// The mode at the requested position.
85        mode: InterpretMode,
86    },
87    /// Get the style at a specific position in a text document.
88    StyleAt {
89        /// The style at the requested position.
90        style: Vec<Option<JsonValue>>,
91    },
92}
93
94/// A request to get the code context of a text document.
95#[derive(Debug, Clone, Deserialize)]
96#[serde(tag = "kind")]
97pub struct InteractCodeContextRequest {
98    /// The path to the text document.
99    pub path: PathBuf,
100    /// The queries to execute.
101    pub query: Vec<Option<InteractCodeContextQuery>>,
102}
103
104impl SemanticRequest for InteractCodeContextRequest {
105    type Response = Vec<Option<InteractCodeContextResponse>>;
106
107    fn request(self, ctx: &mut LocalContext) -> Option<Self::Response> {
108        let mut responses = Vec::new();
109
110        let source = ctx.source_by_path(&self.path).ok()?;
111
112        for query in self.query {
113            responses.push(query.and_then(|query| match query {
114                InteractCodeContextQuery::PathAt { code, inputs: base } => {
115                    let res = eval_path_expr(ctx, &code, base)?;
116                    Some(InteractCodeContextResponse::PathAt(res))
117                }
118                InteractCodeContextQuery::ModeAt { position } => {
119                    let cursor = ctx.to_typst_pos(position, &source)?;
120                    let mode = Self::mode_at(&source, cursor)?;
121                    Some(InteractCodeContextResponse::ModeAt { mode })
122                }
123                InteractCodeContextQuery::StyleAt { position, style } => {
124                    let mut world = ctx.world().clone();
125                    log::info!(
126                        "style at position {position:?} . {style:?} when main is {:?}",
127                        world.main()
128                    );
129                    let cursor = ctx.to_typst_pos(position, &source)?;
130                    let root = LinkedNode::new(source.root());
131                    let mut leaf = root.leaf_at_compat(cursor)?;
132                    log::info!("style at leaf {leaf:?} . {style:?}");
133
134                    if !matches!(leaf.kind(), SyntaxKind::Text | SyntaxKind::MathText) {
135                        return None;
136                    }
137
138                    if matches!(leaf.parent_kind(), Some(SyntaxKind::Raw)) {
139                        leaf = leaf.parent()?.clone();
140                    }
141
142                    let mode = Self::mode_at(&source, cursor);
143                    if !matches!(
144                        mode,
145                        Some(InterpretMode::Code | InterpretMode::Markup | InterpretMode::Math)
146                    ) {
147                        leaf = leaf.parent()?.clone();
148                    }
149                    let mut mapped_source = source.clone();
150                    let (with, offset) = match mode {
151                        Some(InterpretMode::Code) => ("context text.font", 8),
152                        _ => ("#context text.font", 10),
153                    };
154                    let start = leaf.range().start;
155                    mapped_source.edit(leaf.range(), with);
156
157                    let _ = world.map_shadow_by_id(
158                        mapped_source.id(),
159                        Bytes::new(mapped_source.text().as_bytes().to_vec()),
160                    );
161                    world.take_db();
162
163                    let root = LinkedNode::new(mapped_source.root());
164                    let leaf = root.leaf_at_compat(start + offset)?;
165
166                    log::info!("style at new_leaf {leaf:?} . {style:?}");
167
168                    let mut cursor_styles = analyze_expr(&world, &leaf)
169                        .iter()
170                        .filter_map(|s| s.1.clone())
171                        .collect::<Vec<_>>();
172                    cursor_styles.sort_by_key(|x| x.as_slice().len());
173                    log::info!("style at styles {cursor_styles:?} . {style:?}");
174                    let cursor_style = cursor_styles.into_iter().next_back().unwrap_or_default();
175                    let cursor_style = StyleChain::new(&cursor_style);
176
177                    log::info!("style at style {cursor_style:?} . {style:?}");
178
179                    let style = style
180                        .iter()
181                        .map(|style| Self::style_at(cursor_style, style))
182                        .collect();
183                    let _ = world.map_shadow_by_id(
184                        mapped_source.id(),
185                        Bytes::new(source.text().as_bytes().to_vec()),
186                    );
187
188                    Some(InteractCodeContextResponse::StyleAt { style })
189                }
190            }));
191        }
192
193        Some(responses)
194    }
195}
196
197impl InteractCodeContextRequest {
198    fn mode_at(source: &Source, pos: usize) -> Option<InterpretMode> {
199        // Smart special cases that is definitely at markup
200        if pos == 0 || pos >= source.text().len() {
201            return Some(InterpretMode::Markup);
202        }
203
204        // Get mode
205        let root = LinkedNode::new(source.root());
206        Some(interpret_mode_at(root.leaf_at_compat(pos).as_ref()))
207    }
208
209    fn style_at(cursor_style: StyleChain, style: &str) -> Option<JsonValue> {
210        match style {
211            "text.font" => {
212                let font = cursor_style.get_cloned(TextElem::font).into_value();
213                serde_json::to_value(font).ok()
214            }
215            _ => None,
216        }
217    }
218}
219
220fn eval_path_expr(
221    ctx: &mut LocalContext,
222    code: &str,
223    inputs: Dict,
224) -> Option<QueryResult<serde_json::Value>> {
225    let entry = ctx.world().entry_state();
226    let path = if code.starts_with("{") && code.ends_with("}") {
227        let id = entry
228            .select_in_workspace(Path::new("/__path__.typ"))
229            .main()?;
230
231        let inputs = make_sys(&entry, ctx.world().inputs(), inputs);
232        let (inputs, root, dir, name) = match inputs {
233            Some(EvalSysCtx {
234                inputs,
235                root,
236                dir,
237                name,
238            }) => (Some(inputs), Some(root), dir, Some(name)),
239            None => (None, None, None, None),
240        };
241
242        let mut world = ctx.world().task(tinymist_world::TaskInputs {
243            entry: None,
244            inputs,
245        });
246        // todo: bad performance
247        world.take_db();
248        let _ = world.map_shadow_by_id(id, Bytes::from_string(code.to_owned()));
249
250        tinymist_analysis::upstream::with_vm((&world as &dyn World).track(), |vm| {
251            define_val(vm, "join", Value::Func(join::data().into()));
252            for (key, value) in [("root", root), ("dir", dir), ("name", name)] {
253                if let Some(value) = value {
254                    define_val(vm, key, value);
255                }
256            }
257
258            let mut expr = typst::syntax::parse_code(code);
259            let mapper = RangeMapper::new(std::iter::once(0..code.len()))
260                .context("invalid code range mapper")?;
261            expr.synthesize_mapped(id, &mapper)
262                .context("failed to map code span")?;
263
264            let expr = match expr.cast::<ast::Code>() {
265                Some(v) => v,
266                None => bail!(
267                    "code is not a valid code expression: kind={:?}",
268                    expr.kind()
269                ),
270            };
271            match expr.eval(vm) {
272                Ok(value) => serde_json::to_value(value).context_ut("failed to serialize path"),
273                Err(e) => {
274                    let res =
275                        print_diagnostics_to_string(&world, e.iter(), DiagnosticFormat::Human);
276                    let err = res.unwrap_or_else(|e| e);
277                    bail!("failed to evaluate path expression: {err}")
278                }
279            }
280        })
281    } else {
282        PathPattern::new(code)
283            .substitute(&entry)
284            .context_ut("failed to substitute path pattern")
285            .and_then(|path| {
286                serde_json::to_value(path.deref()).context_ut("failed to serialize path")
287            })
288    };
289    Some(path.into())
290}
291
292#[derive(Debug, Clone, Hash)]
293struct EvalSysCtx {
294    inputs: Arc<LazyHash<Dict>>,
295    root: Value,
296    dir: Option<Value>,
297    name: Value,
298}
299
300#[comemo::memoize]
301fn make_sys(entry: &EntryState, base: Arc<LazyHash<Dict>>, inputs: Dict) -> Option<EvalSysCtx> {
302    let root = entry.root();
303    let main = entry.main();
304
305    log::debug!("Check path {main:?} and root {root:?}");
306
307    let (root, main) = root.zip(main)?;
308
309    // Files in packages are not exported
310    if WorkspaceResolver::is_package_file(main) {
311        return None;
312    }
313    // Files without a path are not exported
314    let path = main.vpath().realize(&root).ok()?;
315
316    // todo: handle untitled path
317    if path.strip_prefix("/untitled").is_ok() {
318        return None;
319    }
320
321    let path = path.strip_prefix(&root).ok()?;
322    let dir = path.parent();
323    let file_name = path.file_name().unwrap_or_default();
324
325    let root = Value::Str(root.to_string_lossy().into());
326
327    let dir = dir.map(|d| Value::Str(d.to_string_lossy().into()));
328
329    let name = file_name.to_string_lossy();
330    let name = name.as_ref().strip_suffix(".typ").unwrap_or(name.as_ref());
331    let name = Value::Str(name.into());
332
333    let mut dict = base.as_ref().deref().clone();
334    for (key, value) in inputs {
335        dict.insert(key, value);
336    }
337    dict.insert("root".into(), root.clone());
338    if let Some(dir) = &dir {
339        dict.insert("dir".into(), dir.clone());
340    }
341    dict.insert("name".into(), name.clone());
342
343    Some(EvalSysCtx {
344        inputs: Arc::new(LazyHash::new(dict)),
345        root,
346        dir,
347        name,
348    })
349}
350
351fn define_val(vm: &mut Vm, name: &str, value: Value) {
352    let ident = SyntaxNode::leaf(SyntaxKind::Ident, name);
353    vm.define(ident.cast::<ast::Ident>().unwrap(), value);
354}
355
356#[typst_macros::func(title = "Join function")]
357fn join(args: &mut Args) -> SourceResult<Value> {
358    let pos = args.take().to_pos();
359    let mut res = PathBuf::new();
360    for arg in pos {
361        match arg {
362            Value::Str(s) => res.push(s.as_str()),
363            _ => {
364                return Err(eco_format!("join argument is not a string: {arg:?}")).at(args.span);
365            }
366        };
367    }
368    Ok(Value::Str(res.to_string_lossy().into()))
369}
370
371/// A result of a query.
372#[derive(Debug, Clone, Serialize, Deserialize)]
373#[serde(untagged)]
374pub enum QueryResult<T> {
375    /// A successful result.
376    Success {
377        /// The value of the result.
378        value: T,
379    },
380    /// An error result.
381    Error {
382        /// The error message.
383        error: EcoString,
384    },
385}
386
387impl<T> QueryResult<T> {
388    /// Creates a successful result.
389    pub fn success(value: T) -> Self {
390        Self::Success { value }
391    }
392
393    /// Creates an error result.
394    pub fn error(error: EcoString) -> Self {
395        Self::Error { error }
396    }
397}
398
399impl<T, E: std::error::Error> From<Result<T, E>> for QueryResult<T> {
400    fn from(value: Result<T, E>) -> Self {
401        match value {
402            Ok(value) => QueryResult::success(value),
403            Err(error) => QueryResult::error(eco_format!("{error}")),
404        }
405    }
406}
407
408#[cfg(test)]
409mod tests {
410    use typst::foundations::dict;
411
412    use super::*;
413    use crate::tests::*;
414
415    #[test]
416    fn test() {
417        snapshot_testing("code_context_path_at", &|ctx, path| {
418            let patterns = [
419                "$root/$dir/$name",
420                "$root/$name",
421                "$root/assets",
422                "$root/assets/$name",
423                r#"{ join(root, "x", dir, "y", name) }"#,
424                r#"{ join(root, 1) }"#,
425                r#"{ join(roo, 1) }"#,
426            ];
427            let inp = [
428                dict! {
429                    "x-path-context" => "vscode-paste",
430                    "x-path-input-uri" => "https://huh.io/img.png",
431                    "x-path-input-name" => "img.png",
432                },
433                dict! {
434                    "x-path-context" => "vscode-paste",
435                    "x-path-input-uri" => "https://huh.io/text.md",
436                    "x-path-input-name" => "text.md",
437                },
438            ];
439
440            let cases = patterns
441                .iter()
442                .map(|pat| (*pat, inp[0].clone()))
443                .chain(inp.iter().map(|inp| {
444                    (
445                        r#"{ import "/resolve.typ": resolve; resolve(join, root, dir, name) }"#,
446                        inp.clone(),
447                    )
448                }));
449
450            let result = cases
451                .map(|(code, inputs)| {
452                    let request = InteractCodeContextRequest {
453                        path: path.clone(),
454                        query: vec![Some(InteractCodeContextQuery::PathAt {
455                            code: code.to_string(),
456                            inputs: inputs.clone(),
457                        })],
458                    };
459                    json!({ "code": code, "inputs": inputs, "response": request.request(ctx) })
460                })
461                .collect::<Vec<_>>();
462            assert_snapshot!(JsonRepr::new_redacted(result, &REDACT_LOC));
463        });
464    }
465}
tinymist_query/code_context.rs

tinymist_query/
code_context.rs
