2 files changed, 54 insertions, 43 deletions

diff --git a/dom/base/CustomElementRegistry.cpp b/dom/base/CustomElementRegistry.cpp
index 1c1e3d6ae..ce2e4b04f 100644
--- a/dom/base/CustomElementRegistry.cpp
+++ b/dom/base/CustomElementRegistry.cpp
@@ -799,40 +799,6 @@ CustomElementRegistry::Upgrade(Element* aElement,
 
   MOZ_ASSERT(aElement->IsHTMLElement(aDefinition->mLocalName));
 
-  AutoJSAPI jsapi;
-  if (NS_WARN_IF(!jsapi.Init(mWindow))) {
-    return;
-  }
-
-  JSContext* cx = jsapi.cx();
-
-  JS::Rooted<JSObject*> reflector(cx, aElement->GetWrapper());
-  if (reflector) {
-    Maybe<JSAutoCompartment> ac;
-    JS::Rooted<JSObject*> prototype(cx, aDefinition->mPrototype);
-    if (aElement->NodePrincipal()->SubsumesConsideringDomain(nsContentUtils::ObjectPrincipal(prototype))) {
-      ac.emplace(cx, reflector);
-      if (!JS_WrapObject(cx, &prototype) ||
-          !JS_SetPrototype(cx, reflector, prototype)) {
-        return;
-      }
-    } else {
-      // We want to set the custom prototype in the compartment where it was
-      // registered. We store the prototype from define() without unwrapped,
-      // hence the prototype's compartment is the compartment where it was
-      // registered.
-      // In the case that |reflector| and |prototype| are in different
-      // compartments, this will set the prototype on the |reflector|'s wrapper
-      // and thus only visible in the wrapper's compartment, since we know
-      // reflector's principal does not subsume prototype's in this case.
-      ac.emplace(cx, prototype);
-      if (!JS_WrapObject(cx, &reflector) ||
-          !JS_SetPrototype(cx, reflector, prototype)) {
-        return;
-      }
-    }
-  }
-
   EnqueueLifecycleCallback(nsIDocument::eCreated, aElement, nullptr, aDefinition);
 }
 
diff --git a/dom/base/nsDocument.cpp b/dom/base/nsDocument.cpp
index a0857faf8..602da421a 100644
--- a/dom/base/nsDocument.cpp
+++ b/dom/base/nsDocument.cpp
@@ -5670,19 +5670,64 @@ nsDocument::CustomElementConstructor(JSContext* aCx, unsigned aArgc, JS::Value*
     return true;
   }
 
-  nsDependentAtomString localName(definition->mLocalName);
+  RefPtr<Element> element;
+
+  // We integrate with construction stack and do prototype swizzling here, so
+  // that old upgrade behavior could also share the new upgrade steps.
+  // And this old upgrade will be remove at some point (when everything is
+  // switched to latest custom element spec).
+  nsTArray<RefPtr<nsGenericHTMLElement>>& constructionStack =
+    definition->mConstructionStack;
+  if (constructionStack.Length()) {
+    element = constructionStack.LastElement();
+    NS_ENSURE_TRUE(element != ALEADY_CONSTRUCTED_MARKER, false);
+
+    // Do prototype swizzling if dom reflector exists.
+    JS::Rooted<JSObject*> reflector(aCx, element->GetWrapper());
+    if (reflector) {
+      Maybe<JSAutoCompartment> ac;
+      JS::Rooted<JSObject*> prototype(aCx, definition->mPrototype);
+      if (element->NodePrincipal()->SubsumesConsideringDomain(nsContentUtils::ObjectPrincipal(prototype))) {
+        ac.emplace(aCx, reflector);
+        if (!JS_WrapObject(aCx, &prototype) ||
+            !JS_SetPrototype(aCx, reflector, prototype)) {
+          return false;
+        }
+      } else {
+        // We want to set the custom prototype in the compartment where it was
+        // registered. We store the prototype from define() without unwrapped,
+        // hence the prototype's compartment is the compartment where it was
+        // registered.
+        // In the case that |reflector| and |prototype| are in different
+        // compartments, this will set the prototype on the |reflector|'s wrapper
+        // and thus only visible in the wrapper's compartment, since we know
+        // reflector's principal does not subsume prototype's in this case.
+        ac.emplace(aCx, prototype);
+        if (!JS_WrapObject(aCx, &reflector) ||
+            !JS_SetPrototype(aCx, reflector, prototype)) {
+          return false;
+        }
+      }
 
-  nsCOMPtr<Element> element =
-    document->CreateElem(localName, nullptr, kNameSpaceID_XHTML);
-  NS_ENSURE_TRUE(element, true);
+      // Wrap into current context.
+      if (!JS_WrapObject(aCx, &reflector)) {
+        return false;
+      }
 
-  if (definition->mLocalName != typeAtom) {
-    // This element is a custom element by extension, thus we need to
-    // do some special setup. For non-extended custom elements, this happens
-    // when the element is created.
-    nsContentUtils::SetupCustomElement(element, &elemName);
+      args.rval().setObject(*reflector);
+      return true;
+    }
+  } else {
+    nsDependentAtomString localName(definition->mLocalName);
+    element =
+      document->CreateElem(localName, nullptr, kNameSpaceID_XHTML,
+                           (definition->mLocalName != typeAtom) ? &elemName
+                                                                : nullptr);
+    NS_ENSURE_TRUE(element, false);
   }
 
+  // The prototype setup happens in Element::WrapObject().
+
   nsresult rv = nsContentUtils::WrapNative(aCx, element, element, args.rval());
   NS_ENSURE_SUCCESS(rv, true);