1 files changed, 24 insertions, 29 deletions

diff --git a/components/style/parallel.rs b/components/style/parallel.rs
index ea39a0c15d7..2de33023f5e 100644
--- a/components/style/parallel.rs
+++ b/components/style/parallel.rs
@@ -50,7 +50,7 @@ type NodeList<N> = SmallVec<[SendNode<N>; WORK_UNIT_MAX]>;
 pub fn traverse_dom<E, D>(traversal: &D,
                           root: E,
                           token: PreTraverseToken,
-                          queue: &rayon::ThreadPool)
+                          pool: &rayon::ThreadPool)
     where E: TElement,
           D: DomTraversal<E>,
 {
@@ -82,16 +82,17 @@ pub fn traverse_dom<E, D>(traversal: &D,
     let traversal_data = PerLevelTraversalData {
         current_dom_depth: depth,
     };
-    let tls = ScopedTLS::<D::ThreadLocalContext>::new(queue);
+    let tls = ScopedTLS::<D::ThreadLocalContext>::new(pool);
     let root = root.as_node().opaque();
 
-    queue.install(|| {
+    pool.install(|| {
         rayon::scope(|scope| {
             traverse_nodes(nodes,
                            DispatchMode::TailCall,
                            root,
                            traversal_data,
                            scope,
+                           pool,
                            traversal,
                            &tls);
         });
@@ -144,6 +145,7 @@ fn top_down_dom<'a, 'scope, E, D>(nodes: &'a [SendNode<E::ConcreteNode>],
                                   root: OpaqueNode,
                                   mut traversal_data: PerLevelTraversalData,
                                   scope: &'a rayon::Scope<'scope>,
+                                  pool: &'scope rayon::ThreadPool,
                                   traversal: &'scope D,
                                   tls: &'scope ScopedTLS<'scope, D::ThreadLocalContext>)
     where E: TElement + 'scope,
@@ -179,6 +181,7 @@ fn top_down_dom<'a, 'scope, E, D>(nodes: &'a [SendNode<E::ConcreteNode>],
                                root,
                                traversal_data_copy,
                                scope,
+                               pool,
                                traversal,
                                tls);
             }
@@ -208,6 +211,7 @@ fn top_down_dom<'a, 'scope, E, D>(nodes: &'a [SendNode<E::ConcreteNode>],
                        root,
                        traversal_data,
                        scope,
+                       pool,
                        traversal,
                        tls);
     }
@@ -231,6 +235,7 @@ fn traverse_nodes<'a, 'scope, E, D>(nodes: NodeList<E::ConcreteNode>,
                                     root: OpaqueNode,
                                     traversal_data: PerLevelTraversalData,
                                     scope: &'a rayon::Scope<'scope>,
+                                    pool: &'scope rayon::ThreadPool,
                                     traversal: &'scope D,
                                     tls: &'scope ScopedTLS<'scope, D::ThreadLocalContext>)
     where E: TElement + 'scope,
@@ -238,42 +243,32 @@ fn traverse_nodes<'a, 'scope, E, D>(nodes: NodeList<E::ConcreteNode>,
 {
     debug_assert!(!nodes.is_empty());
 
+    // This is a tail call from the perspective of the caller. However, we only
+    // want to actually dispatch the job as a tail call if there's nothing left
+    // in our local queue. Otherwise we need to return to it to maintain proper
+    // breadth-first ordering.
+    let may_dispatch_tail = mode.is_tail_call() &&
+        !pool.current_thread_has_pending_tasks().unwrap();
+
     // In the common case, our children fit within a single work unit, in which
     // case we can pass the SmallVec directly and avoid extra allocation.
     if nodes.len() <= WORK_UNIT_MAX {
-        if mode.is_tail_call() {
-            // If this is a tail call, bypass rayon and invoke top_down_dom directly.
-            top_down_dom(&nodes, root, traversal_data, scope, traversal, tls);
+        if may_dispatch_tail {
+            top_down_dom(&nodes, root, traversal_data, scope, pool, traversal, tls);
         } else {
-            // The caller isn't done yet. Append to the queue and return synchronously.
             scope.spawn(move |scope| {
                 let nodes = nodes;
-                top_down_dom(&nodes, root, traversal_data, scope, traversal, tls);
+                top_down_dom(&nodes, root, traversal_data, scope, pool, traversal, tls);
             });
         }
     } else {
-        // FIXME(bholley): This should be an ArrayVec.
-        let mut first_chunk: Option<NodeList<E::ConcreteNode>> = None;
         for chunk in nodes.chunks(WORK_UNIT_MAX) {
-            if mode.is_tail_call() && first_chunk.is_none() {
-                first_chunk = Some(chunk.iter().cloned().collect::<NodeList<E::ConcreteNode>>());
-            } else {
-                let boxed = chunk.iter().cloned().collect::<Vec<_>>().into_boxed_slice();
-                let traversal_data_copy = traversal_data.clone();
-                scope.spawn(move |scope| {
-                    let b = boxed;
-                    top_down_dom(&*b, root, traversal_data_copy, scope, traversal, tls)
-                });
-
-            }
-        }
-
-        // If this is a tail call, bypass rayon for the first chunk and invoke top_down_dom
-        // directly.
-        debug_assert_eq!(first_chunk.is_some(), mode.is_tail_call());
-        if let Some(c) = first_chunk {
-            debug_assert_eq!(c.len(), WORK_UNIT_MAX);
-            top_down_dom(&*c, root, traversal_data, scope, traversal, tls);
+            let boxed = chunk.iter().cloned().collect::<Vec<_>>().into_boxed_slice();
+            let traversal_data_copy = traversal_data.clone();
+            scope.spawn(move |scope| {
+                let b = boxed;
+                top_down_dom(&*b, root, traversal_data_copy, scope, pool, traversal, tls)
+            });
         }
     }
 }