Added ref forwarding to Flex and Box

This would with various helpers and hooks that use refs to somehow affect the component,
say rotate it or draw a bounding box

src/Box.tsx

@@ -2,82 +2,84 @@ import React, { useLayoutEffect, useRef, useMemo, useState } from 'react'
 import * as THREE from 'three'
 import Yoga from 'yoga-layout-prebuilt'
 import { ReactThreeFiber, useFrame } from '@react-three/fiber'
+import mergeRefs from 'react-merge-refs'
 
 import { setYogaProperties, rmUndefFromObj } from './util'
 import { boxContext, flexContext, SharedBoxContext } from './context'
 import { R3FlexProps } from './props'
 import { useReflow, useContext } from './hooks'
 
-/**
- * Box container for 3D Objects.
- * For containing Boxes use `<Flex />`.
- */
-export function Box({
-  // Non-flex props
-  children,
-  centerAnchor,
-
-  // flex props
-  flexDirection,
-  flexDir,
-  dir,
-
-  alignContent,
-  alignItems,
-  alignSelf,
-  align,
-
-  justifyContent,
-  justify,
-
-  flexBasis,
-  basis,
-  flexGrow,
-  grow,
-
-  flexShrink,
-  shrink,
-
-  flexWrap,
-  wrap,
-
-  margin,
-  m,
-  marginBottom,
-  marginLeft,
-  marginRight,
-  marginTop,
-  mb,
-  ml,
-  mr,
-  mt,
-
-  padding,
-  p,
-  paddingBottom,
-  paddingLeft,
-  paddingRight,
-  paddingTop,
-  pb,
-  pl,
-  pr,
-  pt,
-
-  height,
-  width,
-
-  maxHeight,
-  maxWidth,
-  minHeight,
-  minWidth,
-
-  // other
-  ...props
-}: {
+export type BoxProps = {
   centerAnchor?: boolean
   children: React.ReactNode | ((width: number, height: number, centerAnchor?: boolean) => React.ReactNode)
 } & R3FlexProps &
-  Omit<ReactThreeFiber.Object3DNode<THREE.Group, typeof THREE.Group>, 'children'>) {
+  Omit<ReactThreeFiber.Object3DNode<THREE.Group, typeof THREE.Group>, 'children'>
+
+function BoxImpl(
+  {
+    // Non-flex props
+    children,
+    centerAnchor,
+
+    // flex props
+    flexDirection,
+    flexDir,
+    dir,
+
+    alignContent,
+    alignItems,
+    alignSelf,
+    align,
+
+    justifyContent,
+    justify,
+
+    flexBasis,
+    basis,
+    flexGrow,
+    grow,
+
+    flexShrink,
+    shrink,
+
+    flexWrap,
+    wrap,
+
+    margin,
+    m,
+    marginBottom,
+    marginLeft,
+    marginRight,
+    marginTop,
+    mb,
+    ml,
+    mr,
+    mt,
+
+    padding,
+    p,
+    paddingBottom,
+    paddingLeft,
+    paddingRight,
+    paddingTop,
+    pb,
+    pl,
+    pr,
+    pt,
+
+    height,
+    width,
+
+    maxHeight,
+    maxWidth,
+    minHeight,
+    minWidth,
+
+    // other
+    ...props
+  }: BoxProps,
+  ref: React.Ref<THREE.Group>
+) {
   // must memoize or the object literal will cause every dependent of flexProps to rerender everytime
   const flexProps: R3FlexProps = useMemo(() => {
     const _flexProps = {
@@ -228,10 +230,18 @@ export function Box({
   const sharedBoxContext = useMemo<SharedBoxContext>(() => ({ node, size, centerAnchor }), [node, size, centerAnchor])
 
   return (
-    <group ref={group} {...props}>
+    <group ref={mergeRefs([group, ref])} {...props}>
       <boxContext.Provider value={sharedBoxContext}>
         {typeof children === 'function' ? children(size[0], size[1], centerAnchor) : children}
       </boxContext.Provider>
     </group>
   )
 }
+
+/**
+ * Box container for 3D Objects.
+ * For containing Boxes use `<Flex />`.
+ */
+export const Box = React.forwardRef<THREE.Group, BoxProps>(BoxImpl)
+
+Box.displayName = 'Box'

src/Flex.tsx

@@ -1,7 +1,8 @@
 import React, { useLayoutEffect, useMemo, useCallback, PropsWithChildren, useRef } from 'react'
 import Yoga, { YogaNode } from 'yoga-layout-prebuilt'
-import { Vector3, Group, Box3, Object3D } from 'three'
+import * as THREE from 'three'
 import { useFrame, useThree, ReactThreeFiber } from '@react-three/fiber'
+import mergeRefs from 'react-merge-refs'
 
 import {
   setYogaProperties,
@@ -33,86 +34,86 @@ export type FlexProps = PropsWithChildren<
     centerAnchor?: boolean
   }> &
     R3FlexProps &
-    Omit<ReactThreeFiber.Object3DNode<THREE.Group, typeof Group>, 'children'>
+    Omit<ReactThreeFiber.Object3DNode<THREE.Group, typeof THREE.Group>, 'children'>
 >
 interface BoxesItem {
   node: YogaNode
-  group: Group
+  group: THREE.Group
   flexProps: R3FlexProps
   centerAnchor: boolean
 }
 
-/**
- * Flex container. Can contain Boxes
- */
-export function Flex({
-  // Non flex props
-  size = [1, 1, 1],
-  yogaDirection = 'ltr',
-  plane = 'xy',
-  children,
-  scaleFactor = 100,
-  onReflow,
-  disableSizeRecalc,
-  centerAnchor: rootCenterAnchor,
-
-  // flex props
-
-  flexDirection,
-  flexDir,
-  dir,
-
-  alignContent,
-  alignItems,
-  alignSelf,
-  align,
-
-  justifyContent,
-  justify,
-
-  flexBasis,
-  basis,
-  flexGrow,
-  grow,
-  flexShrink,
-  shrink,
-
-  flexWrap,
-  wrap,
-
-  margin,
-  m,
-  marginBottom,
-  marginLeft,
-  marginRight,
-  marginTop,
-  mb,
-  ml,
-  mr,
-  mt,
-
-  padding,
-  p,
-  paddingBottom,
-  paddingLeft,
-  paddingRight,
-  paddingTop,
-  pb,
-  pl,
-  pr,
-  pt,
-
-  height,
-  width,
-
-  maxHeight,
-  maxWidth,
-  minHeight,
-  minWidth,
-
-  // other
-  ...props
-}: FlexProps) {
+function FlexImpl(
+  {
+    // Non flex props
+    size = [1, 1, 1],
+    yogaDirection = 'ltr',
+    plane = 'xy',
+    children,
+    scaleFactor = 100,
+    onReflow,
+    disableSizeRecalc,
+    centerAnchor: rootCenterAnchor,
+
+    // flex props
+
+    flexDirection,
+    flexDir,
+    dir,
+
+    alignContent,
+    alignItems,
+    alignSelf,
+    align,
+
+    justifyContent,
+    justify,
+
+    flexBasis,
+    basis,
+    flexGrow,
+    grow,
+    flexShrink,
+    shrink,
+
+    flexWrap,
+    wrap,
+
+    margin,
+    m,
+    marginBottom,
+    marginLeft,
+    marginRight,
+    marginTop,
+    mb,
+    ml,
+    mr,
+    mt,
+
+    padding,
+    p,
+    paddingBottom,
+    paddingLeft,
+    paddingRight,
+    paddingTop,
+    pb,
+    pl,
+    pr,
+    pt,
+
+    height,
+    width,
+
+    maxHeight,
+    maxWidth,
+    minHeight,
+    minWidth,
+
+    // other
+    ...props
+  }: FlexProps,
+  ref: React.Ref<THREE.Group>
+) {
   // must memoize or the object literal will cause every dependent of flexProps to rerender everytime
   const flexProps: R3FlexProps = useMemo(() => {
     const _flexProps = {
@@ -217,12 +218,12 @@ export function Flex({
     wrap,
   ])
 
-  const rootGroup = useRef<Group>()
+  const rootGroup = useRef<THREE.Group>()
 
   // Keeps track of the yoga nodes of the children and the related wrapper groups
   const boxesRef = useRef<BoxesItem[]>([])
   const registerBox = useCallback(
-    (node: YogaNode, group: Group, flexProps: R3FlexProps, centerAnchor: boolean = false) => {
+    (node: YogaNode, group: THREE.Group, flexProps: R3FlexProps, centerAnchor: boolean = false) => {
       const i = boxesRef.current.findIndex((b) => b.node === node)
       if (i !== -1) {
         boxesRef.current.splice(i, 1)
@@ -258,8 +259,8 @@ export function Flex({
   }, [children, flexProps, requestReflow])
 
   // Common variables for reflow
-  const boundingBox = useMemo(() => new Box3(), [])
-  const vec = useMemo(() => new Vector3(), [])
+  const boundingBox = useMemo(() => new THREE.Box3(), [])
+  const vec = useMemo(() => new THREE.Vector3(), [])
   const mainAxis = plane[0] as Axis
   const crossAxis = plane[1] as Axis
   const depthAxis = getDepthAxis(plane)
@@ -356,10 +357,17 @@ export function Flex({
   })
 
   return (
-    <group ref={rootGroup} {...props}>
+    <group ref={mergeRefs([rootGroup, ref])} {...props}>
       <flexContext.Provider value={sharedFlexContext}>
         <boxContext.Provider value={sharedBoxContext}>{children}</boxContext.Provider>
       </flexContext.Provider>
     </group>
   )
 }
+
+/**
+ * Flex container. Can contain Boxes
+ */
+export const Flex = React.forwardRef<THREE.Group, FlexProps>(FlexImpl)
+
+Flex.displayName = 'Flex'