Component slots (#35)

* Support component slots

* Fix for renaming things

* Add the slots I need

* Change custom node to get the slots it needs to rebuild

* Sync focus state if needed

* Pass node around as well

* formatting changes

* fix: combine inputs and inputgrpups slots

---------

Co-authored-by: Clark McCauley <[email protected]>

.vscode/settings.json

@@ -0,0 +1,4 @@
+{
+    "editor.formatOnSave": true,
+    "editor.defaultFormatter": "esbenp.prettier-vscode"
+}

lib/NodeGraphEditor.tsx

@@ -26,26 +26,29 @@ import { ClipboardItem } from './clipboard'
 import { LayoutEngine, useLayoutEngine } from './layout/layout'
 import { GraphProvider, useGraphStore } from './context/GraphContext.tsx'
 import { DeserializeFunc, SerializeFunc } from './types/store.ts'
+import { GraphSlots } from './types/slots.ts'
 import './tailwind.css'
 
 type NodeGraphEditorProps = Omit<FlowProps, 'edges' | 'nodes'> & {
   onSave?: (data: any) => void
   config: GraphConfig
+  slots?: Partial<GraphSlots>
 }
 
 export const NodeGraphEditor = forwardRef<
   NodeGraphHandle,
   NodeGraphEditorProps
 >(
   (
-    { defaultNodes, defaultEdges, ...props }: NodeGraphEditorProps,
+    { defaultNodes, defaultEdges, slots, ...props }: NodeGraphEditorProps,
     ref,
   ): JSX.Element => {
     return (
       <GraphProvider
         config={props.config}
         initialNodes={defaultNodes}
         initialEdges={defaultEdges}
+        slots={slots}
       >
         <ReactFlowProvider>
           <Flow {...props} ref={ref} />
@@ -56,7 +59,7 @@ export const NodeGraphEditor = forwardRef<
 )
 
 type FlowProps = ReactFlowProps & {
-  backgroundStyles?: CSSProperties,
+  backgroundStyles?: CSSProperties
   /**
    * The default layout engine to use when nodes are provided without positions.
    */
@@ -73,18 +76,14 @@ export type NodeGraphHandle = {
 }
 
 const Flow = forwardRef<NodeGraphHandle, FlowProps>(
-  (
-    { backgroundStyles, layoutEngine, ...props }: FlowProps,
-    ref,
-  ) => {
+  ({ backgroundStyles, layoutEngine, ...props }: FlowProps, ref) => {
     const nodeTypes = useNodeTypes()
     const edgeTypes = useMemo(() => defaultEdgeTypes, [])
     const onConnect = useSocketConnect()
     const config = useGraphStore((store) => store.config)
     const { getState } = useStoreApi()
     const { setNodes, setEdges } = useReactFlow()
 
-
     // Handle clipboard events
     useHotkeys(
       config.keybindings.copy,
@@ -99,7 +98,7 @@ const Flow = forwardRef<NodeGraphHandle, FlowProps>(
     const layout = useLayoutEngine()
     const serialize = useGraphStore((store) => store.serialize)
     const deserialize = useGraphStore((store) => store.deserialize)
-    const addNode = useGraphStore((store) =>store.addNode)
+    const addNode = useGraphStore((store) => store.addNode)
     const removeNode = useGraphStore((store) => store.removeNode)
     const addEdge = useGraphStore((store) => store.addEdge)
     const removeEdge = useGraphStore((store) => store.removeEdge)
@@ -113,7 +112,7 @@ const Flow = forwardRef<NodeGraphHandle, FlowProps>(
         addNode,
         removeNode,
         addEdge,
-        removeEdge
+        removeEdge,
       }),
       [serialize],
     )

lib/components/NodeContainer.tsx

@@ -26,10 +26,21 @@ export function NodeContainer({
   children,
 }: NodeContainerProps) {
   const config = useGraphStore((store) => store.config)
+  const slots = useGraphStore((store) => store.slots)
   const nodeConfig = config.getNodeConfig(node.type!)!
   const nodeKindConfig = config.getNodeKindConfig(nodeConfig.kind)
   const [collapsed, toggleCollapsed] = useNodeCollapsed()
 
+  const headerProps = useMemo(
+    () => ({
+      defaultTitle: nodeConfig.name,
+      color: nodeKindConfig.color,
+      collapsed,
+      toggleCollapsed,
+    }),
+    [nodeConfig, nodeKindConfig, collapsed, toggleCollapsed],
+  )
+
   if (collapsed) {
     return (
       <CollapsedNodeContainer
@@ -52,12 +63,7 @@ export function NodeContainer({
         }}
         className={draggable ? undefined : 'nodrag'}
       >
-        <NodeHeader
-          defaultTitle={nodeConfig.name}
-          color={nodeKindConfig.color}
-          collapsed={false}
-          toggleCollapsed={toggleCollapsed}
-        />
+        <slots.header {...headerProps} />
         {children}
       </div>
     )

lib/components/NodeHeader.tsx

@@ -3,7 +3,7 @@ import { useNodeFieldValue } from '../hooks/node'
 import './NodeHeader.css'
 import { GoTriangleDown, GoTriangleRight } from 'react-icons/go'
 
-type NodeHeaderProps = {
+export type NodeHeaderProps = {
   defaultTitle: string
   color: string
   collapsed?: boolean

lib/config.ts

@@ -4,6 +4,8 @@ import {
   getBuiltinInputs,
   InputSlots,
 } from './components/inputs.ts'
+import { NodeBodySlots, NodeFocusState } from './node-builder.tsx'
+import { Node } from '@xyflow/react'
 
 export const ANY = '__any'
 
@@ -172,6 +174,11 @@ type WithType<T, K> = T & {
 
 export type InputProps = BaseInputProps & NodeInputConfig & ValueTypeConfig
 
+export interface CustomNodeProps extends NodeFocusState{
+  node: Node
+  slots: NodeBodySlots
+}
+
 export class GraphConfig {
   readonly valueTypes: ValueTypes = {}
   readonly keybindings: KeyBindings
@@ -182,7 +189,7 @@ export class GraphConfig {
     [key: string]: NodeConfig
   } = {}
   private customNodes: {
-    [key: string]: JSXElementConstructor<any>
+    [key: string]: JSXElementConstructor<CustomNodeProps>
   } = {}
   private inputs: {
     [key: string]: JSXElementConstructor<any>
@@ -237,11 +244,11 @@ export class GraphConfig {
     return this
   }
 
-  registerCustomNode<T>(
+  registerCustomNode(
     name: string,
     type: string,
     kind: string,
-    node: JSXElementConstructor<T>,
+    node: JSXElementConstructor<CustomNodeProps>,
     inputs: NodeInputConfig[],
     outputs: NodeOutputConfig[],
   ) {
@@ -269,7 +276,7 @@ export class GraphConfig {
     this.validate()
   }
 
-  customNode<T>(type: string): JSXElementConstructor<T> {
+  customNode(type: string): JSXElementConstructor<CustomNodeProps> {
     return this.customNodes[type]
   }
 
@@ -364,16 +371,11 @@ export class GraphConfig {
     buildNode: (
       config: GraphConfig,
       node: NodeConfig,
+      type: string,
     ) => JSXElementConstructor<any>,
   ): Record<string, JSXElementConstructor<any>> {
     return Object.entries(this.nodeTypes)
-      .map(([type, node]): [string, JSXElementConstructor<any>] => {
-        if (node.custom) {
-          return [type, this.customNode(type)]
-        } else {
-          return [type, buildNode(this, node)]
-        }
-      })
+      .map(([type, node]): [string, JSXElementConstructor<any>] => [type, buildNode(this, node, type)])
       .reduce(
         (acc: Record<string, JSXElementConstructor<any>>, [type, node]) => {
           acc[type] = node

lib/context/GraphContext.tsx

@@ -5,6 +5,7 @@ import { GraphConfig } from '../config.ts'
 import { useStoreWithEqualityFn } from 'zustand/traditional'
 import { GraphStore } from '../types/store.ts'
 import { addNodeInternals } from '../utilities.ts'
+import { GraphSlots } from '../types/slots.ts'
 
 type ContextType = ReturnType<typeof createGraphStore> | null
 
@@ -31,17 +32,21 @@ export function useGraphApi(): NonNullable<ContextType> {
   return store
 }
 
+type GraphProviderProps = {
+  children: ReactNode
+  config: GraphConfig
+  initialNodes?: Graph.Node[]
+  initialEdges?: Graph.Edge[]
+  slots?: Partial<GraphSlots>
+}
+
 export function GraphProvider({
   children,
   config,
   initialNodes,
   initialEdges,
-}: {
-  children: ReactNode
-  config: GraphConfig
-  initialNodes?: Graph.Node[]
-  initialEdges?: Graph.Edge[]
-}) {
+  slots,
+}: GraphProviderProps) {
   const storeRef = useRef<ContextType | null>(null)
   if (!storeRef.current) {
     const nodes = initialNodes
@@ -50,6 +55,7 @@ export function GraphProvider({
 
     storeRef.current = createGraphStore({
       config,
+      slots,
       nodes,
       edges: initialEdges ?? [],
     })

lib/hooks/config.ts

@@ -1,6 +1,6 @@
 import { JSXElementConstructor, useMemo } from 'react'
 import { GraphConfig, IGraphConfig } from '../config.ts'
-import { buildNode } from '../node-types.tsx'
+import { buildNode } from '../node-builder.tsx'
 import { useGraphStore } from '../context/GraphContext.tsx'
 
 export function useNodeTypes(): Record<string, JSXElementConstructor<any>> {

lib/node-types.tsx → lib/node-builder.tsx

@@ -1,4 +1,4 @@
-import {
+import React, {
   FunctionComponent,
   JSX,
   memo,
@@ -21,6 +21,68 @@ import { NodeContainer } from './components/NodeContainer'
 import { useFocusBlur } from './hooks/focus'
 import { Handle } from './components/Handle.tsx'
 import { InputGroup } from './components/InputGroup.tsx'
+import { useGraphStore } from './index.ts'
+
+export interface NodeFocusState {
+  node: Node
+  isFocused: boolean
+  onFocus: () => void
+  onBlur: () => void
+}
+export interface NodeBodySlots {
+  bodyTop?: React.ComponentType<NodeFocusState>
+  bodyBottom?: React.ComponentType<NodeFocusState>
+  inputs: React.JSX.Element[]
+  outputs: React.JSX.Element[]
+}
+export interface NodeBodyProps extends NodeFocusState {
+  slots: NodeBodySlots
+  isFocused: boolean
+  onFocus: () => void
+  onBlur: () => void
+}
+export function NodeBody({
+  node,
+  slots,
+  isFocused,
+  onBlur,
+  onFocus,
+}: NodeBodyProps) {
+  return (
+    <div
+      style={{
+        padding: '8px 0 12px',
+        display: 'flex',
+        flexDirection: 'column',
+      }}
+    >
+      {slots.bodyTop && (
+        <slots.bodyTop
+          isFocused={isFocused}
+          onBlur={onBlur}
+          onFocus={onFocus}
+          node={node}
+        />
+      )}
+      {slots.outputs}
+      {slots.inputs}
+      {slots.bodyBottom && (
+        <slots.bodyBottom
+          isFocused={isFocused}
+          onBlur={onBlur}
+          onFocus={onFocus}
+          node={node}
+        />
+      )}
+    </div>
+  )
+}
+
+export function NodeWrapper({
+  children,
+}: NodeFocusState & { children: ReactNode }) {
+  return <>{children}</>
+}
 
 /**
  * Determines whether a node component should be re-rendered based
@@ -34,9 +96,11 @@ const isComponentChanged = (a: Node, b: Node) =>
 export function buildNode(
   config: GraphConfig,
   nodeConfig: NodeConfig,
+  type: string,
 ): FunctionComponent<Node> {
   function component(node: Node): ReactElement {
     const [isFocused, onFocus, onBlur] = useFocusBlur()
+    const slots = useGraphStore((store) => store.slots)
 
     function getInputElements(
       inputs: NodeInputConfig[],
@@ -103,20 +167,45 @@ export function buildNode(
       ))
     }, [edgeIds])
 
+    const bodySlots: NodeBodySlots = useMemo(() => {
+      return {
+        bodyTop: slots.bodyTop,
+        bodyBottom: slots.bodyBottom,
+        inputs: inputs.concat(inputGroups),
+        outputs,
+      }
+    }, [slots, inputs, outputs, inputGroups])
+
+    if (nodeConfig.custom) {
+      const CustomNode = config.customNode(type)
+      return (
+        <CustomNode
+          isFocused={isFocused}
+          onFocus={onFocus}
+          onBlur={onBlur}
+          node={node}
+          slots={bodySlots}
+        />
+      )
+    }
+
     return (
-      <NodeContainer draggable={!isFocused} node={node}>
-        <div
-          style={{
-            padding: '8px 0 12px',
-            display: 'flex',
-            flexDirection: 'column',
-          }}
-        >
-          {outputs}
-          {inputs}
-          {inputGroups}
-        </div>
-      </NodeContainer>
+      <slots.wrapper
+        isFocused={isFocused}
+        onFocus={onFocus}
+        onBlur={onBlur}
+        node={node}
+      >
+        <NodeContainer draggable={!isFocused} node={node}>
+          <slots.body
+            slots={bodySlots}
+            isFocused={isFocused}
+            onFocus={onFocus}
+            onBlur={onBlur}
+            node={node}
+          />
+        </NodeContainer>
+      </slots.wrapper>
     )
   }
   return memo(component, isComponentChanged)