From 81fe4479ed3108ecc58cc70f7c2a4dd554b47423 Mon Sep 17 00:00:00 2001
From: Andrew Prifer <andrew.prifer@gmail.com>
Date: Tue, 28 Jun 2022 15:59:03 +0200
Subject: [PATCH] Add r3f custom components (<e.custom>)

---
 .../src/shared/custom-r3f-component/App.tsx   | 38 ++++++++++++
 .../src/shared/custom-r3f-component/index.tsx | 10 ++++
 packages/r3f/src/main/editable.tsx            | 60 +++++++++++++++++--
 3 files changed, 104 insertions(+), 4 deletions(-)
 create mode 100644 packages/playground/src/shared/custom-r3f-component/App.tsx
 create mode 100644 packages/playground/src/shared/custom-r3f-component/index.tsx

diff --git a/packages/playground/src/shared/custom-r3f-component/App.tsx b/packages/playground/src/shared/custom-r3f-component/App.tsx
new file mode 100644
index 000000000..3c810b3ff
--- /dev/null
+++ b/packages/playground/src/shared/custom-r3f-component/App.tsx
@@ -0,0 +1,38 @@
+import {editable as e, SheetProvider} from '@theatre/r3f'
+import {getProject} from '@theatre/core'
+import React from 'react'
+import {Canvas} from '@react-three/fiber'
+
+function App() {
+  return (
+    <div
+      onClick={() => {
+        // return setBgIndex((bgIndex) => (bgIndex + 1) % bgs.length)
+      }}
+      style={{
+        height: '100vh',
+      }}
+    >
+      <Canvas
+        dpr={[1.5, 2]}
+        linear
+        gl={{preserveDrawingBuffer: true}}
+        frameloop="demand"
+      >
+        <SheetProvider sheet={getProject('Space').sheet('Scene')}>
+          <ambientLight intensity={0.75} />
+          <e.custom
+            uniqueName="points"
+            customComponent="points"
+            editableType="mesh"
+          >
+            <torusKnotGeometry args={[1, 0.3, 128, 64]} />
+            <meshNormalMaterial />
+          </e.custom>
+        </SheetProvider>
+      </Canvas>
+    </div>
+  )
+}
+
+export default App
diff --git a/packages/playground/src/shared/custom-r3f-component/index.tsx b/packages/playground/src/shared/custom-r3f-component/index.tsx
new file mode 100644
index 000000000..550726371
--- /dev/null
+++ b/packages/playground/src/shared/custom-r3f-component/index.tsx
@@ -0,0 +1,10 @@
+import React from 'react'
+import ReactDOM from 'react-dom'
+import App from './App'
+import studio from '@theatre/studio'
+import extension from '@theatre/r3f/dist/extension'
+
+studio.extend(extension)
+studio.initialize()
+
+ReactDOM.render(<App />, document.getElementById('root'))
diff --git a/packages/r3f/src/main/editable.tsx b/packages/r3f/src/main/editable.tsx
index 34969e9c6..6af81d36e 100644
--- a/packages/r3f/src/main/editable.tsx
+++ b/packages/r3f/src/main/editable.tsx
@@ -10,16 +10,22 @@ import type {EditableFactoryConfig} from './editableFactoryConfigUtils'
 import {makeStoreKey} from './utils'
 import type {$FixMe} from '../types'
 
-const createEditable = <Keys extends keyof JSX.IntrinsicElements>(
+const createEditable = <
+  // This should really be extends keyof JSX.IntrinsicElements, but that results in a union
+  // that is "too complex to represent".
+  Keys extends string,
+>(
   config: EditableFactoryConfig,
 ) => {
   const editable = <
-    T extends ComponentType<any> | Keys | 'primitive',
+    T extends ComponentType<any> | Keys | 'primitive' | 'custom',
     U extends T extends Keys ? T : Keys,
   >(
     Component: T,
-    type: T extends 'primitive' ? null : U,
+    type: T extends 'primitive' | 'custom' ? null : U,
   ) => {
+    // Since we can't make T conform to keyof JSX.IntrinsicElements (see above) we have to @ts-ignore here.
+    // @ts-ignore
     type Props = Omit<ComponentProps<T>, 'visible'> & {
       uniqueName: string
       visible?: boolean | 'editor'
@@ -30,6 +36,14 @@ const createEditable = <Keys extends keyof JSX.IntrinsicElements>(
             editableType: U
           }
         : {}) &
+      (T extends 'custom'
+        ? {
+            customComponent: keyof JSX.IntrinsicElements | ComponentType<any>
+            editableType: U
+          }
+        : {}) &
+      // Since we can't make T conform to keyof JSX.IntrinsicElements (see above) we have to @ts-ignore here.
+      // @ts-ignore
       RefAttributes<JSX.IntrinsicElements[U]>
 
     return forwardRef(
@@ -38,6 +52,7 @@ const createEditable = <Keys extends keyof JSX.IntrinsicElements>(
           uniqueName,
           visible,
           editableType,
+          customComponent,
           additionalProps,
           objRef,
           ...props
@@ -45,7 +60,10 @@ const createEditable = <Keys extends keyof JSX.IntrinsicElements>(
         ref,
       ) => {
         const actualType = type ?? editableType
+        const ActualComponent = customComponent ?? Component
 
+        // Since we can't make T conform to keyof JSX.IntrinsicElements (see above) we have to @ts-ignore here.
+        // @ts-ignore
         const objectRef = useRef<JSX.IntrinsicElements[U]>()
 
         const sheet = useCurrentSheet()!
@@ -134,7 +152,7 @@ const createEditable = <Keys extends keyof JSX.IntrinsicElements>(
 
         return (
           // @ts-ignore
-          <Component
+          <ActualComponent
             ref={mergeRefs([objectRef, ref])}
             {...props}
             visible={visible !== 'editor' && visible}
@@ -157,17 +175,51 @@ const createEditable = <Keys extends keyof JSX.IntrinsicElements>(
       ]),
     ),
     primitive: editable('primitive', null),
+    custom: editable('custom', null),
   } as unknown as {
     [Property in Keys]: React.ForwardRefExoticComponent<
       React.PropsWithoutRef<
+        // Since we can't make T conform to keyof JSX.IntrinsicElements (see above) we have to @ts-ignore here.
+        // @ts-ignore
         Omit<JSX.IntrinsicElements[Property], 'visible'> & {
           uniqueName: string
           visible?: boolean | 'editor'
           additionalProps?: $FixMe
           objRef?: $FixMe
+          // Since we can't make T conform to keyof JSX.IntrinsicElements (see above) we have to @ts-ignore here.
+          // @ts-ignore
         } & React.RefAttributes<JSX.IntrinsicElements[Property]>
       >
     >
+  } & {
+    primitive: React.ForwardRefExoticComponent<
+      React.PropsWithoutRef<
+        {
+          object: any
+          uniqueName: string
+          visible?: boolean | 'editor'
+          additionalProps?: $FixMe
+          objRef?: $FixMe
+          editableType: keyof JSX.IntrinsicElements
+        } & React.RefAttributes<JSX.IntrinsicElements['primitive']>
+      > & {
+        // Have to reproduce the primitive component's props here because we need to
+        // lift this index type here to the outside to make auto-complete work
+        [props: string]: any
+      }
+    >
+    custom: React.ForwardRefExoticComponent<
+      React.PropsWithoutRef<{
+        uniqueName: string
+        visible?: boolean | 'editor'
+        additionalProps?: $FixMe
+        objRef?: $FixMe
+        editableType: keyof JSX.IntrinsicElements
+        customComponent: keyof JSX.IntrinsicElements | ComponentType<any>
+      }> & {
+        [props: string]: any
+      }
+    >
   }
 
   return Object.assign(editable, extensions)