Implement dynamic scene trees in r3f

This makes it possible to add/remove objects on the fly, do hot-module reloading, change object configs on the fly, and more.

packages/r3f/src/extension/components/EditableProxy.tsx

@@ -43,21 +43,21 @@ const EditableProxy: VFC<EditableProxyProps> = ({storeKey, object}) => {
   useLayoutEffect(() => {
     const originalVisibility = object.visible
 
-    if (editable.visibleOnlyInEditor) {
+    if (editable?.visibleOnlyInEditor) {
       object.visible = true
     }
 
     return () => {
       object.visible = originalVisibility
     }
-  }, [editable.visibleOnlyInEditor, object.visible])
+  }, [editable?.visibleOnlyInEditor, object.visible])
 
   const [hovered, setHovered] = useState(false)
 
   // Helpers
   const scene = useThree((state) => state.scene)
   const helper = useMemo<Helper | undefined>(
-    () => editable.objectConfig.createHelper?.(object),
+    () => editable?.objectConfig.createHelper?.(object),
     [object],
   )
   useEffect(() => {
@@ -92,6 +92,7 @@ const EditableProxy: VFC<EditableProxyProps> = ({storeKey, object}) => {
 
   // subscribe to external changes
   useEffect(() => {
+    if (!editable) return
     const sheetObject = editable.sheetObject
     const objectConfig = editable.objectConfig
 
@@ -114,6 +115,8 @@ const EditableProxy: VFC<EditableProxyProps> = ({storeKey, object}) => {
     }
   }, [editable])
 
+  if (!editable) return null
+
   return (
     <>
       <group

packages/r3f/src/extension/components/ProxyManager.tsx

@@ -58,18 +58,19 @@ const ProxyManager: VFC<ProxyManagerProps> = ({orbitControlsRef}) => {
 
     sceneProxy.traverse((object) => {
       if (object.userData.__editable) {
-        // there are duplicate theatreKeys in the scene, only display one instance in the editor
+        const theatreKey = object.userData.__storeKey
-        if (editableProxies[object.userData.__storeKey]) {
+
+        if (
+          // there are duplicate theatreKeys in the scene, only display one instance in the editor
+          editableProxies[theatreKey] ||
+          // this object has been unmounted
+          !editables[theatreKey]
+        ) {
           object.parent!.remove(object)
         } else {
-          const theatreKey = object.userData.__storeKey
-
           editableProxies[theatreKey] = {
             portal: createPortal(
-              <EditableProxy
+              <EditableProxy storeKey={theatreKey} object={object} />,
-                storeKey={object.userData.__storeKey}
-                object={object}
-              />,
               object.parent!,
             ),
             object: object,

packages/r3f/src/main/editable.tsx

@@ -1,4 +1,11 @@
-import type {ComponentProps, ComponentType, Ref, RefAttributes} from 'react'
+import {
+  ComponentProps,
+  ComponentType,
+  Ref,
+  RefAttributes,
+  useMemo,
+  useState,
+} from 'react'
 import React, {forwardRef, useEffect, useLayoutEffect, useRef} from 'react'
 import {allRegisteredObjects, editorStore} from './store'
 import mergeRefs from 'react-merge-refs'
@@ -7,7 +14,8 @@ import {useCurrentSheet} from './SheetProvider'
 import defaultEditableFactoryConfig from './defaultEditableFactoryConfig'
 import type {EditableFactoryConfig} from './editableFactoryConfigUtils'
 import {makeStoreKey} from './utils'
-import type {$FixMe} from '../types'
+import type {$FixMe, $IntentionalAny} from '../types'
+import type {ISheetObject} from '@theatre/core'
 
 const createEditable = <Keys extends keyof JSX.IntrinsicElements>(
   config: EditableFactoryConfig,
@@ -49,21 +57,18 @@ const createEditable = <Keys extends keyof JSX.IntrinsicElements>(
 
         const sheet = useCurrentSheet()!
 
-        const sheetObject = sheet.object(
+        const [sheetObject, setSheetObject] = useState<
-          theatreKey,
+          undefined | ISheetObject<$FixMe>
-          Object.assign(
+        >(undefined)
-            {
-              ...additionalProps,
-            },
-            // @ts-ignore
-            ...Object.values(config[actualType].props).map(
-              // @ts-ignore
-              (value) => value.type,
-            ),
-          ),
-        )
 
-        const storeKey = makeStoreKey(sheetObject.address)
+        const storeKey = useMemo(
+          () =>
+            makeStoreKey({
+              ...sheet.address,
+              objectKey: theatreKey as $IntentionalAny,
+            }),
+          [sheet, theatreKey],
+        )
 
         const invalidate = useInvalidate()
 
@@ -102,22 +107,53 @@ Then you can use it in your JSX like any other editable component. Note the make
         // create sheet object and add editable to store
         useLayoutEffect(() => {
           if (!sheet) return
+          if (sheetObject) {
+            sheet.object(
+              theatreKey,
+              Object.assign(
+                {
+                  ...additionalProps,
+                },
+                // @ts-ignore
+                ...Object.values(config[actualType].props).map(
+                  // @ts-ignore
+                  (value) => value.type,
+                ),
+              ),
+              {override: true},
+            )
+            return
+          } else {
+            const sheetObject = sheet.object(
+              theatreKey,
+              Object.assign(
+                {
+                  ...additionalProps,
+                },
+                // @ts-ignore
+                ...Object.values(config[actualType].props).map(
+                  // @ts-ignore
+                  (value) => value.type,
+                ),
+              ),
+            )
+            allRegisteredObjects.add(sheetObject)
+            setSheetObject(sheetObject)
 
-          allRegisteredObjects.add(sheetObject)
+            if (objRef)
+              typeof objRef === 'function'
+                ? objRef(sheetObject)
+                : (objRef.current = sheetObject)
 
-          if (objRef)
+            editorStore.getState().addEditable(storeKey, {
-            typeof objRef === 'function'
+              type: actualType,
-              ? objRef(sheetObject)
+              sheetObject,
-              : (objRef.current = sheetObject)
+              visibleOnlyInEditor: visible === 'editor',
-
+              // @ts-ignore
-          editorStore.getState().addEditable(storeKey, {
+              objectConfig: config[actualType],
-            type: actualType,
+            })
-            sheetObject,
+          }
-            visibleOnlyInEditor: visible === 'editor',
+        }, [sheet, storeKey, additionalProps])
-            // @ts-ignore
-            objectConfig: config[actualType],
-          })
-        }, [sheet, storeKey])
 
         // store initial values of props
         useLayoutEffect(() => {
@@ -161,6 +197,9 @@ Then you can use it in your JSX like any other editable component. Note the make
 
           return () => {
             untap()
+            sheetObject.sheet.deleteObject(theatreKey)
+            allRegisteredObjects.delete(sheetObject)
+            editorStore.getState().removeEditable(storeKey)
           }
         }, [sheetObject])
 

packages/r3f/src/main/store.ts

@@ -45,6 +45,7 @@ export type EditorStore = {
   init: (scene: Scene, gl: WebGLRenderer) => void
 
   addEditable: (theatreKey: string, editable: Editable<any>) => void
+  removeEditable: (theatreKey: string) => void
   createSnapshot: () => void
   setSnapshotProxyObject: (
     proxyObject: Object3D | null,
@@ -85,6 +86,16 @@ const config: StateCreator<EditorStore> = (set, get) => {
       }))
     },
 
+    removeEditable: (theatreKey) => {
+      set((state) => {
+        const editables = {...state.editables}
+        delete editables[theatreKey]
+        return {
+          editables,
+        }
+      })
+    },
+
     createSnapshot: () => {
       set((state) => ({
         sceneSnapshot: state.scene?.clone() ?? null,