Move the node finding function to op detail data class

PiperOrigin-RevId: 571051870

frontend/app/components/op_profile/op_profile_data.ts

@@ -1,4 +1,5 @@
 import {Node} from 'org_xprof/frontend/app/common/interfaces/op_profile.jsonpb_decls';
+import {ProfileOrNull} from 'org_xprof/frontend/app/common/interfaces/data_table';
 import * as utils from 'org_xprof/frontend/app/common/utils/utils';
 
 /** An op profile data class. */
@@ -36,4 +37,31 @@ export class OpProfileData {
       this.flopsUtilizationPercent = undefined;
     }
   }
+
+  // TODO: Make this function stateful or move to util.
+  findNode(
+      opProfile: ProfileOrNull, nodeName: string,
+      moduleName: string|undefined): Node|null|undefined {
+    if (!opProfile || !opProfile.byProgram) return null;
+    for (const topLevelNode of opProfile.byProgram.children!) {
+      // Find the program id from HloOpProfile by the selected XLA module.
+      // Assuming that the XLA modules and program ids are the same.
+      if (topLevelNode.name === moduleName) {
+        const node = this.findNodeHelper(topLevelNode.children, nodeName);
+        if (node) return node;
+      }
+    }
+    return null;
+  }
+
+  private findNodeHelper(children: Node[]|null|undefined, name: string): Node
+      |null|undefined {
+    if (!children) return null;
+    for (const node of children) {
+      if (node.name === name) return node;
+      const findChildren = this.findNodeHelper(node.children, name);
+      if (findChildren) return findChildren;
+    }
+    return null;
+  }
 }