From 05489aad2cabe03e09a4fab909f4558132518d5e Mon Sep 17 00:00:00 2001
From: bendtherules <abhasbhattacharya2@gmail.com>
Date: Thu, 8 Aug 2019 11:51:33 +0530
Subject: [PATCH] src/utils - Add getRootFiberNodeFromDOM helper To find root
 Fiber node from overall html document or a particular element. Ideally, this
 should be called with the same container element used in ReactDom.render.

---
 src/mocked-types/index.ts |  8 ++++++++
 src/utils.ts              | 38 ++++++++++++++++++++++++++++++++++++--
 2 files changed, 44 insertions(+), 2 deletions(-)

diff --git a/src/mocked-types/index.ts b/src/mocked-types/index.ts
index 8d779fe..775b975 100644
--- a/src/mocked-types/index.ts
+++ b/src/mocked-types/index.ts
@@ -1,5 +1,13 @@
 import * as React from "react";
 
+export interface FiberNodeDOMContainer extends Element {
+  _reactRootContainer: {
+    _internalRoot: {
+      current: FiberNode | null;
+    };
+  };
+}
+
 export type FiberNode =
   | FiberNodeForComponentClass
   | FiberNodeForFunctionComponent
diff --git a/src/utils.ts b/src/utils.ts
index e6d943b..f275b93 100644
--- a/src/utils.ts
+++ b/src/utils.ts
@@ -3,7 +3,8 @@ import {
   FiberNode,
   FiberNodeisHTMLLike,
   FiberNodeForFunctionComponent,
-  FiberNodeForComponentClass
+  FiberNodeForComponentClass,
+  FiberNodeDOMContainer
 } from "./mocked-types";
 
 function isNodeHtmlLike(node: FiberNode): node is FiberNodeisHTMLLike {
@@ -60,6 +61,38 @@ function isConstructorFunctionComponent(
   return !isConstructorComponentClass(ctr);
 }
 
+function doesElementContainRootFiberNode(
+  element: Element
+): element is FiberNodeDOMContainer {
+  return element.hasOwnProperty("_reactRootContainer");
+}
+
+/**
+ * Util to find root React Fiber node from html DOM tree.
+ * Returns null, if not found.SHould be called after ReactDOM.render is finished.
+ * @param startElement Starting DOM element to seach from.
+ * If not found, it checks inside its child nodes. Defaults to document.body
+ */
+function getRootFiberNodeFromDOM(startElement?: Element): FiberNode | null {
+  if (startElement === undefined) {
+    startElement = document.body;
+  }
+
+  if (doesElementContainRootFiberNode(startElement)) {
+    return startElement._reactRootContainer._internalRoot.current;
+  }
+
+  let returnFiberNode = null;
+  for (const childNode of startElement.children) {
+    returnFiberNode = getRootFiberNodeFromDOM(childNode);
+    if (returnFiberNode !== null) {
+      return returnFiberNode;
+    }
+  }
+
+  return returnFiberNode;
+}
+
 export {
   isNodeHtmlLike,
   isNodeNotHtmlLike,
@@ -67,5 +100,6 @@ export {
   isNodeComponentClass,
   isConstructorHtmlLike,
   isConstructorComponentClass,
-  isConstructorFunctionComponent
+  isConstructorFunctionComponent,
+  getRootFiberNodeFromDOM
 };