fix(mosaic): forward refs from every part that renders through useRender

.changeset/happy-waves-bake.md

@@ -0,0 +1,2 @@
+---
+---

packages/mosaic/docs/headless.md

@@ -42,6 +42,9 @@ Every part that renders a DOM element calls **`useRender`** (from
 Copy `src/primitives/select/select-trigger.tsx`: it is the reference part. Its
 params are documented on `UseRenderParams*` in `utils/use-render.tsx`.
 
+- **Forward the ref.** Every part that renders through `useRender` is a
+  `React.forwardRef` and passes its `ref` to `useRender`;
+  `mosaic/forward-ref-parts` enforces this.
 - **`ownProps satisfies DefaultProps<Tag>`** for authored defaults, not `as`.
 - **`mergeProps(internal, consumer)` — internal first, consumer second.** Event
   handlers chain (internal fires, then consumer), `style` shallow-merges,

packages/mosaic/eslint/config.mjs

@@ -4,7 +4,7 @@ import vitest from '@vitest/eslint-plugin';
 import boundaries from 'eslint-plugin-boundaries';
 import noUnneededEffects from 'eslint-plugin-react-you-might-not-need-an-effect';
 
-import { noDynamicMessageCatalogs, noModelReactState } from './rules.mjs';
+import { forwardRefParts, noDynamicMessageCatalogs, noModelReactState } from './rules.mjs';
 
 const mosaicRoot = path.resolve(import.meta.dirname, '..');
 
@@ -29,6 +29,7 @@ export default [
         rules: {
           'no-model-react-state': noModelReactState,
           'no-dynamic-message-catalogs': noDynamicMessageCatalogs,
+          'forward-ref-parts': forwardRefParts,
         },
       },
     },
@@ -54,6 +55,14 @@ export default [
       'mosaic/no-model-react-state': 'error',
     },
   },
+  {
+    name: 'packages/mosaic - parts',
+    files: ['src/primitives/**/*.tsx', 'src/components/**/*.tsx'],
+    ignores: ['src/**/__tests__/**', 'src/**/*.test.{ts,tsx}'],
+    rules: {
+      'mosaic/forward-ref-parts': 'error',
+    },
+  },
   {
     name: 'packages/mosaic - effects',
     files: ['src/features/**/*.{ts,tsx}'],

packages/mosaic/eslint/import-guards.test.mjs

@@ -27,14 +27,18 @@ const flowFixtures = [
   'sample.messages.ts',
 ].map(file => `${featureDir}/${file}`);
 const testFixture = 'src/hooks/__tests__/sample.test.ts';
+const partFixture = 'src/primitives/sample-part.tsx';
+const partTestFixture = 'src/primitives/__tests__/sample-part.test.tsx';
 const rootEslint = new ESLint({ cwd: root });
 const eslint = new ESLint({
   cwd: mosaicRoot,
   overrideConfig: {
     languageOptions: {
       parserOptions: {
         projectService: {
-          allowDefaultProject: [...flowFixtures, ...layerFixtures, testFixture].map(file => `packages/mosaic/${file}`),
+          allowDefaultProject: [...flowFixtures, ...layerFixtures, testFixture, partFixture, partTestFixture].map(
+            file => `packages/mosaic/${file}`,
+          ),
           maximumDefaultProjectFileMatchCount_THIS_WILL_SLOW_DOWN_LINTING: 32,
         },
       },
@@ -345,6 +349,50 @@ describe('Mosaic layers', { timeout: 30_000 }, () => {
   });
 });
 
+describe('Mosaic parts', { timeout: 30_000 }, () => {
+  async function forwardRefMessages(file, source) {
+    const messages = await lintAt(
+      file,
+      ["import React from 'react';", "import { useRender } from './utils';", source].join('\n'),
+    );
+    expect(messages.find(message => message.fatal)).toBeUndefined();
+    return messages.filter(message => message.ruleId === 'mosaic/forward-ref-parts').map(message => message.message);
+  }
+
+  test.each([
+    "export function Part(props: { render?: never }) { return useRender({ defaultTagName: 'div', props }); }",
+    "export const Part = (props: { render?: never }) => useRender({ defaultTagName: 'div', props });",
+    "export function Part(props: { render?: never }) { return <Composite render={() => useRender({ defaultTagName: 'div', props })} />; }",
+    "export const Part = React.forwardRef<HTMLDivElement, object>(function Part(props) { return useRender({ defaultTagName: 'div', props }); });",
+    "export const Part = React.forwardRef<HTMLDivElement, object>(function Part(props, ref) { const internalRef = React.useRef(null); return useRender({ defaultTagName: 'div', ref: internalRef, props }); });",
+    "export const Part = React.forwardRef<HTMLDivElement, object>(function Part(props, ref) { const internalRef = React.useRef(null); return useRender({ defaultTagName: 'div', ref: [internalRef], props }); });",
+    "export const Part = React.forwardRef<HTMLDivElement, object>(function Part(props, ref) { return useRender({ defaultTagName: 'div', ref: null, props }); });",
+  ])('rejects a part that drops the ref: %s', async source => {
+    expect(await forwardRefMessages(partFixture, source)).toEqual([
+      'Parts that render through useRender must use React.forwardRef and pass the ref to useRender (see primitives/select/select-trigger.tsx).',
+    ]);
+  });
+
+  test.each([
+    "export const Part = React.forwardRef<HTMLButtonElement, object>(function Part(props, ref) { return useRender({ defaultTagName: 'button', ref: [null, ref], props }); });",
+    "export const Part = React.forwardRef<HTMLDivElement, object>(function Part(props, ref) { return useRender({ defaultTagName: 'div', ref, props }); });",
+    "export const Part = React.forwardRef<HTMLDivElement, object>(function Part(props, ref) { return useRender({ defaultTagName: 'div', ref: [null, ref ?? null], props }); });",
+    "export const Part = React.forwardRef<HTMLDivElement, object>(function Part(props: { open?: boolean }, ref) { const finalRef = props.open ? ref : [null, ref]; return useRender({ defaultTagName: 'div', ref: finalRef, props }); });",
+    'export function Provider(props: { children: React.ReactNode }) { return props.children; }',
+  ])('allows parts that forward the ref or render nothing: %s', async source => {
+    expect(await forwardRefMessages(partFixture, source)).toEqual([]);
+  });
+
+  test('ignores test files', async () => {
+    expect(
+      await forwardRefMessages(
+        partTestFixture,
+        "export function Part(props: { render?: never }) { return useRender({ defaultTagName: 'div', props }); }",
+      ),
+    ).toEqual([]);
+  });
+});
+
 describe('Mosaic tests', { timeout: 30_000 }, () => {
   async function vitestRules(source) {
     const messages = await lintAt(testFixture, ["import { describe, expect, test } from 'vitest';", source].join('\n'));
@@ -450,6 +498,7 @@ describe('Mosaic package configuration', { timeout: 30_000 }, () => {
       'boundaries/dependencies',
       'mosaic/no-model-react-state',
       'mosaic/no-dynamic-message-catalogs',
+      'mosaic/forward-ref-parts',
       '@typescript-eslint/no-import-type-side-effects',
       ...Object.keys(noUnneededEffects.configs.strict.rules),
     ]) {

packages/mosaic/eslint/rules.mjs

@@ -77,3 +77,194 @@ export const noDynamicMessageCatalogs = {
     };
   },
 };
+
+function isForwardRefCall(node) {
+  if (node?.type !== 'CallExpression') {
+    return false;
+  }
+  const callee = node.callee;
+  if (callee.type === 'Identifier') {
+    return callee.name === 'forwardRef';
+  }
+  return (
+    callee.type === 'MemberExpression' &&
+    callee.object.type === 'Identifier' &&
+    callee.object.name === 'React' &&
+    propertyName(callee.property, callee.computed) === 'forwardRef'
+  );
+}
+
+function unwrapExpression(node) {
+  let current = node;
+  while (current && ['TSAsExpression', 'TSSatisfiesExpression', 'TSNonNullExpression'].includes(current.type)) {
+    current = current.expression;
+  }
+  return current;
+}
+
+function isFunction(node) {
+  return ['FunctionDeclaration', 'FunctionExpression', 'ArrowFunctionExpression'].includes(node?.type);
+}
+
+function useRenderCalls(fn, visitorKeys) {
+  const calls = [];
+  const visit = node => {
+    if (!node || typeof node.type !== 'string') {
+      return;
+    }
+    if (node.type === 'CallExpression' && node.callee.type === 'Identifier' && node.callee.name === 'useRender') {
+      calls.push(node);
+    }
+    for (const key of visitorKeys[node.type] ?? []) {
+      const child = node[key];
+      if (Array.isArray(child)) {
+        child.forEach(visit);
+      } else {
+        visit(child);
+      }
+    }
+  };
+  visit(fn);
+  return calls;
+}
+
+function includesRef(value, refName, scope) {
+  const node = unwrapExpression(value);
+  if (node?.type === 'Identifier') {
+    if (node.name === refName) {
+      return true;
+    }
+    const variable = findVariable(scope, node.name);
+    const definition = variable?.defs.length === 1 ? variable.defs[0] : undefined;
+    return definition?.node.type === 'VariableDeclarator' && includesRef(definition.node.init, refName, scope);
+  }
+  if (node?.type === 'LogicalExpression') {
+    return includesRef(node.left, refName, scope);
+  }
+  if (node?.type === 'ConditionalExpression') {
+    return includesRef(node.consequent, refName, scope) && includesRef(node.alternate, refName, scope);
+  }
+  return node?.type === 'ArrayExpression' && node.elements.some(element => includesRef(element, refName, scope));
+}
+
+function findVariable(scope, name) {
+  for (let current = scope; current; current = current.upper) {
+    const variable = current.set.get(name);
+    if (variable) {
+      return variable;
+    }
+  }
+  return undefined;
+}
+
+function passesRef(call, refName, scope) {
+  const [options] = call.arguments;
+  if (options?.type !== 'ObjectExpression') {
+    return true;
+  }
+  return options.properties.some(
+    property =>
+      property.type === 'Property' &&
+      propertyName(property.key, property.computed) === 'ref' &&
+      includesRef(property.value, refName, scope),
+  );
+}
+
+export const forwardRefParts = {
+  meta: {
+    type: 'problem',
+    schema: [],
+    messages: {
+      forwardRef:
+        'Parts that render through useRender must use React.forwardRef and pass the ref to useRender (see primitives/select/select-trigger.tsx).',
+    },
+  },
+  create(context) {
+    const isComponentName = name => /^[A-Z]/.test(name);
+
+    const check = (reportNode, init) => {
+      const expression = unwrapExpression(init);
+      const wrapped = isForwardRefCall(expression);
+      const fn = wrapped ? unwrapExpression(expression.arguments[0]) : expression;
+      if (!isFunction(fn)) {
+        return;
+      }
+      const calls = useRenderCalls(fn, context.sourceCode.visitorKeys);
+      if (calls.length === 0) {
+        return;
+      }
+      const refName = fn.params[1]?.type === 'Identifier' ? fn.params[1].name : undefined;
+      if (!wrapped || calls.some(call => !passesRef(call, refName, context.sourceCode.getScope(call)))) {
+        context.report({ node: reportNode, messageId: 'forwardRef' });
+      }
+    };
+
+    const checkDeclaration = declaration => {
+      if (declaration?.type === 'FunctionDeclaration' && declaration.id && isComponentName(declaration.id.name)) {
+        check(declaration.id, declaration);
+      }
+      if (declaration?.type === 'VariableDeclaration') {
+        for (const declarator of declaration.declarations) {
+          if (declarator.id.type === 'Identifier' && isComponentName(declarator.id.name) && declarator.init) {
+            check(declarator.id, declarator.init);
+          }
+        }
+      }
+    };
+
+    return {
+      Program(program) {
+        const localDeclarations = new Map();
+        for (const statement of program.body) {
+          const declaration = statement.type === 'ExportNamedDeclaration' ? statement.declaration : statement;
+          if (declaration?.type === 'FunctionDeclaration' && declaration.id) {
+            localDeclarations.set(declaration.id.name, declaration);
+          }
+          if (declaration?.type === 'VariableDeclaration') {
+            for (const declarator of declaration.declarations) {
+              if (declarator.id.type === 'Identifier') {
+                localDeclarations.set(declarator.id.name, { type: 'VariableDeclaration', declarations: [declarator] });
+              }
+            }
+          }
+        }
+
+        const checked = new Set();
+        const checkOnce = declaration => {
+          if (declaration && !checked.has(declaration)) {
+            checked.add(declaration);
+            checkDeclaration(declaration);
+          }
+        };
+
+        for (const statement of program.body) {
+          if (statement.type === 'ExportNamedDeclaration') {
+            if (statement.declaration?.type === 'FunctionDeclaration') {
+              checkOnce(localDeclarations.get(statement.declaration.id?.name));
+            } else if (statement.declaration?.type === 'VariableDeclaration') {
+              for (const declarator of statement.declaration.declarations) {
+                if (declarator.id.type === 'Identifier') {
+                  checkOnce(localDeclarations.get(declarator.id.name));
+                }
+              }
+            } else if (!statement.source) {
+              for (const specifier of statement.specifiers) {
+                if (specifier.local.type === 'Identifier' && isComponentName(specifier.exported.name ?? '')) {
+                  checkOnce(localDeclarations.get(specifier.local.name));
+                }
+              }
+            }
+          }
+          if (statement.type === 'ExportDefaultDeclaration') {
+            const declaration = statement.declaration;
+            if (declaration.type === 'Identifier') {
+              checkOnce(localDeclarations.get(declaration.name));
+            } else if (declaration.type === 'FunctionDeclaration' && declaration.id) {
+              checkOnce(localDeclarations.get(declaration.id.name));
+            }
+          }
+        }
+      },
+    };
+  },
+};

packages/mosaic/src/components/menu/menu.test.tsx

@@ -371,4 +371,19 @@ describe('Mosaic Menu', () => {
     );
     expect(ref.current).toBe(screen.getByRole('button'));
   });
+
+  it('forwards a ref to the separator', () => {
+    const ref = React.createRef<HTMLDivElement>();
+    render(
+      <Menu.Root defaultOpen>
+        <Menu.Trigger />
+        <Menu.Popup>
+          <Menu.Item>Cut</Menu.Item>
+          <Menu.Separator ref={ref} />
+          <Menu.Item>Copy</Menu.Item>
+        </Menu.Popup>
+      </Menu.Root>,
+    );
+    expect(ref.current).toBe(screen.getByRole('separator'));
+  });
 });

packages/mosaic/src/components/menu/menu.tsx

@@ -226,13 +226,17 @@ export const MenuItem = React.forwardRef<HTMLButtonElement, MenuItemProps>(funct
 export type MenuSeparatorProps = Omit<PrimitiveMenuSeparatorProps, 'className' | 'style'> & MosaicStyleProps;
 
 /** A full-bleed divider between groups of items. */
-export function MenuSeparator({ xstyle, ...rest }: MenuSeparatorProps): React.ReactElement {
+export const MenuSeparator = React.forwardRef<HTMLDivElement, MenuSeparatorProps>(function MosaicMenuSeparator(
+  { xstyle, ...rest },
+  ref,
+) {
   return (
     <Primitive.Separator
+      ref={ref}
       {...mergeStyleProps(themeProps('menu-separator'), stylex.props(reset.base, slots.separator.base, xstyle), rest)}
     />
   );
-}
+});
 
 export const Menu = {
   // Renders no element of its own; re-exported unchanged so the state provider sits

packages/mosaic/src/primitives/accordion/accordion-header.tsx

@@ -1,17 +1,22 @@
 'use client';
 
+import React from 'react';
+
 import { type ComponentProps, mergeProps, useRender } from '../utils';
 
 export type AccordionHeaderProps = ComponentProps<'h3'>;
 
-export function AccordionHeader(props: AccordionHeaderProps) {
-  const { render, ...otherProps } = props;
+export const AccordionHeader = React.forwardRef<HTMLHeadingElement, AccordionHeaderProps>(
+  function AccordionHeader(props, ref) {
+    const { render, ...otherProps } = props;
 
-  const defaultProps = {};
+    const defaultProps = {};
 
-  return useRender({
-    defaultTagName: 'h3',
-    render,
-    props: mergeProps<'h3'>(defaultProps, otherProps),
-  });
-}
+    return useRender({
+      defaultTagName: 'h3',
+      ref,
+      render,
+      props: mergeProps<'h3'>(defaultProps, otherProps),
+    });
+  },
+);