Testing use-indexed-children with yoga

yoga.tsx

import * as React from 'react'
import * as Yoga from 'yoga-layout-prebuilt'
import { useIndexedChildren, useIndexPath } from 'use-indexed-children'
import useConstant from 'use-constant'

type Node = ReturnType<typeof Yoga.Node.create>

const LayoutContext = React.createContext<{ root: Node; node: Node } | null>(
  null
)

function useLayout({
  width,
  height,
}: {
  width: number | 'fill'
  height: number | 'fill'
}) {
  const node = useConstant(() => Yoga.Node.create())
  const firstRender = React.useRef(true)
  let layoutContext = React.useContext(LayoutContext)

  const setSize = () => {
    if (width === 'fill') {
      node.setFlex(1)
    } else {
      node.setWidth(width)
    }

    if (height === 'fill') {
      node.setFlex(1)
    } else {
      node.setHeight(height)
    }
  }

  if (layoutContext === null) {
    if (firstRender.current) {
      setSize()
    }

    React.useLayoutEffect(() => {
      if (firstRender.current) {
        firstRender.current = false
      }
    }, [])

    layoutContext = { root: node, node }
  } else {
    const { index } = useIndexPath()

    if (firstRender.current) {
      setSize()
      layoutContext.root.insertChild(node, index)
    }

    React.useLayoutEffect(() => {
      if (firstRender.current) {
        firstRender.current = false
      } else {
        layoutContext!.root.insertChild(node, index)
      }

      return () => {
        layoutContext!.root.removeChild(node)
      }
    }, [index])

    layoutContext = { root: layoutContext.root, node }
  }

  React.useEffect(() => {
    setSize()
  }, [width, height])

  return layoutContext
}

function useComputedLayout({
  width,
  height,
}: {
  width: number
  height: number
}) {
  const { root, node } = useLayout({ width, height })

  return React.useMemo(() => {
    root.calculateLayout(
      root.getWidth().value,
      root.getHeight().value,
      Yoga.DIRECTION_LTR
    )

    return node.getComputedLayout()
  }, [node])
}

function Row({
  children,
  columns,
  rows,
}: {
  children: React.ReactNode
  columns: number
  rows: number
}) {
  const root = useLayout({ width: columns, height: rows })

  return (
    <LayoutContext.Provider value={root}>
      {useIndexedChildren(children)}
    </LayoutContext.Provider>
  )
}

function Space() {
  useLayout({ width: 'fill', height: 'fill' })

  return null
}

function Rectangle({ width, height }: { width: number; height: number }) {
  const layout = useComputedLayout({ width, height })

  return (
    <div
      style={{
        width: layout.width,
        height: layout.height,
        transform: `translate(${layout.left}px, ${layout.top}px)`,
        backgroundColor: 'red',
      }}
    />
  )
}

export default function App() {
  return (
    <div style={{ width: 100, height: 100, background: 'pink' }}>
      <Row columns={100} rows={100}>
        <Space />
        <Rectangle width={20} height={20} />
        <Space />
      </Row>
    </div>
  )
}