[react-testing] descendants optimizations for faster react element tr…

…ee updates
Shopify · Apr 7, 2021 · 0674baa · 0674baa
1 parent 26f1a44
commit 0674baa
Show file tree

Hide file tree

Showing 4 changed files with 92 additions and 201 deletions.
diff --git a/packages/react-testing/src/element.ts b/packages/react-testing/src/element.ts
@@ -45,6 +45,24 @@ export class Element<Props> implements Node<Props> {
     return this.elementChildren;
   }
 
+  private get elementDescendants() {
+    if (!this.descendantsCache) {
+      this.descendantsCache = getDescendants(this);
+    }
+
+    return this.descendantsCache;
+  }
+
+  private get elementChildren() {
+    if (!this.elementChildrenCache) {
+      this.elementChildrenCache = this.allChildren.filter(
+        element => typeof element !== 'string',
+      ) as Element<unknown>[];
+    }
+
+    return this.elementChildrenCache;
+  }
+
   get descendants() {
     return this.elementDescendants;
   }
@@ -71,23 +89,14 @@ export class Element<Props> implements Node<Props> {
     return domNodes[0] || null;
   }
 
-  private readonly elementChildren: Element<unknown>[];
-  private readonly elementDescendants: Element<unknown>[];
+  private descendantsCache: Element<unknown>[] | null = null;
+  private elementChildrenCache: Element<unknown>[] | null = null;
 
   constructor(
     private readonly tree: Tree<Props>,
     private readonly allChildren: (Element<unknown> | string)[],
-    allDescendants: (Element<unknown> | string)[],
     public readonly root: Root,
-  ) {
-    this.elementChildren = allChildren.filter(
-      element => typeof element !== 'string',
-    ) as Element<unknown>[];
-
-    this.elementDescendants = allDescendants.filter(
-      element => typeof element !== 'string',
-    ) as Element<unknown>[];
-  }
+  ) {}
 
   data(key: string): string | undefined {
     return this.props[key.startsWith('data-') ? key : `data-${key}`];
@@ -252,3 +261,18 @@ function equalSubset(subset: object, full: object) {
     key => key in full && full[key] === subset[key],
   );
 }
+
+function getDescendants(element: any) {
+  const descendants: Element<unknown>[] = [];
+  // eslint-disable-next-line @typescript-eslint/prefer-for-of
+  for (let i = 0; i < element.allChildren.length; i++) {
+    const child = element.allChildren[i];
+    if (typeof child !== 'string') {
+      descendants.push(child);
+      // eslint-disable-next-line prefer-spread
+      descendants.push.apply(descendants, child.elementDescendants);
+    }
+  }
+
+  return descendants;
+}
diff --git a/packages/react-testing/src/root.tsx b/packages/react-testing/src/root.tsx
@@ -253,10 +253,12 @@ export class Root<Props> implements Node<Props> {
     if (this.wrapper == null) {
       this.root = null;
     } else {
-      const topElement = flatten(
-        ((this.wrapper as unknown) as ReactInstance)._reactInternalFiber,
+      const topElement = fiberToElement(
+        findCurrentFiberUsingSlowPath(
+          ((this.wrapper as unknown) as ReactInstance)._reactInternalFiber,
+        ),
         this,
-      )[0];
+      );
 
       this.root = this.resolveRoot(topElement as any) as any;
     }
@@ -284,52 +286,40 @@ function defaultRender(element: React.ReactElement<unknown>) {
   return element;
 }
 
-function flatten(
-  element: Fiber,
+function fiberToElement(
+  node: Fiber,
   root: Root<unknown>,
-): (Element<unknown> | string)[] {
-  const node: Fiber = findCurrentFiberUsingSlowPath(element);
-
+): Element<unknown> | string {
   if (node.tag === Tag.HostText) {
-    return [node.memoizedProps as string];
-  }
-
-  const props = {...((node.memoizedProps as any) || {})};
-  const {children, descendants} = childrenToTree(node.child, root);
-
-  return [
-    new Element(
-      {
-        tag: node.tag,
-        type: node.type,
-        props,
-        instance: node.stateNode,
-      },
-      children,
-      descendants,
-      root,
-    ),
-    ...descendants,
-  ];
+    return node.memoizedProps as string;
+  }
+
+  const props = node.memoizedProps as any;
+  const children = childrenToTree(node.child, root);
+
+  return new Element(
+    {
+      tag: node.tag,
+      type: node.type,
+      props,
+      instance: node.stateNode,
+    },
+    children,
+    root,
+  );
 }
 
 function childrenToTree(fiber: Fiber | null, root: Root<unknown>) {
   let currentFiber = fiber;
   const children: (string | Element<unknown>)[] = [];
-  const descendants: (string | Element<unknown>)[] = [];
 
   while (currentFiber != null) {
-    const result = flatten(currentFiber, root);
-
-    if (result.length > 0) {
-      children.push(result[0]);
-      descendants.push(...result);
-    }
-
+    const result = fiberToElement(currentFiber, root);
+    children.push(result);
     currentFiber = currentFiber.sibling;
   }
 
-  return {children, descendants};
+  return children;
 }
 
 function isPromise<T>(promise: T | Promise<T>): promise is Promise<T> {