From ebddb72d714cd130fcc1ce0bb4dbdba4667b312b Mon Sep 17 00:00:00 2001
From: Stefan Schmutz <stefan.schmutz@3dit.de>
Date: Sun, 7 Jan 2024 00:58:40 +0100
Subject: [PATCH] fix: fix array of reused primitives

---
 packages/fiber/src/core/renderer.ts         | 10 +++-
 packages/fiber/tests/core/renderer.test.tsx | 51 +++++++++++++++++++++
 2 files changed, 59 insertions(+), 2 deletions(-)

diff --git a/packages/fiber/src/core/renderer.ts b/packages/fiber/src/core/renderer.ts
index 3963e7fcff..2d74c7ae5a 100644
--- a/packages/fiber/src/core/renderer.ts
+++ b/packages/fiber/src/core/renderer.ts
@@ -32,6 +32,7 @@ export type LocalState = {
   previousAttach: any
   memoizedProps: { [key: string]: any }
   autoRemovedBeforeAppend?: boolean
+  isCreatedBeforeAttach?: boolean
 }
 
 export type AttachFnType = (parent: Instance, self: Instance) => () => void
@@ -96,7 +97,7 @@ function createRenderer<TCanvas>(_roots: Map<TCanvas, Root>, _getEventPriority?:
     if (type === 'primitive') {
       if (props.object === undefined) throw new Error("R3F: Primitives without 'object' are invalid!")
       const object = props.object as Instance
-      instance = prepare<Instance>(object, { type, root, attach, primitive: true })
+      instance = prepare<Instance>(object, { type, root, attach, primitive: true, isCreatedBeforeAttach: true })
     } else {
       const target = catalogue[name]
       if (!target) {
@@ -149,6 +150,7 @@ function createRenderer<TCanvas>(_roots: Map<TCanvas, Root>, _getEventPriority?:
       if (!added) parentInstance.__r3f?.objects.push(child)
       if (!child.__r3f) prepare(child, {})
       child.__r3f.parent = parentInstance
+      child.__r3f.isCreatedBeforeAttach = false
       updateInstance(child)
       invalidateInstance(child)
     }
@@ -175,6 +177,7 @@ function createRenderer<TCanvas>(_roots: Map<TCanvas, Root>, _getEventPriority?:
       if (!added) parentInstance.__r3f?.objects.push(child)
       if (!child.__r3f) prepare(child, {})
       child.__r3f.parent = parentInstance
+      child.__r3f.isCreatedBeforeAttach = false
       updateInstance(child)
       invalidateInstance(child)
     }
@@ -222,7 +225,10 @@ function createRenderer<TCanvas>(_roots: Map<TCanvas, Root>, _getEventPriority?:
       }
 
       // Remove references
-      delete (child as Partial<Instance>).__r3f
+      // only remove primitive reference if it is not freshly created
+      if (!(child as Partial<Instance>).__r3f?.isCreatedBeforeAttach) {
+        delete (child as Partial<Instance>).__r3f
+      }
 
       // Dispose item whenever the reconciler feels like it
       if (shouldDispose && child.dispose && child.type !== 'Scene') {
diff --git a/packages/fiber/tests/core/renderer.test.tsx b/packages/fiber/tests/core/renderer.test.tsx
index 1a9e93be3f..cf271536e8 100644
--- a/packages/fiber/tests/core/renderer.test.tsx
+++ b/packages/fiber/tests/core/renderer.test.tsx
@@ -1015,4 +1015,55 @@ describe('renderer', () => {
     expect(meshDispose).toBeCalledTimes(1)
     expect(primitiveDispose).not.toBeCalled()
   })
+
+  it('should not clear reattached primitive', async () => {
+    const primitive1 = new THREE.Mesh()
+    const primitive2 = new THREE.Mesh()
+    /* This test still fails 
+    await act(async () =>
+      root.render(
+        <group>
+          <primitive object={primitive1} dispose={null} />
+        </group>,
+      ),
+    )
+
+    expect((primitive1 as any).__r3f?.type).toBe('primitive')
+
+    await act(async () =>
+      root.render(
+        <group>
+          <primitive object={primitive2} dispose={null} />
+          <primitive object={primitive1} dispose={null} />
+        </group>,
+      ),
+    )
+    expect((primitive1 as any).__r3f?.type).toBe('primitive')
+    expect((primitive2 as any).__r3f?.type).toBe('primitive')
+    */
+
+    // Initialize a list of primitives
+    await act(async () =>
+      root.render(
+        <group>
+          <primitive key={'a'} object={primitive1} name="p2" dispose={null} />
+        </group>,
+      ),
+    )
+
+    expect((primitive1 as any).__r3f?.type).toBe('primitive')
+
+    // New list of primitives while reusing primitive object
+    await act(async () =>
+      root.render(
+        <group>
+          <primitive key={'b'} object={primitive1} name="p1" dispose={null} />
+          <primitive key={'c'} object={primitive2} name="p2" dispose={null} />
+        </group>,
+      ),
+    )
+
+    expect((primitive1 as any).__r3f?.type).toBe('primitive')
+    expect((primitive2 as any).__r3f?.type).toBe('primitive')
+  })
 })