[WIP] Run iterations in parallelparallel-iterations

4 files changed, 107 insertions, 44 deletions

diff --git a/crates/typst/src/engine.rs b/crates/typst/src/engine.rs
index 2e2525b2..8997a9a1 100644
--- a/crates/typst/src/engine.rs
+++ b/crates/typst/src/engine.rs
@@ -128,13 +128,13 @@ pub struct Sink {
     /// because the introspector is not yet ready. We first ignore that and
     /// proceed with empty content and only if the error remains by the end
     /// of the last iteration, we promote it.
-    delayed: EcoVec<SourceDiagnostic>,
+    pub(crate) delayed: EcoVec<SourceDiagnostic>,
     /// Warnings emitted during iteration.
-    warnings: EcoVec<SourceDiagnostic>,
+    pub(crate) warnings: EcoVec<SourceDiagnostic>,
     /// Hashes of all warning's spans and messages for warning deduplication.
     warnings_set: HashSet<u128>,
     /// A sequence of traced values for a span.
-    values: EcoVec<(Value, Option<Styles>)>,
+    pub(crate) values: EcoVec<(Value, Option<Styles>)>,
 }
 
 impl Sink {
@@ -186,7 +186,7 @@ impl Sink {
     }
 
     /// Extend from another sink.
-    fn extend(
+    pub fn extend(
         &mut self,
         delayed: EcoVec<SourceDiagnostic>,
         warnings: EcoVec<SourceDiagnostic>,
diff --git a/crates/typst/src/introspection/introspector.rs b/crates/typst/src/introspection/introspector.rs
index fe59cb00..bd49e43f 100644
--- a/crates/typst/src/introspection/introspector.rs
+++ b/crates/typst/src/introspection/introspector.rs
@@ -38,21 +38,25 @@ pub struct Introspector {
 }
 
 impl Introspector {
-    /// Applies new frames in-place, reusing the existing allocations.
+    /// Build a new introspector for a page collection.
     #[typst_macros::time(name = "introspect")]
-    pub fn rebuild(&mut self, pages: &[Page]) {
-        self.pages = pages.len();
-        self.elems.clear();
-        self.labels.clear();
-        self.keys.clear();
-        self.page_numberings.clear();
-        self.queries.clear();
+    pub fn new(pages: &[Page]) -> Self {
+        let mut this = Self {
+            pages: pages.len(),
+            elems: IndexMap::new(),
+            labels: HashMap::new(),
+            keys: HashMap::new(),
+            page_numberings: Vec::with_capacity(pages.len()),
+            queries: QueryCache::default(),
+        };
 
         for (i, page) in pages.iter().enumerate() {
             let page_nr = NonZeroUsize::new(1 + i).unwrap();
-            self.extract(&page.frame, page_nr, Transform::identity());
-            self.page_numberings.push(page.numbering.clone());
+            this.extract(&page.frame, page_nr, Transform::identity());
+            this.page_numberings.push(page.numbering.clone());
         }
+
+        this
     }
 
     /// Extract metadata from a frame.
@@ -328,10 +332,6 @@ impl QueryCache {
     fn insert(&self, hash: u128, output: EcoVec<Content>) {
         self.0.write().unwrap().insert(hash, output);
     }
-
-    fn clear(&mut self) {
-        self.0.get_mut().unwrap().clear();
-    }
 }
 
 impl Clone for QueryCache {
diff --git a/crates/typst/src/lib.rs b/crates/typst/src/lib.rs
index 50575d12..2ce1cebc 100644
--- a/crates/typst/src/lib.rs
+++ b/crates/typst/src/lib.rs
@@ -58,15 +58,17 @@ pub use typst_utils as utils;
 
 use std::collections::HashSet;
 use std::ops::{Deref, Range};
+use std::sync::{Arc, Mutex};
 
 use comemo::{Track, Tracked, Validate};
 use ecow::{EcoString, EcoVec};
+use rayon::iter::{IndexedParallelIterator, IntoParallelRefIterator, ParallelIterator};
 use typst_timing::{timed, TimingScope};
 
 use crate::diag::{warning, FileResult, SourceDiagnostic, SourceResult, Warned};
 use crate::engine::{Engine, Route, Sink, Traced};
 use crate::foundations::{
-    Array, Bytes, Datetime, Dict, Module, Scope, StyleChain, Styles, Value,
+    Array, Bytes, Content, Datetime, Dict, Module, Scope, StyleChain, Styles, Value,
 };
 use crate::introspection::Introspector;
 use crate::layout::{Alignment, Dir};
@@ -105,8 +107,8 @@ fn compile_inner(
     traced: Tracked<Traced>,
     sink: &mut Sink,
 ) -> SourceResult<Document> {
-    let library = world.library();
-    let styles = StyleChain::new(&library.styles);
+    // Yes, this is a hack. But it's ok, this is just an experimental branch.
+    static INTROSPECTORS: Mutex<Vec<Arc<Introspector>>> = Mutex::new(Vec::new());
 
     // First evaluate the main source file into a module.
     let content = crate::eval::eval(
@@ -118,36 +120,57 @@ fn compile_inner(
     )?
     .content();
 
+    // If we have introspectors from previous compilations, we can speed up
+    // compilation with a small speculative trick: Since quite often the
+    // iterations we go through will be exactly the same ones as previously with
+    // no relevant introspector changes, we can run all iterations in parallel
+    // using the previous introspectors.
+    let mut speculation = INTROSPECTORS
+        .lock()
+        .unwrap()
+        .clone()
+        .par_iter()
+        .enumerate()
+        .map(|(iter, introspector)| {
+            let mut sink = Sink::new();
+            (layout(world, &introspector, traced, &mut sink, &content, iter), sink)
+        })
+        .collect::<Vec<_>>()
+        .into_iter();
+
     let mut iter = 0;
-    let mut document = Document::default();
+    let mut introspector = Arc::new(Introspector::new(&[]));
+    let mut introspectors = vec![];
 
     // Relayout until all introspections stabilize.
     // If that doesn't happen within five attempts, we give up.
-    loop {
-        // The name of the iterations for timing scopes.
-        const ITER_NAMES: &[&str] =
-            &["layout (1)", "layout (2)", "layout (3)", "layout (4)", "layout (5)"];
-        let _scope = TimingScope::new(ITER_NAMES[iter], None);
+    let document = loop {
+        // Save this iteration's introspector.
+        introspectors.push(introspector.clone());
 
         // Clear delayed errors.
         sink.delayed();
 
-        let constraint = <Introspector as Validate>::Constraint::new();
-        let mut engine = Engine {
-            world,
-            introspector: document.introspector.track_with(&constraint),
-            traced,
-            sink: sink.track_mut(),
-            route: Route::default(),
+        // Try to use a speculative result if it was valid for the
+        // real current introspector. If not, layout as usual.
+        let (result, constraint) = if let Some((pair, subsink)) =
+            speculation.next().filter(|((_, constraint), _)| {
+                timed!("check speculation", introspector.validate(&constraint))
+            }) {
+            sink.extend(subsink.delayed, subsink.warnings, subsink.values);
+            pair
+        } else {
+            layout(world, &introspector, traced, sink, &content, iter)
         };
 
-        // Layout!
-        document = content.layout_document(&mut engine, styles)?;
-        document.introspector.rebuild(&document.pages);
+        let document = result?;
+        introspector = document.introspector.clone();
+
+        // Bump the iteration.
         iter += 1;
 
-        if timed!("check stabilized", document.introspector.validate(&constraint)) {
-            break;
+        if timed!("check stabilized", introspector.validate(&constraint)) {
+            break document;
         }
 
         if iter >= 5 {
@@ -155,9 +178,12 @@ fn compile_inner(
                 Span::detached(), "layout did not converge within 5 attempts";
                 hint: "check if any states or queries are updating themselves"
             ));
-            break;
+            break document;
         }
-    }
+    };
+
+    // Save the introspectors for next time.
+    *INTROSPECTORS.lock().unwrap() = introspectors;
 
     // Promote delayed errors.
     let delayed = sink.delayed();
@@ -168,6 +194,39 @@ fn compile_inner(
     Ok(document)
 }
 
+/// Run a single layout iteration with the given introspector.
+///
+/// Returns the document and the introspector constraints for it.
+fn layout(
+    world: Tracked<dyn World + '_>,
+    introspector: &Introspector,
+    traced: Tracked<Traced>,
+    sink: &mut Sink,
+    content: &Content,
+    iter: usize,
+) -> (SourceResult<Document>, <Introspector as Validate>::Constraint) {
+    // The name of the iterations for timing scopes.
+    const ITERS: &[&str] =
+        &["layout (1)", "layout (2)", "layout (3)", "layout (4)", "layout (5)"];
+    let _scope = TimingScope::new(ITERS[iter], None);
+
+    let constraint = <Introspector as Validate>::Constraint::new();
+    let mut engine = Engine {
+        world,
+        introspector: introspector.track_with(&constraint),
+        traced,
+        sink: sink.track_mut(),
+        route: Route::default(),
+    };
+
+    let library = world.library();
+    let styles = StyleChain::new(&library.styles);
+
+    // Layout!
+    let document = content.layout_document(&mut engine, styles);
+    (document, constraint)
+}
+
 /// Deduplicate diagnostics.
 fn deduplicate(mut diags: EcoVec<SourceDiagnostic>) -> EcoVec<SourceDiagnostic> {
     let mut unique = HashSet::new();
@@ -189,7 +248,7 @@ fn deduplicate(mut diags: EcoVec<SourceDiagnostic>) -> EcoVec<SourceDiagnostic>
 /// information on when something can change. For example, fonts typically don't
 /// change and can thus even be cached across multiple compilations (for
 /// long-running applications like `typst watch`). Source files on the other
-/// hand can change and should thus be cleared after each compilation. Advanced
+/// hand can change and might thus be cleared after each compilation. Advanced
 /// clients like language servers can also retain the source files and
 /// [edit](Source::edit) them in-place to benefit from better incremental
 /// performance.
diff --git a/crates/typst/src/model/document.rs b/crates/typst/src/model/document.rs
index 77044112..62a202c7 100644
--- a/crates/typst/src/model/document.rs
+++ b/crates/typst/src/model/document.rs
@@ -1,3 +1,5 @@
+use std::sync::Arc;
+
 use ecow::EcoString;
 
 use crate::diag::{bail, HintedStrResult, SourceResult};
@@ -107,13 +109,15 @@ impl Packed<DocumentElem> {
             pages.extend(result?.finalize(engine, &mut page_counter)?);
         }
 
+        let introspector = Introspector::new(&pages);
+
         Ok(Document {
             pages,
             title: DocumentElem::title_in(styles).map(|content| content.plain_text()),
             author: DocumentElem::author_in(styles).0,
             keywords: DocumentElem::keywords_in(styles).0,
             date: DocumentElem::date_in(styles),
-            introspector: Introspector::default(),
+            introspector: Arc::new(introspector),
         })
     }
 }
@@ -154,7 +158,7 @@ pub struct Document {
     /// The document's creation date.
     pub date: Smart<Option<Datetime>>,
     /// Provides the ability to execute queries on the document.
-    pub introspector: Introspector,
+    pub introspector: Arc<Introspector>,
 }
 
 #[cfg(test)]