Merge pull request #391 from pmndrs/feat/select-exclude

feat(Select): add exclude prop

Author: kvvasuu

docs/selection.mdx

@@ -17,6 +17,8 @@ when things have to be grouped. These two components take care of it:
 <Select
   children: JSX.Element | JSX.Element[]
   enabled?: boolean
+  /** Keeps this subtree out of the selection, even when an ancestor Select is enabled. */
+  exclude?: boolean
 >
 ```
 
@@ -36,7 +38,7 @@ instance `Outline`) will acknowledge it.
 ```
 
 Selection can be nested and group multiple object, higher up selection take
-precence over lower ones. The following for instance will select everything.
+precedence over lower ones. The following for instance will select everything.
 Remove the outmost `enabled` and only the two mesh group is selected. You can
 flip the selections or bind them to interactions and state.
 
@@ -51,3 +53,16 @@ flip the selections or bind them to interactions and state.
   </Select>
 </Select>
 ```
+
+To keep part of an enabled subtree out of the selection, wrap it in a `Select`
+with `exclude`. Selects nested inside an excluded subtree can still select
+their own objects.
+
+```jsx
+<Select enabled>
+  <mesh /> {/* selected */}
+  <Select exclude>
+    <mesh /> {/* not selected */}
+  </Select>
+</Select>
+```

src/Selection.tsx

@@ -2,8 +2,11 @@ import { type ThreeElements } from '@react-three/fiber'
 import {
   createContext,
   use,
+  useCallback,
   useEffect,
+  useLayoutEffect,
   useMemo,
+  useReducer,
   useRef,
   useState,
   type Dispatch,
@@ -19,10 +22,23 @@ export type Api = {
 }
 export type SelectApi = Omit<ThreeElements['group'], 'ref'> & {
   enabled?: boolean
+  /** Keeps this subtree out of the selection, even when an ancestor Select is enabled. */
+  exclude?: boolean
 }
 
 export const selectionContext = /* @__PURE__ */ createContext<Api | null>(null)
 
+// Lets a nested Select ask its ancestors to re-traverse (e.g. when `exclude` toggles).
+const refreshContext = /* @__PURE__ */ createContext<(() => void) | null>(null)
+
+// How many Selects claim each object, so one letting go doesn't drop another's claim.
+const claims = new WeakMap<Object3D, number>()
+const excluded = new WeakSet<Object3D>()
+
+function updateClaims(objects: Object3D[], delta: number): void {
+  for (const o of objects) claims.set(o, (claims.get(o) ?? 0) + delta)
+}
+
 export function Selection({ children, enabled = true }: { enabled?: boolean; children: ReactNode }) {
   const [selected, select] = useState<Object3D[]>([])
   const value = useMemo(() => ({ selected, select, enabled }), [selected, select, enabled])
@@ -35,52 +51,76 @@ function isSelectable(object: Object3D): boolean {
   return !!(o.isMesh || o.isLine || o.isPoints)
 }
 
-export function Select({ enabled = false, children, ...props }: SelectApi) {
+function collect(object: Object3D, out: Object3D[]): void {
+  if (excluded.has(object)) return
+  if (isSelectable(object)) out.push(object)
+  for (const child of object.children) collect(child, out)
+}
+
+export function Select({ enabled = false, exclude = false, children, ...props }: SelectApi) {
   const group = useRef<Group>(null!)
   // Stable, unlike the context value - avoids retriggering off our own write.
   const select = use(selectionContext)?.select
+  const refreshParent = use(refreshContext)
   const claimed = useRef<Object3D[]>([])
+  const [version, bump] = useReducer((v: number) => v + 1, 0)
+  const refresh = useCallback(() => {
+    bump()
+    refreshParent?.()
+  }, [refreshParent])
+
+  useLayoutEffect(() => {
+    if (!exclude) return
+    const g = group.current
+    excluded.add(g)
+    refreshParent?.()
+    return () => {
+      excluded.delete(g)
+      refreshParent?.()
+    }
+  }, [exclude, refreshParent])
 
   useEffect(() => {
     if (!select) return
 
     const current: Object3D[] = []
-    if (enabled) {
-      group.current.traverse((o) => {
-        if (isSelectable(o)) current.push(o)
-      })
-    }
+    if (enabled && !exclude) for (const child of group.current.children) collect(child, current)
 
     const previouslyClaimed = claimed.current
     claimed.current = current
+    updateClaims(previouslyClaimed, -1)
+    updateClaims(current, 1)
 
     select((prev) => {
       const prevSet = new Set(prev)
-      const currentSet = new Set(current)
       const toAdd = current.filter((o) => !prevSet.has(o))
-      const toRemove = previouslyClaimed.filter((o) => !currentSet.has(o) && prevSet.has(o))
+      const toRemove = previouslyClaimed.filter((o) => prevSet.has(o) && !claims.get(o))
       if (!toAdd.length && !toRemove.length) return prev
       const toRemoveSet = toRemove.length ? new Set(toRemove) : null
       const kept = toRemoveSet ? prev.filter((o) => !toRemoveSet.has(o)) : prev
       return toAdd.length ? [...kept, ...toAdd] : kept
     })
-  }, [enabled, children, select])
+  }, [enabled, exclude, children, select, version])
 
   // Separate from the effect above so unmount cleanup doesn't fire on every enabled/children change.
   useEffect(() => {
     return () => {
       if (!select || !claimed.current.length) return
-      const stillClaimed = new Set(claimed.current)
+      const released = new Set(claimed.current)
+      updateClaims(claimed.current, -1)
+      claimed.current = []
       select((prev) => {
-        const next = prev.filter((o) => !stillClaimed.has(o))
+        const next = prev.filter((o) => !released.has(o) || claims.get(o))
         return next.length !== prev.length ? next : prev
       })
     }
   }, [select])
 
   return (
-    <group ref={group} {...props}>
-      {children}
-    </group>
+    <refreshContext.Provider value={refresh}>
+      <group ref={group} {...props}>
+        {children}
+      </group>
+    </refreshContext.Provider>
   )
 }

src/tests/Selection.test.tsx

@@ -2,7 +2,7 @@ import * as React from 'react'
 import * as THREE from 'three'
 import { afterEach, describe, expect, it } from 'vitest'
 import { Select, Selection, selectionContext } from '../Selection'
-import { flush, root } from './test-utils'
+import { flush, root, strict } from './test-utils'
 
 afterEach(async () => {
   await React.act(async () => {
@@ -321,4 +321,169 @@ describe('Selection/Select', () => {
     const afterSecond = updaters[1](afterFirst)
     expect(afterSecond).toBe(afterFirst)
   })
+
+  it('keeps an object selected via the inner Select after the outer one disables', async () => {
+    const snapshots: THREE.Object3D[][] = []
+    const meshRef = React.createRef<THREE.Mesh>()
+
+    const render = (outerEnabled: boolean) =>
+      root.render(
+        <Selection>
+          <Capture onSnapshot={(s) => snapshots.push(s)} />
+          <Select enabled={outerEnabled}>
+            <Select enabled>
+              <mesh ref={meshRef}>
+                <boxGeometry />
+                <meshBasicMaterial />
+              </mesh>
+            </Select>
+          </Select>
+        </Selection>
+      )
+
+    await React.act(async () => render(true))
+    await flush()
+    await flush()
+    expect(snapshots[snapshots.length - 1]).toContain(meshRef.current)
+
+    await React.act(async () => render(false))
+    await flush()
+    await flush()
+    expect(snapshots[snapshots.length - 1]).toContain(meshRef.current)
+  })
+
+  it('keeps claims balanced under StrictMode double-invoked effects', async () => {
+    const snapshots: THREE.Object3D[][] = []
+    const meshRef = React.createRef<THREE.Mesh>()
+
+    const render = (mounted: boolean, outerEnabled: boolean) =>
+      root.render(
+        strict(
+          <Selection>
+            <Capture onSnapshot={(s) => snapshots.push(s)} />
+            {mounted && (
+              <Select enabled={outerEnabled}>
+                <Select enabled>
+                  <mesh ref={meshRef} />
+                </Select>
+              </Select>
+            )}
+          </Selection>
+        )
+      )
+
+    await React.act(async () => render(true, true))
+    await flush()
+    await flush()
+    const mesh = meshRef.current
+    expect(snapshots[snapshots.length - 1]).toEqual([mesh])
+
+    await React.act(async () => render(true, false))
+    await flush()
+    await flush()
+    expect(snapshots[snapshots.length - 1]).toEqual([mesh])
+
+    await React.act(async () => render(false, false))
+    await flush()
+    await flush()
+    expect(snapshots[snapshots.length - 1]).toEqual([])
+  })
+
+  describe('exclude', () => {
+    it('keeps an excluded subtree out of an enabled ancestor Select (#154)', async () => {
+      const snapshots: THREE.Object3D[][] = []
+      const keptRef = React.createRef<THREE.Mesh>()
+      const excludedRef = React.createRef<THREE.Mesh>()
+
+      await React.act(async () =>
+        root.render(
+          <Selection>
+            <Capture onSnapshot={(s) => snapshots.push(s)} />
+            <Select enabled>
+              <mesh ref={keptRef} />
+              <Select exclude>
+                <group>
+                  <mesh ref={excludedRef} />
+                </group>
+              </Select>
+            </Select>
+          </Selection>
+        )
+      )
+      await flush()
+      await flush()
+
+      const last = snapshots[snapshots.length - 1]
+      expect(last).toContain(keptRef.current)
+      expect(last).not.toContain(excludedRef.current)
+    })
+
+    it('reacts to exclude toggling inside a nested component, without the ancestor re-rendering', async () => {
+      const snapshots: THREE.Object3D[][] = []
+      const meshRef = React.createRef<THREE.Mesh>()
+      let setExclude: (value: boolean) => void = () => {}
+
+      function Item() {
+        const [exclude, set] = React.useState(false)
+        setExclude = set
+        return (
+          <Select exclude={exclude}>
+            <mesh ref={meshRef} />
+          </Select>
+        )
+      }
+
+      await React.act(async () =>
+        root.render(
+          <Selection>
+            <Capture onSnapshot={(s) => snapshots.push(s)} />
+            <Select enabled>
+              <Item />
+            </Select>
+          </Selection>
+        )
+      )
+      await flush()
+      await flush()
+      expect(snapshots[snapshots.length - 1]).toContain(meshRef.current)
+
+      await React.act(async () => setExclude(true))
+      await flush()
+      await flush()
+      expect(snapshots[snapshots.length - 1]).not.toContain(meshRef.current)
+
+      await React.act(async () => setExclude(false))
+      await flush()
+      await flush()
+      expect(snapshots[snapshots.length - 1]).toContain(meshRef.current)
+    })
+
+    it('still lets an enabled Select inside an excluded subtree select its own objects', async () => {
+      const snapshots: THREE.Object3D[][] = []
+      const excludedRef = React.createRef<THREE.Mesh>()
+      const reselectedRef = React.createRef<THREE.Mesh>()
+
+      await React.act(async () =>
+        root.render(
+          <Selection>
+            <Capture onSnapshot={(s) => snapshots.push(s)} />
+            <Select enabled>
+              <Select exclude>
+                <mesh ref={excludedRef} />
+                <Select enabled>
+                  <mesh ref={reselectedRef} />
+                </Select>
+              </Select>
+            </Select>
+          </Selection>
+        )
+      )
+      await flush()
+      await flush()
+
+      const last = snapshots[snapshots.length - 1]
+      expect(last).not.toContain(excludedRef.current)
+      expect(last).toContain(reselectedRef.current)
+    })
+  })
 })