Segments.tsx

/* eslint-disable @typescript-eslint/no-namespace */
import * as THREE from 'three';
import * as React from 'react';
import mergeRefs from 'react-merge-refs';
import { extend, useFrame, ReactThreeFiber, useThree, ThreeEvent } from '@react-three/fiber';
import { Line2, LineSegmentsGeometry, LineMaterial } from 'three-stdlib';
import { InterleavedBufferAttribute } from 'three';
import { Instance } from '@react-three/fiber/dist/declarations/src/core/renderer';

type SegmentsProps = {
    limit: number;
    lineWidth: number;
    renderOrder: number;
    children: React.ReactNode;
};

type Api = {
    subscribe: (ref: React.RefObject<SegmentObject2>) => void;
};

type SegmentRef = React.RefObject<SegmentObject2>;

type SegmentProps = Omit<JSX.IntrinsicElements['segmentObject'], 'start' | 'end' | 'color'> & {
    start: ReactThreeFiber.Vector3;
    end: ReactThreeFiber.Vector3;
    color?: ReactThreeFiber.Color;
};

const context = React.createContext<Api>(null!);

const Segments = React.forwardRef<Line2, SegmentsProps>((props, forwardedRef) => {
    React.useMemo(() => extend({ SegmentObject2 }), []);

    const { limit = 1000, lineWidth = 1.0, children, ...rest } = props;
    const [segments, setSegments] = React.useState<Array<SegmentRef>>([]);

    const [line] = React.useState(() => new Line2());
    const [material] = React.useState(() => new LineMaterial());
    const [geometry] = React.useState(() => new LineSegmentsGeometry());

    const [positions] = React.useState<number[]>(() => Array(limit * 6).fill(0));
    const [colors] = React.useState<number[]>(() => Array(limit * 6).fill(0));

    const size = useThree((state) => state.size);

    const reset = React.useRef(false);

    const api: Api = React.useMemo(
        () => ({
            subscribe: (ref: React.RefObject<SegmentObject2>) => {
                setSegments((segments) => [...segments, ref]);

                return () => {
                    setSegments((segments) => segments.filter((item) => item.current !== ref.current));
                    reset.current = true;
                };
            },
        }),
        []
    );

    React.useEffect(() => {
        geometry.setPositions(positions);
        geometry.setColors(colors);
    }, [geometry, positions, colors]);

    useFrame(() => {
        const positions = (geometry.attributes['instanceStart'] as InterleavedBufferAttribute).data.array as Float32Array;
        const colors = (geometry.attributes['instanceColorStart'] as InterleavedBufferAttribute).data.array as Float32Array;

        let updated = false;

        for (let i = 0; i < limit; i++) {
            const segment = segments[i]?.current;

            if (reset.current) {
                positions[i * 6 + 0] = segment?.start?.x ?? 0;
                positions[i * 6 + 1] = segment?.start?.y ?? 0;
                positions[i * 6 + 2] = segment?.start?.z ?? 0;

                positions[i * 6 + 3] = segment?.end?.x ?? 0;
                positions[i * 6 + 4] = segment?.end?.y ?? 0;
                positions[i * 6 + 5] = segment?.end?.z ?? 0;

                colors[i * 6 + 0] = segment?.color.r ?? 0;
                colors[i * 6 + 1] = segment?.color.g ?? 0;
                colors[i * 6 + 2] = segment?.color.b ?? 0;

                colors[i * 6 + 3] = segment?.color.r ?? 0;
                colors[i * 6 + 4] = segment?.color.g ?? 0;
                colors[i * 6 + 5] = segment?.color.b ?? 0;

                updated = true;
            }

            if (segment && segment.needsUpdate) {
                positions[i * 6 + 0] = segment.start.x;
                positions[i * 6 + 1] = segment.start.y;
                positions[i * 6 + 2] = segment.start.z;

                positions[i * 6 + 3] = segment.end.x;
                positions[i * 6 + 4] = segment.end.y;
                positions[i * 6 + 5] = segment.end.z;

                colors[i * 6 + 0] = segment.color.r;
                colors[i * 6 + 1] = segment.color.g;
                colors[i * 6 + 2] = segment.color.b;

                colors[i * 6 + 3] = segment.color.r;
                colors[i * 6 + 4] = segment.color.g;
                colors[i * 6 + 5] = segment.color.b;

                updated = true;
                segment.needsUpdate = false;
            }
        }

        if (updated) {
            geometry.attributes['instanceStart'].needsUpdate = true;
            geometry.attributes['instanceColorStart'].needsUpdate = true;
            geometry.computeBoundingSphere();
            geometry.computeBoundingBox();
        }

        if (reset) {
            reset.current = false;
        }
    });

    const getHandlers = (e: ThreeEvent<Event>) => {
        const segment = segments[e.faceIndex!]?.current;

        if (!segment) {
            return undefined;
        }

        const instance = (segment as unknown as Instance).__r3f;
        return instance.handlers;
    };

    return (
        <primitive
            renderOrder={props.renderOrder}
            object={line}
            ref={forwardedRef}
            onClick={(e: ThreeEvent<MouseEvent>) => getHandlers(e)?.onClick?.(e)}
            onContextMenu={(e: ThreeEvent<PointerEvent>) => getHandlers(e)?.onContextMenu?.(e)}
            onDoubleClick={(e: ThreeEvent<PointerEvent>) => getHandlers(e)?.onDoubleClick?.(e)}
            onWheel={(e: ThreeEvent<WheelEvent>) => getHandlers(e)?.onWheel?.(e)}
            onPointerDown={(e: ThreeEvent<PointerEvent>) => getHandlers(e)?.onPointerDown?.(e)}
            onPointerUp={(e: ThreeEvent<PointerEvent>) => getHandlers(e)?.onPointerUp?.(e)}
            onPointerEnter={(e: ThreeEvent<PointerEvent>) => getHandlers(e)?.onPointerEnter?.(e)}
            onPointerLeave={(e: ThreeEvent<PointerEvent>) => getHandlers(e)?.onPointerLeave?.(e)}
            onPointerMove={(e: ThreeEvent<PointerEvent>) => getHandlers(e)?.onPointerMove?.(e)}
            onPointerCancel={(e: ThreeEvent<PointerEvent>) => getHandlers(e)?.onPointerCancel?.(e)}
        >
            <primitive object={geometry} attach="geometry" />
            <primitive
                object={material}
                attach="material"
                vertexColors={true}
                resolution={[size.width, size.height]}
                transparent={true}
                depthWrite={false}
                linewidth={lineWidth}
                {...rest}
            />
            <context.Provider value={api}>{children}</context.Provider>
        </primitive>
    );
});

declare global {
    namespace JSX {
        interface IntrinsicElements {
            segmentObject2: ReactThreeFiber.Object3DNode<SegmentObject2, typeof SegmentObject2>;
        }
    }
}

export class SegmentObject2 {
    color: THREE.Color;
    start: THREE.Vector3;
    end: THREE.Vector3;
    needsUpdate: boolean;

    constructor() {
        this.color = new THREE.Color('white');
        this.start = new THREE.Vector3(0, 0, 0);
        this.end = new THREE.Vector3(0, 0, 0);
        this.needsUpdate = false;
    }
}

const normPos = (pos: SegmentProps['start']): SegmentObject2['start'] =>
    pos instanceof THREE.Vector3 ? pos : new THREE.Vector3(...(typeof pos === 'number' ? [pos, pos, pos] : pos));

const Segment = React.forwardRef<SegmentObject2, SegmentProps>(
    (
        {
            color,
            start,
            end,
            onClick,
            onContextMenu,
            onDoubleClick,
            onWheel,
            onPointerDown,
            onPointerUp,
            onPointerEnter,
            onPointerLeave,
            onPointerMove,
            onPointerCancel,
        },
        forwardedRef
    ) => {
        const api = React.useContext<Api>(context);
        if (!api) throw new Error('Segment must used inside Segments component.');

        const ref = React.useRef<SegmentObject2>(null);

        React.useLayoutEffect(() => api.subscribe(ref), []);

        return (
            <segmentObject2
                ref={mergeRefs([ref, forwardedRef])}
                color={color}
                start={normPos(start)}
                end={normPos(end)}
                onClick={onClick}
                onContextMenu={onContextMenu}
                onDoubleClick={onDoubleClick}
                onWheel={onWheel}
                onPointerDown={onPointerDown}
                onPointerUp={onPointerUp}
                onPointerEnter={onPointerEnter}
                onPointerLeave={onPointerLeave}
                onPointerMove={onPointerMove}
                onPointerCancel={onPointerCancel}
            />
        );
    }
);

export { Segments, Segment };