fix: memoize HMR, strict check constructors

example/package.json

@@ -7,26 +7,27 @@
     "serve": "vite preview"
   },
   "dependencies": {
-    "@react-spring/core": "^9.4.4",
-    "@react-spring/three": "^9.4.4",
-    "@react-three/drei": "^9.1.2",
+    "@react-spring/core": "^9.6.1",
+    "@react-spring/three": "^9.6.1",
+    "@react-three/drei": "^9.56.24",
     "@use-gesture/react": "latest",
-    "react": "^18.0.0",
-    "react-dom": "^18.0.0",
-    "react-merge-refs": "^1.1.0",
+    "@vitejs/plugin-react": "^3.1.0",
+    "react": "^18.2.0",
+    "react-dom": "^18.2.0",
+    "react-merge-refs": "^2.0.1",
     "react-use-refs": "^1.0.1",
-    "styled-components": "^5.3.5",
-    "three": "^0.139.2",
-    "three-stdlib": "^2.9.1",
+    "styled-components": "^5.3.6",
+    "three": "^0.149.0",
+    "three-stdlib": "^2.21.8",
     "use-error-boundary": "^2.0.6",
-    "wouter": "^2.8.0-alpha.2",
-    "zustand": "^3.7.1"
+    "wouter": "^2.10.0",
+    "zustand": "^4.3.3"
   },
   "devDependencies": {
-    "@types/react": "^17.0.43",
-    "@types/styled-components": "^5.1.24",
+    "@types/react": "^18.0.28",
+    "@types/styled-components": "^5.1.26",
     "@vitejs/plugin-react-refresh": "^1.3.6",
-    "typescript": "^4.6.3",
-    "vite": "^2.9.1"
+    "typescript": "^4.9.5",
+    "vite": "^4.1.1"
   }
 }

example/src/demos/ContextMenuOverride.tsx

@@ -12,7 +12,8 @@ export default function App() {
       <ambientLight />
       <pointLight position={[10, 10, 10]} />
       <mesh
-        position={[0, 0, 0]}
+        scale={[2, 2, 2]}
+        position={[1, 0, 0]}
         onContextMenu={(ev) => {
           ev.nativeEvent.preventDefault()
           set((value) => !value)

example/vite.config.ts

@@ -1,9 +1,9 @@
 import { defineConfig } from 'vite'
-import reactRefresh from '@vitejs/plugin-react-refresh'
+import react from '@vitejs/plugin-react'
 
 export default defineConfig({
   optimizeDeps: {
     exclude: ['@react-three/fiber'],
   },
-  plugins: [reactRefresh()],
+  plugins: [react()],
 })

packages/fiber/src/core/renderer.ts

@@ -21,6 +21,7 @@ export type Root = { fiber: Reconciler.FiberRoot; store: UseBoundStore<RootState
 
 export type LocalState = {
   type: string
+  fiber: Reconciler.Fiber | null
   root: UseBoundStore<RootState>
   // objects and parent are used when children are added with `attach` instead of being added to the Object3D scene graph
   objects: Instance[]
@@ -30,7 +31,6 @@ export type LocalState = {
   handlers: Partial<EventHandlers>
   attach?: AttachType
   previousAttach: any
-  memoizedProps: { [key: string]: any }
   autoRemovedBeforeAppend?: boolean
 }
 
@@ -89,6 +89,8 @@ function createRenderer<TCanvas>(_roots: Map<TCanvas, Root>, _getEventPriority?:
     type: string,
     { args = [], attach, ...props }: InstanceProps,
     root: UseBoundStore<RootState>,
+    _hostContext: null,
+    fiber: Reconciler.Fiber,
   ) {
     let name = `${type[0].toUpperCase()}${type.slice(1)}`
     let instance: Instance
@@ -114,8 +116,7 @@ function createRenderer<TCanvas>(_roots: Map<TCanvas, Root>, _getEventPriority?:
         type,
         root,
         attach,
-        // Save args in case we need to reconstruct later for HMR
-        memoizedProps: { args },
+        fiber,
       })
     }
 
@@ -130,6 +131,7 @@ function createRenderer<TCanvas>(_roots: Map<TCanvas, Root>, _getEventPriority?:
     // why it passes "true" here
     // There is no reason to apply props to injects
     if (name !== 'inject') applyProps(instance, props)
+    instance.__r3f.fiber = fiber
     return instance
   }
 
@@ -226,7 +228,6 @@ function createRenderer<TCanvas>(_roots: Map<TCanvas, Root>, _getEventPriority?:
         delete ((child as Partial<Instance>).__r3f as Partial<LocalState>).root
         delete ((child as Partial<Instance>).__r3f as Partial<LocalState>).objects
         delete ((child as Partial<Instance>).__r3f as Partial<LocalState>).handlers
-        delete ((child as Partial<Instance>).__r3f as Partial<LocalState>).memoizedProps
         if (!isPrimitive) delete (child as Partial<Instance>).__r3f
       }
 
@@ -254,7 +255,7 @@ function createRenderer<TCanvas>(_roots: Map<TCanvas, Root>, _getEventPriority?:
     const parent = instance.__r3f?.parent
     if (!parent) return
 
-    const newInstance = createInstance(type, newProps, instance.__r3f.root)
+    const newInstance = createInstance(type, newProps, instance.__r3f.root, null, fiber)
 
     // https://github.com/pmndrs/react-three-fiber/issues/1348
     // When args change the instance has to be re-constructed, which then

packages/fiber/src/core/utils.ts

@@ -56,13 +56,13 @@ export class ErrorBoundary extends React.Component<
 }
 
 export const DEFAULT = '__default'
+export const DEFAULTS = new Map()
 
 export type DiffSet = {
-  memoized: { [key: string]: any }
   changes: [key: string, value: unknown, isEvent: boolean, keys: string[]][]
 }
 
-export const isDiffSet = (def: any): def is DiffSet => def && !!(def as DiffSet).memoized && !!(def as DiffSet).changes
+export const isDiffSet = (def: any): def is DiffSet => def && !!(def as DiffSet).changes
 export type ClassConstructor = { new (): void }
 
 export type ObjectMap = {
@@ -152,11 +152,11 @@ export function prepare<T = THREE.Object3D>(object: T, state?: Partial<LocalStat
       type: '',
       root: null as unknown as UseBoundStore<RootState>,
       previousAttach: null,
-      memoizedProps: {},
       eventCount: 0,
       handlers: {},
       objects: [],
       parent: null,
+      fiber: null,
       ...state,
     }
   }
@@ -241,11 +241,7 @@ export function diffProps(
     }
   })
 
-  const memoized: { [key: string]: any } = { ...props }
-  if (localState.memoizedProps && localState.memoizedProps.args) memoized.args = localState.memoizedProps.args
-  if (localState.memoizedProps && localState.memoizedProps.attach) memoized.attach = localState.memoizedProps.attach
-
-  return { memoized, changes }
+  return { changes }
 }
 
 // This function applies a set of changes to the instance
@@ -254,12 +250,9 @@ export function applyProps(instance: Instance, data: InstanceProps | DiffSet) {
   const localState = (instance.__r3f ?? {}) as LocalState
   const root = localState.root
   const rootState = root?.getState?.() ?? {}
-  const { memoized, changes } = isDiffSet(data) ? data : diffProps(instance, data)
+  const { changes } = isDiffSet(data) ? data : diffProps(instance, data)
   const prevHandlers = localState.eventCount
 
-  // Prepare memoized props
-  if (instance.__r3f) instance.__r3f.memoizedProps = memoized
-
   for (let i = 0; i < changes.length; i++) {
     let [key, value, isEvent, keys] = changes[i]
     let currentInstance = instance
@@ -284,11 +277,13 @@ export function applyProps(instance: Instance, data: InstanceProps | DiffSet) {
     if (value === DEFAULT + 'remove') {
       if (currentInstance.constructor) {
         // create a blank slate of the instance and copy the particular parameter.
-        // @ts-ignore
-        const defaultClassCall = new currentInstance.constructor(...(currentInstance.__r3f.memoizedProps.args ?? []))
-        value = defaultClassCall[key]
-        // destroy the instance
-        if (defaultClassCall.dispose) defaultClassCall.dispose()
+        let ctor = DEFAULTS.get(currentInstance.constructor)
+        if (!ctor) {
+          // @ts-ignore
+          ctor = new currentInstance.constructor()
+          DEFAULTS.set(currentInstance.constructor, ctor)
+        }
+        value = ctor[key]
       } else {
         // instance does not have constructor, just set it to 0
         value = 0
@@ -313,7 +308,7 @@ export function applyProps(instance: Instance, data: InstanceProps | DiffSet) {
         targetProp.copy &&
         value &&
         (value as ClassConstructor).constructor &&
-        targetProp.constructor.name === (value as ClassConstructor).constructor.name
+        targetProp.constructor === (value as ClassConstructor).constructor
       ) {
         targetProp.copy(value)
       }

packages/test-renderer/src/createTestInstance.ts

@@ -2,7 +2,7 @@ import { Object3D } from 'three'
 
 import type { MockInstance, MockScene, Obj, TestInstanceChildOpts } from './types/internal'
 
-import { expectOne, matchProps, findAll } from './helpers/testInstance'
+import { expectOne, matchProps, findAll, getMemoizedProps } from './helpers/testInstance'
 
 export class ReactThreeTestInstance<TInstance extends Object3D = Object3D> {
   _fiber: MockInstance
@@ -20,7 +20,7 @@ export class ReactThreeTestInstance<TInstance extends Object3D = Object3D> {
   }
 
   public get props(): Obj {
-    return this._fiber.__r3f.memoizedProps
+    return getMemoizedProps(this._fiber)
   }
 
   public get parent(): ReactThreeTestInstance | null {

packages/test-renderer/src/helpers/testInstance.ts

@@ -1,5 +1,34 @@
 import { ReactThreeTestInstance } from '../createTestInstance'
-import type { Obj } from '../types/internal'
+import type { MockInstance, Obj } from '../types/internal'
+
+const INTERNAL_PROPS = [
+  // Exclude react-internal props
+  'children',
+  'key',
+  'ref',
+  // Exclude "correct" JSX props from breaking snapshots
+  // TODO: remove in v9
+  // https://github.com/pmndrs/react-three-fiber/pull/2757#issuecomment-1430362907
+  'attach',
+  'object',
+]
+
+export function getMemoizedProps(instance: MockInstance): Record<string, unknown> {
+  const props: Record<string, unknown> = {}
+
+  // Gets only instance props from its Fiber
+  for (const fiber of [instance.__r3f.fiber, instance.__r3f.fiber?.alternate]) {
+    if (fiber) {
+      for (const key in fiber.memoizedProps) {
+        if (!INTERNAL_PROPS.includes(key)) props[key] ??= fiber.memoizedProps[key]
+      }
+    }
+  }
+
+  props.args ??= []
+
+  return props
+}
 
 export const expectOne = <TItem>(items: TItem[], msg: string) => {
   if (items.length === 1) {

packages/test-renderer/src/helpers/tree.ts

@@ -1,6 +1,7 @@
 import type { TreeNode, Tree } from '../types/public'
 import type { MockSceneChild, MockScene } from '../types/internal'
 import { lowerCaseFirstLetter } from './strings'
+import { getMemoizedProps } from './testInstance'
 
 const treeObjectFactory = (
   type: TreeNode['type'],
@@ -16,7 +17,7 @@ const toTreeBranch = (obj: MockSceneChild[]): TreeNode[] =>
   obj.map((child) => {
     return treeObjectFactory(
       lowerCaseFirstLetter(child.type || child.constructor.name),
-      { ...child.__r3f.memoizedProps },
+      { ...getMemoizedProps(child) },
       toTreeBranch([...(child.children || []), ...child.__r3f.objects]),
     )
   })
@@ -25,7 +26,7 @@ export const toTree = (root: MockScene): Tree =>
   root.children.map((obj) =>
     treeObjectFactory(
       lowerCaseFirstLetter(obj.type),
-      { ...obj.__r3f.memoizedProps },
+      { ...getMemoizedProps(obj) },
       toTreeBranch([...(obj.children as MockSceneChild[]), ...(obj.__r3f.objects as MockSceneChild[])]),
     ),
   )