feat: RuleEnhancer & unwrapped POJOs rule

packages/eslint-plugin/src/configs.ts

@@ -1,16 +1,20 @@
 import type { TSESLint } from '@typescript-eslint/utils';
 import { integerDivision } from './rules/integerDivision.ts';
+import { unwrappedPojos } from './rules/unwrappedPojos.ts';
 
 export const rules = {
   'integer-division': integerDivision,
+  'unwrapped-pojo': unwrappedPojos,
 } as const;
 
 type Rules = Record<`typegpu/${keyof typeof rules}`, TSESLint.FlatConfig.RuleEntry>;
 
 export const recommendedRules: Rules = {
   'typegpu/integer-division': 'warn',
+  'typegpu/unwrapped-pojo': 'warn',
 };
 
 export const allRules: Rules = {
   'typegpu/integer-division': 'error',
+  'typegpu/unwrapped-pojo': 'error',
 };

packages/eslint-plugin/src/enhanceRule.ts

@@ -0,0 +1,81 @@
+import type { RuleContext, RuleListener } from '@typescript-eslint/utils/ts-eslint';
+
+export type RuleEnhancer<TState> = (context: RuleContext<string, unknown[]>) => {
+  visitors: RuleListener;
+  state: TState;
+};
+
+type State<TMap extends Record<string, RuleEnhancer<unknown>>> = {
+  [K in keyof TMap]: TMap[K] extends RuleEnhancer<infer S> ? S : never;
+};
+
+/**
+ * Allows enhancing rule code with additional context provided by RuleEnhancers (reusable node visitors collecting data).
+ * @param enhancers a record of RuleEnhancers
+ * @param rule a visitor with an additional `state` argument that allows access to the enhancers' data
+ * @returns a resulting `(context: Context) => RuleListener` function
+ *
+ * @example
+ * // inside of `createRule`
+ *   create: enhanceRule({ metadata: metadataTrackingEnhancer }, (context, state) => {
+ *     const { metadata } = state;
+ *
+ *     return {
+ *       ObjectExpression(node) {
+ *         if (metadata.shouldReport()) {
+ *           context.report({ node, messageId: 'error' });
+ *         }
+ *       },
+ *     };
+ */
+export function enhanceRule<
+  TMap extends Record<string, RuleEnhancer<unknown>>,
+  Context extends RuleContext<string, unknown[]>,
+>(enhancers: TMap, rule: (context: Context, state: State<TMap>) => RuleListener) {
+  return (context: Context) => {
+    const enhancerVisitors: RuleListener[] = [];
+    const combinedState: Record<string, unknown> = {};
+
+    for (const [key, enhancer] of Object.entries(enhancers)) {
+      const initializedEnhancer = enhancer(context);
+      enhancerVisitors.push(initializedEnhancer.visitors);
+      combinedState[key] = initializedEnhancer.state;
+    }
+
+    const initializedRule = rule(context, combinedState as State<TMap>);
+
+    return mergeVisitors([...enhancerVisitors, initializedRule]);
+  };
+}
+
+/**
+ * Merges all passed visitors into one visitor.
+ * Retains visitor order:
+ * - on node enter, visitors are called in `visitorsList` order,
+ * - on node exit, visitors are called in reversed order.
+ */
+function mergeVisitors(visitors: RuleListener[]): RuleListener {
+  const merged: RuleListener = {};
+
+  const allKeys = new Set(visitors.flatMap((v) => Object.keys(v)));
+
+  for (const key of allKeys) {
+    const listeners = visitors.map((v) => v[key]).filter((fn) => fn !== undefined);
+
+    if (listeners.length === 0) {
+      continue;
+    }
+
+    // Reverse order if node is an exit node
+    if (key.endsWith(':exit')) {
+      listeners.reverse();
+    }
+
+    merged[key] = (...args: unknown[]) => {
+      // biome-ignore lint/suspicious/useIterableCallbackReturn: those functions return void
+      listeners.forEach((fn) => (fn as (...args: unknown[]) => void)(...args));
+    };
+  }
+
+  return merged;
+}

packages/eslint-plugin/src/enhancers/directiveTracking.ts

@@ -0,0 +1,70 @@
+import type { TSESTree } from '@typescript-eslint/utils';
+import type { RuleListener } from '@typescript-eslint/utils/ts-eslint';
+import type { RuleEnhancer } from '../enhanceRule.ts';
+
+export type DirectiveData = {
+  insideUseGpu: () => boolean;
+};
+
+/**
+ * A RuleEnhancer that tracks whether the current node is inside a 'use gpu' function.
+ *
+ * @privateRemarks
+ * Should the need arise, the API could be updated to expose:
+ * - a list of directives of the current function,
+ * - directives of other visited functions,
+ * - top level directives.
+ */
+export const directiveTracking: RuleEnhancer<DirectiveData> = () => {
+  const stack: string[][] = [];
+
+  const visitors: RuleListener = {
+    FunctionDeclaration(node) {
+      stack.push(getDirectives(node));
+    },
+    FunctionExpression(node) {
+      stack.push(getDirectives(node));
+    },
+    ArrowFunctionExpression(node) {
+      stack.push(getDirectives(node));
+    },
+
+    'FunctionDeclaration:exit'() {
+      stack.pop();
+    },
+    'FunctionExpression:exit'() {
+      stack.pop();
+    },
+    'ArrowFunctionExpression:exit'() {
+      stack.pop();
+    },
+  };
+
+  return {
+    visitors,
+    state: { insideUseGpu: () => (stack.at(-1) ?? []).includes('use gpu') },
+  };
+};
+
+function getDirectives(
+  node:
+    | TSESTree.FunctionDeclaration
+    | TSESTree.FunctionExpression
+    | TSESTree.ArrowFunctionExpression,
+): string[] {
+  const body = node.body;
+  if (body.type !== 'BlockStatement') {
+    return [];
+  }
+
+  const directives: string[] = [];
+  for (const statement of body.body) {
+    if (statement.type === 'ExpressionStatement' && statement.directive) {
+      directives.push(statement.directive);
+    } else {
+      break;
+    }
+  }
+
+  return directives;
+}

packages/eslint-plugin/src/rules/unwrappedPojos.ts

@@ -0,0 +1,48 @@
+import { enhanceRule } from '../enhanceRule.ts';
+import { directiveTracking } from '../enhancers/directiveTracking.ts';
+import { createRule } from '../ruleCreator.ts';
+
+export const unwrappedPojos = createRule({
+  name: 'unwrapped-pojo',
+  meta: {
+    type: 'problem',
+    docs: {
+      description: `Wrap Plain Old JavaScript Objects with schemas.`,
+    },
+    messages: {
+      unwrappedPojo:
+        '{{snippet}} is a POJO that is not wrapped in a schema. To allow WGSL resolution, wrap it in a schema call. You only need to wrap the outermost object.',
+    },
+    schema: [],
+  },
+  defaultOptions: [],
+
+  create: enhanceRule({ directives: directiveTracking }, (context, state) => {
+    const { directives } = state;
+
+    return {
+      ObjectExpression(node) {
+        if (!directives.insideUseGpu()) {
+          return;
+        }
+        if (node.parent?.type === 'Property') {
+          // a part of a bigger struct
+          return;
+        }
+        if (node.parent?.type === 'CallExpression') {
+          // wrapped in a schema call
+          return;
+        }
+        if (node.parent?.type === 'ReturnStatement') {
+          // likely inferred (shelled fn or shell-less entry) so we cannot report
+          return;
+        }
+        context.report({
+          node,
+          messageId: 'unwrappedPojo',
+          data: { snippet: context.sourceCode.getText(node) },
+        });
+      },
+    };
+  }),
+});

packages/eslint-plugin/tests/unwrappedPojos.test.ts

@@ -0,0 +1,76 @@
+import { describe } from 'vitest';
+import { ruleTester } from './ruleTester.ts';
+import { unwrappedPojos } from '../src/rules/unwrappedPojos.ts';
+
+describe('unwrappedPojos', () => {
+  ruleTester.run('unwrappedPojos', unwrappedPojos, {
+    valid: [
+      // correctly wrapped
+      "function func() { 'use gpu'; const wrapped = Schema({ a: 1 }); }",
+      "const func = function() { 'use gpu'; const wrapped = Schema({ a: 1 }); }",
+      "() => { 'use gpu'; const wrapped = Schema({ a: 1 }); }",
+
+      // not inside 'use gpu' function
+      'const pojo = { a: 1 };',
+      'function func() { const unwrapped = { a: 1 }; }',
+      'const func = function () { const unwrapped = { a: 1 }; }',
+      '() => { const unwrapped = { a: 1 }; }',
+      'function func() { return { a: 1 }; }',
+      'const func = function () { return { a: 1 }; }',
+      '() => { return { a: 1 }; }',
+
+      // return from 'use gpu' function
+      "function func() { 'use gpu'; return { a: 1 }; }",
+      "const func = function() { 'use gpu'; return { a: 1 }; }",
+      "() => { 'use gpu'; return { a: 1 }; }",
+      "() => { 'use gpu'; return { a: { b: 1 } }; }",
+    ],
+    invalid: [
+      {
+        code: "function func() { 'use gpu'; const unwrapped = { a: 1 }; }",
+        errors: [
+          {
+            messageId: 'unwrappedPojo',
+            data: { snippet: '{ a: 1 }' },
+          },
+        ],
+      },
+      {
+        code: "const func = function() { 'use gpu'; const unwrapped = { a: 1 }; }",
+        errors: [
+          {
+            messageId: 'unwrappedPojo',
+            data: { snippet: '{ a: 1 }' },
+          },
+        ],
+      },
+      {
+        code: "() => { 'use gpu'; const unwrapped = { a: 1 }; }",
+        errors: [
+          {
+            messageId: 'unwrappedPojo',
+            data: { snippet: '{ a: 1 }' },
+          },
+        ],
+      },
+      {
+        code: "function func() { 'unknown directive'; 'use gpu'; const unwrapped = { a: 1 }; }",
+        errors: [
+          {
+            messageId: 'unwrappedPojo',
+            data: { snippet: '{ a: 1 }' },
+          },
+        ],
+      },
+      {
+        code: "() => { 'use gpu'; const unwrapped = { a: { b: 1 } }; }",
+        errors: [
+          {
+            messageId: 'unwrappedPojo',
+            data: { snippet: '{ a: { b: 1 } }' },
+          },
+        ],
+      },
+    ],
+  });
+});