chaiNNer-org · joeyballentine · May 23, 2024 · May 23, 2024
diff --git a/src/common/util.ts b/src/common/util.ts
@@ -105,6 +105,28 @@ export const cacheLast = <K extends NonNullable<unknown>, T>(
     };
 };
 
+export interface CacheOptions {
+    readonly maxSize?: number;
+}
+export const cached = <K, T extends NonNullable<unknown>>(
+    fn: (arg: K) => T,
+    { maxSize = 100 }: CacheOptions = {}
+): ((arg: K) => T) => {
+    const cache = new Map<K, T>();
+    return (arg: K): T => {
+        let c = cache.get(arg);
+        if (c === undefined) {
+            if (cache.size >= maxSize && cache.size > 0) {
+                const firstKey = cache.keys().next().value as K;
+                cache.delete(firstKey);
+            }
+            c = fn(arg);
+            cache.set(arg, c);
+        }
+        return c;
+    };
+};
+
 export const debounce = (fn: () => void, delay: number): (() => void) => {
     let id: NodeJS.Timeout | undefined;
     return () => {

diff --git a/src/renderer/components/Handle.tsx b/src/renderer/components/Handle.tsx
@@ -1,12 +1,13 @@
-import { Box, Tooltip, chakra } from '@chakra-ui/react';
-import React, { memo } from 'react';
+import { Box, SystemStyleObject, Tooltip, chakra } from '@chakra-ui/react';
+import React, { memo, useMemo } from 'react';
 import { Connection, Position, Handle as RFHandle } from 'reactflow';
 import { useContext } from 'use-context-selector';
 import { Validity } from '../../common/Validity';
 import { FakeNodeContext } from '../contexts/FakeExampleContext';
 import { createConicGradient } from '../helpers/colorTools';
 import { noContextMenu } from '../hooks/useContextMenu';
 import { useIsCollapsedNode } from '../hooks/useIsCollapsedNode';
+import { useSettings } from '../hooks/useSettings';
 import { Markdown } from './Markdown';
 
 export type HandleType = 'input' | 'output';
@@ -92,17 +93,20 @@ const Div = chakra('div', {
     baseStyle: {},
 });
 
+const getComputedColor = (color: string) =>
+    getComputedStyle(document.documentElement).getPropertyValue(color);
+
 export interface HandleProps {
     id: string;
     type: HandleType;
     validity: Validity;
     isValidConnection: (connection: Readonly<Connection>) => boolean;
     handleColors: readonly string[];
-    connectedColor: string | undefined;
+    connectedColor: readonly string[] | undefined;
     isIterated: boolean;
 }
 
-export const Handle = memo(
+const VisibleHandle = memo(
     ({
         id,
         type,
@@ -112,11 +116,38 @@ export const Handle = memo(
         connectedColor,
         isIterated,
     }: HandleProps) => {
-        const isCollapsed = useIsCollapsedNode();
-        if (isCollapsed) return null;
+        const { theme } = useSettings();
+        const connectedBg = useMemo(() => {
+            return getComputedColor('--connection-color');
+            // eslint-disable-next-line react-hooks/exhaustive-deps
+        }, [theme]);
 
-        const isConnected = !!connectedColor;
-        const connectedBg = 'var(--connection-color)';
+        const handleStyle = useMemo((): SystemStyleObject => {
+            const isConnected = !!connectedColor;
+            if (!isConnected) {
+                return {
+                    borderWidth: '0px',
+                    borderColor: 'transparent',
+                    background: createConicGradient(handleColors),
+                };
+            }
+
+            const size = 16;
+            const rectSize = 12;
+            const offset = (size - rectSize) / 2;
+            const radius = isIterated ? 1 : rectSize / 2;
+
+            const svg = `<svg xmlns='http://www.w3.org/2000/svg' viewBox='0 0 ${size} ${size}'>
+                <rect x='${offset}' y='${offset}' width='${rectSize}' height='${rectSize}' rx='${radius}' fill='${connectedBg}' />
+            </svg>`;
+            const bgImage = `url("data:image/svg+xml,${encodeURIComponent(svg)}")`;
+
+            return {
+                borderWidth: '0px',
+                borderColor: 'transparent',
+                background: `${bgImage}, ${createConicGradient(connectedColor)}`,
+            };
+        }, [connectedColor, handleColors, connectedBg, isIterated]);
 
         return (
             <Div
@@ -146,14 +177,13 @@ export const Handle = memo(
                 sx={{
                     width: '16px',
                     height: '16px',
-                    borderWidth: isConnected ? '2px' : '0px',
-                    borderColor: isConnected ? connectedColor : 'transparent',
                     transition: '0.15s ease-in-out',
-                    background: isConnected ? connectedBg : createConicGradient(handleColors),
+                    transitionProperty: 'width, height, margin, opacity, filter',
                     boxShadow: `${type === 'input' ? '+' : '-'}2px 2px 2px #00000014`,
                     filter: validity.isValid ? undefined : 'grayscale(100%)',
                     opacity: validity.isValid ? 1 : 0.3,
                     position: 'relative',
+                    ...handleStyle,
                 }}
                 type={type}
                 validity={validity}
@@ -162,3 +192,28 @@ export const Handle = memo(
         );
     }
 );
+export const Handle = memo(
+    ({
+        id,
+        type,
+        validity,
+        isValidConnection,
+        handleColors,
+        connectedColor,
+        isIterated,
+    }: HandleProps) => {
+        const isCollapsed = useIsCollapsedNode();
+        if (isCollapsed) return null;
+        return (
+            <VisibleHandle
+                connectedColor={connectedColor}
+                handleColors={handleColors}
+                id={id}
+                isIterated={isIterated}
+                isValidConnection={isValidConnection}
+                type={type}
+                validity={validity}
+            />
+        );
+    }
+);
diff --git a/src/renderer/components/inputs/InputContainer.tsx b/src/renderer/components/inputs/InputContainer.tsx
@@ -9,8 +9,7 @@ import { assertNever, stringifyTargetHandle } from '../../../common/util';
 import { VALID, invalid } from '../../../common/Validity';
 import { GlobalVolatileContext } from '../../contexts/GlobalNodeState';
 import { InputContext } from '../../contexts/InputContext';
-import { defaultColor } from '../../helpers/accentColors';
-import { useSourceTypeColor } from '../../hooks/useSourceTypeColor';
+import { useSourceTypeColors } from '../../hooks/useSourceTypeColor';
 import { useTypeColor } from '../../hooks/useTypeColor';
 import { Handle } from '../Handle';
 import { Markdown } from '../Markdown';
@@ -66,8 +65,7 @@ export const InputHandle = memo(
         }, [connectingFrom, id, targetHandle, isValidConnection]);
 
         const handleColors = useTypeColor(connectableType);
-
-        const sourceTypeColor = useSourceTypeColor(targetHandle);
+        const sourceTypeColor = useSourceTypeColors(targetHandle);
 
         return (
             <HStack
@@ -81,12 +79,7 @@ export const InputHandle = memo(
                     position="absolute"
                 >
                     <Handle
-                        connectedColor={
-                            isConnected
-                                ? (sourceTypeColor === defaultColor ? null : sourceTypeColor) ??
-                                  handleColors[0]
-                                : undefined
-                        }
+                        connectedColor={isConnected ? sourceTypeColor ?? handleColors : undefined}
                         handleColors={handleColors}
                         id={targetHandle}
                         isIterated={isIterated}

diff --git a/src/renderer/components/node/CollapsedHandles.tsx b/src/renderer/components/node/CollapsedHandles.tsx
@@ -7,9 +7,9 @@ import { Output } from '../../../common/common-types';
 import { FunctionDefinition } from '../../../common/types/function';
 import { stringifySourceHandle, stringifyTargetHandle } from '../../../common/util';
 import { BackendContext } from '../../contexts/BackendContext';
-import { createConicGradient } from '../../helpers/colorTools';
+import { defaultColor } from '../../helpers/accentColors';
 import { NodeState } from '../../helpers/nodeState';
-import { useSourceTypeColor } from '../../hooks/useSourceTypeColor';
+import { useSourceTypeColors } from '../../hooks/useSourceTypeColor';
 import { useTypeColor } from '../../hooks/useTypeColor';
 
 interface InputHandleProps {
@@ -18,7 +18,7 @@ interface InputHandleProps {
 }
 
 const InputHandle = memo(({ isIterated, targetHandle }: InputHandleProps) => {
-    const sourceTypeColor = useSourceTypeColor(targetHandle);
+    const sourceTypeColor = useSourceTypeColors(targetHandle)?.[0] ?? defaultColor;
 
     return (
         <Box
@@ -72,7 +72,7 @@ const OutputHandle = memo(
                     isConnectable={false}
                     position={Position.Right}
                     style={{
-                        borderColor: createConicGradient(handleColors),
+                        borderColor: handleColors[0],
                         borderRadius: isIterated ? '10%' : '50%',
                     }}
                     type="source"

diff --git a/src/renderer/components/outputs/OutputContainer.tsx b/src/renderer/components/outputs/OutputContainer.tsx
@@ -61,7 +61,7 @@ export const OutputHandle = memo(
                 right="-6px"
             >
                 <Handle
-                    connectedColor={isConnected ? handleColors[0] : undefined}
+                    connectedColor={isConnected ? handleColors : undefined}
                     handleColors={handleColors}
                     id={sourceHandle}
                     isIterated={isIterated}

diff --git a/src/renderer/helpers/accentColors.ts b/src/renderer/helpers/accentColors.ts
@@ -9,53 +9,46 @@ import {
 import { CategoryMap } from '../../common/CategoryMap';
 import { CategoryId } from '../../common/common-types';
 import { getChainnerScope } from '../../common/types/chainner-scope';
-
-export const defaultColor = '#718096';
+import { cached } from '../../common/util';
 
 const getComputedColor = (color: string) =>
     getComputedStyle(document.documentElement).getPropertyValue(color);
 
-const colorList = () => {
+const resolveName = cached((name: string): Type => {
     const scope = getChainnerScope();
+    return evaluate(new NamedExpression(name), scope);
+});
+
+const colorList = () => {
     return [
-        {
-            type: evaluate(new NamedExpression('Directory'), scope),
-            color: getComputedColor('--type-color-directory'),
-        },
-        {
-            type: evaluate(new NamedExpression('Image'), scope),
-            color: getComputedColor('--type-color-image'),
-        },
+        { type: resolveName('Directory'), color: getComputedColor('--type-color-directory') },
+        { type: resolveName('Image'), color: getComputedColor('--type-color-image') },
         { type: NumberType.instance, color: getComputedColor('--type-color-number') },
         { type: StringType.instance, color: getComputedColor('--type-color-string') },
-        {
-            type: evaluate(new NamedExpression('bool'), scope),
-            color: getComputedColor('--type-color-bool'),
-        },
-        {
-            type: evaluate(new NamedExpression('Color'), scope),
-            color: getComputedColor('--type-color-color'),
-        },
-        {
-            type: evaluate(new NamedExpression('PyTorchModel'), scope),
-            color: getComputedColor('--type-color-torch'),
-        },
-        {
-            type: evaluate(new NamedExpression('OnnxModel'), scope),
-            color: getComputedColor('--type-color-onnx'),
-        },
-        {
-            type: evaluate(new NamedExpression('NcnnNetwork'), scope),
-            color: getComputedColor('--type-color-ncnn'),
-        },
+        { type: resolveName('bool'), color: getComputedColor('--type-color-bool') },
+        { type: resolveName('Color'), color: getComputedColor('--type-color-color') },
+        { type: resolveName('PyTorchModel'), color: getComputedColor('--type-color-torch') },
+        { type: resolveName('OnnxModel'), color: getComputedColor('--type-color-onnx') },
+        { type: resolveName('NcnnNetwork'), color: getComputedColor('--type-color-ncnn') },
     ];
 };
+// eslint-disable-next-line @typescript-eslint/no-unused-vars
+const cachedColorList = cached((_theme: string) => colorList());
 
+export const defaultColor = '#718096';
 const defaultColorList = [defaultColor] as const;
 
-export const getTypeAccentColors = (inputType: Type): readonly [string, ...string[]] => {
+export const getTypeAccentColors = (
+    inputType: Type,
+    theme?: string
+): readonly [string, ...string[]] => {
+    if (inputType.underlying === 'never') {
+        // never is common enough to warrant a special optimization
+        return defaultColorList;
+    }
+
     const colors: string[] = [];
-    const allColors = colorList();
+    const allColors = theme ? cachedColorList(theme) : colorList();
     for (const { type, color } of allColors) {
         if (!isDisjointWith(type, inputType)) {
             colors.push(color);

diff --git a/src/renderer/hooks/useSourceTypeColor.ts b/src/renderer/hooks/useSourceTypeColor.ts
@@ -1,3 +1,4 @@
+import { NeverType } from '@chainner/navi';
 import { useMemo } from 'react';
 import { Node, useReactFlow } from 'reactflow';
 import { useContext } from 'use-context-selector';
@@ -7,7 +8,7 @@ import { BackendContext } from '../contexts/BackendContext';
 import { GlobalVolatileContext } from '../contexts/GlobalNodeState';
 import { useTypeColor } from './useTypeColor';
 
-export const useSourceTypeColor = (targetHandle: string) => {
+export const useSourceTypeColors = (targetHandle: string) => {
     const { functionDefinitions } = useContext(BackendContext);
     const { edgeChanges, typeState } = useContext(GlobalVolatileContext);
     const { getEdges, getNode } = useReactFlow();
@@ -34,7 +35,7 @@ export const useSourceTypeColor = (targetHandle: string) => {
         }
     }, [sourceHandle, functionDefinitions, typeState, getNode]);
 
-    const sourceTypeColor = useTypeColor(sourceType);
+    const colors = useTypeColor(sourceType ?? NeverType.instance);
 
-    return sourceTypeColor[0];
+    return sourceType ? colors : undefined;
 };
diff --git a/src/renderer/hooks/useTypeColor.ts b/src/renderer/hooks/useTypeColor.ts
@@ -1,10 +1,9 @@
 import { Type } from '@chainner/navi';
 import { useMemo } from 'react';
-import { defaultColor, getTypeAccentColors } from '../helpers/accentColors';
+import { getTypeAccentColors } from '../helpers/accentColors';
 import { useSettings } from './useSettings';
 
-export const useTypeColor = (type: Type | undefined) => {
+export const useTypeColor = (type: Type) => {
     const { theme } = useSettings();
-    // eslint-disable-next-line react-hooks/exhaustive-deps
-    return useMemo(() => (type ? getTypeAccentColors(type) : [defaultColor]), [type, theme]);
+    return useMemo(() => getTypeAccentColors(type, theme), [type, theme]);
 };