1 files changed, 28 insertions, 22 deletions

diff --git a/dom/base/CustomElementRegistry.cpp b/dom/base/CustomElementRegistry.cpp
index 75477dbaa..6c643d861 100644
--- a/dom/base/CustomElementRegistry.cpp
+++ b/dom/base/CustomElementRegistry.cpp
@@ -78,12 +78,21 @@ CustomElementCallback::CustomElementCallback(Element* aThisObject,
 {
 }
 
+//-----------------------------------------------------
+// CustomElementData
+
 CustomElementData::CustomElementData(nsIAtom* aType)
-  : mType(aType),
-    mCurrentCallback(-1),
-    mElementIsBeingCreated(false),
-    mCreatedCallbackInvoked(true),
-    mAssociatedMicroTask(-1)
+  : CustomElementData(aType, CustomElementData::State::eUndefined)
+{
+}
+
+CustomElementData::CustomElementData(nsIAtom* aType, State aState)
+  : mType(aType)
+  , mCurrentCallback(-1)
+  , mElementIsBeingCreated(false)
+  , mCreatedCallbackInvoked(true)
+  , mAssociatedMicroTask(-1)
+  , mState(aState)
 {
 }
 
@@ -99,6 +108,9 @@ CustomElementData::RunCallbackQueue()
   mCurrentCallback = -1;
 }
 
+//-----------------------------------------------------
+// CustomElementRegistry
+
 // Only needed for refcounted objects.
 NS_IMPL_CYCLE_COLLECTION_CLASS(CustomElementRegistry)
 
@@ -305,10 +317,15 @@ CustomElementRegistry::SetupCustomElement(Element* aElement,
     aElement->SetAttr(kNameSpaceID_None, nsGkAtoms::is, *aTypeExtension, true);
   }
 
-  CustomElementDefinition* data = LookupCustomElementDefinition(
+  // SetupCustomElement() should be called with an element that don't have
+  // CustomElementData setup, if not we will hit the assertion in
+  // SetCustomElementData().
+  aElement->SetCustomElementData(new CustomElementData(typeAtom));
+
+  CustomElementDefinition* definition = LookupCustomElementDefinition(
     aElement->NodeInfo()->LocalName(), aTypeExtension);
 
-  if (!data) {
+  if (!definition) {
     // The type extension doesn't exist in the registry,
     // thus we don't need to enqueue callback or adjust
     // the "is" attribute, but it is possibly an upgrade candidate.
@@ -316,7 +333,7 @@ CustomElementRegistry::SetupCustomElement(Element* aElement,
     return;
   }
 
-  if (data->mLocalName != tagAtom) {
+  if (definition->mLocalName != tagAtom) {
     // The element doesn't match the local name for the
     // definition, thus the element isn't a custom element
     // and we don't need to do anything more.
@@ -325,7 +342,7 @@ CustomElementRegistry::SetupCustomElement(Element* aElement,
 
   // Enqueuing the created callback will set the CustomElementData on the
   // element, causing prototype swizzling to occur in Element::WrapObject.
-  EnqueueLifecycleCallback(nsIDocument::eCreated, aElement, nullptr, data);
+  EnqueueLifecycleCallback(nsIDocument::eCreated, aElement, nullptr, definition);
 }
 
 void
@@ -334,7 +351,8 @@ CustomElementRegistry::EnqueueLifecycleCallback(nsIDocument::ElementCallbackType
                                                 LifecycleCallbackArgs* aArgs,
                                                 CustomElementDefinition* aDefinition)
 {
-  CustomElementData* elementData = aCustomElement->GetCustomElementData();
+  RefPtr<CustomElementData> elementData = aCustomElement->GetCustomElementData();
+  MOZ_ASSERT(elementData, "CustomElementData should exist");
 
   // Let DEFINITION be ELEMENT's definition
   CustomElementDefinition* definition = aDefinition;
@@ -354,16 +372,6 @@ CustomElementRegistry::EnqueueLifecycleCallback(nsIDocument::ElementCallbackType
     }
   }
 
-  if (!elementData) {
-    // Create the custom element data the first time
-    // that we try to enqueue a callback.
-    elementData = new CustomElementData(definition->mType);
-    // aCustomElement takes ownership of elementData
-    aCustomElement->SetCustomElementData(elementData);
-    MOZ_ASSERT(aType == nsIDocument::eCreated,
-               "First callback should be the created callback");
-  }
-
   // Let CALLBACK be the callback associated with the key NAME in CALLBACKS.
   CallbackFunction* func = nullptr;
   switch (aType) {
@@ -881,8 +889,6 @@ CustomElementRegistry::Upgrade(Element* aElement,
     }
   } // Leave prototype's compartment.
 
-  // Enqueuing the created callback will set the CustomElementData on the
-  // element, causing prototype swizzling to occur in Element::WrapObject.
   EnqueueLifecycleCallback(nsIDocument::eCreated, aElement, nullptr, aDefinition);
 }