@finos/legend-query-builder 4.15.3 → 4.15.4

package/src/stores/explorer/QueryFunctionsExplorerState.ts

@@ -16,12 +16,16 @@
 
 import {
   type PackageableElement,
-  ConcreteFunctionDefinition,
+  type FunctionAnalysisInfo,
+  type PureModel,
   Package,
   Unit,
   ROOT_PACKAGE_NAME,
+  buildFunctionAnalysisInfoFromConcreteFunctionDefinition,
+  getOrCreateGraphPackage,
 } from '@finos/legend-graph';
 import {
+  ActionState,
   addUniqueEntry,
   guaranteeNonNullable,
   isNonNullable,
@@ -47,16 +51,12 @@ export class QueryBuilderFunctionsExplorerTreeNodeData implements TreeNodeData {
   label: string;
   childrenIds: string[] = [];
   isOpen?: boolean | undefined;
-  packageableElement: PackageableElement;
+  package?: PackageableElement;
+  functionAnalysisInfo?: FunctionAnalysisInfo | undefined;
 
-  constructor(
-    id: string,
-    label: string,
-    packageableElement: PackageableElement,
-  ) {
+  constructor(id: string, label: string) {
     this.id = id;
     this.label = label;
-    this.packageableElement = packageableElement;
   }
 }
 
@@ -73,29 +73,38 @@ const getValidDisplayablePackageSet = (
   }
 };
 
-export const generateFunctionsExplorerTreeNodeData = (
+const generateFunctionsExplorerTreeNodeDataFromPackage = (
   queryBuilderState: QueryBuilderState,
-  element: PackageableElement,
+  element: Package,
   rootPackageName: ROOT_PACKAGE_NAME,
 ): QueryBuilderFunctionsExplorerTreeNodeData => ({
   id: element.path,
   label: element.name,
-  childrenIds:
-    element instanceof Package
-      ? element.children
-          .filter((child) => !(child instanceof Unit))
-          .filter(
-            (child) =>
-              (child instanceof Package &&
-                getValidDisplayablePackageSet(
-                  queryBuilderState,
-                  rootPackageName,
-                ).has(child)) ||
-              child instanceof ConcreteFunctionDefinition,
-          )
-          .map((child) => child.path)
-      : [],
-  packageableElement: element,
+  childrenIds: element.children
+    .filter((child) => !(child instanceof Unit))
+    .filter(
+      (child) =>
+        child instanceof Package &&
+        getValidDisplayablePackageSet(queryBuilderState, rootPackageName).has(
+          child,
+        ),
+    )
+    .map((child) => child.path)
+    .concat(
+      queryBuilderState.functionsExplorerState.packagePathToFunctionInfoMap
+        ?.get(element.path)
+        ?.map((info) => info.functionPath) ?? [],
+    ),
+  package: element,
+});
+
+export const generateFunctionsExplorerTreeNodeDataFromFunctionAnalysisInfo = (
+  functionAnalysisInfo: FunctionAnalysisInfo,
+): QueryBuilderFunctionsExplorerTreeNodeData => ({
+  id: functionAnalysisInfo.functionPath,
+  label: functionAnalysisInfo.name,
+  childrenIds: [],
+  functionAnalysisInfo: functionAnalysisInfo,
 });
 
 const generateFunctionsExplorerTreeNodeChilrdren = (
@@ -104,34 +113,48 @@ const generateFunctionsExplorerTreeNodeChilrdren = (
   data: TreeData<QueryBuilderFunctionsExplorerTreeNodeData>,
   rootPackageName = ROOT_PACKAGE_NAME.MAIN,
 ): void => {
+  const functionInfoMap =
+    rootPackageName === ROOT_PACKAGE_NAME.MAIN
+      ? queryBuilderState.functionsExplorerState.functionInfoMap
+      : queryBuilderState.functionsExplorerState.dependencyFunctionInfoMap;
   const validDisplayablePackageSet = getValidDisplayablePackageSet(
     queryBuilderState,
     rootPackageName,
   );
-  node.childrenIds = (node.packageableElement as Package).children
-    .filter((child) => !(child instanceof Unit))
-    .filter(
-      (child) =>
-        // avoid displaying empty packages
-        (child instanceof Package && validDisplayablePackageSet.has(child)) ||
-        child instanceof ConcreteFunctionDefinition,
-    )
-    .map((child) => child.path);
-  (node.packageableElement as Package).children
-    .filter((child) => !(child instanceof Unit))
+  const qualifiedExtraFunctionPaths =
+    queryBuilderState.functionsExplorerState.packagePathToFunctionInfoMap
+      ?.get(node.id)
+      ?.map((info) => info.functionPath);
+
+  const childrenPackages = (node.package as Package).children
     .filter(
       (child) =>
         // avoid displaying empty packages
-        (child instanceof Package && validDisplayablePackageSet.has(child)) ||
-        child instanceof ConcreteFunctionDefinition,
-    )
-    .map((child) =>
-      generateFunctionsExplorerTreeNodeData(
-        queryBuilderState,
-        child,
-        rootPackageName,
-      ),
+        child instanceof Package && validDisplayablePackageSet.has(child),
     )
+    .map((child) => child as Package);
+
+  node.childrenIds = childrenPackages
+    .map((child) => child.path)
+    .concat(qualifiedExtraFunctionPaths ?? []);
+
+  const childNodesFromPackage = childrenPackages.map((child) =>
+    generateFunctionsExplorerTreeNodeDataFromPackage(
+      queryBuilderState,
+      child,
+      rootPackageName,
+    ),
+  );
+
+  const childNodesFromFunction = qualifiedExtraFunctionPaths
+    ?.map((path) => functionInfoMap?.get(path))
+    .filter(isNonNullable)
+    .map((info) =>
+      generateFunctionsExplorerTreeNodeDataFromFunctionAnalysisInfo(info),
+    );
+
+  childNodesFromPackage
+    .concat(childNodesFromFunction ?? [])
     .forEach((childNode) => {
       const currentNode = data.nodes.get(childNode.id);
       if (currentNode) {
@@ -157,14 +180,20 @@ export const getFunctionsExplorerTreeData = (
   switch (rootPackageName) {
     case ROOT_PACKAGE_NAME.PROJECT_DEPENDENCY_ROOT:
       if (
-        queryBuilderState.graphManagerState.graph.dependencyManager.functions
-          .length === 0
+        !queryBuilderState.functionsExplorerState.dependencyFunctionInfoMap ||
+        Array.from(
+          queryBuilderState.functionsExplorerState.dependencyFunctionInfoMap,
+        ).length === 0
       ) {
         return { rootIds, nodes };
       }
       break;
     default:
-      if (queryBuilderState.graphManagerState.graph.ownFunctions.length === 0) {
+      if (
+        !queryBuilderState.functionsExplorerState.functionInfoMap ||
+        Array.from(queryBuilderState.functionsExplorerState.functionInfoMap)
+          .length === 0
+      ) {
         return { rootIds, nodes };
       }
   }
@@ -172,22 +201,19 @@ export const getFunctionsExplorerTreeData = (
   roots.forEach((root) => {
     root.children
       .slice()
-      .filter((child) => !(child instanceof Unit))
       .filter(
         (child) =>
           child instanceof Package && validDisplayablePackageSet.has(child),
       )
+      .map((child) => child as Package)
       .sort((a, b) => a.name.localeCompare(b.name))
-      .sort(
-        (a, b) =>
-          (b instanceof Package ? 1 : 0) - (a instanceof Package ? 1 : 0),
-      )
       .forEach((childPackage) => {
-        const childTreeNodeData = generateFunctionsExplorerTreeNodeData(
-          queryBuilderState,
-          childPackage,
-          rootPackageName,
-        );
+        const childTreeNodeData =
+          generateFunctionsExplorerTreeNodeDataFromPackage(
+            queryBuilderState,
+            childPackage,
+            rootPackageName,
+          );
         addUniqueEntry(rootIds, childTreeNodeData.id);
         nodes.set(childTreeNodeData.id, childTreeNodeData);
       });
@@ -211,16 +237,14 @@ export const getFunctionsExplorerTreeNodeChildren = (
     .map((id) => data.nodes.get(id))
     .filter(isNonNullable)
     .sort(compareLabelFn)
-    .sort(
-      (a, b) =>
-        (b.packageableElement instanceof Package ? 1 : 0) -
-        (a.packageableElement instanceof Package ? 1 : 0),
-    );
+    .sort((a, b) => (b.package ? 1 : 0) - (a.package ? 1 : 0));
 };
 
 const getAllPackagesFromElement = (element: PackageableElement): Package[] => {
   if (element.package) {
-    return [element.package].concat(getAllPackagesFromElement(element.package));
+    return (element instanceof Package ? [element] : []).concat(
+      [element.package].concat(getAllPackagesFromElement(element.package)),
+    );
   }
   return [];
 };
@@ -228,30 +252,36 @@ const getAllPackagesFromElement = (element: PackageableElement): Package[] => {
 export class QueryFunctionExplorerState {
   readonly uuid = uuid();
   queryFunctionsState: QueryFunctionsExplorerState;
-  concreteFunctionDefinition: ConcreteFunctionDefinition;
+  functionAnalysisInfo: FunctionAnalysisInfo;
 
   constructor(
     queryFunctionsState: QueryFunctionsExplorerState,
-    concreteFunctionDefinition: ConcreteFunctionDefinition,
+    functionAnalysisInfo: FunctionAnalysisInfo,
   ) {
     makeObservable(this, {
-      concreteFunctionDefinition: observable,
+      functionAnalysisInfo: observable,
     });
     this.queryFunctionsState = queryFunctionsState;
-    this.concreteFunctionDefinition = concreteFunctionDefinition;
+    this.functionAnalysisInfo = functionAnalysisInfo;
   }
 }
 
 export class QueryFunctionsExplorerState {
+  readonly initState = ActionState.create();
+
   queryBuilderState: QueryBuilderState;
   treeData?: TreeData<QueryBuilderFunctionsExplorerTreeNodeData> | undefined;
   dependencyTreeData?:
     | TreeData<QueryBuilderFunctionsExplorerTreeNodeData>
     | undefined;
+  _functionGraph: PureModel;
   functionExplorerStates: QueryFunctionExplorerState[] = [];
   dependencyFunctionExplorerStates: QueryFunctionExplorerState[] = [];
   displayablePackagesSet: Set<Package> = new Set<Package>();
   dependencyDisplayablePackagesSet: Set<Package> = new Set<Package>();
+  functionInfoMap?: Map<string, FunctionAnalysisInfo>;
+  dependencyFunctionInfoMap?: Map<string, FunctionAnalysisInfo>;
+  packagePathToFunctionInfoMap?: Map<string, FunctionAnalysisInfo[]>;
 
   constructor(queryBuilderState: QueryBuilderState) {
     makeObservable(this, {
@@ -259,13 +289,22 @@ export class QueryFunctionsExplorerState {
       dependencyFunctionExplorerStates: observable.ref,
       treeData: observable.ref,
       dependencyTreeData: observable.ref,
+      _functionGraph: observable,
+      functionInfoMap: observable,
+      dependencyFunctionInfoMap: observable,
+      packagePathToFunctionInfoMap: observable,
+      setFunctionInfoMap: action,
+      setDependencyFunctionInfoMap: action,
       setTreeData: action,
+      setPackagePathToFunctionInfoMap: action,
       setDependencyTreeData: action,
       refreshTree: action,
       onTreeNodeSelect: action,
+      initializeTreeData: action,
     });
     this.queryBuilderState = queryBuilderState;
-    this.initializeTreeData();
+    this._functionGraph =
+      this.queryBuilderState.graphManagerState.createNewGraph();
   }
 
   getTreeData(
@@ -279,18 +318,48 @@ export class QueryFunctionsExplorerState {
     }
   }
 
+  setFunctionInfoMap(info: Map<string, FunctionAnalysisInfo>) {
+    this.functionInfoMap = info;
+  }
+
+  setDependencyFunctionInfoMap(info: Map<string, FunctionAnalysisInfo>) {
+    this.dependencyFunctionInfoMap = info;
+  }
+
+  setPackagePathToFunctionInfoMap(map: Map<string, FunctionAnalysisInfo[]>) {
+    this.packagePathToFunctionInfoMap = map;
+  }
+
   async initializeDisplayablePackagesSet(): Promise<void> {
-    this.queryBuilderState.graphManagerState.graph.ownFunctions
-      .map((f) => getAllPackagesFromElement(f))
-      .flat()
-      .forEach((pkg) => this.displayablePackagesSet.add(pkg));
+    if (this.functionInfoMap) {
+      Array.from(this.functionInfoMap.values())
+        .map((info) =>
+          getOrCreateGraphPackage(
+            this._functionGraph,
+            info.packagePath,
+            undefined,
+          ),
+        )
+        .map((f) => getAllPackagesFromElement(f))
+        .flat()
+        .forEach((pkg) => this.displayablePackagesSet.add(pkg));
+    }
   }
 
   async initializeDependencyDisplayablePackagesSet(): Promise<void> {
-    this.queryBuilderState.graphManagerState.graph.dependencyManager.functions
-      .map((f) => getAllPackagesFromElement(f))
-      .flat()
-      .forEach((pkg) => this.dependencyDisplayablePackagesSet.add(pkg));
+    if (this.dependencyFunctionInfoMap) {
+      Array.from(this.dependencyFunctionInfoMap.values())
+        .map((info) =>
+          getOrCreateGraphPackage(
+            this._functionGraph,
+            info.packagePath,
+            undefined,
+          ),
+        )
+        .map((f) => getAllPackagesFromElement(f))
+        .flat()
+        .forEach((pkg) => this.dependencyDisplayablePackagesSet.add(pkg));
+    }
   }
 
   setTreeData(
@@ -327,7 +396,7 @@ export class QueryFunctionsExplorerState {
     data: TreeData<QueryBuilderFunctionsExplorerTreeNodeData>,
     rootPackageName = ROOT_PACKAGE_NAME.MAIN,
   ): void => {
-    if (node.packageableElement instanceof Package) {
+    if (node.package) {
       if (node.childrenIds.length) {
         node.isOpen = !node.isOpen;
         generateFunctionsExplorerTreeNodeChilrdren(
@@ -347,41 +416,105 @@ export class QueryFunctionsExplorerState {
     }
   };
 
+  initializeFunctionInfoMap(): void {
+    const functionInfoMap = new Map<string, FunctionAnalysisInfo>();
+    const dependencyFunctionInfoMap = new Map<string, FunctionAnalysisInfo>();
+    const functionInfos =
+      buildFunctionAnalysisInfoFromConcreteFunctionDefinition(
+        this.queryBuilderState.graphManagerState.graph.ownFunctions,
+        this._functionGraph,
+      );
+    functionInfos.forEach((info) =>
+      functionInfoMap.set(info.functionPath, info),
+    );
+    if (
+      this.queryBuilderState.graphManagerState.graph.dependencyManager
+        .hasDependencies
+    ) {
+      const dependencyFunctions =
+        this.queryBuilderState.graphManagerState.graph.dependencyManager
+          .functions;
+      const dependencyFunctionInfos =
+        buildFunctionAnalysisInfoFromConcreteFunctionDefinition(
+          dependencyFunctions,
+          this._functionGraph,
+        );
+      dependencyFunctionInfos.forEach((info) =>
+        dependencyFunctionInfoMap.set(info.functionPath, info),
+      );
+    }
+    const queryBuilderFunctionAnalysisInfo =
+      this.queryBuilderState.buildFunctionAnalysisInfo();
+    if (queryBuilderFunctionAnalysisInfo) {
+      Array.from(
+        queryBuilderFunctionAnalysisInfo.functionInfoMap.entries(),
+      ).forEach(([path, info]) => functionInfoMap.set(path, info));
+      Array.from(
+        queryBuilderFunctionAnalysisInfo.dependencyFunctionInfoMap.entries(),
+      ).forEach(([path, info]) => {
+        dependencyFunctionInfoMap.set(path, info);
+      });
+    }
+    const packagePathToFunctionInfoMap = new Map<
+      string,
+      FunctionAnalysisInfo[]
+    >();
+    Array.from(functionInfoMap.values())
+      .concat(Array.from(dependencyFunctionInfoMap.values()))
+      .forEach((info) => {
+        const curr = packagePathToFunctionInfoMap.get(info.packagePath);
+        if (curr) {
+          packagePathToFunctionInfoMap.set(info.packagePath, [...curr, info]);
+        } else {
+          packagePathToFunctionInfoMap.set(info.packagePath, [info]);
+        }
+      });
+    this.setPackagePathToFunctionInfoMap(packagePathToFunctionInfoMap);
+    this.setFunctionInfoMap(functionInfoMap);
+    this.setDependencyFunctionInfoMap(dependencyFunctionInfoMap);
+  }
+
   initializeTreeData(): void {
+    if (!this.initState.isInInitialState) {
+      return;
+    }
+
+    this.initState.inProgress();
+    this.initializeFunctionInfoMap();
     this.initializeDisplayablePackagesSet()
       .catch(noop())
       .finally(() => {
         this.setTreeData(
           getFunctionsExplorerTreeData(
-            [this.queryBuilderState.graphManagerState.graph.root],
+            [this._functionGraph.root],
             this.queryBuilderState,
+            ROOT_PACKAGE_NAME.MAIN,
           ),
         );
-        this.functionExplorerStates =
-          this.queryBuilderState.graphManagerState.graph.ownFunctions.map(
-            (f) => new QueryFunctionExplorerState(this, f),
-          );
+        this.functionExplorerStates = this.functionInfoMap
+          ? Array.from(this.functionInfoMap.values()).map(
+              (info) => new QueryFunctionExplorerState(this, info),
+            )
+          : [];
       });
-    if (
-      this.queryBuilderState.graphManagerState.graph.dependencyManager
-        .hasDependencies
-    ) {
+    if (this.dependencyFunctionInfoMap) {
       this.initializeDependencyDisplayablePackagesSet()
         .catch(noop())
         .finally(() => {
           this.setDependencyTreeData(
             getFunctionsExplorerTreeData(
-              this.queryBuilderState.graphManagerState.graph.dependencyManager
-                .roots,
+              [this._functionGraph.root],
               this.queryBuilderState,
               ROOT_PACKAGE_NAME.PROJECT_DEPENDENCY_ROOT,
             ),
           );
-          this.dependencyFunctionExplorerStates =
-            this.queryBuilderState.graphManagerState.graph.dependencyManager.functions.map(
-              (f) => new QueryFunctionExplorerState(this, f),
-            );
+          this.dependencyFunctionExplorerStates = this.dependencyFunctionInfoMap
+            ? Array.from(this.dependencyFunctionInfoMap.values()).map(
+                (info) => new QueryFunctionExplorerState(this, info),
+              )
+            : [];
         });
     }
+    this.initState.pass();
   }
 }