From ef2cb915086496fb1c708b0ba8a25957ac4b7a4b Mon Sep 17 00:00:00 2001
From: Robert O'Callahan <robert@ocallahan.org>
Date: Mon, 10 Jan 2011 02:55:00 -0800
Subject: [PATCH] Bug 624279. Setup print mode for non-reftest-wait documents
 as soon as 'load' fires. This ensures that the initial reflow will happen in
 print mode. r=dbaron, a=test

This patch also ensures that the root element at the time of 'load' firing
is the element we always look for test metadata on --- even if the document
is removed from the window, or the root element is removed or replaced by
another root element, etc.
---
 layout/tools/reftest/reftest.js | 23 ++++++++++++-----------
 1 file changed, 12 insertions(+), 11 deletions(-)

diff --git a/layout/tools/reftest/reftest.js b/layout/tools/reftest/reftest.js
index 201dc5954121..bd04897ea55f 100644
--- a/layout/tools/reftest/reftest.js
+++ b/layout/tools/reftest/reftest.js
@@ -969,9 +969,7 @@ const STATE_WAITING_FOR_REFTEST_WAIT_REMOVAL = 1;
 const STATE_WAITING_TO_FINISH = 2;
 const STATE_COMPLETED = 3;
 
-function WaitForTestEnd() {
-    var currentDoc = gBrowser.contentDocument;
-    var contentRootElement = currentDoc ? currentDoc.documentElement : null;
+function WaitForTestEnd(contentRootElement) {
     var stopAfterPaintReceived = false;
     var state = STATE_WAITING_TO_FIRE_INVALIDATE_EVENT;
 
@@ -1176,23 +1174,20 @@ function OnDocumentLoad(event)
 
     var contentRootElement = currentDoc ? currentDoc.documentElement : null;
     setupZoom(contentRootElement);
+    var inPrintMode = false;
 
     function AfterOnLoadScripts() {
-        if (doPrintMode(contentRootElement)) {
-            LogInfo("AfterOnLoadScripts setting up print mode");
-            setupPrintMode();
-        }
-
         // Take a snapshot now. We need to do this before we check whether
         // we should wait, since this might trigger dispatching of
         // MozPaintWait events and make shouldWaitForExplicitPaintWaiters() true
         // below.
         InitCurrentCanvasWithSnapshot();
 
-        if (shouldWaitForExplicitPaintWaiters()) {
+        if (shouldWaitForExplicitPaintWaiters() ||
+            (!inPrintMode && doPrintMode(contentRootElement))) {
             LogInfo("AfterOnLoadScripts belatedly entering WaitForTestEnd");
             // Go into reftest-wait mode belatedly.
-            WaitForTestEnd();
+            WaitForTestEnd(contentRootElement);
         } else {
             RecordResult();
         }
@@ -1204,8 +1199,14 @@ function OnDocumentLoad(event)
         // unsuppressed, after the onload event has finished dispatching.
         gFailureReason = "timed out waiting for test to complete (trying to get into WaitForTestEnd)";
         LogInfo("OnDocumentLoad triggering WaitForTestEnd");
-        setTimeout(WaitForTestEnd, 0);
+        setTimeout(WaitForTestEnd, 0, contentRootElement);
     } else {
+        if (doPrintMode(contentRootElement)) {
+            LogInfo("OnDocumentLoad setting up print mode");
+            setupPrintMode();
+            inPrintMode = true;
+        }
+
         // Since we can't use a bubbling-phase load listener from chrome,
         // this is a capturing phase listener.  So do setTimeout twice, the
         // first to get us after the onload has fired in the content, and