@needle-tools/engine/src/engine-components/DragControlsConstraints.ts

Needle Engine is a web-based runtime for 3D apps. It runs on your machine for development with great integrations into editors like Unity or Blender - and can be deployed onto any device! It is flexible, extensible and networking and XR are built-in.

import { Box3, Euler, Object3D, Plane, Quaternion, Ray, Vector3 } from "three";

import type { Context } from "../engine/engine_setup.js";

import { GameObject } from "./Component.js";
import type { IGameObject } from "../engine/engine_types.js";

// #region Interfaces
/**
 * Context passed to {@link IDragConstraint.init} at drag start and whenever the drag context
 * resets (e.g. transitioning from multi-touch back to single-pointer).
 * Provides a snapshot of the dragged object and attachment data a constraint needs
 * to initialize or re-initialize its locked reference values.
 */
export interface IDragConstraintContext {
    /** The object being dragged. Cast to IGameObject for world-space properties. */
    readonly gameObject: Object3D;
    /** Drag attachment point in the dragged object's local space. */
    readonly hitPointInLocalSpace: Vector3;
    /** Surface normal at drag attachment in the dragged object's local space. */
    readonly hitNormalInLocalSpace: Vector3;
    /**
     * Local-space bounding box at scale=1. Non-null only for multi-touch drags
     * where scale-aware constraints (e.g. keeping the bounds bottom grounded) are needed.
     */
    readonly boundsAtScaleOne: Box3 | null;
    /**
     * World scale of the object captured at the start of the current drag phase.
     * Non-null alongside boundsAtScaleOne for multi-touch; null for single-pointer.
     */
    readonly initialWorldScale: Vector3 | null;
}

/**
 * Contract for drag constraints applied after a follow object's position is resolved each frame.
 * Implement this interface to add custom position/rotation/scale restrictions.
 */
export interface IDragConstraint {
    /**
     * Called once at drag start and again whenever the drag context resets
     * (e.g. multi-touch → single-pointer transition). Capture any object snapshot
     * (position, rotation, scale) you need to hold fixed during the drag.
     * Constraints that need no dynamic initialization may omit this method.
     */
    init?(context: IDragConstraintContext): void;
    /** Modifies followObject in-place. Invoked after position is resolved each frame. */
    apply(followObject: GameObject): void;
}
// #endregion


// #region Constraints
/** Snaps the follow object's world position to a uniform grid. Resolution ≤ 0 is a no-op.
 *  For XZ-plane mode, use {@link GrabPointPlaneConstraint} instead (it handles both
 *  plane-clamping and optional grid-snapping on the plane). */
export class GridSnapConstraint implements IDragConstraint {
    constructor(public snapGridResolution: number) {}
    apply(followObject: GameObject): void {
        const r = this.snapGridResolution;
        if (r <= 0) return;
        const wp = followObject.worldPosition;
        wp.x = Math.round(wp.x / r) * r;
        wp.y = Math.round(wp.y / r) * r;
        wp.z = Math.round(wp.z / r) * r;
        followObject.worldPosition = wp;
        followObject.updateMatrix();
    }
}

/** Projects the grabbed attachment point back onto a plane after position is resolved,
 *  with optional grid snapping on that plane.
 *  Owned by the strategy that requires plane-clamping (e.g. {@link XZPlaneDragStrategy}). */
export class GrabPointPlaneConstraint implements IDragConstraint {
    /** Plane to project the grabbed point onto. Typically the strategy's fixed drag plane. */
    public readonly plane: Plane;
    /** Attachment point in dragged-object local space. Set once at drag start; the handler
     *  mutates this Vector3 in-place so updates are reflected automatically. */
    public hitPointInLocalSpace: Vector3 | null = null;
    /** Grid snap resolution on the plane. 0 = projection only, no snapping. */
    public snapResolution: number = 0;

    constructor(plane: Plane) {
        this.plane = plane;
    }

    init(ctx: IDragConstraintContext): void {
        this.hitPointInLocalSpace = ctx.hitPointInLocalSpace;
    }

    apply(followObject: GameObject): void {
        if (!this.hitPointInLocalSpace) return;

        // Compute the grabbed point in world space.
        const grabbedWP = this.hitPointInLocalSpace.clone();
        (followObject as unknown as Object3D).localToWorld(grabbedWP);

        // Preserve the constant pivot↔grab-point offset.
        const followWP = followObject.worldPosition;
        const ox = followWP.x - grabbedWP.x;
        const oy = followWP.y - grabbedWP.y;
        const oz = followWP.z - grabbedWP.z;

        // Optionally snap the grabbed point to the grid first.
        let snapped = false;
        const r = this.snapResolution;
        if (r > 0) {
            grabbedWP.x = Math.round(grabbedWP.x / r) * r;
            grabbedWP.y = Math.round(grabbedWP.y / r) * r;
            grabbedWP.z = Math.round(grabbedWP.z / r) * r;
            snapped = true;
        }

        // Project back onto the plane so the grabbed point stays on the drag surface.
        const bx = grabbedWP.x, by = grabbedWP.y, bz = grabbedWP.z;
        this.plane.projectPoint(grabbedWP, grabbedWP);
        const projectionMoved = Math.abs(grabbedWP.x - bx) + Math.abs(grabbedWP.y - by) + Math.abs(grabbedWP.z - bz) > 1e-5;

        if (snapped || projectionMoved) {
            // Restore follow-object offset.
            followWP.set(grabbedWP.x + ox, grabbedWP.y + oy, grabbedWP.z + oz);
            followObject.worldPosition = followWP;
            followObject.updateMatrix();
        }
    }
}

/** Locks the follow object's world rotation to the quaternion captured at drag-start. */
export class KeepRotationConstraint implements IDragConstraint {
    private readonly _savedQuat: Quaternion = new Quaternion();
    init(ctx: IDragConstraintContext): void {
        this._savedQuat.copy((ctx.gameObject as unknown as IGameObject).worldQuaternion);
    }
    apply(followObject: GameObject): void {
        followObject.worldQuaternion = this._savedQuat;
        followObject.updateMatrix();
    }
}

/** Locks the follow object's world scale to the value captured at drag-start. */
export class KeepScaleConstraint implements IDragConstraint {
    private readonly _savedScale: Vector3 = new Vector3(1, 1, 1);
    init(ctx: IDragConstraintContext): void {
        this._savedScale.copy((ctx.gameObject as unknown as IGameObject).worldScale);
    }
    apply(followObject: GameObject): void {
        followObject.worldScale = this._savedScale;
        followObject.updateMatrix();
    }
}

/**
 * Clamps the world scale of the dragged object to a [min, max] range.
 *
 * When `relativeToInitialScale` is `true` (default), `min` and `max` are treated as
 * multipliers of the object's world scale captured at drag start — for example, `min=0.01`
 * means the object cannot shrink below 1% of its original size.
 *
 * When `relativeToInitialScale` is `false`, `min` and `max` are applied directly as
 * absolute world-space scale values to each axis independently.
 *
 * @example Prevent negative / near-zero scale with default relative mode:
 * ```ts
 * new ScaleLimitConstraint(0.01, 10000, true, object.worldScale.clone());
 * ```
 */
export class ScaleLimitConstraint implements IDragConstraint {
    /**
     * @param min                    Lower bound. Relative mode: fraction of initial scale. Raw mode: absolute world-scale per axis.
     * @param max                    Upper bound. Relative mode: fraction of initial scale. Raw mode: absolute world-scale per axis.
     * @param relativeToInitialScale When `true`, clamp is relative to `initialWorldScale`. When `false`, each axis is clamped independently.
     * @param initialWorldScale      Reference to the object's world scale captured at drag start. Required for relative mode.
     */
    private readonly _initialWorldScale: Vector3 = new Vector3(1, 1, 1);

    constructor(
        public min: number,
        public max: number,
        public relativeToInitialScale: boolean = true,
    ) {}

    init(ctx: IDragConstraintContext): void {
        const ws = ctx.initialWorldScale ?? (ctx.gameObject as unknown as IGameObject).worldScale;
        this._initialWorldScale.copy(ws);
    }

    apply(followObject: GameObject): void {
        const ws = followObject.worldScale;
        if (this.relativeToInitialScale) {
            // Derive the uniform scale ratio from the magnitude of the current world scale
            // relative to the initial world scale magnitude.
            const initLen = this._initialWorldScale.length();
            if (initLen < 1e-10) return;
            const ratio = ws.length() / initLen;
            const clamped = Math.max(this.min, Math.min(this.max, ratio));
            if (Math.abs(clamped - ratio) > 1e-9) {
                ws.copy(this._initialWorldScale).multiplyScalar(clamped);
                followObject.worldScale = ws;
                followObject.updateMatrix();
            }
        } else {
            // Absolute per-axis clamp.
            const cx = Math.max(this.min, Math.min(this.max, ws.x));
            const cy = Math.max(this.min, Math.min(this.max, ws.y));
            const cz = Math.max(this.min, Math.min(this.max, ws.z));
            if (cx !== ws.x || cy !== ws.y || cz !== ws.z) {
                ws.set(cx, cy, cz);
                followObject.worldScale = ws;
                followObject.updateMatrix();
            }
        }
    }
}

/**
 * Flags controlling which rotation axes are frozen by {@link FixedRotationAxesConstraint}.
 * Values can be combined with the bitwise OR operator (`|`).
 * @example Freeze X and Z: `RotationAxis.X | RotationAxis.Z`
 */
export enum RotationAxis {
    X = 1,
    Y = 2,
    Z = 4,
}

/**
 * Freezes individual rotation axes (X, Y, and/or Z) of the dragged object.
 * The locked axis values are captured once at construction time and restored every frame.
 *
 * Set {@link useLocalSpace} to `true` to lock axes in the object's local space;
 * leave it `false` (default) to lock axes in world space.
 *
 * @example Lock Y-axis rotation in world space:
 * ```ts
 * const c = new FixedRotationAxesConstraint(RotationAxis.Y, false);
 * // init is called automatically by DragControls, or call manually: c.init(ctx);
 * ```
 */
export class FixedRotationAxesConstraint implements IDragConstraint {
    private readonly _startEuler = new Euler();
    private readonly _eulerCache = new Euler();

    /**
     * @param frozenAxes       Bitfield of {@link RotationAxis} values indicating which axes to lock.
     * @param useLocalSpace    When `true`, axes are locked in the object's local space; otherwise world space.
     */
    constructor(
        public frozenAxes: RotationAxis,
        public useLocalSpace: boolean = false,
    ) {}

    init(ctx: IDragConstraintContext): void {
        const q = this.useLocalSpace
            ? ctx.gameObject.quaternion
            : (ctx.gameObject as unknown as IGameObject).worldQuaternion;
        this._startEuler.setFromQuaternion(q, 'XYZ');
    }

    apply(followObject: GameObject): void {
        const frozenAxes = this.frozenAxes;
        if ((frozenAxes as number) === 0) return;

        // Read the current rotation in the chosen space.
        const current = this.useLocalSpace
            ? (followObject as unknown as Object3D).quaternion
            : followObject.worldQuaternion;

        this._eulerCache.setFromQuaternion(current, 'XYZ');

        // Overwrite the locked components with their start values.
        if (frozenAxes & RotationAxis.X) this._eulerCache.x = this._startEuler.x;
        if (frozenAxes & RotationAxis.Y) this._eulerCache.y = this._startEuler.y;
        if (frozenAxes & RotationAxis.Z) this._eulerCache.z = this._startEuler.z;

        // Recompose and write back.
        _tmpQuat.setFromEuler(this._eulerCache);

        if (this.useLocalSpace) {
            (followObject as unknown as Object3D).quaternion.copy(_tmpQuat);
        } else {
            followObject.worldQuaternion = _tmpQuat;
        }
        followObject.updateMatrix();
    }
}

/** Shared scratch quaternion for {@link FixedRotationAxesConstraint}. */
const _tmpQuat = new Quaternion();


/**
 * Constrains a dragged object to rotate only around a fixed world-space axis.
 * Uses swing-twist decomposition to extract only the twist component around
 * the given axis from the delta rotation since drag start, discarding any
 * perpendicular swing that would tilt the object.
 *
 * This is the correct constraint for XZ-plane dragging when the parent is
 * rotated — it restricts the object to yaw around the plane normal regardless
 * of how the plane is oriented in world space.
 */
export class AxisRotationConstraint implements IDragConstraint {
    private readonly _startQuat: Quaternion = new Quaternion();
    private readonly _axis: Vector3;

    constructor(axis: Vector3) {
        this._axis = axis.clone().normalize();
    }

    init(ctx: IDragConstraintContext): void {
        this._startQuat.copy((ctx.gameObject as unknown as IGameObject).worldQuaternion);
    }

    apply(followObject: GameObject): void {
        const q = followObject.worldQuaternion;

        // q_delta = q_current * q_start⁻¹  (motion applied since drag start)
        _axisQ0.copy(this._startQuat).invert();
        _axisQ1.multiplyQuaternions(q, _axisQ0);

        // Swing-twist decomposition: project the vector part of q_delta onto the twist axis.
        const nx = this._axis.x, ny = this._axis.y, nz = this._axis.z;
        const dot = _axisQ1.x * nx + _axisQ1.y * ny + _axisQ1.z * nz;
        const px = dot * nx, py = dot * ny, pz = dot * nz;
        const qw = _axisQ1.w;

        // twist = normalize(px, py, pz, qw)
        const len = Math.sqrt(px * px + py * py + pz * pz + qw * qw);
        if (len < 1e-10) {
            // Degenerate: delta is a pure swing ~180° perpendicular to the axis.
            // No twist component — snap back to start orientation.
            followObject.worldQuaternion = this._startQuat;
        } else {
            const inv = 1 / len;
            _axisQ0.set(px * inv, py * inv, pz * inv, qw * inv);
            // q_result = twist * q_start
            _axisQ1.multiplyQuaternions(_axisQ0, this._startQuat);
            followObject.worldQuaternion = _axisQ1;
        }
        followObject.updateMatrix();
    }
}

/** Scratch quaternions for {@link AxisRotationConstraint}. */
const _axisQ0 = new Quaternion();
const _axisQ1 = new Quaternion();

/** Locks the follow object's signed distance from a plane to a value set at drag start.
 *  Owned by {@link XZPlaneDragStrategy} to prevent slow Y drift in the side-view fallback path.
 *  Immune to grab-point projection errors, stale matrixWorld, and epsilon guards.
 *
 *  When {@link boundsBottomSignedDistFromPivot} is non-zero (set by {@link MultiTouchDragHandler}
 *  for XZPlane scaling), the effective locked height is shifted each frame so the **bottom of the
 *  object's bounds** stays at the height captured at drag start instead of the pivot. */
export class PlaneHeightLockConstraint implements IDragConstraint {
    private _lockedHeight: number = 0;
    private _boundsBottomSignedDistFromPivot: number = 0;
    /** Current pinch-scale ratio, updated each frame by {@link MultiTouchDragHandler}. */
    public currentScaleRatio: number = 1;

    constructor(private readonly plane: Plane) {}

    init(ctx: IDragConstraintContext): void {
        const wp = (ctx.gameObject as unknown as IGameObject).worldPosition;
        this._lockedHeight = this.plane.normal.dot(wp) + this.plane.constant;
        // When bounds and initial scale are provided (two-pointer drag), configure the
        // constraint to keep the bounds bottom at a constant height while pinch-scaling.
        // bounds.min.y is in local space at scale=1; multiplied by initialWorldScaleY it gives
        // the world-space signed-distance offset from pivot to the bottom of the bounds.
        this._boundsBottomSignedDistFromPivot = (ctx.boundsAtScaleOne !== null && ctx.initialWorldScale !== null)
            ? ctx.boundsAtScaleOne.min.y * ctx.initialWorldScale.y
            : 0;
        this.currentScaleRatio = 1;
    }

    apply(followObject: GameObject): void {
        // When scaling (_boundsBottomSignedDistFromPivot != 0), shift the target height so the
        // bounds bottom stays fixed rather than the pivot. At ratio == 1 the correction is 0.
        const targetHeight = this._lockedHeight
            + this._boundsBottomSignedDistFromPivot * (1 - this.currentScaleRatio);
        const wp = followObject.worldPosition;
        const delta = targetHeight - (this.plane.normal.dot(wp) + this.plane.constant);
        if (Math.abs(delta) > 1e-9) {
            wp.addScaledVector(this.plane.normal, delta);
            followObject.worldPosition = wp;
            followObject.updateMatrix();
        }
    }
}
// #endregion


/**
 * Used by {@link SnapToSurfacesDragStrategy} for two-pointer (multi-touch) drags.
 * Casts a downward ray each frame from the follow object's world position and adjusts
 * it so the dragged object's bounding-box contact face rests on the detected surface.
 *
 * Only injected into the constraint pipeline for multi-touch; the single-pointer path
 * continues to use {@link SnapToSurfacesDragStrategy.update} (pointer ray + drag plane
 * intersection) unchanged.
 */
export class SnapToSurfaceConstraint implements IDragConstraint {
    private _gameObject: Object3D | null = null;
    private _boundsAtScaleOne: Box3 | null = null;

    // Per-frame scratch vectors — avoids per-frame allocations.
    private readonly _t1 = new Vector3();
    private readonly _t2 = new Vector3();
    private readonly _t3 = new Vector3();
    private readonly _t4 = new Vector3();

    constructor(private readonly _context: Context) {}

    init(ctx: IDragConstraintContext): void {
        this._gameObject = ctx.gameObject;
        this._boundsAtScaleOne = ctx.boundsAtScaleOne;
    }

    apply(followObject: GameObject): void {
        if (!this._gameObject || !this._boundsAtScaleOne) return;

        // Cast a downward ray from 0.5 units above the follow object's current world position.
        // The small upward offset avoids self-intersection when the object sits on a surface.
        // World-Y down is used; it finds the nearest surface below the follow object.
        const currentWP = followObject.worldPosition;
        this._t1.set(currentWP.x, currentWP.y + 0.5, currentWP.z);
        const ray = new Ray(this._t1, _snapToSurfaceDown);

        const hits = this._context.physics.raycastFromRay(ray, {
            testObject: (o: Object3D) => o !== (followObject as unknown as Object3D) && o !== this._gameObject
        });

        if (hits.length === 0) return;
        const hit = hits[0];
        if (!hit.face) return;

        // Compute world-space surface normal.
        this._t2.copy(hit.normal ?? hit.face.normal).applyQuaternion(hit.object.worldQuaternion);
        const worldNormal = this._t2;

        // Compute the bounds contact face in local-scale-1 space — mirrors the anchor
        // calculation in SnapToSurfacesDragStrategy.update() (needsAnchorUpdate block).
        this._boundsAtScaleOne.getCenter(this._t3);
        this._boundsAtScaleOne.getSize(this._t4);
        // Contact face: move center toward the surface by half the extent along worldNormal.
        this._t3.sub(this._t4.multiplyScalar(0.5).multiply(worldNormal));
        const contactAlongNormal = this._t3.dot(worldNormal);

        // Anchor in local-scale-1 space. For multi-touch hitPointInLocalSpace is (0,0,0), so
        // the anchor simplifies to: worldNormal * contactAlongNormal.
        this._t3.copy(worldNormal).multiplyScalar(contactAlongNormal);

        // Transform anchor to world space via the follow object's current matrix.
        // This correctly accounts for the object's actual rotation and pinch-scale.
        this._t4.copy(this._t3);
        followObject.localToWorld(this._t4);   // _t4 is now anchorWP

        // Shift follow object so anchorWP coincides with hit.point:
        //   newFollowWP = currentWP + (hit.point − anchorWP)
        this._t1.copy(hit.point).sub(this._t4).add(currentWP);
        followObject.worldPosition = this._t1;
        followObject.updateMatrix();
    }
}

/** Shared downward direction vector for {@link SnapToSurfaceConstraint}. */
const _snapToSurfaceDown = new Vector3(0, -1, 0);


// #region Utilities
/**
 * Runs the constraint pipeline — applies each constraint to followObject in order.
 * Shared by {@link DragPointerHandler} and {@link MultiTouchDragHandler}.
 */
export function applyFollowObjectConstraints(
    followObject: GameObject,
    constraints: readonly IDragConstraint[],
): void {
    for (const c of constraints) c.apply(followObject);
}
// #endregion