[react-testing] improved performance by making descendants calculatio…

…ns lazy for element children

packages/react-testing/src/element.ts

@@ -45,6 +45,14 @@ export class Element<Props> implements Node<Props> {
     return this.elementChildren;
   }
 
+  private get elementDescendants() {
+    if (!this.descendantsCache) {
+      this.descendantsCache = getDescendants(this);
+    }
+
+    return this.descendantsCache;
+  }
+
   get descendants() {
     return this.elementDescendants;
   }
@@ -72,21 +80,19 @@ export class Element<Props> implements Node<Props> {
   }
 
   private readonly elementChildren: Element<unknown>[];
-  private readonly elementDescendants: Element<unknown>[];
+  private descendantsCache: Element<unknown>[] | null;
 
   constructor(
     private readonly tree: Tree<Props>,
     private readonly allChildren: (Element<unknown> | string)[],
-    allDescendants: (Element<unknown> | string)[],
+    elementDescendants: Element<unknown>[] | null,
     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>[];
+    this.descendantsCache = elementDescendants;
   }
 
   data(key: string): string | undefined {
@@ -252,3 +258,14 @@ function equalSubset(subset: object, full: object) {
     key => key in full && full[key] === subset[key],
   );
 }
+
+function getDescendants(element) {
+  const descendants: Element<unknown>[] = [];
+  // eslint-disable-next-line @typescript-eslint/prefer-for-of
+  for (let i = 0; i < element.elementChildren.length; i++) {
+    const child = element.elementChildren[i];
+    descendants.push(child, ...getDescendants(child));
+  }
+
+  return descendants;
+}

packages/react-testing/src/root.tsx

@@ -253,10 +253,10 @@ export class Root<Props> implements Node<Props> {
     if (this.wrapper == null) {
       this.root = null;
     } else {
-      const topElement = flatten(
+      const topElement = fiberToElement(
         ((this.wrapper as unknown) as ReactInstance)._reactInternalFiber,
         this,
-      )[0];
+      );
 
       this.root = this.resolveRoot(topElement as any) as any;
     }
@@ -284,52 +284,49 @@ function defaultRender(element: React.ReactElement<unknown>) {
   return element;
 }
 
-function flatten(
+function fiberToElement(
   element: Fiber,
   root: Root<unknown>,
-): (Element<unknown> | string)[] {
+  descendants = [],
+  rootLevel = true,
+): Element<unknown> | string {
   const node: Fiber = findCurrentFiberUsingSlowPath(element);
 
   if (node.tag === Tag.HostText) {
-    return [node.memoizedProps as string];
+    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,
-  ];
+  const children = childrenToTree(node.child, root, descendants);
+
+  return new Element(
+    {
+      tag: node.tag,
+      type: node.type,
+      props,
+      instance: node.stateNode,
+    },
+    children,
+    rootLevel ? descendants : null,
+    root,
+  );
 }
 
-function childrenToTree(fiber: Fiber | null, root: Root<unknown>) {
+function childrenToTree(fiber: Fiber | null, root: Root<unknown>, descendants) {
   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, descendants, false);
+    children.push(result);
+    if (typeof result !== 'string') {
+      descendants.push(result);
     }
 
     currentFiber = currentFiber.sibling;
   }
 
-  return {children, descendants};
+  return children;
 }
 
 function isPromise<T>(promise: T | Promise<T>): promise is Promise<T> {

packages/react-testing/src/tests/element.test.tsx

@@ -122,7 +122,7 @@ describe('Element', () => {
       const element = new Element(
         defaultTree,
         [divOne, 'Some text', divTwo],
-        [divOne, 'Some text', divTwo],
+        [divOne, divTwo],
         defaultRoot,
       );
 
@@ -141,17 +141,6 @@ describe('Element', () => {
 
       expect(element).toHaveProperty('descendants', [divOne, divTwo]);
     });
-
-    it('does not return string descendants', () => {
-      const element = new Element(
-        defaultTree,
-        [divOne],
-        [divOne, 'Some text', divTwo],
-        defaultRoot,
-      );
-
-      expect(element).toHaveProperty('descendants', [divOne, divTwo]);
-    });
   });
 
   describe('#domNodes', () => {
@@ -271,7 +260,7 @@ describe('Element', () => {
       const element = new Element(
         {...defaultTree, tag: Tag.FunctionComponent, type: DummyComponent},
         [elementChild, childTextTwo],
-        [elementChild, childTextTwo, descendantText],
+        [elementChild],
         defaultRoot,
       );
 
@@ -282,7 +271,7 @@ describe('Element', () => {
       const element = new Element(
         {...defaultTree, tag: Tag.HostPortal, type: DummyComponent},
         ['Hello world'],
-        ['Hello world'],
+        [],
         defaultRoot,
       );
 
@@ -310,15 +299,14 @@ describe('Element', () => {
     it('concatenates the HTML contents of all child elements and child text', () => {
       const childHtml = 'foo ';
       const childText = 'bar';
-      const descendantText = ' baz?';
 
       const elementChild = new Element(defaultTree, [], [], defaultRoot);
       jest.spyOn(elementChild, 'text').mockImplementation(() => childHtml);
 
       const element = new Element(
         {...defaultTree, tag: Tag.FunctionComponent, type: DummyComponent},
         [elementChild, childText],
-        [elementChild, childText, descendantText],
+        [elementChild],
         defaultRoot,
       );
 
@@ -329,7 +317,7 @@ describe('Element', () => {
       const element = new Element(
         {...defaultTree, tag: Tag.HostPortal, type: DummyComponent},
         ['Hello world'],
-        ['Hello world'],
+        [],
         defaultRoot,
       );