Added support for multiple sheets in r3f

packages/plugin-r3f/src/Wrapper.tsx

@@ -1,12 +1,26 @@
-import React, {useLayoutEffect} from 'react'
+import React, {
+  createContext,
+  useContext,
+  useLayoutEffect,
+  useState,
+} from 'react'
 import {useThree} from '@react-three/fiber'
 import type {ISheet} from '@theatre/core'
 import {bindToCanvas} from './store'
 
+const ctx = createContext<{sheet: ISheet | undefined} | undefined>(undefined)
+
+export const useWrapperContext = ():
+  | {sheet: ISheet | undefined}
+  | undefined => {
+  return useContext(ctx)
+}
+
 const Wrapper: React.FC<{
   getSheet: () => ISheet
 }> = (props) => {
   const {scene, gl} = useThree((s) => ({scene: s.scene, gl: s.gl}))
+  const [sheet, setSheet] = useState<ISheet | undefined>(undefined)
 
   useLayoutEffect(() => {
     const sheet = props.getSheet()
@@ -15,10 +29,11 @@ const Wrapper: React.FC<{
         `getSheet() in <Wrapper getSheet={getSheet}> has returned an invalid value`,
       )
     }
-    bindToCanvas({sheet, gl, scene})
+    setSheet(sheet)
+    bindToCanvas({gl, scene})
   }, [scene, gl])
 
-  return <>{props.children}</>
+  return <ctx.Provider value={{sheet}}>{props.children}</ctx.Provider>
 }
 
 export default Wrapper

packages/plugin-r3f/src/components/editable.tsx

@@ -18,6 +18,7 @@ import mergeRefs from 'react-merge-refs'
 import type {$FixMe} from '@theatre/shared/utils/types'
 import type {ISheetObject} from '@theatre/core'
 import useInvalidate from './useInvalidate'
+import {useWrapperContext} from '../Wrapper'
 
 interface Elements {
   group: Group
@@ -50,7 +51,14 @@ const editable = <
     ({uniqueName, visible, editableType, ...props}: Props, ref) => {
       const objectRef = useRef<Elements[U]>()
 
-      const sheet = useEditorStore((state) => state.sheet)
+      const wrapperContext = useWrapperContext()
+      if (!wrapperContext) {
+        throw new Error(
+          `Editable components must be a descendent of a <Wrapper>`,
+        )
+      }
+
+      const {sheet} = wrapperContext
 
       const [sheetObject, setSheetObject] = useState<
         undefined | ISheetObject<$FixMe>

packages/plugin-r3f/src/components/useSelected.tsx

@@ -1,6 +1,5 @@
 import {useLayoutEffect, useRef, useState} from 'react'
-import {useEditorStore} from '../store'
-import shallow from 'zustand/shallow'
+import {allRegisteredObjects} from '../store'
 import studio from '@theatre/studio'
 import type {ISheetObject} from '@theatre/core'
 
@@ -9,13 +8,12 @@ export function useSelected(): undefined | string {
   const stateRef = useRef(state)
   stateRef.current = state
 
-  const sheet = useEditorStore((state) => state.sheet, shallow)
-
   useLayoutEffect(() => {
     const setFromStudio = (selection: typeof studio.selection) => {
       const item = selection.find(
         (s): s is ISheetObject =>
-          s.type === 'Theatre_SheetObject_PublicAPI' && s.sheet === sheet,
+          s.type === 'Theatre_SheetObject_PublicAPI' &&
+          allRegisteredObjects.has(s),
       )
       if (!item) {
         set(undefined)
@@ -25,17 +23,15 @@ export function useSelected(): undefined | string {
     }
     setFromStudio(studio.selection)
     return studio.onSelectionChange(setFromStudio)
-  }, [sheet])
+  }, [])
 
   return state
 }
 
 export function getSelected(): undefined | string {
-  const sheet = useEditorStore.getState().sheet
-  if (!sheet) return undefined
   const item = studio.selection.find(
     (s): s is ISheetObject =>
-      s.type === 'Theatre_SheetObject_PublicAPI' && s.sheet === sheet,
+      s.type === 'Theatre_SheetObject_PublicAPI' && allRegisteredObjects.has(s),
   )
   if (!item) {
     return undefined

packages/plugin-r3f/src/store.ts

@@ -2,7 +2,7 @@ import type {StateCreator} from 'zustand'
 import create from 'zustand'
 import type {Object3D, Scene, WebGLRenderer} from 'three'
 import {Group} from 'three'
-import type {ISheet, ISheetObject} from '@theatre/core'
+import type {ISheetObject} from '@theatre/core'
 import {types} from '@theatre/core'
 
 export type EditableType =
@@ -134,7 +134,6 @@ export interface EditableState {
 }
 
 export type EditorStore = {
-  sheet: ISheet | null
   sheetObjects: {[uniqueName in string]?: BaseSheetObjectType}
   scene: Scene | null
   gl: WebGLRenderer | null
@@ -150,7 +149,6 @@ export type EditorStore = {
     scene: Scene,
     gl: WebGLRenderer,
     allowImplicitInstancing: boolean,
-    sheet: ISheet,
   ) => void
 
   addEditable: <T extends EditableType>(type: T, uniqueName: string) => void
@@ -178,12 +176,11 @@ const config: StateCreator<EditorStore> = (set, get) => {
     editablesSnapshot: null,
     initialEditorCamera: {},
 
-    init: (scene, gl, allowImplicitInstancing, sheet) => {
+    init: (scene, gl, allowImplicitInstancing) => {
       set({
         scene,
         gl,
         allowImplicitInstancing,
-        sheet,
       })
     },
 
@@ -266,17 +263,15 @@ export const useEditorStore = create<EditorStore>(config)
 
 export type BindFunction = (options: {
   allowImplicitInstancing?: boolean
-  sheet: ISheet
   gl: WebGLRenderer
   scene: Scene
 }) => void
 
 export const bindToCanvas: BindFunction = ({
   allowImplicitInstancing = false,
-  sheet,
   gl,
   scene,
 }) => {
   const init = useEditorStore.getState().init
-  init(scene, gl, allowImplicitInstancing, sheet)
+  init(scene, gl, allowImplicitInstancing)
 }