feat: expose store creator api and support pass custom store to canvas

packages/fiber/package.json

@@ -52,7 +52,7 @@
     "react-use-measure": "^2.1.1",
     "scheduler": "^0.21.0",
     "suspend-react": "^0.1.3",
-    "zustand": "^3.7.1"
+    "zustand": "^4.4.1"
   },
   "peerDependencies": {
     "expo": ">=43.0",

packages/fiber/src/core/events.ts

@@ -1,7 +1,7 @@
 import * as THREE from 'three'
 import { ContinuousEventPriority, DiscreteEventPriority, DefaultEventPriority } from 'react-reconciler/constants'
 import { getRootState } from './utils'
-import type { UseBoundStore } from 'zustand'
+import type { StoreApi, UseBoundStore } from 'zustand'
 import type { Instance } from './renderer'
 import type { RootState } from './store'
 import type { Properties } from '../three-types'
@@ -152,7 +152,7 @@ function releaseInternalPointerCapture(
   }
 }
 
-export function removeInteractivity(store: UseBoundStore<RootState>, object: THREE.Object3D) {
+export function removeInteractivity(store: UseBoundStore<StoreApi<RootState>>, 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)
@@ -168,7 +168,7 @@ export function removeInteractivity(store: UseBoundStore<RootState>, object: THR
   })
 }
 
-export function createEvents(store: UseBoundStore<RootState>) {
+export function createEvents(store: UseBoundStore<StoreApi<RootState>>) {
   /** Calculates delta */
   function calculateDistance(event: DomEvent) {
     const { internal } = store.getState()

packages/fiber/src/core/hooks.tsx

@@ -52,7 +52,7 @@ export function useThree<T = RootState>(
   selector: StateSelector<RootState, T> = (state) => state as unknown as T,
   equalityFn?: EqualityChecker<T>,
 ) {
-  return useStore()(selector, equalityFn)
+  return useStore()(selector, equalityFn as any)
 }
 
 /**

packages/fiber/src/core/index.tsx

@@ -1,12 +1,13 @@
 import * as THREE from 'three'
 import * as React from 'react'
 import { ConcurrentRoot } from 'react-reconciler/constants'
-import create, { UseBoundStore } from 'zustand'
+import create, { StoreApi, UseBoundStore } from 'zustand'
 
 import * as ReactThreeFiber from '../three-types'
 import {
   Renderer,
   createStore,
+  three,
   isRenderer,
   context,
   RootState,
@@ -99,7 +100,7 @@ export type RenderProps<TCanvas extends Canvas> = {
     manual?: boolean
   }
   /** An R3F event manager to manage elements' pointer events */
-  events?: (store: UseBoundStore<RootState>) => EventManager<HTMLElement>
+  events?: (store: UseBoundStore<StoreApi<RootState>>) => 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 */
@@ -121,7 +122,7 @@ const createRendererInstance = <TCanvas extends Canvas>(gl: GLProps, canvas: TCa
 
 export type ReconcilerRoot<TCanvas extends Canvas> = {
   configure: (config?: RenderProps<TCanvas>) => ReconcilerRoot<TCanvas>
-  render: (element: React.ReactNode) => UseBoundStore<RootState>
+  render: (element: React.ReactNode) => UseBoundStore<StoreApi<RootState>>
   unmount: () => void
 }
 
@@ -147,7 +148,10 @@ function computeInitialSize(canvas: Canvas, defaultSize?: Size): Size {
   return { width: 0, height: 0, top: 0, left: 0 }
 }
 
-function createRoot<TCanvas extends Canvas>(canvas: TCanvas): ReconcilerRoot<TCanvas> {
+function createRoot<TCanvas extends Canvas>(
+  canvas: TCanvas,
+  store?: UseBoundStore<StoreApi<RootState>>,
+): ReconcilerRoot<TCanvas> {
   // Check against mistaken use of createRoot
   const prevRoot = roots.get(canvas)
   const prevFiber = prevRoot?.fiber
@@ -166,12 +170,13 @@ function createRoot<TCanvas extends Canvas>(canvas: TCanvas): ReconcilerRoot<TCa
         console.error
 
   // Create store
-  const store = prevStore || createStore(invalidate, advance)
+  const actualStore = store || prevStore || createStore(invalidate, advance)
   // Create renderer
   const fiber =
-    prevFiber || reconciler.createContainer(store, ConcurrentRoot, null, false, null, '', logRecoverableError, null)
+    prevFiber ||
+    reconciler.createContainer(actualStore, ConcurrentRoot, null, false, null, '', logRecoverableError, null)
   // Map it
-  if (!prevRoot) roots.set(canvas, { fiber, store })
+  if (!prevRoot) roots.set(canvas, { fiber, store: actualStore })
 
   // Locals
   let onCreated: ((state: RootState) => void) | undefined
@@ -199,7 +204,7 @@ function createRoot<TCanvas extends Canvas>(canvas: TCanvas): ReconcilerRoot<TCa
         onPointerMissed,
       } = props
 
-      let state = store.getState()
+      let state = actualStore.getState()
 
       // Set up renderer (one time only!)
       let gl = state.gl
@@ -255,14 +260,14 @@ function createRoot<TCanvas extends Canvas>(canvas: TCanvas): ReconcilerRoot<TCa
       if (!state.xr) {
         // Handle frame behavior in WebXR
         const handleXRFrame = (timestamp: number, frame?: _XRFrame) => {
-          const state = store.getState()
+          const state = actualStore.getState()
           if (state.frameloop === 'never') return
           advance(timestamp, true, state, frame)
         }
 
         // Toggle render switching on session
         const handleSessionChange = () => {
-          const state = store.getState()
+          const state = actualStore.getState()
           state.gl.xr.enabled = state.gl.xr.isPresenting
 
           state.gl.xr.setAnimationLoop(state.gl.xr.isPresenting ? handleXRFrame : null)
@@ -272,12 +277,12 @@ function createRoot<TCanvas extends Canvas>(canvas: TCanvas): ReconcilerRoot<TCa
         // WebXR session manager
         const xr = {
           connect() {
-            const gl = store.getState().gl
+            const gl = actualStore.getState().gl
             gl.xr.addEventListener('sessionstart', handleSessionChange)
             gl.xr.addEventListener('sessionend', handleSessionChange)
           },
           disconnect() {
-            const gl = store.getState().gl
+            const gl = actualStore.getState().gl
             gl.xr.removeEventListener('sessionstart', handleSessionChange)
             gl.xr.removeEventListener('sessionend', handleSessionChange)
           },
@@ -341,7 +346,7 @@ function createRoot<TCanvas extends Canvas>(canvas: TCanvas): ReconcilerRoot<TCa
       if (glConfig && !is.fun(glConfig) && !isRenderer(glConfig) && !is.equ(glConfig, gl, shallowLoose))
         applyProps(gl as any, glConfig as any)
       // Store events internally
-      if (events && !state.events.handlers) state.set({ events: events(store) })
+      if (events && !state.events.handlers) state.set({ events: events(actualStore) })
       // Check size, allow it to take on container bounds initially
       const size = computeInitialSize(canvas, propsSize)
       if (!is.equ(size, state.size, shallowLoose)) {
@@ -368,12 +373,12 @@ function createRoot<TCanvas extends Canvas>(canvas: TCanvas): ReconcilerRoot<TCa
       if (!configured) this.configure()
 
       reconciler.updateContainer(
-        <Provider store={store} children={children} onCreated={onCreated} rootElement={canvas} />,
+        <Provider store={actualStore} children={children} onCreated={onCreated} rootElement={canvas} />,
         fiber,
         null,
         () => undefined,
       )
-      return store
+      return actualStore
     },
     unmount() {
       unmountComponentAtNode(canvas)
@@ -385,7 +390,7 @@ function render<TCanvas extends Canvas>(
   children: React.ReactNode,
   canvas: TCanvas,
   config: RenderProps<TCanvas>,
-): UseBoundStore<RootState> {
+): UseBoundStore<StoreApi<RootState>> {
   console.warn('R3F.render is no longer supported in React 18. Use createRoot instead!')
   const root = createRoot(canvas)
   root.configure(config)
@@ -399,7 +404,7 @@ function Provider<TCanvas extends Canvas>({
   rootElement,
 }: {
   onCreated?: (state: RootState) => void
-  store: UseBoundStore<RootState>
+  store: UseBoundStore<StoreApi<RootState>>
   children: React.ReactNode
   rootElement: TCanvas
 }) {
@@ -414,6 +419,7 @@ function Provider<TCanvas extends Canvas>({
     if (!store.getState().events.connected) state.events.connect?.(rootElement)
     // eslint-disable-next-line react-hooks/exhaustive-deps
   }, [])
+
   return <context.Provider value={store}>{children}</context.Provider>
 }
 
@@ -603,6 +609,8 @@ export {
   addTail,
   flushGlobalEffects,
   getRootState,
+  createStore,
+  three,
   act,
   buildGraph,
   roots as _roots,

packages/fiber/src/core/renderer.ts

@@ -1,5 +1,5 @@
 import * as THREE from 'three'
-import { UseBoundStore } from 'zustand'
+import { StoreApi, UseBoundStore } from 'zustand'
 import Reconciler from 'react-reconciler'
 import { unstable_IdlePriority as idlePriority, unstable_scheduleCallback as scheduleCallback } from 'scheduler'
 import { DefaultEventPriority } from 'react-reconciler/constants'
@@ -18,11 +18,11 @@ import {
 import { RootState } from './store'
 import { EventHandlers, removeInteractivity } from './events'
 
-export type Root = { fiber: Reconciler.FiberRoot; store: UseBoundStore<RootState> }
+export type Root = { fiber: Reconciler.FiberRoot; store: UseBoundStore<StoreApi<RootState>> }
 
 export type LocalState = {
   type: string
-  root: UseBoundStore<RootState>
+  root: UseBoundStore<StoreApi<RootState>>
   // objects and parent are used when children are added with `attach` instead of being added to the Object3D scene graph
   objects: Instance[]
   parent: Instance | null
@@ -41,7 +41,7 @@ export type AttachType = string | AttachFnType
 interface HostConfig {
   type: string
   props: InstanceProps
-  container: UseBoundStore<RootState>
+  container: UseBoundStore<StoreApi<RootState>>
   instance: Instance
   textInstance: void
   suspenseInstance: Instance
@@ -89,7 +89,7 @@ function createRenderer<TCanvas>(_roots: Map<TCanvas, Root>, _getEventPriority?:
   function createInstance(
     type: string,
     { args = [], attach, ...props }: InstanceProps,
-    root: UseBoundStore<RootState>,
+    root: UseBoundStore<StoreApi<RootState>>,
   ) {
     let name = `${type[0].toUpperCase()}${type.slice(1)}`
     let instance: Instance