Fix search for accessibility view

Summary:
Bug: On search, the accessibility views didn't used to expand.Look at the following video to understand the bug

{F153135892}

Reviewed By: danielbuechele

Differential Revision: D14407511

fbshipit-source-id: 011bf31b5bd10f5a7dad3d5d30703ec3bdeaf9c1

src/plugins/layout/layout2/ProxyArchiveClient.js

@@ -25,6 +25,7 @@ function constructSearchResultTree(
   isMatch: boolean,
   children: Array<SearchResultTree>,
   AXMode: boolean,
+  AXNode: ?Element,
 ): SearchResultTree {
   let searchResult = {
     id: node.id,
@@ -32,9 +33,8 @@ function constructSearchResultTree(
     hasChildren: children.length > 0,
     children: children.length > 0 ? children : null,
     element: node,
-    axElement: null,
+    axElement: AXNode,
   };
-  searchResult[`${propsForPersistedState(AXMode).ELEMENT}`] = node;
   return searchResult;
 }
 
@@ -64,7 +64,15 @@ export function searchNodes(
   }
 
   if (match || children.length > 0) {
-    return cloneDeep(constructSearchResultTree(node, match, children, AXMode));
+    return cloneDeep(
+      constructSearchResultTree(
+        node,
+        match,
+        children,
+        AXMode,
+        AXMode ? state.AXelements[node.id] : null,
+      ),
+    );
   }
   return null;
 }

src/plugins/layout/layout2/Search.js

@@ -73,21 +73,29 @@ export default class Search extends Component<Props, State> {
     });
 
     if (!query) {
-      this.displaySearchResults({query: '', results: null});
+      this.displaySearchResults(
+        {query: '', results: null},
+        this.props.inAXMode,
+      );
     } else {
       this.props.client
         .call('getSearchResults', {query, axEnabled: this.props.inAXMode})
-        .then(response => this.displaySearchResults(response));
+        .then(response =>
+          this.displaySearchResults(response, this.props.inAXMode),
+        );
     }
   }
 
-  displaySearchResults({
+  displaySearchResults(
+    {
       results,
       query,
     }: {
       results: ?SearchResultTree,
       query: string,
-  }) {
+    },
+    axMode: boolean,
+  ) {
     this.setState({
       outstandingSearchQuery:
         query === this.state.outstandingSearchQuery
@@ -106,10 +114,30 @@ export default class Search extends Component<Props, State> {
           expanded: element.children.some(c => searchResultIDs.has(c)),
         },
       }),
-      {...this.props.persistedState.elements},
+      this.props.persistedState.elements,
     );
 
-    this.props.setPersistedState({elements});
+    let {AXelements} = this.props.persistedState;
+    if (axMode) {
+      AXelements = searchResults.reduce(
+        (acc: ElementMap, {axElement}: SearchResultTree) => {
+          if (!axElement) {
+            return acc;
+          }
+          return {
+            ...acc,
+            [axElement.id]: {
+              ...axElement,
+              // expand all search results, that we have have children for
+              expanded: axElement.children.some(c => searchResultIDs.has(c)),
+            },
+          };
+        },
+        this.props.persistedState.AXelements,
+      );
+    }
+
+    this.props.setPersistedState({elements, AXelements});
 
     this.props.onSearchResults({
       matches: new Set(