Default Layout

Set and use a default layout if one or more of the nodes does not have a position.
clarkmcc · Jan 27, 2024 · 86eadd6 · 86eadd6
1 parent e4e947c
commit 86eadd6
Show file tree

Hide file tree

Showing 3 changed files with 57 additions and 15 deletions.
diff --git a/lib/NodeGraphEditor.tsx b/lib/NodeGraphEditor.tsx
@@ -7,6 +7,7 @@ import {
   ReactFlowProps,
   ReactFlowProvider,
   useEdgesState,
+  useNodesInitialized,
   useNodesState,
   useReactFlow,
   useStoreApi,
@@ -23,24 +24,50 @@ import {
   useMemo,
   JSX,
   CSSProperties,
+  useState,
+  useEffect,
 } from 'react'
 import { defaultEdgeTypes } from './edge-types'
 import { IGraphConfig } from './config'
 import { useSocketConnect } from './hooks/connect'
 import { useHotkeys } from 'react-hotkeys-hook'
 import { ClipboardItem } from './clipboard'
-import { LayoutEngine, useLayoutEngine } from './layout/layout'
+import { LayoutEngine, getLayoutFunction, useLayoutEngine } from './layout/layout'
+
+const options = {
+  includeHiddenNodes: false,
+};
+
+export default function useDefaultLayout(layoutFn:((nodes: Node[], edges:Edge[]) => Node[])| undefined) {
+  const { getNodes, getEdges } = useReactFlow();
+  const nodesInitialized = useNodesInitialized(options);
+  const [layoutedNodes, setLayoutedNodes] = useState(getNodes());
+
+  useEffect(() => {
+    if (nodesInitialized && typeof layoutFn == "function") {
+      const found = !!getNodes().find((node) => node.position == undefined);
+      if (found) {
+        setLayoutedNodes(layoutFn(getNodes(), getEdges()));
+      } else {
+        setLayoutedNodes(getNodes());
+      }
+    }
+  }, [nodesInitialized]);
+
+  return layoutedNodes;
+}
 
 type NodeGraphEditorProps = Omit<FlowProps, 'edges' | 'nodes'> & {
-  onSave?: (data: any) => void
+  onSave?: (data: any) => void,
+  defaultLayout?: LayoutEngine,
 }
 
 export const NodeGraphEditor = forwardRef<
   NodeGraphHandle,
   NodeGraphEditorProps
 >(
   (
-    { defaultNodes, defaultEdges, ...props }: NodeGraphEditorProps,
+    { defaultNodes, defaultEdges, defaultLayout, ...props }: NodeGraphEditorProps,
     ref,
   ): JSX.Element => {
     const [nodes, , onNodesChange] = useNodesState(defaultNodes ?? [])
@@ -54,6 +81,7 @@ export const NodeGraphEditor = forwardRef<
           edges={edges}
           onNodesChange={onNodesChange}
           onEdgesChange={onEdgesChange}
+          defaultLayout={defaultLayout}
         />
       </ReactFlowProvider>
     )
@@ -81,15 +109,16 @@ export const ExampleNodeGraphEditor = forwardRef<
 })
 
 type FlowProps = ReactFlowProps & {
-  backgroundStyles?: CSSProperties
+  backgroundStyles?: CSSProperties,
+  defaultLayout?: LayoutEngine
 }
 export type NodeGraphHandle = {
   layout: () => void
 }
 
 const Flow = forwardRef<NodeGraphHandle, FlowProps>(
   (
-    { defaultNodes, defaultEdges, backgroundStyles, ...props }: FlowProps,
+    { backgroundStyles, defaultLayout, ...props }: FlowProps,
     ref,
   ) => {
     const nodeTypes = useNodeTypes()
@@ -98,6 +127,12 @@ const Flow = forwardRef<NodeGraphHandle, FlowProps>(
     const [config] = useGraphConfig()
     const { getState } = useStoreApi()
     const { setNodes, setEdges } = useReactFlow()
+    const layoutedNodes = useDefaultLayout(getLayoutFunction(defaultLayout))
+    useEffect(() => {
+      if (layoutedNodes && defaultLayout !== undefined) {
+        setNodes(layoutedNodes)
+      }
+    }, [layoutedNodes])
 
     // Handle clipboard events
     useHotkeys(

diff --git a/lib/NodeGraphEditorInputGroups.stories.tsx b/lib/NodeGraphEditorInputGroups.stories.tsx
@@ -11,6 +11,7 @@ import { useBuildGraphConfig } from './hooks/config.ts'
 import { InputProps } from './config.ts'
 import { Wheel } from '@uiw/react-color'
 import { useNodeFieldValue } from './hooks/node.ts'
+import { LayoutEngine } from './layout/layout.ts'
 
 const meta = {
   title: 'Node Graph Editor',
@@ -201,7 +202,7 @@ const meta = {
     )
     return (
       <GraphConfigProvider defaultConfig={config}>
-        <NodeGraphEditor defaultNodes={nodes} defaultEdges={edges}>
+        <NodeGraphEditor defaultNodes={nodes} defaultEdges={edges} defaultLayout={LayoutEngine.Dagre}>
           <Background color="#52525b" variant={BackgroundVariant.Dots} />
         </NodeGraphEditor>
       </GraphConfigProvider>
@@ -228,30 +229,30 @@ export const InputGroups: Story = {
       {
         id: '1',
         type: 'bsdf',
-        position: { x: 350, y: 100 },
+        // position: { x: 350, y: 100 },
         data: {
           __inputGroupsExpanded: ['Specular'],
         },
       },
       {
         id: '2',
         type: 'number',
-        position: { x: 100, y: 100 },
+        // position: { x: 100, y: 100 },
         data: {},
       },
       {
         id: '3',
         type: 'number',
-        position: { x: 100, y: 200 },
+        // position: { x: 100, y: 200 },
         data: {},
       },
       {
         id: '4',
         type: 'color',
-        position: { x: 100, y: 300 },
+        // position: { x: 100, y: 300 },
         data: {},
       },
-    ],
+    ] as Node[],
     edges: [
       {
         id: 'e1',

diff --git a/lib/layout/layout.ts b/lib/layout/layout.ts
@@ -1,6 +1,6 @@
 import { useCallback } from 'react'
 import { computeDagreLayout } from './dagre'
-import { Instance, Node, useReactFlow } from '@xyflow/react'
+import { Edge, Instance, Node, useReactFlow } from '@xyflow/react'
 
 export enum LayoutEngine {
   Dagre,
@@ -20,10 +20,16 @@ function computeLayout(
   getNodes: Instance.GetNodes<any>,
   getEdges: Instance.GetEdges<any>,
 ): Node[] {
+  const layoutFn = getLayoutFunction(engine)
+  if (!layoutFn) throw new Error(`Unknown layout engine ${engine}`)
+  return layoutFn(getNodes(), getEdges())
+}
+
+export function getLayoutFunction(
+  engine: LayoutEngine | undefined,
+): ((nodes: Node[], edges: Edge[]) => Node[]) | undefined {
   switch (engine) {
     case LayoutEngine.Dagre:
-      return computeDagreLayout(getNodes(), getEdges())
-    default:
-      throw new Error(`Unknown layout engine ${engine}`)
+      return computeDagreLayout
   }
 }