Make `typst-timing` WASM-compatible (#5689)

1 files changed, 217 insertions, 155 deletions

diff --git a/crates/typst-timing/src/lib.rs b/crates/typst-timing/src/lib.rs
index b4653170..6da2cdf0 100644
--- a/crates/typst-timing/src/lib.rs
+++ b/crates/typst-timing/src/lib.rs
@@ -1,149 +1,13 @@
 //! Performance timing for Typst.
 
-#![cfg_attr(target_arch = "wasm32", allow(dead_code, unused_variables))]
-
-use std::hash::Hash;
 use std::io::Write;
 use std::num::NonZeroU64;
-use std::sync::atomic::AtomicBool;
-use std::sync::atomic::Ordering::Relaxed;
-use std::thread::ThreadId;
-use std::time::{Duration, SystemTime};
+use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
 
 use parking_lot::Mutex;
 use serde::ser::SerializeSeq;
 use serde::{Serialize, Serializer};
 
-/// Whether the timer is enabled. Defaults to `false`.
-static ENABLED: AtomicBool = AtomicBool::new(false);
-
-/// The global event recorder.
-static RECORDER: Mutex<Recorder> = Mutex::new(Recorder::new());
-
-/// The recorder of events.
-struct Recorder {
-    /// The events that have been recorded.
-    events: Vec<Event>,
-    /// The discriminator of the next event.
-    discriminator: u64,
-}
-
-impl Recorder {
-    /// Create a new recorder.
-    const fn new() -> Self {
-        Self { events: Vec::new(), discriminator: 0 }
-    }
-}
-
-/// An event that has been recorded.
-#[derive(Clone, Copy, Eq, PartialEq, Hash)]
-struct Event {
-    /// Whether this is a start or end event.
-    kind: EventKind,
-    /// The start time of this event.
-    timestamp: SystemTime,
-    /// The discriminator of this event.
-    id: u64,
-    /// The name of this event.
-    name: &'static str,
-    /// The raw value of the span of code that this event was recorded in.
-    span: Option<NonZeroU64>,
-    /// The thread ID of this event.
-    thread_id: ThreadId,
-}
-
-/// Whether an event marks the start or end of a scope.
-#[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
-enum EventKind {
-    Start,
-    End,
-}
-
-/// Enable the timer.
-#[inline]
-pub fn enable() {
-    // We only need atomicity and no synchronization of other
-    // operations, so `Relaxed` is fine.
-    ENABLED.store(true, Relaxed);
-}
-
-/// Whether the timer is enabled.
-#[inline]
-pub fn is_enabled() -> bool {
-    ENABLED.load(Relaxed)
-}
-
-/// Clears the recorded events.
-#[inline]
-pub fn clear() {
-    RECORDER.lock().events.clear();
-}
-
-/// A scope that records an event when it is dropped.
-pub struct TimingScope {
-    name: &'static str,
-    span: Option<NonZeroU64>,
-    id: u64,
-    thread_id: ThreadId,
-}
-
-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> {
-        #[cfg(not(target_arch = "wasm32"))]
-        if is_enabled() {
-            return Some(Self::new_impl(name, span));
-        }
-        None
-    }
-
-    /// Create a new scope without checking if timing is enabled.
-    fn new_impl(name: &'static str, span: Option<NonZeroU64>) -> Self {
-        let timestamp = SystemTime::now();
-        let thread_id = std::thread::current().id();
-
-        let mut recorder = RECORDER.lock();
-        let id = recorder.discriminator;
-        recorder.discriminator += 1;
-        recorder.events.push(Event {
-            kind: EventKind::Start,
-            timestamp,
-            id,
-            name,
-            span,
-            thread_id,
-        });
-
-        Self { name, span, id, thread_id }
-    }
-}
-
-impl Drop for TimingScope {
-    fn drop(&mut self) {
-        let event = Event {
-            kind: EventKind::End,
-            timestamp: SystemTime::now(),
-            id: self.id,
-            name: self.name,
-            span: self.span,
-            thread_id: self.thread_id,
-        };
-
-        RECORDER.lock().events.push(event);
-    }
-}
-
 /// Creates a timing scope around an expression.
 ///
 /// The output of the expression is returned.
@@ -179,6 +43,46 @@ macro_rules! timed {
     }};
 }
 
+thread_local! {
+    /// Data that is initialized once per thread.
+    static THREAD_DATA: ThreadData = ThreadData {
+        id: {
+            // We only need atomicity and no synchronization of other
+            // operations, so `Relaxed` is fine.
+            static COUNTER: AtomicU64 = AtomicU64::new(1);
+            COUNTER.fetch_add(1, Ordering::Relaxed)
+        },
+        #[cfg(all(target_arch = "wasm32", feature = "wasm"))]
+        timer: WasmTimer::new(),
+    };
+}
+
+/// Whether the timer is enabled. Defaults to `false`.
+static ENABLED: AtomicBool = AtomicBool::new(false);
+
+/// The list of collected events.
+static EVENTS: Mutex<Vec<Event>> = Mutex::new(Vec::new());
+
+/// Enable the timer.
+#[inline]
+pub fn enable() {
+    // We only need atomicity and no synchronization of other
+    // operations, so `Relaxed` is fine.
+    ENABLED.store(true, Ordering::Relaxed);
+}
+
+/// Whether the timer is enabled.
+#[inline]
+pub fn is_enabled() -> bool {
+    ENABLED.load(Ordering::Relaxed)
+}
+
+/// Clears the recorded events.
+#[inline]
+pub fn clear() {
+    EVENTS.lock().clear();
+}
+
 /// Export data as JSON for Chrome's tracing tool.
 ///
 /// The `source` function is called for each span to get the source code
@@ -205,19 +109,15 @@ pub fn export_json<W: Write>(
         line: u32,
     }
 
-    let recorder = RECORDER.lock();
-    let run_start = recorder
-        .events
-        .first()
-        .map(|event| event.timestamp)
-        .unwrap_or_else(SystemTime::now);
+    let lock = EVENTS.lock();
+    let events = lock.as_slice();
 
     let mut serializer = serde_json::Serializer::new(writer);
     let mut seq = serializer
-        .serialize_seq(Some(recorder.events.len()))
+        .serialize_seq(Some(events.len()))
         .map_err(|e| format!("failed to serialize events: {e}"))?;
 
-    for event in recorder.events.iter() {
+    for event in events.iter() {
         seq.serialize_element(&Entry {
             name: event.name,
             cat: "typst",
@@ -225,17 +125,9 @@ pub fn export_json<W: Write>(
                 EventKind::Start => "B",
                 EventKind::End => "E",
             },
-            ts: event
-                .timestamp
-                .duration_since(run_start)
-                .unwrap_or(Duration::ZERO)
-                .as_nanos() as f64
-                / 1_000.0,
+            ts: event.timestamp.micros_since(events[0].timestamp),
             pid: 1,
-            tid: unsafe {
-                // Safety: `thread_id` is a `ThreadId` which is a `u64`.
-                std::mem::transmute_copy(&event.thread_id)
-            },
+            tid: event.thread_id,
             args: event.span.map(&mut source).map(|(file, line)| Args { file, line }),
         })
         .map_err(|e| format!("failed to serialize event: {e}"))?;
@@ -245,3 +137,173 @@ pub fn export_json<W: Write>(
 
     Ok(())
 }
+
+/// A scope that records an event when it is dropped.
+pub struct TimingScope {
+    name: &'static str,
+    span: Option<NonZeroU64>,
+    thread_id: u64,
+}
+
+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> {
+        if is_enabled() {
+            return Some(Self::new_impl(name, span));
+        }
+        None
+    }
+
+    /// Create a new scope without checking if timing is enabled.
+    fn new_impl(name: &'static str, span: Option<NonZeroU64>) -> Self {
+        let (thread_id, timestamp) =
+            THREAD_DATA.with(|data| (data.id, Timestamp::now_with(data)));
+        EVENTS.lock().push(Event {
+            kind: EventKind::Start,
+            timestamp,
+            name,
+            span,
+            thread_id,
+        });
+        Self { name, span, thread_id }
+    }
+}
+
+impl Drop for TimingScope {
+    fn drop(&mut self) {
+        let timestamp = Timestamp::now();
+        EVENTS.lock().push(Event {
+            kind: EventKind::End,
+            timestamp,
+            name: self.name,
+            span: self.span,
+            thread_id: self.thread_id,
+        });
+    }
+}
+
+/// An event that has been recorded.
+struct Event {
+    /// Whether this is a start or end event.
+    kind: EventKind,
+    /// The time at which this event occurred.
+    timestamp: Timestamp,
+    /// The name of this event.
+    name: &'static str,
+    /// The raw value of the span of code that this event was recorded in.
+    span: Option<NonZeroU64>,
+    /// The thread ID of this event.
+    thread_id: u64,
+}
+
+/// Whether an event marks the start or end of a scope.
+#[derive(Debug, Clone, Copy, Eq, PartialEq)]
+enum EventKind {
+    Start,
+    End,
+}
+
+/// A cross-platform way to get the current time.
+#[derive(Copy, Clone)]
+struct Timestamp {
+    #[cfg(not(target_arch = "wasm32"))]
+    inner: std::time::SystemTime,
+    #[cfg(target_arch = "wasm32")]
+    inner: f64,
+}
+
+impl Timestamp {
+    fn now() -> Self {
+        #[cfg(target_arch = "wasm32")]
+        return THREAD_DATA.with(Self::now_with);
+
+        #[cfg(not(target_arch = "wasm32"))]
+        Self { inner: std::time::SystemTime::now() }
+    }
+
+    #[allow(unused_variables)]
+    fn now_with(data: &ThreadData) -> Self {
+        #[cfg(all(target_arch = "wasm32", feature = "wasm"))]
+        return Self { inner: data.timer.now() };
+
+        #[cfg(all(target_arch = "wasm32", not(feature = "wasm")))]
+        return Self { inner: 0.0 };
+
+        #[cfg(not(target_arch = "wasm32"))]
+        Self::now()
+    }
+
+    fn micros_since(self, start: Self) -> f64 {
+        #[cfg(target_arch = "wasm32")]
+        return (self.inner - start.inner) * 1000.0;
+
+        #[cfg(not(target_arch = "wasm32"))]
+        (self
+            .inner
+            .duration_since(start.inner)
+            .unwrap_or(std::time::Duration::ZERO)
+            .as_nanos() as f64
+            / 1_000.0)
+    }
+}
+
+/// Per-thread data.
+struct ThreadData {
+    /// The thread's ID.
+    ///
+    /// In contrast to `std::thread::current().id()`, this is wasm-compatible
+    /// and also a bit cheaper to access because the std version does a bit more
+    /// stuff (including cloning an `Arc`).
+    id: u64,
+    /// A way to get the time from WebAssembly.
+    #[cfg(all(target_arch = "wasm32", feature = "wasm"))]
+    timer: WasmTimer,
+}
+
+/// A way to get the time from WebAssembly.
+#[cfg(all(target_arch = "wasm32", feature = "wasm"))]
+struct WasmTimer {
+    /// The cached JS performance handle for the thread.
+    perf: web_sys::Performance,
+    /// The cached JS time origin.
+    time_origin: f64,
+}
+
+#[cfg(all(target_arch = "wasm32", feature = "wasm"))]
+impl WasmTimer {
+    fn new() -> Self {
+        // Retrieve `performance` from global object, either the window or
+        // globalThis.
+        let perf = web_sys::window()
+            .and_then(|window| window.performance())
+            .or_else(|| {
+                use web_sys::wasm_bindgen::JsCast;
+                web_sys::js_sys::global()
+                    .dyn_into::<web_sys::WorkerGlobalScope>()
+                    .ok()
+                    .and_then(|scope| scope.performance())
+            })
+            .expect("failed to get JS performance handle");
+
+        // Every thread gets its own time origin. To make the results consistent
+        // across threads, we need to add this to each `now()` call.
+        let time_origin = perf.time_origin();
+
+        Self { perf, time_origin }
+    }
+
+    fn now(&self) -> f64 {
+        self.time_origin + self.perf.now()
+    }
+}