feat: quickfix for data location

EXTENSION.md

@@ -154,6 +154,12 @@ When no license is specified on a contract, the `solidity(1878)` warning is rais
 
 ![Add pragma solidity](https://raw.githubusercontent.com/NomicFoundation/hardhat-vscode/main/docs/gifs/add-pragma.gif "Add pragma solidity")
 
+#### Specifying data location for variables
+
+Some types require you to specify a data location (memory, storage, calldata), depending on where they are defined. The available code actions allow the user to add, change or remove data locations depending on the error being raised.
+
+![Data location quickfix](https://raw.githubusercontent.com/NomicFoundation/hardhat-vscode/main/docs/gifs/data-location.gif "Specify data location")
+
 ### Commands
 
 #### Compile project

server/src/compilerDiagnostics/compilerDiagnostics.ts

@@ -8,6 +8,7 @@ import { MarkContractAbstract } from "./diagnostics/MarkContractAbstract";
 import { SpecifyVisibility } from "./diagnostics/SpecifyVisibility";
 import { SpecifyCompilerVersion } from "./diagnostics/SpecifyCompilerVersion";
 import { CompilerDiagnostic } from "./types";
+import { SpecifyDataLocation } from "./diagnostics/SpecifyDataLocation";
 
 export const compilerDiagnostics: { [key: string]: CompilerDiagnostic } = [
   new AddOverrideSpecifier(),
@@ -19,4 +20,5 @@ export const compilerDiagnostics: { [key: string]: CompilerDiagnostic } = [
   new MarkContractAbstract(),
   new SpecifyVisibility(),
   new SpecifyCompilerVersion(),
+  new SpecifyDataLocation(),
 ].reduce((acc, item) => ({ ...acc, [item.code]: item }), {});

server/src/compilerDiagnostics/diagnostics/SpecifyDataLocation.ts

@@ -0,0 +1,152 @@
+/* eslint-disable @typescript-eslint/no-non-null-assertion */
+import {
+  CodeAction,
+  CodeActionKind,
+  Diagnostic,
+  Range,
+} from "vscode-languageserver/node";
+import { TextDocument } from "vscode-languageserver-textdocument";
+import _ from "lodash";
+import { ResolveActionsContext } from "../types";
+import { SolcError, ServerState } from "../../types";
+import { passThroughConversion } from "../conversions/passThroughConversion";
+
+const MUST_SPECIFY_LOCATION_REGEX = /must be(.*)for/;
+const CANNOT_SPECIFY_REGEX = /can only be specified for/;
+const DATA_LOCATION_REGEX = /storage|memory|calldata/g;
+const DATA_LOCATION_WHITESPACE_REGEX = /(storage|memory|calldata)(\s|$)/g;
+
+export class SpecifyDataLocation {
+  public code = "6651";
+  public blocks: string[] = [];
+
+  public fromHardhatCompilerError(
+    document: TextDocument,
+    error: SolcError
+  ): Diagnostic {
+    return passThroughConversion(document, error);
+  }
+
+  public resolveActions(
+    serverState: ServerState,
+    diagnostic: Diagnostic,
+    context: ResolveActionsContext
+  ): CodeAction[] {
+    const { document, uri } = context;
+
+    if (MUST_SPECIFY_LOCATION_REGEX.test(diagnostic.message)) {
+      // A variable declaration didn't specify data location and it should (e.g. arrays)
+      // Build a list of valid locations and sort by custom criteria (most frequent first)
+      const allowedLocations = this._getAllowedLocationsFromMessage(
+        diagnostic.message
+      ).sort(
+        (a, b) =>
+          this._getLocationSortWeight(a) - this._getLocationSortWeight(b)
+      );
+
+      return allowedLocations.map((location) =>
+        this._buildAddLocationAction(location, document, uri, diagnostic.range)
+      );
+    } else if (CANNOT_SPECIFY_REGEX.test(diagnostic.message)) {
+      // A variable declaration specified data location and it shouldn't (e.g. an integer)
+      return [this._buildRemoveLocationAction(document, uri, diagnostic.range)];
+    } else {
+      throw new Error(
+        `Unexpected ${this.code} error message: ${diagnostic.message}`
+      );
+    }
+  }
+
+  private _getAllowedLocationsFromMessage(message: string) {
+    const allowedLocationSubstring = message.match(MUST_SPECIFY_LOCATION_REGEX);
+
+    return allowedLocationSubstring![1].match(DATA_LOCATION_REGEX) ?? [];
+  }
+
+  private _buildAddLocationAction(
+    location: string,
+    document: TextDocument,
+    uri: string,
+    range: Range
+  ): CodeAction {
+    // Find the declaration text, e.g. 'uint[] foo'
+    const variableDeclaration = document.getText(range);
+
+    // Remove any potential existing data location
+    const normalizedVariableDeclaration =
+      this._removeDataLocationFromDeclaration(variableDeclaration);
+
+    // Split into array of words by whitespaces
+    const tokens = normalizedVariableDeclaration.split(/\s+/);
+
+    // Add data location after type
+    const type = tokens.shift();
+    const newDeclaration = [type, location, ...tokens].join(" ");
+
+    return {
+      title: `Specify '${location}' as data location`,
+      kind: CodeActionKind.QuickFix,
+      isPreferred: false,
+      edit: {
+        changes: {
+          [uri]: [
+            {
+              range,
+              newText: newDeclaration,
+            },
+          ],
+        },
+      },
+    };
+  }
+
+  private _buildRemoveLocationAction(
+    document: TextDocument,
+    uri: string,
+    range: Range
+  ): CodeAction {
+    // Find the declaration text, e.g. 'uint[] foo'
+    const variableDeclaration = document.getText(range);
+
+    // Remove any potential existing data location
+    const normalizedVariableDeclaration =
+      this._removeDataLocationFromDeclaration(variableDeclaration);
+
+    // Replace multiple whitespaces with single ones
+    const newDeclaration = normalizedVariableDeclaration.replace(/\s+/g, " ");
+
+    return {
+      title: `Remove specified data location`,
+      kind: CodeActionKind.QuickFix,
+      isPreferred: false,
+      edit: {
+        changes: {
+          [uri]: [
+            {
+              range,
+              newText: newDeclaration,
+            },
+          ],
+        },
+      },
+    };
+  }
+
+  private _removeDataLocationFromDeclaration(declaration: string): string {
+    return declaration.replace(DATA_LOCATION_WHITESPACE_REGEX, "");
+  }
+
+  // We want the most frequent options to be at the top
+  private _getLocationSortWeight(location: string) {
+    switch (location) {
+      case "memory":
+        return 0;
+      case "storage":
+        return 1;
+      case "calldata":
+        return 2;
+      default:
+        return 3;
+    }
+  }
+}

server/src/services/validation/validate.ts

@@ -143,30 +143,30 @@ export async function validate(
 function sendResults(
   serverState: ServerState,
   change: TextDocumentChangeEvent<TextDocument>,
-  completeMessage: ValidationResult,
+  validationResult: ValidationResult,
   openDocuments: OpenDocuments,
   projectBasePath: string
 ) {
-  switch (completeMessage.status) {
+  switch (validationResult.status) {
     case "JOB_COMPLETION_ERROR":
-      jobCompletionErrorFail(serverState, change, completeMessage);
+      jobCompletionErrorFail(serverState, change, validationResult);
       break;
     case "VALIDATION_FAIL":
-      validationFail(serverState, change, completeMessage);
+      validationFail(serverState, change, validationResult);
       break;
     case "VALIDATION_PASS":
-      validationPass(serverState, change, completeMessage, openDocuments);
+      validationPass(serverState, change, validationResult, openDocuments);
       break;
     case "BUILD_INPUT_ERROR":
       handleBuildInputError(
         serverState,
         change,
-        completeMessage,
+        validationResult,
         projectBasePath
       );
       break;
     default:
-      assertUnknownMessageStatus(completeMessage);
+      assertUnknownMessageStatus(validationResult);
       break;
   }
 }

test/integration/projects/main/contracts/codeactions/SpecifyDataLocation.sol

@@ -0,0 +1,16 @@
+// SPDX-License-Identifier: MIT
+pragma solidity >=0.7.0 <0.9.0;
+
+contract Foo {
+  // Constructor quickfixes
+  constructor(uint256[] p1, string calldata p2) {}
+
+  // Function parameters and return value quickfixes
+  function foo(uint256[] storage p3, string p4)
+    public
+    returns (bytes, string storage)
+  {
+    // Variable quickfixes
+    uint256 memory singleUint;
+  }
+}

test/integration/projects/main/contracts/codeactions/SpecifyDataLocation_fixed.sol

@@ -0,0 +1,16 @@
+// SPDX-License-Identifier: MIT
+pragma solidity >=0.7.0 <0.9.0;
+
+contract Foo {
+  // Constructor quickfixes
+  constructor(uint256[] memory p1, string memory p2) {}
+
+  // Function parameters and return value quickfixes
+  function foo(uint256[] memory p3, string memory p4)
+    public
+    returns (bytes memory, string memory )
+  {
+    // Variable quickfixes
+    uint256 singleUint;
+  }
+}

test/integration/tests/codeactions/specifyDataLocation.test.ts

@@ -0,0 +1,138 @@
+import * as vscode from "vscode";
+import { checkOrWaitDiagnostic } from "../../helpers/assertions";
+import { getTestContractUri } from "../../helpers/getTestContract";
+import {
+  applyQuickfix,
+  compareWithFile,
+  goToPosition,
+  openFileInEditor,
+  openQuickfixMenu,
+} from "../../helpers/editor";
+
+suite("codeactions - specify data location", function () {
+  test("add, change and remove multiple data locations", async () => {
+    const uri = getTestContractUri(
+      "main/contracts/codeactions/SpecifyDataLocation.sol"
+    );
+    const editor = await openFileInEditor(uri);
+
+    // 1st constructor param
+    await checkOrWaitDiagnostic(
+      uri,
+      new vscode.Range(new vscode.Position(5, 14), new vscode.Position(5, 26)),
+      vscode.DiagnosticSeverity.Error,
+      "solidity",
+      'Data location must be "storage" or "memory" for constructor parameter'
+    );
+
+    goToPosition(new vscode.Position(5, 14));
+
+    await openQuickfixMenu();
+    await applyQuickfix(0);
+
+    // 2nd constructor param
+
+    await checkOrWaitDiagnostic(
+      uri,
+      new vscode.Range(new vscode.Position(5, 35), new vscode.Position(5, 53)),
+      vscode.DiagnosticSeverity.Error,
+      "solidity",
+      'Data location must be "storage" or "memory" for constructor parameter'
+    );
+
+    goToPosition(new vscode.Position(5, 35));
+
+    await openQuickfixMenu();
+    await applyQuickfix(0);
+
+    // 1st function param
+
+    await checkOrWaitDiagnostic(
+      uri,
+      new vscode.Range(new vscode.Position(8, 15), new vscode.Position(8, 35)),
+      vscode.DiagnosticSeverity.Error,
+      "solidity",
+      'Data location must be "memory" or "calldata" for parameter in function'
+    );
+
+    goToPosition(new vscode.Position(8, 15));
+
+    await openQuickfixMenu();
+    await applyQuickfix(0);
+
+    // 2nd function param
+
+    await checkOrWaitDiagnostic(
+      uri,
+      new vscode.Range(new vscode.Position(8, 36), new vscode.Position(8, 45)),
+      vscode.DiagnosticSeverity.Error,
+      "solidity",
+      'Data location must be "memory" or "calldata" for parameter in function'
+    );
+
+    goToPosition(new vscode.Position(8, 36));
+
+    await openQuickfixMenu();
+    await applyQuickfix(0);
+
+    // 1st function return
+
+    await checkOrWaitDiagnostic(
+      uri,
+      new vscode.Range(
+        new vscode.Position(10, 13),
+        new vscode.Position(10, 18)
+      ),
+      vscode.DiagnosticSeverity.Error,
+      "solidity",
+      'Data location must be "memory" or "calldata" for return parameter'
+    );
+
+    goToPosition(new vscode.Position(10, 13));
+
+    await openQuickfixMenu();
+    await applyQuickfix(0);
+
+    // 2nd function return
+
+    await checkOrWaitDiagnostic(
+      uri,
+      new vscode.Range(
+        new vscode.Position(10, 27),
+        new vscode.Position(10, 41)
+      ),
+      vscode.DiagnosticSeverity.Error,
+      "solidity",
+      'Data location must be "memory" or "calldata" for return parameter'
+    );
+
+    goToPosition(new vscode.Position(10, 27));
+
+    await openQuickfixMenu();
+    await applyQuickfix(0);
+
+    // variable
+
+    await checkOrWaitDiagnostic(
+      uri,
+      new vscode.Range(new vscode.Position(13, 4), new vscode.Position(13, 29)),
+      vscode.DiagnosticSeverity.Error,
+      "solidity",
+      "Data location can only be specified for array, struct or mapping types"
+    );
+
+    goToPosition(new vscode.Position(13, 4));
+
+    await openQuickfixMenu();
+    await applyQuickfix(0);
+
+    // Assert with final file
+
+    compareWithFile(
+      editor,
+      getTestContractUri(
+        "main/contracts/codeactions/SpecifyDataLocation_fixed.sol"
+      )
+    );
+  });
+});