refactor(types)!: upgrade Zustand to v4 (#2558)

Author: CodyJasonBennett

packages/fiber/package.json

@@ -49,7 +49,7 @@
     "react-use-measure": "^2.1.1",
     "scheduler": "^0.21.0",
     "suspend-react": "^0.0.8",
-    "zustand": "^3.7.1"
+    "zustand": "^4.1.2"
   },
   "peerDependencies": {
     "expo": ">=46.0",

packages/fiber/src/core/events.ts

@@ -1,9 +1,8 @@
 import * as THREE from 'three'
 import { ContinuousEventPriority, DiscreteEventPriority, DefaultEventPriority } from 'react-reconciler/constants'
 import { getRootState } from './utils'
-import type { UseBoundStore } from 'zustand'
 import type { Instance } from './renderer'
-import type { RootState } from './store'
+import type { RootState, RootStore } from './store'
 
 export interface Intersection extends THREE.Intersection {
   /** The event source (the object which registered the handler) */
@@ -147,7 +146,7 @@ function releaseInternalPointerCapture(
   }
 }
 
-export function removeInteractivity(store: UseBoundStore<RootState>, object: THREE.Object3D) {
+export function removeInteractivity(store: RootStore, object: THREE.Object3D) {
   const { internal } = store.getState()
   // Removes every trace of an object from the data store
   internal.interaction = internal.interaction.filter((o) => o !== object)
@@ -163,7 +162,7 @@ export function removeInteractivity(store: UseBoundStore<RootState>, object: THR
   })
 }
 
-export function createEvents(store: UseBoundStore<RootState>) {
+export function createEvents(store: RootStore) {
   /** Calculates delta */
   function calculateDistance(event: DomEvent) {
     const { internal } = store.getState()

packages/fiber/src/core/hooks.tsx

@@ -1,6 +1,5 @@
 import * as THREE from 'three'
 import * as React from 'react'
-import { StateSelector, EqualityChecker } from 'zustand'
 import { GLTF } from 'three/examples/jsm/loaders/GLTFLoader'
 import { suspend, preload, clear } from 'suspend-react'
 import { context, RootState, RenderCallback, StageTypes } from './store'
@@ -49,8 +48,8 @@ export function useStore() {
  * @see https://docs.pmnd.rs/react-three-fiber/api/hooks#usethree
  */
 export function useThree<T = RootState>(
-  selector: StateSelector<RootState, T> = (state) => state as unknown as T,
-  equalityFn?: EqualityChecker<T>,
+  selector: (state: RootState) => T = (state) => state as unknown as T,
+  equalityFn?: <T>(state: T, newState: T) => boolean,
 ) {
   return useStore()(selector, equalityFn)
 }

packages/fiber/src/core/index.tsx

@@ -1,7 +1,7 @@
 import * as THREE from 'three'
 import * as React from 'react'
 import { ConcurrentRoot } from 'react-reconciler/constants'
-import create, { StoreApi, UseBoundStore } from 'zustand'
+import create from 'zustand'
 
 import { ThreeElement } from '../three-types'
 import {
@@ -18,6 +18,7 @@ import {
   Subscription,
   FrameloopLegacy,
   Frameloop,
+  RootStore,
 } from './store'
 import { reconciler, extend, Root } from './renderer'
 import { createLoop, addEffect, addAfterEffect, addTail } from './loop'
@@ -97,7 +98,7 @@ export type RenderProps<TCanvas extends Element> = {
     manual?: boolean
   }
   /** An R3F event manager to manage elements' pointer events */
-  events?: (store: UseBoundStore<RootState>) => EventManager<HTMLElement>
+  events?: (store: RootStore) => EventManager<HTMLElement>
   /** Callback after the canvas has rendered (but not yet committed) */
   onCreated?: (state: RootState) => void
   /** Response for pointer clicks that have missed any target */
@@ -122,7 +123,7 @@ const createRendererInstance = <TElement extends Element>(gl: GLProps, canvas: T
     })
 }
 
-const createStages = (stages: Stage[] | undefined, store: UseBoundStore<RootState, StoreApi<RootState>>) => {
+const createStages = (stages: Stage[] | undefined, store: RootStore) => {
   const state = store.getState()
   let subscribers: Subscription[]
   let subscription: Subscription
@@ -157,7 +158,7 @@ const createStages = (stages: Stage[] | undefined, store: UseBoundStore<RootStat
 
 export type ReconcilerRoot<TCanvas extends Element> = {
   configure: (config?: RenderProps<TCanvas>) => ReconcilerRoot<TCanvas>
-  render: (element: React.ReactNode) => UseBoundStore<RootState>
+  render: (element: React.ReactNode) => RootStore
   unmount: () => void
 }
 
@@ -366,7 +367,7 @@ function render<TCanvas extends Element>(
   children: React.ReactNode,
   canvas: TCanvas,
   config: RenderProps<TCanvas>,
-): UseBoundStore<RootState> {
+): RootStore {
   console.warn('R3F.render is no longer supported in React 18. Use createRoot instead!')
   const root = createRoot(canvas)
   root.configure(config)
@@ -380,7 +381,7 @@ function Provider<TElement extends Element>({
   rootElement,
 }: {
   onCreated?: (state: RootState) => void
-  store: UseBoundStore<RootState>
+  store: RootStore
   children: React.ReactNode
   rootElement: TElement
   parent?: React.MutableRefObject<TElement | undefined>

packages/fiber/src/core/renderer.ts

@@ -1,14 +1,13 @@
 import * as THREE from 'three'
-import { UseBoundStore } from 'zustand'
 import Reconciler from 'react-reconciler'
 import { unstable_IdlePriority as idlePriority, unstable_scheduleCallback as scheduleCallback } from 'scheduler'
 import { is, diffProps, applyProps, invalidateInstance, attach, detach, prepare } from './utils'
-import { RootState } from './store'
+import { RootState, RootStore } from './store'
 import { removeInteractivity, getEventPriority, EventHandlers } from './events'
 
 export interface Root {
   fiber: Reconciler.FiberRoot
-  store: UseBoundStore<RootState>
+  store: RootStore
 }
 
 export type AttachFnType<O = any> = (parent: any, self: O) => () => void
@@ -31,7 +30,7 @@ export interface InstanceProps<T = any> {
 }
 
 export interface Instance<O = any> {
-  root: UseBoundStore<RootState>
+  root: RootStore
   type: string
   parent: Instance | null
   children: Instance[]
@@ -47,7 +46,7 @@ export interface Instance<O = any> {
 interface HostConfig {
   type: string
   props: Instance['props']
-  container: UseBoundStore<RootState>
+  container: RootStore
   instance: Instance
   textInstance: void
   suspenseInstance: Instance
@@ -63,11 +62,7 @@ interface HostConfig {
 const catalogue: Catalogue = {}
 const extend = (objects: Partial<Catalogue>): void => void Object.assign(catalogue, objects)
 
-function createInstance(
-  type: string,
-  props: HostConfig['props'],
-  root: UseBoundStore<RootState>,
-): HostConfig['instance'] {
+function createInstance(type: string, props: HostConfig['props'], root: RootStore): HostConfig['instance'] {
   // Get target from catalogue
   const name = `${type[0].toUpperCase()}${type.slice(1)}`
   const target = catalogue[name]

packages/fiber/src/core/stages.ts

@@ -1,14 +1,12 @@
 import { MutableRefObject } from 'react'
-import { StoreApi, UseBoundStore } from 'zustand'
-import { RootState } from './store'
+import { RootState, RootStore } from './store'
 
 export interface UpdateCallback {
   (state: RootState, delta: number, frame?: XRFrame): void
 }
 
 export type UpdateCallbackRef = MutableRefObject<UpdateCallback>
-type Store = UseBoundStore<RootState, StoreApi<RootState>>
-export type UpdateSubscription = { ref: UpdateCallbackRef; store: Store }
+export type UpdateSubscription = { ref: UpdateCallbackRef; store: RootStore }
 
 export type FixedStageOptions = { fixedStep?: number; maxSubsteps?: number }
 export type FixedStageProps = { fixedStep: number; maxSubsteps: number; accumulator: number; alpha: number }
@@ -48,7 +46,7 @@ export class Stage {
    * @param store - The store to be used with the callback execution.
    * @returns A function to remove the subscription.
    */
-  add(ref: UpdateCallbackRef, store: Store) {
+  add(ref: UpdateCallbackRef, store: RootStore) {
     this.subscribers.push({ ref, store })
 
     return () => {

packages/fiber/src/core/store.ts

@@ -1,6 +1,6 @@
 import * as THREE from 'three'
 import * as React from 'react'
-import create, { GetState, SetState, StoreApi, UseBoundStore } from 'zustand'
+import create, { StoreApi, UseBoundStore } from 'zustand'
 import { DomEvent, EventManager, PointerCaptureTarget, ThreeEvent } from './events'
 import { calculateDpr, Camera, isOrthographicCamera, prepare, updateCamera } from './utils'
 import { FixedStage, Stage } from './stages'
@@ -28,7 +28,7 @@ export interface Intersection extends THREE.Intersection {
 export type Subscription = {
   ref: React.MutableRefObject<RenderCallback>
   priority: number
-  store: UseBoundStore<RootState, StoreApi<RootState>>
+  store: RootStore
 }
 
 export type Dpr = number | [min: number, max: number]
@@ -89,18 +89,14 @@ export type InternalState = {
   render: 'auto' | 'manual'
   /** The max delta time between two frames. */
   maxDelta: number
-  subscribe: (
-    callback: React.MutableRefObject<RenderCallback>,
-    priority: number,
-    store: UseBoundStore<RootState, StoreApi<RootState>>,
-  ) => () => void
+  subscribe: (callback: React.MutableRefObject<RenderCallback>, priority: number, store: RootStore) => () => void
 }
 
 export type RootState = {
   /** Set current state */
-  set: SetState<RootState>
+  set: StoreApi<RootState>['setState']
   /** Get current state */
-  get: GetState<RootState>
+  get: StoreApi<RootState>['getState']
   /** The instance of the renderer */
   gl: THREE.WebGLRenderer
   /** Default camera */
@@ -156,17 +152,19 @@ export type RootState = {
   /** When the canvas was clicked but nothing was hit */
   onPointerMissed?: (event: MouseEvent) => void
   /** If this state model is layerd (via createPortal) then this contains the previous layer */
-  previousRoot?: UseBoundStore<RootState, StoreApi<RootState>>
+  previousRoot?: RootStore
   /** Internals */
   internal: InternalState
 }
 
-const context = React.createContext<UseBoundStore<RootState>>(null!)
+export type RootStore = UseBoundStore<StoreApi<RootState>>
+
+const context = React.createContext<RootStore>(null!)
 
 const createStore = (
   invalidate: (state?: RootState, frames?: number) => void,
   advance: (timestamp: number, runGlobalEffects?: boolean, state?: RootState, frame?: XRFrame) => void,
-): UseBoundStore<RootState> => {
+): RootStore => {
   const rootStore = create<RootState>((set, get) => {
     const position = new THREE.Vector3()
     const defaultTarget = new THREE.Vector3()
@@ -311,11 +309,7 @@ const createStore = (
         render: 'auto',
         maxDelta: 1 / 10,
         priority: 0,
-        subscribe: (
-          ref: React.MutableRefObject<RenderCallback>,
-          priority: number,
-          store: UseBoundStore<RootState, StoreApi<RootState>>,
-        ) => {
+        subscribe: (ref: React.MutableRefObject<RenderCallback>, priority: number, store: RootStore) => {
           const state = get()
           const internal = state.internal
           // If this subscription was given a priority, it takes rendering into its own hands

packages/fiber/src/core/utils.ts

@@ -1,9 +1,8 @@
 import * as THREE from 'three'
 import * as React from 'react'
 import type { Fiber } from 'react-reconciler'
-import type { UseBoundStore } from 'zustand'
 import type { EventHandlers } from './events'
-import type { Dpr, RootState, Size } from './store'
+import type { Dpr, RootState, RootStore, Size } from './store'
 import type { ConstructorRepresentation, Instance } from './renderer'
 
 export type Camera = THREE.OrthographicCamera | THREE.PerspectiveCamera
@@ -154,12 +153,7 @@ export function getInstanceProps<T = any>(queue: Fiber['pendingProps']): Instanc
 }
 
 // Each object in the scene carries a small LocalState descriptor
-export function prepare<T = any>(
-  target: T,
-  root: UseBoundStore<RootState>,
-  type: string,
-  props: Instance<T>['props'],
-): Instance<T> {
+export function prepare<T = any>(target: T, root: RootStore, type: string, props: Instance<T>['props']): Instance<T> {
   const object = target as unknown as Instance['object']
 
   // Create instance descriptor

packages/fiber/src/index.tsx

@@ -17,6 +17,7 @@ export type {
   RenderCallback,
   Performance,
   RootState,
+  RootStore,
 } from './core/store'
 export type { ThreeEvent, Events, EventManager, ComputeFunction } from './core/events'
 export type { ObjectMap, Camera } from './core/utils'

packages/fiber/src/native.tsx

@@ -17,6 +17,7 @@ export type {
   RenderCallback,
   Performance,
   RootState,
+  RootStore,
 } from './core/store'
 export type { ThreeEvent, Events, EventManager, ComputeFunction } from './core/events'
 export type { ObjectMap, Camera } from './core/utils'

packages/fiber/src/native/events.ts

@@ -1,5 +1,4 @@
-import { UseBoundStore } from 'zustand'
-import { RootState } from '../core/store'
+import { RootState, RootStore } from '../core/store'
 import { createEvents, DomEvent, EventManager, Events } from '../core/events'
 import { GestureResponderEvent } from 'react-native'
 /* eslint-disable import/default, import/no-named-as-default, import/no-named-as-default-member */
@@ -31,7 +30,7 @@ const DOM_EVENTS = {
 }
 
 /** Default R3F event manager for react-native */
-export function createTouchEvents(store: UseBoundStore<RootState>): EventManager<HTMLElement> {
+export function createTouchEvents(store: RootStore): EventManager<HTMLElement> {
   const { handlePointer } = createEvents(store)
 
   const handleTouch = (event: GestureResponderEvent, name: keyof typeof EVENTS) => {

packages/fiber/src/web/events.ts

@@ -1,5 +1,4 @@
-import { UseBoundStore } from 'zustand'
-import { RootState } from '../core/store'
+import { RootState, RootStore } from '../core/store'
 import { EventManager, Events, createEvents, DomEvent } from '../core/events'
 
 const DOM_EVENTS = {
@@ -16,7 +15,7 @@ const DOM_EVENTS = {
 } as const
 
 /** Default R3F event manager for web */
-export function createPointerEvents(store: UseBoundStore<RootState>): EventManager<HTMLElement> {
+export function createPointerEvents(store: RootStore): EventManager<HTMLElement> {
   const { handlePointer } = createEvents(store)
 
   return {

packages/fiber/tests/index.test.tsx

@@ -1,7 +1,6 @@
 import * as React from 'react'
 import * as THREE from 'three'
-import { ReconcilerRoot, createRoot, act, useFrame, useThree, createPortal, RootState } from '../src/index'
-import { UseBoundStore } from 'zustand'
+import { ReconcilerRoot, createRoot, act, useFrame, useThree, createPortal, RootState, RootStore } from '../src/index'
 import { privateKeys } from '../src/core/store'
 
 let root: ReconcilerRoot<HTMLCanvasElement> = null!
@@ -26,35 +25,35 @@ describe('createRoot', () => {
   })
 
   it('should handle an performance changing functions', async () => {
-    let state: UseBoundStore<RootState> = null!
+    let store: RootStore = null!
     await act(async () => {
-      state = root.configure({ dpr: [1, 2], performance: { min: 0.2 } }).render(<group />)
+      store = root.configure({ dpr: [1, 2], performance: { min: 0.2 } }).render(<group />)
     })
 
-    expect(state.getState().viewport.initialDpr).toEqual(window.devicePixelRatio)
-    expect(state.getState().performance.min).toEqual(0.2)
-    expect(state.getState().performance.current).toEqual(1)
+    expect(store.getState().viewport.initialDpr).toEqual(window.devicePixelRatio)
+    expect(store.getState().performance.min).toEqual(0.2)
+    expect(store.getState().performance.current).toEqual(1)
 
     await act(async () => {
-      state.getState().setDpr(0.1)
+      store.getState().setDpr(0.1)
     })
 
-    expect(state.getState().viewport.dpr).toEqual(0.1)
+    expect(store.getState().viewport.dpr).toEqual(0.1)
 
     jest.useFakeTimers()
 
     await act(async () => {
-      state.getState().performance.regress()
+      store.getState().performance.regress()
       jest.advanceTimersByTime(100)
     })
 
-    expect(state.getState().performance.current).toEqual(0.2)
+    expect(store.getState().performance.current).toEqual(0.2)
 
     await act(async () => {
       jest.advanceTimersByTime(200)
     })
 
-    expect(state.getState().performance.current).toEqual(1)
+    expect(store.getState().performance.current).toEqual(1)
 
     jest.useRealTimers()
   })