Arbitrary arguments for typst_macros::time

Cargo.lock

@@ -3247,6 +3247,7 @@ dependencies = [
 name = "typst-timing"
 version = "0.13.1"
 dependencies = [
+ "indexmap 2.7.1",
  "parking_lot",
  "serde",
  "serde_json",

crates/typst-macros/src/lib.rs

@@ -333,9 +333,18 @@ pub fn derive_cast(item: BoundaryStream) -> BoundaryStream {
 /// invocations of the function and store them in a global map. The map can be
 /// accessed through the `typst_trace::RECORDER` static.
 ///
-/// You can also specify the span of the function invocation:
+/// You can also specify the following arguments (in that order):
-/// - `#[time(span = ..)]` to record the span, which will be used for the
+/// - the name of the event type using `#[time(name = ..)]`, by default uses the
-///   `EventKey`.
+///   function's name.
+/// - the span of the function declaration using `#[time(span = ..)]`.
+/// - the span of the callsite using `#[time(callsite = ..)]`.
+/// - the name of the function being called using `#[time(func = ..)]`.
+/// - any extra arguments, which must be at the end of the argument list, they
+///   can be of four types:
+///   - serialize arguments `#[time(<name> = <value>)]`
+///   - debug arguments `#[time(<name> = ?<value>)]`
+///   - display arguments `#[time(<name> = #<value>)]`
+///   - span arguments `#[time(<name> = $<value>)]`
 ///
 /// By default, all tracing is omitted using the `wasm32` target flag.
 /// This is done to avoid bloating the web app, which doesn't need tracing.

crates/typst-macros/src/time.rs

@@ -1,9 +1,9 @@
 use proc_macro2::TokenStream;
-use quote::{quote, ToTokens};
+use quote::{format_ident, quote, ToTokens};
 use syn::parse::{Parse, ParseStream};
 use syn::{parse_quote, Result};
 
-use crate::util::{kw, parse_key_value, parse_string};
+use crate::util::{eat_comma, kw, parse_key_value, parse_string};
 
 /// Expand the `#[time(..)]` macro.
 pub fn time(stream: TokenStream, item: syn::ItemFn) -> Result<TokenStream> {
@@ -13,10 +13,11 @@ pub fn time(stream: TokenStream, item: syn::ItemFn) -> Result<TokenStream> {
 
 /// The `..` in `#[time(..)]`.
 pub struct Meta {
+    pub name: Option<String>,
     pub span: Option<syn::Expr>,
     pub callsite: Option<syn::Expr>,
-    pub name: Option<String>,
     pub func: Option<syn::Expr>,
+    pub extras: Vec<(String, Mode, syn::Expr)>,
 }
 
 impl Parse for Meta {
@@ -26,41 +27,92 @@ impl Parse for Meta {
             span: parse_key_value::<kw::span, syn::Expr>(input)?,
             callsite: parse_key_value::<kw::callsite, syn::Expr>(input)?,
             func: parse_key_value::<kw::func, syn::Expr>(input)?,
+            extras: {
+                let mut pairs = Vec::new();
+                while input.peek(syn::Ident) {
+                    let key: syn::Ident = input.parse()?;
+                    let _: syn::Token![=] = input.parse()?;
+
+                    // Get the mode of this extra argument.
+                    let mode = Mode::parse(input)?;
+
+                    let value = input.parse()?;
+                    eat_comma(input);
+
+                    pairs.push((key.to_string(), mode, value));
+                }
+                pairs
+            },
         })
     }
 }
 
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+
+pub enum Mode {
+    Span,
+    Serialize,
+    Debug,
+    Display,
+}
+
+impl Parse for Mode {
+    fn parse(input: ParseStream) -> Result<Self> {
+        if input.peek(syn::Token![$]) {
+            input.parse::<syn::Token![$]>()?;
+            Ok(Self::Span)
+        } else if input.peek(syn::Token![?]) {
+            input.parse::<syn::Token![?]>()?;
+            Ok(Self::Debug)
+        } else if input.peek(syn::Token![#]) {
+            input.parse::<syn::Token![#]>()?;
+            Ok(Self::Display)
+        } else {
+            Ok(Self::Serialize)
+        }
+    }
+}
+
 fn create(meta: Meta, mut item: syn::ItemFn) -> Result<TokenStream> {
     let name = meta.name.unwrap_or_else(|| item.sig.ident.to_string());
-
+    let mut extras = Vec::new();
-    let func = match meta.func {
+    if let Some(func) = &meta.func {
-        Some(func) => {
+        extras.push(quote! { .with_func(#func) });
-            if meta.callsite.is_none() {
-                bail!(func, "the `func` argument can only be used with a callsite")
     }
 
-            quote! { Some(#func.into()) }
+    if let Some(span) = &meta.span {
+        extras.push(quote! { .with_span(#span.into_raw()) });
     }
-        None => quote! { None },
+
+    if let Some(callsite) = &meta.callsite {
+        extras.push(quote! { .with_callsite(#callsite.into_raw()) });
+    }
+
+    for (key, mode, value) in &meta.extras {
+        let (method, transform) = match mode {
+            Mode::Span => {
+                (format_ident!("with_named_span"), Some(quote! { .into_raw() }))
+            }
+            Mode::Debug => (format_ident!("with_debug"), None),
+            Mode::Display => (format_ident!("with_display"), None),
+            Mode::Serialize => (format_ident!("with_arg"), None),
         };
 
-    let construct = match (meta.span.as_ref(), meta.callsite.as_ref()) {
+        extras.push(quote! { .#method(#key, (#value) #transform) });
-        (Some(span), Some(callsite)) => quote! {
+        if matches!(mode, Mode::Serialize) {
-            with_callsite(#name, Some(#span.into_raw()), Some(#callsite.into_raw()), #func)
+            let error_msg = format!("failed to serialize {key}");
-        },
+            extras.push(quote! { .expect(#error_msg) })
-        (Some(span), None) => quote! {
+        }
-            with_span(#name, Some(#span.into_raw()))
-        },
-        (None, Some(expr)) => {
-            bail!(expr, "cannot have a callsite span without a main span")
     }
-        (None, None) => quote! { new(#name) },
-    };
 
     item.block.stmts.insert(
         0,
         parse_quote! {
-            let __scope = ::typst_timing::TimingScope::#construct;
+            let __scope = ::typst_timing::TimingScope::new(#name).map(|__scope| {
+                __scope
+                    #(#extras)*
+                    .build()
+            });
         },
     );
 

crates/typst-timing/Cargo.toml

@@ -13,6 +13,7 @@ keywords = { workspace = true }
 readme = { workspace = true }
 
 [dependencies]
+indexmap = { workspace = true }
 parking_lot = { workspace = true }
 serde = { workspace = true }
 serde_json = { workspace = true }

crates/typst-timing/src/lib.rs

@@ -1,9 +1,14 @@
 //! Performance timing for Typst.
 
+use std::borrow::Cow;
+use std::fmt::Display;
 use std::io::Write;
 use std::num::NonZeroU64;
+use std::ops::Not;
 use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
+use std::sync::Arc;
 
+use indexmap::IndexMap;
 use parking_lot::Mutex;
 use serde::ser::SerializeSeq;
 use serde::{Serialize, Serializer};
@@ -100,19 +105,8 @@ pub fn export_json<W: Write>(
         ts: f64,
         pid: u64,
         tid: u64,
-        args: Option<Args<'a>>,
-    }
-
-    #[derive(Serialize)]
-    struct Args<'a> {
-        file: String,
-        line: u32,
         #[serde(skip_serializing_if = "Option::is_none")]
-        function_name: Option<&'a str>,
+        args: Option<IndexMap<Cow<'a, str>, Cow<'a, serde_json::Value>>>,
-        #[serde(skip_serializing_if = "Option::is_none")]
-        callsite_file: Option<String>,
-        #[serde(skip_serializing_if = "Option::is_none")]
-        callsite_line: Option<u32>,
     }
 
     let lock = EVENTS.lock();
@@ -124,6 +118,16 @@ pub fn export_json<W: Write>(
         .map_err(|e| format!("failed to serialize events: {e}"))?;
 
     for event in events.iter() {
+        let mut args = IndexMap::new();
+        if let Some(func) = event.func.as_ref() {
+            args.insert("func".into(), Cow::Owned(serde_json::json!(func)));
+        }
+
+        for (key, arg) in event.arguments.iter() {
+            arg.to_json(&mut source, key, &mut args)
+                .map_err(|e| format!("failed to serialize event argument: {e}"))?;
+        }
+
         seq.serialize_element(&Entry {
             name: event.name,
             cat: "typst",
@@ -134,21 +138,7 @@ pub fn export_json<W: Write>(
             ts: event.timestamp.micros_since(events[0].timestamp),
             pid: 1,
             tid: event.thread_id,
-            args: event.span.map(&mut source).map(|(file, line)| {
+            args: args.is_empty().not().then_some(args),
-                let (callsite_file, callsite_line) = match event.callsite.map(&mut source)
-                {
-                    Some((a, b)) => (Some(a), Some(b)),
-                    None => (None, None),
-                };
-
-                Args {
-                    file,
-                    line,
-                    callsite_file,
-                    callsite_line,
-                    function_name: event.func.as_deref(),
-                }
-            }),
         })
         .map_err(|e| format!("failed to serialize event: {e}"))?;
     }
@@ -159,87 +149,155 @@ pub fn export_json<W: Write>(
 }
 
 /// A scope that records an event when it is dropped.
+#[must_use]
 pub struct TimingScope {
     name: &'static str,
-    span: Option<NonZeroU64>,
-    callsite: Option<NonZeroU64>,
     func: Option<String>,
-    thread_id: u64,
+    args: IndexMap<&'static str, EventArgument>,
 }
 
 impl TimingScope {
     /// Create a new scope if timing is enabled.
     #[inline]
     pub fn new(name: &'static str) -> Option<Self> {
-        Self::with_span(name, None)
-    }
-
-    /// Create a new scope with a span if timing is enabled.
-    ///
-    /// The span is a raw number because `typst-timing` can't depend on
-    /// `typst-syntax` (or else `typst-syntax` couldn't depend on
-    /// `typst-timing`).
-    #[inline]
-    pub fn with_span(name: &'static str, span: Option<NonZeroU64>) -> Option<Self> {
-        Self::with_callsite(name, span, None, None)
-    }
-
-    /// Create a new scope with a span if timing is enabled.
-    ///
-    /// The span is a raw number because `typst-timing` can't depend on
-    /// `typst-syntax` (or else `typst-syntax` couldn't depend on
-    /// `typst-timing`).
-    #[inline]
-    pub fn with_callsite(
-        name: &'static str,
-        span: Option<NonZeroU64>,
-        callsite: Option<NonZeroU64>,
-        func: Option<String>,
-    ) -> Option<Self> {
         if is_enabled() {
-            return Some(Self::new_impl(name, span, callsite, func));
+            Some(Self { name, func: None, args: IndexMap::new() })
-        }
+        } else {
             None
         }
+    }
+
+    pub fn with_func(mut self, func: impl ToString) -> Self {
+        self.func = Some(func.to_string());
+        self
+    }
+
+    pub fn with_span(mut self, span: NonZeroU64) -> Self {
+        self.args.insert("span", EventArgument::Span(span));
+        self
+    }
+
+    pub fn with_callsite(mut self, callsite: NonZeroU64) -> Self {
+        self.args.insert("callsite", EventArgument::Span(callsite));
+        self
+    }
+
+    pub fn with_named_span(mut self, name: &'static str, span: NonZeroU64) -> Self {
+        self.args.insert(name, EventArgument::Span(span));
+        self
+    }
+
+    pub fn with_display(self, name: &'static str, value: impl Display) -> Self {
+        self.with_arg(name, value.to_string())
+            .expect("failed to serialize display value")
+    }
+
+    pub fn with_debug(self, name: &'static str, value: impl std::fmt::Debug) -> Self {
+        self.with_arg(name, format!("{value:?}"))
+            .expect("failed to serialize debug value")
+    }
+
+    pub fn with_arg(
+        mut self,
+        arg: &'static str,
+        value: impl Serialize,
+    ) -> Result<Self, serde_json::Error> {
+        let value = serde_json::to_value(value)?;
+        self.args.insert(arg, EventArgument::Value(value));
+        Ok(self)
+    }
 
     /// Create a new scope without checking if timing is enabled.
-    fn new_impl(
+    pub fn build(self) -> TimingScopeGuard {
-        name: &'static str,
-        span: Option<NonZeroU64>,
-        callsite: Option<NonZeroU64>,
-        func: Option<String>,
-    ) -> Self {
         let (thread_id, timestamp) =
             THREAD_DATA.with(|data| (data.id, Timestamp::now_with(data)));
-        EVENTS.lock().push(Event {
+        let event = Event {
             kind: EventKind::Start,
             timestamp,
-            name,
+            name: self.name,
-            span,
+            func: self.func.clone(),
-            callsite,
-            func: func.clone(),
             thread_id,
-        });
+            arguments: Arc::new(self.args),
-        Self { name, span, callsite: None, thread_id, func }
+        };
+        EVENTS.lock().push(event.clone());
+        TimingScopeGuard { scope: Some(event) }
     }
 }
 
-impl Drop for TimingScope {
+pub struct TimingScopeGuard {
+    scope: Option<Event>,
+}
+
+impl Drop for TimingScopeGuard {
     fn drop(&mut self) {
         let timestamp = Timestamp::now();
-        EVENTS.lock().push(Event {
+
-            kind: EventKind::End,
+        let mut scope = self.scope.take().expect("scope already dropped");
-            timestamp,
+        scope.timestamp = timestamp;
-            name: self.name,
+        scope.kind = EventKind::End;
-            span: self.span,
+
-            callsite: self.callsite,
+        EVENTS.lock().push(scope);
-            thread_id: self.thread_id,
+    }
-            func: std::mem::take(&mut self.func),
+}
-        });
+
+enum EventArgument {
+    Span(NonZeroU64),
+    Value(serde_json::Value),
+}
+
+impl EventArgument {
+    fn to_json<'a>(
+        &'a self,
+        mut source: impl FnMut(NonZeroU64) -> (String, u32),
+        key: &'static str,
+        out: &mut IndexMap<Cow<'static, str>, Cow<'a, serde_json::Value>>,
+    ) -> Result<(), serde_json::Error> {
+        match self {
+            EventArgument::Span(span) => {
+                let (file, line) = source(*span);
+                // Insert file and line information for the span
+                if key == "span" {
+                    out.insert("file".into(), Cow::Owned(serde_json::json!(file)));
+                    out.insert("line".into(), Cow::Owned(serde_json::json!(line)));
+
+                    return Ok(());
+                }
+
+                // Small optimization for callsite
+                if key == "callsite" {
+                    out.insert(
+                        "callsite_file".into(),
+                        Cow::Owned(serde_json::json!(file)),
+                    );
+                    out.insert(
+                        "callsite_line".into(),
+                        Cow::Owned(serde_json::json!(line)),
+                    );
+
+                    return Ok(());
+                }
+
+                out.insert(
+                    format!("{key}_file").into(),
+                    Cow::Owned(serde_json::json!(file)),
+                );
+
+                out.insert(
+                    format!("{key}_line").into(),
+                    Cow::Owned(serde_json::json!(line)),
+                );
+            }
+            EventArgument::Value(value) => {
+                out.insert(key.into(), Cow::Borrowed(value));
+            }
+        }
+
+        Ok(())
     }
 }
 
 /// An event that has been recorded.
+#[derive(Clone)]
 struct Event {
     /// Whether this is a start or end event.
     kind: EventKind,
@@ -247,10 +305,8 @@ struct Event {
     timestamp: Timestamp,
     /// The name of this event.
     name: &'static str,
-    /// The raw value of the span of code that this event was recorded in.
+    /// The additional arguments of this event.
-    span: Option<NonZeroU64>,
+    arguments: Arc<IndexMap<&'static str, EventArgument>>,
-    /// The raw value of the callsite span of the code that this event was recorded in.
-    callsite: Option<NonZeroU64>,
     /// The function being called (if any).
     func: Option<String>,
     /// The thread ID of this event.