2 files changed, 40 insertions, 27 deletions

diff --git a/tests/src/benches.rs b/tests/src/benches.rs
index 21fdbc7a..b1e77692 100644
--- a/tests/src/benches.rs
+++ b/tests/src/benches.rs
@@ -80,17 +80,17 @@ fn bench_typeset(iai: &mut Iai) {
     )
     .unwrap();
     let content = module.content();
-    iai.run(|| typst::model::typeset(world.track(), &content));
+    iai.run(|| typst::model::typeset(world.track(), tracer.track_mut(), &content));
 }
 
 fn bench_compile(iai: &mut Iai) {
     let world = BenchWorld::new();
-    iai.run(|| typst::compile(&world, &world.source));
+    iai.run(|| typst::compile(&world));
 }
 
 fn bench_render(iai: &mut Iai) {
     let world = BenchWorld::new();
-    let document = typst::compile(&world, &world.source).unwrap();
+    let document = typst::compile(&world).unwrap();
     iai.run(|| typst::export::render(&document.pages[0], 1.0, Color::WHITE))
 }
 
@@ -124,6 +124,18 @@ impl World for BenchWorld {
         &self.library
     }
 
+    fn main(&self) -> &Source {
+        &self.source
+    }
+
+    fn resolve(&self, path: &Path) -> FileResult<SourceId> {
+        Err(FileError::NotFound(path.into()))
+    }
+
+    fn source(&self, _: SourceId) -> &Source {
+        &self.source
+    }
+
     fn book(&self) -> &Prehashed<FontBook> {
         &self.book
     }
@@ -135,12 +147,4 @@ impl World for BenchWorld {
     fn file(&self, path: &Path) -> FileResult<Buffer> {
         Err(FileError::NotFound(path.into()))
     }
-
-    fn resolve(&self, path: &Path) -> FileResult<SourceId> {
-        Err(FileError::NotFound(path.into()))
-    }
-
-    fn source(&self, _: SourceId) -> &Source {
-        unimplemented!()
-    }
 }
diff --git a/tests/src/tests.rs b/tests/src/tests.rs
index 1fb40204..9e6d94a3 100644
--- a/tests/src/tests.rs
+++ b/tests/src/tests.rs
@@ -207,6 +207,7 @@ struct TestWorld {
     fonts: Vec<Font>,
     paths: RefCell<HashMap<PathBuf, PathSlot>>,
     sources: FrozenVec<Box<Source>>,
+    main: SourceId,
 }
 
 #[derive(Default)]
@@ -236,6 +237,7 @@ impl TestWorld {
             fonts,
             paths: RefCell::default(),
             sources: FrozenVec::new(),
+            main: SourceId::detached(),
         }
     }
 }
@@ -249,19 +251,8 @@ impl World for TestWorld {
         &self.library
     }
 
-    fn book(&self) -> &Prehashed<FontBook> {
-        &self.book
-    }
-
-    fn font(&self, id: usize) -> Option<Font> {
-        Some(self.fonts[id].clone())
-    }
-
-    fn file(&self, path: &Path) -> FileResult<Buffer> {
-        self.slot(path)
-            .buffer
-            .get_or_init(|| read(path).map(Buffer::from))
-            .clone()
+    fn main(&self) -> &Source {
+        self.source(self.main)
     }
 
     fn resolve(&self, path: &Path) -> FileResult<SourceId> {
@@ -278,20 +269,38 @@ impl World for TestWorld {
     fn source(&self, id: SourceId) -> &Source {
         &self.sources[id.into_u16() as usize]
     }
+
+    fn book(&self) -> &Prehashed<FontBook> {
+        &self.book
+    }
+
+    fn font(&self, id: usize) -> Option<Font> {
+        Some(self.fonts[id].clone())
+    }
+
+    fn file(&self, path: &Path) -> FileResult<Buffer> {
+        self.slot(path)
+            .buffer
+            .get_or_init(|| read(path).map(Buffer::from))
+            .clone()
+    }
 }
 
 impl TestWorld {
     fn set(&mut self, path: &Path, text: String) -> SourceId {
         let slot = self.slot(path);
-        if let Some(&Ok(id)) = slot.source.get() {
+        let id = if let Some(&Ok(id)) = slot.source.get() {
             drop(slot);
             self.sources.as_mut()[id.into_u16() as usize].replace(text);
             id
         } else {
             let id = self.insert(path, text);
             slot.source.set(Ok(id)).unwrap();
+            drop(slot);
             id
-        }
+        };
+        self.main = id;
+        id
     }
 
     fn slot(&self, path: &Path) -> RefMut<PathSlot> {
@@ -448,7 +457,7 @@ fn test_part(
         println!("Model:\n{:#?}\n", module.content());
     }
 
-    let (mut frames, errors) = match typst::compile(world, source) {
+    let (mut frames, errors) = match typst::compile(world) {
         Ok(document) => (document.pages, vec![]),
         Err(errors) => (vec![], *errors),
     };