feat(cli): replace gas report

.github/workflows/build.yml

@@ -29,9 +29,3 @@ jobs:
 
       - name: Outdated files, run `pnpm build` and commit them
         uses: ./.github/actions/require-empty-diff
-
-      - name: Generate gas reports
-        run: pnpm gas-report
-
-      - name: Outdated files, run `pnpm gas-report` and commit them
-        uses: ./.github/actions/require-empty-diff

.github/workflows/release.yml

@@ -30,12 +30,6 @@ jobs:
       - name: Build
         uses: ./.github/actions/build
 
-      - name: Generate gas reports
-        run: pnpm gas-report
-
-      - name: Outdated files, run `pnpm gas-report` and commit them
-        uses: ./.github/actions/require-empty-diff
-
       - name: Set deployment token
         run: npm config set '//registry.npmjs.org/:_authToken' "${NPM_TOKEN}"
         env:
@@ -64,12 +58,6 @@ jobs:
       - name: Build
         uses: ./.github/actions/build
 
-      - name: Generate gas reports
-        run: pnpm gas-report
-
-      - name: Outdated files, run `pnpm gas-report` and commit them
-        uses: ./.github/actions/require-empty-diff
-
       - name: Set deployment token
         run: npm config set '//registry.npmjs.org/:_authToken' "${NPM_TOKEN}"
         env:

.github/workflows/test.yml

@@ -23,6 +23,12 @@ jobs:
       - name: Run tests
         run: pnpm test
 
+      - name: Generate gas reports
+        run: pnpm gas-report
+
+      - name: Outdated files, run `pnpm gas-report` and commit them
+        uses: ./.github/actions/require-empty-diff
+
       - name: Lint
         run: pnpm lint
 

packages/cli/package.json

@@ -51,6 +51,7 @@
     "nice-grpc-web": "^2.0.1",
     "openurl": "^1.1.1",
     "path": "^0.12.7",
+    "strip-ansi": "^7.1.0",
     "table": "^6.8.1",
     "throttle-debounce": "^5.0.0",
     "typechain": "^8.1.1",

packages/cli/src/commands/gas-report.ts

@@ -1,23 +1,27 @@
 import type { CommandModule } from "yargs";
-import { readFileSync, writeFileSync, rmSync } from "fs";
+import { readFileSync, writeFileSync } from "fs";
 import { execa } from "execa";
 import chalk from "chalk";
 import { table, getBorderCharacters } from "table";
+import stripAnsi from "strip-ansi";
 
 /**
  * Print the gas report to the console, save it to a file and compare it to a previous gas report if provided.
- * Requires forge to be installed, and gas test files including `// !gasreport` comments, like this:
+ *
+ * Requires foundry to be installed. Inherit from GasReporter and use startGasReport/endGasReport to measure gas used.
  *
  * ```solidity
- * contract GasTest is DSTestPlus {
+ * contract MyContractTest is Test, GasReporter {
  *   function testBuffer() public pure {
- *     // !gasreport allocate a buffer
+ *     startGasReport("allocate a buffer");
  *     Buffer buffer = Buffer_.allocate(32);
+ *     endGasReport();
  *
  *     bytes32 value = keccak256("some data");
  *
- *     // !gasreport append 32 bytes to a buffer
+ *     startGasReport("append 32 bytes to a buffer");
  *     buffer.append(value);
+ *     endGasReport();
  *  }
  * }
  * ```
@@ -30,48 +34,35 @@ type Options = {
 };
 
 type GasReportEntry = {
-  source: string;
+  file: string;
+  test: string;
   name: string;
-  functionCall: string;
   gasUsed: number;
   prevGasUsed?: number;
 };
 
 type GasReport = GasReportEntry[];
 
-const tempFileSuffix = "MudGasReport";
-
 const commandModule: CommandModule<Options, Options> = {
   command: "gas-report",
 
   describe: "Create a gas report",
 
   builder(yargs) {
     return yargs.options({
-      path: { type: "array", string: true, default: ["Gas.t.sol"], desc: "File containing the gas tests" },
       save: { type: "string", desc: "Save the gas report to a file" },
       compare: { type: "string", desc: "Compare to an existing gas report" },
     });
   },
 
-  async handler({ path: files, save, compare }) {
-    const validFiles = files.filter((file) => file.endsWith(".t.sol"));
-    const tempFiles = await Promise.all(validFiles.map((file) => createGasReport(file)));
-
-    process.once("SIGINT", () => {
-      console.log("caught sigint, deleting temp files");
-      tempFiles.forEach((file) => rmSync(file));
-    });
-
+  async handler({ save, compare }) {
     let gasReport: GasReport;
     try {
       gasReport = await runGasReport();
-    } catch {
+    } catch (error) {
+      console.error(error);
       setTimeout(() => process.exit());
       return;
-    } finally {
-      // Delete the temporary files
-      tempFiles.forEach((file) => rmSync(file));
     }
 
     // If this gas report should be compared to an existing one, load the existing one
@@ -80,9 +71,7 @@ const commandModule: CommandModule<Options, Options> = {
         const compareGasReport: GasReport = JSON.parse(readFileSync(compare, "utf8"));
         // Merge the previous gas report with the new one
         gasReport = gasReport.map((entry) => {
-          const prevEntry = compareGasReport.find(
-            (e) => e.name === entry.name && e.functionCall === entry.functionCall
-          );
+          const prevEntry = compareGasReport.find((e) => e.file === entry.file && e.name === entry.name);
           return { ...entry, prevGasUsed: prevEntry?.gasUsed };
         });
       } catch {
@@ -103,88 +92,73 @@ const commandModule: CommandModule<Options, Options> = {
 
 export default commandModule;
 
-async function createGasReport(filename: string): Promise<string> {
-  console.log("Creating gas report for", chalk.bold(filename));
-
-  // Parse the given test file, and add gas reporting wherever requested by a `// !gasreport` comment
-  const fileContents = readFileSync(filename, "utf8");
-  let newFile = fileContents;
-
-  // Use a regex to find first line of each function
-  const functionRegex = new RegExp(/function (.*){/g);
-  // Insert a line to declare the _gasreport variable at the start of each function
-  let functionMatch;
-  while ((functionMatch = functionRegex.exec(fileContents)) !== null) {
-    const functionSignature = functionMatch[0];
-    newFile = newFile.replace(functionSignature, `${functionSignature}\nuint256 _gasreport;`);
-  }
-
-  // A gasreport comment has a name (written after the comment) and a function call (written on the next line)
-  // Create a regex to extract both the name and the function call
-  const regex = new RegExp(/\/\/ !gasreport (.*)\n(.*)/g);
-
-  // Apply the regex and loop through the matches,
-  // and create a new file with the gasreport comments replaced by the gas report
-  let match;
-  while ((match = regex.exec(fileContents)) !== null) {
-    const name = match[1];
-    const functionCall = match[2].trim();
-
-    newFile = newFile.replace(
-      match[0],
-      `
-_gasreport = gasleft();
-${functionCall}
-_gasreport = _gasreport - gasleft();
-console.log("GAS REPORT(${filename}): ${name} [${functionCall.replaceAll('"', '\\"')}]:", _gasreport);`
-    );
-  }
-
-  // Remove all occurrences of `pure` with `view`
-  newFile = newFile.replace(/pure/g, "view");
-
-  // Write the new file to disk (temporarily)
-  // Create the temp file by replacing the previous file name with MudGasReport
-  const tempFileName = filename.replace(/\.t\.sol$/, `${tempFileSuffix}.t.sol`);
-  writeFileSync(tempFileName, newFile);
-
-  return tempFileName;
-}
-
 async function runGasReport(): Promise<GasReport> {
   console.log("Running gas report");
   const gasReport: GasReport = [];
 
   // Extract the logs from the child process
-  let logs = "";
+  let stdout: string;
   try {
     // Run the generated file using forge
-    const child = execa("forge", ["test", "--match-path", `*${tempFileSuffix}*`, "-vvv"], {
+    const child = execa("forge", ["test", "-vvv"], {
       stdio: ["inherit", "pipe", "inherit"],
+      env: { GAS_REPORTER_ENABLED: "true" },
     });
-    logs = (await child).stdout;
+    stdout = (await child).stdout;
   } catch (error: any) {
     console.log(error.stdout ?? error);
     console.log(chalk.red("\n-----------\nError while running the gas report (see above)"));
     throw error;
   }
 
   // Extract the gas reports from the logs
+  const lines = stdout.split("\n").map(stripAnsi);
+  const gasReportPattern = /^\s*GAS REPORT: (\d+) (.*)$/;
+  const testFunctionPattern = /^\[(?:PASS|FAIL).*\] (\w+)\(\)/;
+  const testFilePattern = /^Running \d+ tests? for (.*):(.*)$/;
+
+  function nearestLine(pattern: RegExp, startIndex: number = lines.length - 1): number {
+    for (let i = startIndex; i >= 0; i--) {
+      const line = lines[i];
+      if (pattern.test(line)) {
+        return i;
+      }
+    }
+    return -1;
+  }
+
+  for (let i = 0; i < lines.length; i++) {
+    const matches = lines[i].match(gasReportPattern);
+    if (!matches) continue;
+
+    const gasUsed = parseInt(matches[1]);
+    const name = matches[2];
+
+    const testFunctionLineIndex = nearestLine(testFunctionPattern, i);
+    if (testFunctionLineIndex === -1) {
+      throw new Error("Could not find nearest test function, did `forge test` output change?");
+    }
+    const testFileLineIndex = nearestLine(testFilePattern, testFunctionLineIndex);
+    if (testFileLineIndex === -1) {
+      throw new Error("Could not find nearest test filename, did `forge test` output change?");
+    }
+
+    const functionMatches = lines[testFunctionLineIndex].match(testFunctionPattern);
+    if (!functionMatches) {
+      throw new Error("Could not parse test function name, did `forge test` output change?");
+    }
+    const fileMatches = lines[testFileLineIndex].match(testFilePattern);
+    if (!fileMatches) {
+      throw new Error("Could not parse test filename, did `forge test` output change?");
+    }
+
+    const test = functionMatches[1];
+    const file = fileMatches[1];
 
-  // Create a regex to find all lines starting with `GAS REPORT:` and extract the name, function call and gas used
-  const gasReportRegex = new RegExp(/GAS REPORT\((.*)\): (.*) \[(.*)\]: (.*)/g);
-
-  // Loop through the matches and print the gas report
-  let gasReportMatch;
-  while ((gasReportMatch = gasReportRegex.exec(logs)) !== null) {
-    const source = gasReportMatch[1];
-    const name = gasReportMatch[2];
-    const functionCall = gasReportMatch[3].replace(";", "");
-    const gasUsed = parseInt(gasReportMatch[4]);
-    gasReport.push({ source, name, functionCall, gasUsed });
+    gasReport.push({ file, test, name, gasUsed });
   }
 
-  gasReport.sort((a, b) => a.source.localeCompare(b.source));
+  gasReport.sort((a, b) => a.file.localeCompare(b.file));
 
   return gasReport;
 }
@@ -193,9 +167,9 @@ function printGasReport(gasReport: GasReport, compare?: string) {
   if (compare) console.log(chalk.bold(`Gas report compared to ${compare}`));
 
   const headers = [
-    chalk.bold("Source"),
+    chalk.bold("File"),
+    chalk.bold("Test"),
     chalk.bold("Name"),
-    chalk.bold("Function call"),
     chalk.bold("Gas used"),
     ...(compare ? [chalk.bold("Prev gas used"), chalk.bold("Difference")] : []),
   ];
@@ -205,7 +179,7 @@ function printGasReport(gasReport: GasReport, compare?: string) {
     const diffEntry = diff > 0 ? chalk.red(`+${diff}`) : diff < 0 ? chalk.green(`${diff}`) : diff;
     const compareColumns = compare ? [entry.prevGasUsed ?? "n/a", diffEntry] : [];
     const gasUsedEntry = diff > 0 ? chalk.red(entry.gasUsed) : diff < 0 ? chalk.green(entry.gasUsed) : entry.gasUsed;
-    return [entry.source, entry.name, entry.functionCall, gasUsedEntry, ...compareColumns];
+    return [entry.file, entry.test, entry.name, gasUsedEntry, ...compareColumns];
   });
 
   const rows = [headers, ...values];

packages/std-contracts/src/test/GasReporter.sol

@@ -0,0 +1,62 @@
+// SPDX-License-Identifier: MIT
+pragma solidity >=0.8.0;
+
+import { Test, console } from "forge-std/Test.sol";
+
+contract GasReporter is Test {
+  string private __currentGasReportName;
+  uint256 private __currentGasReportValue = gasleft();
+  mapping(string => uint256) private __gasReports;
+  string[] private __gasReportNames;
+
+  function startGasReport(string memory name) internal {
+    if (!vm.envOr("GAS_REPORTER_ENABLED", false)) return;
+    require(
+      bytes(__currentGasReportName).length == 0,
+      string.concat(
+        'gas report "',
+        __currentGasReportName,
+        '" is already running and only one report can be run at a time'
+      )
+    );
+    require(__gasReports[name] == 0, string.concat('gas report "', name, '" already used for this test'));
+    __currentGasReportName = name;
+    vm.pauseGasMetering();
+    __currentGasReportValue = gasleft();
+    vm.resumeGasMetering();
+  }
+
+  function endGasReport() internal {
+    uint256 gas = gasleft();
+    if (!vm.envOr("GAS_REPORTER_ENABLED", false)) return;
+    // subtract 160 gas used by the GasReporter itself
+    // add 1 gas so we can later check that this is set
+    uint256 gasUsed = __currentGasReportValue - gas - 160 + 1;
+    require(gasUsed > 0, "gas report didn't use gas");
+    __gasReports[__currentGasReportName] = gasUsed;
+    __gasReportNames.push(__currentGasReportName);
+    printGasReport(__currentGasReportName);
+    __currentGasReportName = "";
+  }
+
+  modifier gasReport(string memory name) {
+    startGasReport(name);
+    _;
+    endGasReport();
+  }
+
+  function getGasUsed(string memory name) internal view returns (uint256) {
+    require(__gasReports[name] > 0, string.concat('gas report "', name, '" not found'));
+    return __gasReports[name];
+  }
+
+  function printGasReport(string memory name) internal view {
+    console.log(string.concat("GAS REPORT: ", vm.toString(__gasReports[name]), " ", name));
+  }
+
+  function printGasReports() internal view {
+    for (uint256 i = 0; i < __gasReportNames.length; i++) {
+      printGasReport(__gasReportNames[i]);
+    }
+  }
+}