From aa0f742e1a52d4f6565cfcb9503f4729841af475 Mon Sep 17 00:00:00 2001
From: Cody Bennett <hi@codyb.co>
Date: Mon, 24 Feb 2025 11:48:21 -0600
Subject: [PATCH] fix(Merged): restores recursive implementation (#2366)

---
 .storybook/stories/Merged.stories.tsx | 30 ++++++++++
 src/core/Instances.tsx                | 85 ++++++++++++++++++---------
 2 files changed, 87 insertions(+), 28 deletions(-)
 create mode 100644 .storybook/stories/Merged.stories.tsx

diff --git a/.storybook/stories/Merged.stories.tsx b/.storybook/stories/Merged.stories.tsx
new file mode 100644
index 000000000..a554f8228
--- /dev/null
+++ b/.storybook/stories/Merged.stories.tsx
@@ -0,0 +1,30 @@
+import * as React from 'react'
+import { Meta, StoryObj } from '@storybook/react'
+
+import { Setup } from '../Setup'
+
+import { useGLTF, Merged, Instance } from '../../src'
+
+export default {
+  title: 'Performance/Merged',
+  component: Merged,
+  decorators: [
+    (Story) => (
+      <Setup>
+        <Story />
+      </Setup>
+    ),
+  ],
+} satisfies Meta<typeof Merged>
+
+type Story = StoryObj<typeof Merged>
+
+function Scene() {
+  const { nodes } = useGLTF('suzanne.glb', true)
+  return <Merged meshes={nodes}>{({ Suzanne }) => <Suzanne />}</Merged>
+}
+
+export const DefaultStory = {
+  render: (args) => <Scene {...args} />,
+  name: 'Default',
+} satisfies Story
diff --git a/src/core/Instances.tsx b/src/core/Instances.tsx
index 5c0a19510..ab9e01849 100644
--- a/src/core/Instances.tsx
+++ b/src/core/Instances.tsx
@@ -218,39 +218,68 @@ export const Instances: ForwardRefComponent<InstancesProps, THREE.InstancedMesh>
   )
 })
 
-export interface MergedProps extends InstancesProps {
-  meshes: THREE.Mesh[]
-  children: React.ReactNode
+export interface MergedProps extends Omit<InstancesProps, 'children'> {
+  meshes: THREE.Mesh[] | Record<string, THREE.Object3D>
+  children: (
+    ...instances: [React.FC<InstanceProps> & Record<string, React.FC<InstanceProps>>, ...React.FC<InstanceProps>[]]
+  ) => React.ReactNode
 }
 
-export const Merged: ForwardRefComponent<any, THREE.Group> = React.forwardRef<THREE.Group, any>(function Merged(
-  { meshes, children, ...rest },
-  ref
-) {
-  const instances: React.FC[] = []
-
-  if (Array.isArray(meshes)) {
-    for (const mesh of meshes) {
-      if (mesh?.isMesh) {
-        instances.push((props) => (
-          <Instances key={mesh.geometry.uuid} geometry={mesh.geometry} material={mesh.material} {...rest} {...props} />
-        ))
-      }
-    }
-  } else if (meshes != null && typeof meshes === 'object') {
-    for (const key in meshes) {
-      const mesh = meshes[key]
-      if (mesh?.isMesh) {
-        instances.push((props) => (
-          <Instances key={mesh.geometry.uuid} geometry={mesh.geometry} material={mesh.material} {...rest} {...props} />
-        ))
-      }
-    }
-  }
+// TODO: make this non-recursive and type-safe
+export const Merged: ForwardRefComponent<MergedProps, THREE.Group> = /* @__PURE__ */ React.forwardRef<
+  THREE.Group,
+  MergedProps
+>(function Merged({ meshes, children, ...props }, ref) {
+  const isArray = Array.isArray(meshes)
+  // Filter out meshes from collections, which may contain non-meshes
+  // @ts-expect-error
+  if (!isArray) for (const key of Object.keys(meshes)) if (!meshes[key].isMesh) delete meshes[key]
+
+  const render = (args) =>
+    isArray
+      ? // @ts-expect-error
+        children(...args)
+      : children(
+          // @ts-expect-error
+          Object.keys(meshes)
+            // @ts-expect-error
+            .filter((key) => meshes[key].isMesh)
+            .reduce((acc, key, i) => ({ ...acc, [key]: args[i] }), {})
+        )
+
+  // @ts-expect-error
+  const components = (isArray ? meshes : Object.values(meshes)).map(({ geometry, material }) => (
+    <Instances key={geometry.uuid} geometry={geometry} material={material} {...props} />
+  ))
 
-  return <group ref={ref}>{children(instances)}</group>
+  return <group ref={ref}>{renderRecursive(render, components)}</group>
 })
 
+// https://github.com/jamesplease/react-composer
+function renderRecursive(
+  render: Function,
+  components: Array<React.ReactElement<{ children: any }> | Function>,
+  results: unknown[] = []
+): React.ReactElement {
+  // Once components is exhausted, we can render out the results array.
+  if (!components[0]) {
+    return render(results)
+  }
+
+  // Continue recursion for remaining items.
+  // results.concat([value]) ensures [...results, value] instead of [...results, ...value]
+  function nextRender(value) {
+    return renderRecursive(render, components.slice(1), results.concat([value]))
+  }
+
+  // Each props.components entry is either an element or function [element factory]
+  return typeof components[0] === 'function'
+    ? // When it is a function, produce an element by invoking it with "render component values".
+      components[0]({ results, render: nextRender })
+    : // When it is an element, enhance the element's props with the render prop.
+      React.cloneElement(components[0], { children: nextRender })
+}
+
 /** Idea and implementation for global instances and instanced attributes by
 /*  Matias Gonzalez Fernandez https://x.com/matiNotFound
 /*  and Paul Henschel https://x.com/0xca0a