gizmo.tsx

import React, { useCallback, useEffect, useRef, useState, useMemo } from "react";
import * as three from "three"
import { useThree } from 'react-three-fiber';
import { Lines } from "./lines";
import { ScaleHandles } from "./scaleHandles";
import { MoveHandles } from "./moveHandles";
import { IntersectionPlane } from "./intersectionPlane";
import { useDrag } from '../../../hooks/useDrag';
import { useHover } from "../../../hooks/useHover";
import { intersectObjectWithRay } from "./intersection";

interface Props {
    children: React.ReactElement<three.Object3D>;
    selected: boolean;
    onDragStart?: () => void;
    onDragEnd?: () => void;
    onPositionChanged?: (position: three.Vector3) => void;
    position?: three.Vector3;
}

export const Gizmo = (props: Props) => {

    const innerRef = useRef<three.Object3D>();

    const { invalidate, raycaster } = useThree();

    const offset = new three.Vector3();

    const [dragging, setDragging] = useState(false);
    const [axis, setAxis] = useState<"X" | "Y" | "Z" | "XY" | "YZ" | "XZ" | "XYZ" | "E">("XZ");
    const [mode, setMode] = useState<"translate" | "scale" | "rotate">("translate");

    const getPosition = useCallback(() => {
        return innerRef.current?.position;
    }, []);

    const setPosition = useCallback((position: three.Vector3) => {
        innerRef.current?.position.copy(position);
        invalidate();

        if (props.onPositionChanged) props.onPositionChanged(position);
    }, []);

    const bindDrag = useDrag({
        onDrag: () => {

            if (!planeRef.current) return;

            const planeIntersect = intersectObjectWithRay(planeRef.current, raycaster, true);

            if (planeIntersect) {
                planeIntersect.point.sub(offset);

                // snap to whole numbers
                planeIntersect.point.setX(Math.floor(planeIntersect.point.x));
                planeIntersect.point.setZ(Math.floor(planeIntersect.point.z));

                setPosition(planeIntersect.point.clone());
            }
        },
        onDragStart: () => {

            if (!planeRef.current) return;

            const intersection = intersectObjectWithRay(planeRef.current, raycaster, true);

            if (intersection) {
                // you'll never drag from the center, so determine the offset
                offset.copy(intersection.point);

                if (innerRef.current) {
                    offset.sub(innerRef.current.position);
                }
            }

            setDragging(true);

            if (props.onDragStart) props.onDragStart();
        },
        onDragEnd: () => {

            setDragging(false);

            if (props.onDragEnd) props.onDragEnd();
        }
    });

    const [bindHover, hovered] = useHover(false);

    useEffect(() => {
        setMode("translate");
        setAxis("XZ");
    }, [hovered]);

    const planeRef = useRef<three.Mesh>();

    const box = useMemo(() => {

        if (!innerRef.current) {
            return null;
        }

        const box = new three.Box3();
        box.setFromObject(innerRef.current);

        const worldToLocal = new three.Matrix4();
        worldToLocal.getInverse(innerRef.current.matrixWorld);
        box.applyMatrix4(worldToLocal);

        return box;

    }, [innerRef.current]);

    return (
        <>
            <group ref={innerRef}>
                <group {...bindDrag} {...bindHover}>
                    {props.children}
                </group>
                <Lines box={box} />
                <ScaleHandles
                    box={box}
                    dragging={dragging}
                    selected={props.selected}
                />
                <MoveHandles
                    box={box}
                    radius={0.05}
                    height={0.1}
                    radialSegments={30}
                    insetScale={1.6}
                    distance={0.3}
                    dragging={dragging}
                    selected={props.selected}
                    plane={planeRef.current}
                    setAxis={setAxis}
                    setMode={setMode}
                    getPosition={getPosition}
                    setPosition={setPosition}
                    onDragStart={props.onDragStart}
                    onDragEnd={props.onDragEnd}
                />
            </group>
            <IntersectionPlane
                ref={planeRef}
                space="world"
                mode={mode}
                axis={axis}
                object={innerRef.current}
            />
        </>
    );
}

intersectionPlane.tsx

/* eslint-disable react/prop-types */
/* eslint-disable react/display-name */
import * as three from "three";
import { useFrame } from 'react-three-fiber';
import React, { useImperativeHandle, useRef } from "react";

interface Props {
    space: "local" | "world";
    mode: "translate" | "scale" | "rotate";
    axis: "X" | "Y" | "Z" | "XY" | "YZ" | "XZ" | "XYZ" | "E";
    object: three.Object3D | undefined;
}

export const IntersectionPlane = React.forwardRef<three.Mesh | undefined, Props>((props, ref) => {

    const internalRef = useRef<three.Mesh>();
    useImperativeHandle(ref, () => internalRef.current, []);

    const worldPosition = new three.Vector3();
    const worldQuaternion = new three.Quaternion();
    const worldScale = new three.Vector3();
    const cameraPosition = new three.Vector3();
    const cameraQuaternion = new three.Quaternion();
    const cameraScale = new three.Vector3();
    const eye = new three.Vector3();

    const unitX = new three.Vector3(1, 0, 0);
    const unitY = new three.Vector3(0, 1, 0);
    const unitZ = new three.Vector3(0, 0, 1);

    const tempVector = new three.Vector3();
    const dirVector = new three.Vector3();
    const alignVector = new three.Vector3();
    const tempMatrix = new three.Matrix4();
    const identityQuaternion = new three.Quaternion();

    useFrame(({ camera }) => {

        const current = internalRef.current;

        if (!current || !props.object) {
            return;
        }

        props.object.matrixWorld.decompose(worldPosition, worldQuaternion, worldScale)
        camera.matrixWorld.decompose(cameraPosition, cameraQuaternion, cameraScale);
        eye.copy(cameraPosition).sub(worldPosition).normalize();

        current.position.copy(worldPosition);

        if (props.mode === 'scale') props.space = 'local'; // scale always oriented to local rotation

        unitX.set(1, 0, 0).applyQuaternion(props.space === "local" ? worldQuaternion : identityQuaternion);
        unitY.set(0, 1, 0).applyQuaternion(props.space === "local" ? worldQuaternion : identityQuaternion);
        unitZ.set(0, 0, 1).applyQuaternion(props.space === "local" ? worldQuaternion : identityQuaternion);

        // Align the plane for current transform mode, axis and space.
        alignVector.copy(unitY);

        switch (props.mode) {

            case 'translate':
            case 'scale':
                switch (props.axis) {

                    case 'X':
                        alignVector.copy(eye).cross(unitX);
                        dirVector.copy(unitX).cross(alignVector);
                        break;
                    case 'Y':
                        alignVector.copy(eye).cross(unitY);
                        dirVector.copy(unitY).cross(alignVector);
                        break;
                    case 'Z':
                        alignVector.copy(eye).cross(unitZ);
                        dirVector.copy(unitZ).cross(alignVector);
                        break;
                    case 'XY':
                        dirVector.copy(unitZ);
                        break;
                    case 'YZ':
                        dirVector.copy(unitX);
                        break;
                    case 'XZ':
                        alignVector.copy(unitZ);
                        dirVector.copy(unitY);
                        break;
                    case 'XYZ':
                    case 'E':
                        dirVector.set(0, 0, 0);
                        break;

                }

                break;
            case 'rotate':
            default:
                // special case for rotate
                dirVector.set(0, 0, 0);
        }

        if (dirVector.length() === 0) {

            // If in rotate mode, make the plane parallel to camera
            current.quaternion.copy(cameraQuaternion);

        } else {

            tempMatrix.lookAt(tempVector.set(0, 0, 0), dirVector, alignVector);
            current.quaternion.setFromRotationMatrix(tempMatrix);

        }
    }
    );

    return (
        <mesh ref={internalRef}>
            {/* <planeBufferGeometry args={[100000, 100000, 2, 2]} /> */}
            <planeBufferGeometry args={[500, 500, 2, 2]} />
            <meshBasicMaterial attach="material" color="red" wireframe={true} side={three.DoubleSide} toneMapped={true} />
        </mesh>
    )
});

moveHandles.tsx

/* eslint-disable react/display-name */
/* eslint-disable react/prop-types */
import React, { useEffect, useMemo, useRef } from "react";
import * as three from "three"
import { useFrame, useThree } from 'react-three-fiber';
import { Cone } from "drei";
import { useHover } from "../../../hooks/useHover";
import { useDrag } from "../../../hooks/useDrag";
import { intersectObjectWithRay } from "./intersection";

interface MoveHandleProps {
    radius: number;
    height: number;
    radialSegments: number;
    insetScale: number;
    origin: three.Vector3;
    position: three.Vector3;
    scale: three.Vector3;
    distance: number;
    plane: three.Mesh | undefined;
    setMode: React.Dispatch<React.SetStateAction<"translate" | "rotate" | "scale">>;
    setAxis: React.Dispatch<React.SetStateAction<"X" | "Y" | "Z" | "XY" | "YZ" | "XZ" | "XYZ" | "E">>;
    getPosition: () => three.Vector3 | undefined;
    setPosition: (position: three.Vector3) => void;
    onDragStart?: () => void;
    onDragEnd?: () => void;
}

const MoveHandle = React.forwardRef<three.Group, MoveHandleProps>((props, ref) => {

    const { raycaster } = useThree();

    const localRef = useRef<three.Object3D>();
    const translateRef = useRef<three.Vector3>(new three.Vector3());

    const worldPosition = new three.Vector3();
    const worldQuaternion = new three.Quaternion();
    const worldScale = new three.Vector3();

    const distanceStartScale = new three.Vector3(0, props.distance, 0);
    const startScale = new three.Vector3(1, 1, 1);

    useFrame(({ camera }) => {

        if (localRef.current) {

            localRef.current.matrixWorld.decompose(worldPosition, worldQuaternion, worldScale);

            const factor = camera.type === "OrthographicCamera"
                ? (camera.top - camera.bottom) / camera.zoom
                : worldPosition.distanceTo(camera.position) * Math.min(1.9 * Math.tan(Math.PI * camera.fov / 360) / camera.zoom, 7);

            const scaledDistance = distanceStartScale.clone().multiplyScalar(factor / 7).multiply(props.scale);
            const scale = startScale.clone().multiplyScalar(factor / 7).multiply(props.scale);

            translateRef.current.copy(props.position).sub(props.origin).add(scaledDistance);
            localRef.current.position.copy(props.origin).add(translateRef.current);
            localRef.current.scale.copy(scale);
        }
    });

    const bindDrag = useDrag({
        onDrag: () => {

            if (!props.plane) return;

            const intersection = intersectObjectWithRay(props.plane, raycaster, true);

            if (intersection) {
                intersection.point.sub(translateRef.current);

                const position = props.getPosition();
                if (position) {
                    intersection.point.setX(position.x);
                    intersection.point.setZ(position.z);
                }

                props.setPosition(intersection.point.clone());
            }
        },
        onDragStart: () => {
            if (props.onDragStart) props.onDragStart();
        },
        onDragEnd: () => {
            if (props.onDragEnd) props.onDragEnd();
        }
    });

    const [bindHover, hovered] = useHover(false);

    useEffect(() => {
        props.setMode("translate");
        props.setAxis("Y");
    }, [hovered])

    return (
        <group ref={ref} >
            <group {...bindHover} {...bindDrag} ref={localRef}>
                <Cone
                    args={[props.radius, props.height, props.radialSegments]}
                    position={[0, 0.01, 0]}
                    renderOrder={2}
                    material-color="black"
                    material-depthTest={false}
                    material-depthWrite={false}
                    material-blending={1}
                    material-transparent={true}
                />
                <Cone
                    args={[props.radius / props.insetScale, props.height / props.insetScale, props.radialSegments]}
                    renderOrder={3}
                    material-color="red"
                    material-depthTest={false}
                    material-depthWrite={false}
                    material-blending={1}
                    material-transparent={true}
                    visible={hovered}
                />
            </group>
        </group>
    );
});

interface Props {
    box: three.Box3 | null;
    radius: number;
    height: number;
    radialSegments: number;
    insetScale: number;
    distance: number;
    selected: boolean;
    dragging: boolean;
    plane: three.Mesh | undefined;
    setMode: React.Dispatch<React.SetStateAction<"translate" | "rotate" | "scale">>;
    setAxis: React.Dispatch<React.SetStateAction<"X" | "Y" | "Z" | "XY" | "YZ" | "XZ" | "XYZ" | "E">>;
    getPosition: () => three.Vector3 | undefined;
    setPosition: (position: three.Vector3) => void;
    onDragStart?: () => void;
    onDragEnd?: () => void;
}

export const MoveHandles = (props: Props) => {

    const upDirection = new three.Vector3(1, 1, 1);
    const downDirection = new three.Vector3(1, -1, 1);

    const moveUpHandleRef = useRef<three.Group>(null);
    const moveDownHandleRef = useRef<three.Group>(null);

    const points = useMemo(() => {

        if (!props.box) {
            return null;
        }

        const min = props.box.min;
        const max = props.box.max;

        const centerX = (min.x + max.x) / 2;
        const centerY = (min.y + max.y) / 2;
        const centerZ = (min.z + max.z) / 2;

        const center = new three.Vector3(centerX, centerY, centerZ);
        const midDown = new three.Vector3(centerX, min.y, centerZ);
        const midUp = new three.Vector3(centerX, max.y, centerZ);

        return { center, midDown, midUp };

    }, [props.box]);

    useFrame(({ camera }) => {
        const above = camera.rotation.x < 0;

        if (moveUpHandleRef.current) {
            moveUpHandleRef.current.visible = above;
        }

        if (moveDownHandleRef.current) {
            moveDownHandleRef.current.visible = !above;
        }
    });

    if (!points) {
        return null;
    }

    return (
        <group visible={props.selected && !props.dragging}>
            <MoveHandle
                ref={moveUpHandleRef}
                radius={props.radius}
                height={props.height}
                radialSegments={props.radialSegments}
                insetScale={props.insetScale}
                distance={props.distance}
                origin={points.center}
                position={points.midUp}
                scale={upDirection}
                plane={props.plane}
                setMode={props.setMode}
                setAxis={props.setAxis}
                getPosition={props.getPosition}
                setPosition={props.setPosition}
                onDragStart={props.onDragStart}
                onDragEnd={props.onDragEnd}
            />
            <MoveHandle
                ref={moveDownHandleRef}
                radius={props.radius}
                height={props.height}
                radialSegments={props.radialSegments}
                insetScale={props.insetScale}
                distance={props.distance}
                origin={points.center}
                position={points.midDown}
                scale={downDirection}
                plane={props.plane}
                setMode={props.setMode}
                setAxis={props.setAxis}
                getPosition={props.getPosition}
                setPosition={props.setPosition}
                onDragStart={props.onDragStart}
                onDragEnd={props.onDragEnd}
            />
        </group>
    );
}