add expand all, root node ids props to layer hierarchy

Author: softmarshmallow

editor/components/editor/editor-hierarchy-layers/editor-layer-hierarchy-tree.tsx

@@ -1,42 +1,77 @@
-import React, { useCallback, useMemo, useState } from "react";
+import React, { useCallback, useEffect, useMemo, useState } from "react";
 import { TreeView } from "@editor-ui/editor";
 import {
   LayerRow,
   IconContainer,
   LayerIcon,
 } from "./editor-layer-hierarchy-item";
-import { useEditorState, useWorkspace } from "core/states";
+import {
+  FigmaReflectRepository,
+  useEditorState,
+  useWorkspace,
+} from "core/states";
 import { useDispatch } from "core/dispatch";
 import {
   flattenNodeTree,
   FlattenedDisplayItemNode,
 } from "./editor-layer-heriarchy-controller";
+import type { ReflectSceneNode } from "@design-sdk/figma-node";
 
 // TODO:
 // - add navigate context menu
 // - add go to main component
 // - add reveal and focus to selected layers
 
-export function DesignLayerHierarchy() {
+/**
+ *
+ * @param props.rootNodeIDs - root node ids to display @default: null
+ * @param props.expandAll - expand all nodes by default @default: false
+ * @returns
+ */
+export function DesignLayerHierarchy({
+  rootNodeIDs = null,
+  expandAll = false,
+}: {
+  rootNodeIDs?: string[];
+  expandAll?: boolean;
+}) {
   const [state] = useEditorState();
+  const { selectedNodes, selectedPage, design } = state;
   const { highlightLayer, highlightedLayer } = useWorkspace();
   const dispatch = useDispatch();
 
-  const { selectedNodes, selectedPage, design } = state;
-
   const [expands, setExpands] = useState<string[]>(state?.selectedNodes ?? []);
 
-  const root = selectedPage
-    ? design.pages.find((p) => p.id == selectedPage).children
-    : [design?.input?.entry];
+  // get the root nodes (if the rootNodeIDs is not specified, use the selected page's children)
+  let roots: ReflectSceneNode[] = [];
+  if (rootNodeIDs?.length > 0) {
+    roots = rootNodeIDs.reduce((acc, item) => {
+      acc.push(findUnder(item, design));
+      return acc;
+    }, []);
+  } else {
+    roots = selectedPage
+      ? design.pages.find((p) => p.id == selectedPage).children
+      : [design?.input?.entry];
+  }
 
   const layers: FlattenedDisplayItemNode[][] = useMemo(() => {
-    return root
-      ? root
+    return roots
+      ? roots
           .filter(Boolean)
           .map((layer) => flattenNodeTree(layer, selectedNodes, expands))
       : [];
-  }, [root, state?.selectedNodes, expands]);
+  }, [roots, state?.selectedNodes, expands]);
+
+  useEffect(() => {
+    if (expandAll) {
+      const ids = layers.reduce((acc, item) => {
+        acc.push(...item.map((i) => i.id));
+        return acc;
+      }, [] as string[]);
+      setExpands(ids);
+    }
+  }, [layers, expandAll]);
 
   const renderItem = useCallback(
     ({
@@ -96,3 +131,20 @@ export function DesignLayerHierarchy() {
   );
   //
 }
+
+/**
+ * This only supports root frame at the moment.
+ * partof: nodeQ
+ * @param node
+ * @param design
+ * @returns
+ */
+function findUnder(node: string, design: FigmaReflectRepository) {
+  for (const page of design.pages) {
+    for (const frame of page.children) {
+      if (frame.id === node) {
+        return frame;
+      }
+    }
+  }
+}