From 41e1cd948bf82d602a274358a7c520374c9b5f08 Mon Sep 17 00:00:00 2001
From: Michael Dougall <6801309+itsdouges@users.noreply.github.com>
Date: Fri, 7 Jun 2024 08:12:00 +1000
Subject: [PATCH] chore: make root optional in container spike

---
 packages/react/src/container.tsx | 20 ++++++++++++++++++--
 packages/react/src/context.tsx   |  5 +++++
 2 files changed, 23 insertions(+), 2 deletions(-)

diff --git a/packages/react/src/container.tsx b/packages/react/src/container.tsx
index a57f3d18..32838e52 100644
--- a/packages/react/src/container.tsx
+++ b/packages/react/src/container.tsx
@@ -1,7 +1,7 @@
 import { EventHandlers } from '@react-three/fiber/dist/declarations/src/core/events'
 import { forwardRef, ReactNode, RefAttributes, useEffect, useMemo, useRef } from 'react'
 import { Object3D } from 'three'
-import { ParentProvider, useParent } from './context.js'
+import { ParentProvider, useParent, useParentOptional } from './context.js'
 import { AddHandlers, usePropertySignals } from './utilts.js'
 import {
   ContainerProperties as BaseContainerProperties,
@@ -11,6 +11,7 @@ import {
   initialize,
 } from '@pmndrs/uikit/internals'
 import { ComponentInternals, useComponentInternals } from './ref.js'
+import { Root } from './root.js'
 
 export type ContainerProperties = {
   name?: string
@@ -18,7 +19,7 @@ export type ContainerProperties = {
 } & BaseContainerProperties &
   EventHandlers
 
-export const Container: (
+const ContainerInner: (
   props: ContainerProperties & RefAttributes<ComponentInternals<BaseContainerProperties & EventHandlers>>,
 ) => ReactNode = forwardRef((properties, ref) => {
   const parent = useParent()
@@ -57,3 +58,18 @@ export const Container: (
     </AddHandlers>
   )
 })
+
+export const Container: (
+  props: ContainerProperties & RefAttributes<ComponentInternals<BaseContainerProperties & EventHandlers>>,
+) => ReactNode = forwardRef((properties, ref) => {
+  const parent = useParentOptional()
+  if (parent) {
+    return <ContainerInner {...properties} ref={ref} />
+  }
+
+  return (
+    <Root>
+      <ContainerInner {...properties} ref={ref} />
+    </Root>
+  )
+})
diff --git a/packages/react/src/context.tsx b/packages/react/src/context.tsx
index 347f4a07..55c6e278 100644
--- a/packages/react/src/context.tsx
+++ b/packages/react/src/context.tsx
@@ -11,4 +11,9 @@ export function useParent(): ParentContext {
   return parent
 }
 
+export function useParentOptional(): ParentContext {
+  const parent = useContext(ParentContext)
+  return parent
+}
+
 export const ParentProvider = ParentContext.Provider