Support STYLO_THREADS=1

1 files changed, 14 insertions, 6 deletions

diff --git a/components/layout_thread_2020/lib.rs b/components/layout_thread_2020/lib.rs
index 45202d51ad9..35b4c92c21b 100644
--- a/components/layout_thread_2020/lib.rs
+++ b/components/layout_thread_2020/lib.rs
@@ -1076,14 +1076,18 @@ impl LayoutThread {
         };
 
         let rayon_pool = STYLE_THREAD_POOL.pool();
-        let rayon_pool = rayon_pool.as_ref().unwrap();
+        let rayon_pool = rayon_pool.as_ref();
 
         let box_tree = if token.should_traverse() {
-            driver::traverse_dom(&traversal, token, Some(rayon_pool));
+            driver::traverse_dom(&traversal, token, rayon_pool);
 
             let root_node = document.root_element().unwrap().as_node();
-            let box_tree =
-                rayon_pool.install(|| BoxTreeRoot::construct(traversal.context(), root_node));
+            let build_box_tree = || BoxTreeRoot::construct(traversal.context(), root_node);
+            let box_tree = if let Some(pool) = rayon_pool {
+                pool.install(build_box_tree)
+            } else {
+                build_box_tree()
+            };
             Some(box_tree)
         } else {
             None
@@ -1096,8 +1100,12 @@ impl LayoutThread {
                 self.viewport_size.width.to_f32_px(),
                 self.viewport_size.height.to_f32_px(),
             );
-            let fragment_tree =
-                rayon_pool.install(|| box_tree.layout(&layout_context, viewport_size));
+            let run_layout = || box_tree.layout(&layout_context, viewport_size);
+            let fragment_tree = if let Some(pool) = rayon_pool {
+                pool.install(run_layout)
+            } else {
+                run_layout()
+            };
             *self.box_tree_root.borrow_mut() = Some(box_tree);
             *self.fragment_tree_root.borrow_mut() = Some(fragment_tree);
         }