feat: Lint rule for invalid assignment (#2235)

Author: aleksanderkatan

packages/eslint-plugin/README.md

@@ -43,6 +43,7 @@ export default defineConfig([
 | Name                                                                   | Description                                                                                 | 🚨 | ⚠️ |
 | :--------------------------------------------------------------------- | :------------------------------------------------------------------------------------------ | :- | :- |
 | [no-integer-division](docs/rules/no-integer-division.md)               | Disallow division incorporating numbers wrapped in 'u32' and 'i32'                          |    | ⭐  |
+| [no-invalid-assignment](docs/rules/no-invalid-assignment.md)           | Disallow assignments that will generate invalid WGSL                                        | ⭐  |    |
 | [no-math](docs/rules/no-math.md)                                       | Disallow usage of JavaScript 'Math' methods inside 'use gpu' functions                      |    | ⭐  |
 | [no-uninitialized-variables](docs/rules/no-uninitialized-variables.md) | Disallow variable declarations without initializers inside 'use gpu' functions              | ⭐  |    |
 | [no-unwrapped-objects](docs/rules/no-unwrapped-objects.md)             | Disallow unwrapped Plain Old JavaScript Objects inside 'use gpu' functions (except returns) | ⭐  |    |

packages/eslint-plugin/docs/rules/no-invalid-assignment.md

@@ -0,0 +1,54 @@
+# typegpu/no-invalid-assignment
+
+📝 Disallow assignments that will generate invalid WGSL.
+
+🚨 This rule is enabled in the ⭐ `recommended` config.
+
+<!-- end auto-generated rule header -->
+
+## Rule details
+
+Examples of **incorrect** code for this rule:
+
+```ts
+const fn = (a) => { 
+  'use gpu'; 
+  a = 1;
+}
+```
+```ts
+const fn = (a) => { 
+  'use gpu'; 
+  a.prop++;
+}
+```
+```ts
+let a;
+const fn = () => { 
+  'use gpu'; 
+  a = 1;
+}
+```
+
+Examples of **correct** code for this rule:
+
+```ts
+const fn = () => {
+  'use gpu';
+  const ref = d.ref(0);
+  other(ref);
+};
+
+const other = (ref: d.ref<number>) => {
+  'use gpu';
+  ref.$ = 1;
+};
+```
+```ts
+const privateVar = tgpu.privateVar(d.u32);
+const fn = () => { 
+  'use gpu'; 
+  privateVar.$ = 1;
+}
+```
+

packages/eslint-plugin/src/configs.ts

@@ -3,12 +3,14 @@ import { noIntegerDivision } from './rules/noIntegerDivision.ts';
 import { noUnwrappedObjects } from './rules/noUnwrappedObjects.ts';
 import { noMath } from './rules/noMath.ts';
 import { noUninitializedVariables } from './rules/noUninitializedVariables.ts';
+import { noInvalidAssignment } from './rules/noInvalidAssignment.ts';
 
 export const rules = {
   'no-integer-division': noIntegerDivision,
   'no-unwrapped-objects': noUnwrappedObjects,
   'no-uninitialized-variables': noUninitializedVariables,
   'no-math': noMath,
+  'no-invalid-assignment': noInvalidAssignment,
 } as const;
 
 type Rules = Record<`typegpu/${keyof typeof rules}`, TSESLint.FlatConfig.RuleEntry>;
@@ -18,11 +20,13 @@ export const recommendedRules: Rules = {
   'typegpu/no-unwrapped-objects': 'error',
   'typegpu/no-uninitialized-variables': 'error',
   'typegpu/no-math': 'warn',
+  'typegpu/no-invalid-assignment': 'error',
 };
 
 export const allRules: Rules = {
   'typegpu/no-integer-division': 'error',
   'typegpu/no-unwrapped-objects': 'error',
   'typegpu/no-uninitialized-variables': 'error',
   'typegpu/no-math': 'error',
+  'typegpu/no-invalid-assignment': 'error',
 };

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

@@ -2,8 +2,13 @@ import type { TSESTree } from '@typescript-eslint/utils';
 import type { RuleListener } from '@typescript-eslint/utils/ts-eslint';
 import type { RuleEnhancer } from '../enhanceRule.ts';
 
+export type FunctionNode =
+  | TSESTree.FunctionDeclaration
+  | TSESTree.FunctionExpression
+  | TSESTree.ArrowFunctionExpression;
+
 export type DirectiveData = {
-  insideUseGpu: () => boolean;
+  getEnclosingTypegpuFunction: () => FunctionNode | undefined;
 };
 
 /**
@@ -16,17 +21,17 @@ export type DirectiveData = {
  * - top level directives.
  */
 export const directiveTracking: RuleEnhancer<DirectiveData> = () => {
-  const stack: string[][] = [];
+  const stack: { node: FunctionNode; directives: string[] }[] = [];
 
   const visitors: RuleListener = {
     FunctionDeclaration(node) {
-      stack.push(getDirectives(node));
+      stack.push({ node, directives: getDirectives(node) });
     },
     FunctionExpression(node) {
-      stack.push(getDirectives(node));
+      stack.push({ node, directives: getDirectives(node) });
     },
     ArrowFunctionExpression(node) {
-      stack.push(getDirectives(node));
+      stack.push({ node, directives: getDirectives(node) });
     },
 
     'FunctionDeclaration:exit'() {
@@ -42,7 +47,15 @@ export const directiveTracking: RuleEnhancer<DirectiveData> = () => {
 
   return {
     visitors,
-    state: { insideUseGpu: () => (stack.at(-1) ?? []).includes('use gpu') },
+    state: {
+      getEnclosingTypegpuFunction: () => {
+        const current = stack.at(-1);
+        if (current && current.directives.includes('use gpu')) {
+          return current.node;
+        }
+        return undefined;
+      },
+    },
   };
 };
 

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

@@ -0,0 +1,90 @@
+import { ASTUtils, type TSESTree } from '@typescript-eslint/utils';
+import { createRule } from '../ruleCreator.ts';
+import { enhanceRule } from '../enhanceRule.ts';
+import { directiveTracking } from '../enhancers/directiveTracking.ts';
+import type { RuleContext } from '@typescript-eslint/utils/ts-eslint';
+
+export const noInvalidAssignment = createRule({
+  name: 'no-invalid-assignment',
+  meta: {
+    type: 'problem',
+    docs: {
+      description: `Disallow assignments that will generate invalid WGSL`,
+    },
+    messages: {
+      parameterAssignment:
+        "Cannot assign to '{{snippet}}' since WGSL parameters are immutable. If you're using d.ref, please either use '.$' or disable this rule",
+      jsAssignment:
+        "Cannot assign to '{{snippet}}' since it is a JS variable defined outside of the current TypeGPU function's scope. Use buffers, workgroup variables or local variables instead",
+    },
+    schema: [],
+  },
+  defaultOptions: [],
+
+  create: enhanceRule({ directives: directiveTracking }, (context, state) => {
+    const { directives } = state;
+
+    return {
+      UpdateExpression(node) {
+        const enclosingFn = directives.getEnclosingTypegpuFunction();
+        validateAssignment(context, node, enclosingFn, node.argument);
+      },
+
+      AssignmentExpression(node) {
+        const enclosingFn = directives.getEnclosingTypegpuFunction();
+        validateAssignment(context, node, enclosingFn, node.left);
+      },
+    };
+  }),
+});
+
+function validateAssignment(
+  context: Readonly<RuleContext<'parameterAssignment' | 'jsAssignment', []>>,
+  node: TSESTree.Node,
+  enclosingFn: TSESTree.Node | undefined,
+  leftNode: TSESTree.Node,
+) {
+  if (!enclosingFn) {
+    return;
+  }
+
+  // follow the member expression chain
+  let assignee = leftNode;
+  while (assignee.type === 'MemberExpression') {
+    if (assignee.property.type === 'Identifier' && assignee.property.name === '$') {
+      // a dollar was used so we assume this assignment is fine
+      return;
+    }
+    assignee = assignee.object;
+  }
+  if (assignee.type !== 'Identifier') {
+    return;
+  }
+
+  // look for a scope that defines the variable
+  const variable = ASTUtils.findVariable(context.sourceCode.getScope(assignee), assignee);
+  // defs is an array because there may be multiple definitions with `var`
+  const def = variable?.defs[0];
+
+  // check if variable is global or was defined outside of current function by checking ranges
+  // NOTE: if the variable is an outer function parameter, then the enclosingFn range will be encompassed by node range
+  if (
+    !def ||
+    (def && (def.node.range[0] < enclosingFn.range[0] || enclosingFn.range[1] < def.node.range[1]))
+  ) {
+    context.report({
+      messageId: 'jsAssignment',
+      node,
+      data: { snippet: context.sourceCode.getText(leftNode) },
+    });
+    return;
+  }
+
+  if (def.type === 'Parameter') {
+    context.report({
+      messageId: 'parameterAssignment',
+      node,
+      data: { snippet: context.sourceCode.getText(leftNode) },
+    });
+  }
+}

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

@@ -24,7 +24,7 @@ export const noMath = createRule({
 
     return {
       CallExpression(node) {
-        if (!directives.insideUseGpu()) {
+        if (!directives.getEnclosingTypegpuFunction()) {
           return;
         }
 

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

@@ -21,7 +21,7 @@ export const noUninitializedVariables = createRule({
 
     return {
       VariableDeclarator(node) {
-        if (!directives.insideUseGpu()) {
+        if (!directives.getEnclosingTypegpuFunction()) {
           return;
         }
         if (node.parent?.parent?.type === 'ForOfStatement') {

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

@@ -22,7 +22,7 @@ export const noUnwrappedObjects = createRule({
 
     return {
       ObjectExpression(node) {
-        if (!directives.insideUseGpu()) {
+        if (!directives.getEnclosingTypegpuFunction()) {
           return;
         }
         let parent = getNonTransparentParent(node);