implement variables in functions

src/components/graph/FunctionNode.tsx

@@ -1,4 +1,4 @@
-import React, { memo, useCallback, useContext, useMemo } from 'react'
+import React, { memo, useCallback, useContext, useMemo, useRef } from 'react'
 import CodeMirror from '@uiw/react-codemirror'
 import { javascript } from '@codemirror/lang-javascript'
 import FlowContext from '../../lib/flowContext'
@@ -14,7 +14,8 @@ import { InputHandle, OutputHandle } from './NodeHandles'
 import { NodeFooterEdgeIndicators } from './NodeFooterEdgeIndicators'
 import { NodeHeaderTitle } from './NodeHeaderTitle'
 import { NodeMenubar } from './NodeMenubar'
-import { Code2 } from 'lucide-react'
+import { MenubarItem, MenubarSub, MenubarSubContent, MenubarSubTrigger } from '../ui/menubar'
+import { Code2, Variable } from 'lucide-react'
 
 type Props = {
     id: string
@@ -28,6 +29,17 @@ export const FunctionNode = memo(function FunctionNode({ id, data, width, height
     const { theme } = useTheme()
     const ctx = useContext(FlowContext)
     const setNodes = ctx?.setNodes
+    const editorRef = useRef<unknown>(null)
+
+    const edges = ctx?.edges ?? []
+    const nodes = ctx?.nodes ?? []
+    const incomingEdges = useMemo(() => edges.filter((e: any) => e.target === id), [edges, id])
+    const incomingIds = useMemo(() => incomingEdges.map((e: any) => e.source).sort(), [incomingEdges])
+    const connectedVariableNodes = useMemo(
+        () => (nodes as any[]).filter((n: any) => incomingIds.includes(n.id) && n.type === 'variable'),
+        [nodes, incomingIds]
+    )
+    const hasConnectedVariables = connectedVariableNodes.length > 0
 
     const onChange = useCallback(
         (val: string) => {
@@ -40,6 +52,34 @@ export const FunctionNode = memo(function FunctionNode({ id, data, width, height
         [id, setNodes]
     )
 
+    const insertAt = useCallback(
+        (insertText: string, mode: 'prepend' | 'append' | 'cursor') => {
+            const ref = editorRef.current as { view: { state: { doc: { length: number }; selection: { main: { from: number } } }; dispatch: (arg: { changes: { from: number; to: number; insert: string } }) => void } } | null
+            if (ref?.view) {
+                const view = ref.view
+                const doc = view.state.doc
+                const len = doc.length
+                let from: number
+                if (mode === 'prepend') from = 0
+                else if (mode === 'append') from = len
+                else from = view.state.selection.main.from
+                view.dispatch({ changes: { from, to: from, insert: insertText } })
+                onChange(view.state.doc.toString())
+                return
+            }
+            if (mode === 'prepend') onChange(insertText + bodyValue)
+            else onChange(bodyValue + insertText)
+        },
+        [onChange, bodyValue]
+    )
+
+    const insertVariableAtCursor = useCallback(
+        (variableNode: any) => {
+            insertAt(variableNode.id, 'cursor')
+        },
+        [insertAt]
+    )
+
     const extensions = useMemo(() => [javascript()], [])
     const [editorHeight, editorContainerRef] = useResizeHeight(120)
     const dimensions =
@@ -53,10 +93,37 @@ export const FunctionNode = memo(function FunctionNode({ id, data, width, height
 
             <BaseNodeContent>
                 <div className="shrink-0 w-full">
-                    <NodeMenubar nodeId={id} nodeType="function" />
+                    <NodeMenubar
+                        nodeId={id}
+                        nodeType="function"
+                        editInputsContent={
+                            hasConnectedVariables ? (
+                                <>
+                                    {connectedVariableNodes.map((n: any) => (
+                                        <MenubarSub key={n.id}>
+                                            <MenubarSubTrigger className="text-xs flex items-center gap-2">
+                                                <Variable className="size-3.5 shrink-0" />
+                                                {n.id}
+                                            </MenubarSubTrigger>
+                                            <MenubarSubContent>
+                                                <MenubarItem
+                                                    className="text-xs"
+                                                    onClick={() => insertVariableAtCursor(n)}
+                                                >
+                                                    Insert at cursor
+                                                </MenubarItem>
+                                            </MenubarSubContent>
+                                        </MenubarSub>
+                                    ))}
+                                </>
+                            ) : undefined
+                        }
+                    />
                 </div>
                 <div ref={editorContainerRef} className="min-h-0 flex-1 w-full nodrag nopan overflow-hidden border-t border-input">
                     <CodeMirror
+                        // @ts-expect-error ref is { view, state, editor }; package ref type not in our node_modules
+                        ref={editorRef}
                         value={bodyValue}
                         height={`${editorHeight}px`}
                         theme={theme}
@@ -72,9 +139,6 @@ export const FunctionNode = memo(function FunctionNode({ id, data, width, height
                 <NodeFooterEdgeIndicators nodeId={id} nodeType="function">
                     {bodyValue ? `${bodyValue.length} chars` : 'none'}
                 </NodeFooterEdgeIndicators>
-                <p className="text-[10px] text-muted-foreground mt-1">
-                    In config: <code className="rounded bg-muted px-0.5">{'{{ x | '}{id}{' }}'}</code> or <code>{'{{ x | '}{id}{'(a, b, key=val) }}'}</code>. Use <code>function(num, x, y, kwargs) &#123; ... &#125;</code> — kwargs has keyword args.
-                </p>
             </BaseNodeFooter>
         </BaseNode>
     )