From 4385c5a4c0e2d5550c041acc4386ae7fc1cb4b7e Mon Sep 17 00:00:00 2001
From: alvarius <alvarius@lattice.xyz>
Date: Tue, 3 Oct 2023 15:23:51 +0100
Subject: [PATCH] feat(gas-reporter): allow gas-reporter to parse stdin (#1688)

---
 .changeset/silver-dolls-shave.md | 12 ++++++++
 packages/gas-report/ts/index.ts  | 49 ++++++++++++++++++++++++--------
 2 files changed, 49 insertions(+), 12 deletions(-)
 create mode 100644 .changeset/silver-dolls-shave.md

diff --git a/.changeset/silver-dolls-shave.md b/.changeset/silver-dolls-shave.md
new file mode 100644
index 0000000000..4cb44e330f
--- /dev/null
+++ b/.changeset/silver-dolls-shave.md
@@ -0,0 +1,12 @@
+---
+"@latticexyz/gas-report": minor
+---
+
+Allow the `gas-report` CLI to parse logs via `stdin`, so it can be used with custom test commands (e.g. `mud test`).
+
+Usage:
+
+```sh
+# replace `forge test -vvv` with the custom test command
+GAS_REPORTER_ENABLED=true forge test -vvv | pnpm gas-report --stdin
+```
diff --git a/packages/gas-report/ts/index.ts b/packages/gas-report/ts/index.ts
index 9fea453b2d..35b644f66c 100644
--- a/packages/gas-report/ts/index.ts
+++ b/packages/gas-report/ts/index.ts
@@ -31,6 +31,7 @@ type Options = {
   path: string[];
   save?: string;
   compare?: string;
+  stdin?: boolean;
 };
 
 type GasReportEntry = {
@@ -52,13 +53,17 @@ const commandModule: CommandModule<Options, Options> = {
     return yargs.options({
       save: { type: "string", desc: "Save the gas report to a file" },
       compare: { type: "string", desc: "Compare to an existing gas report" },
+      stdin: {
+        type: "boolean",
+        desc: "Parse the gas report logs from stdin instead of running an internal test command",
+      },
     });
   },
 
-  async handler({ save, compare }) {
+  async handler(options) {
     let gasReport: GasReport;
     try {
-      gasReport = await runGasReport();
+      gasReport = await runGasReport(options);
     } catch (error) {
       console.error(error);
       setTimeout(() => process.exit());
@@ -66,6 +71,7 @@ const commandModule: CommandModule<Options, Options> = {
     }
 
     // If this gas report should be compared to an existing one, load the existing one
+    let { compare } = options;
     if (compare) {
       try {
         const compareGasReport: GasReport = JSON.parse(readFileSync(compare, "utf8"));
@@ -84,7 +90,7 @@ const commandModule: CommandModule<Options, Options> = {
     printGasReport(gasReport, compare);
 
     // Save gas report to file if requested
-    if (save) saveGasReport(gasReport, save);
+    if (options.save) saveGasReport(gasReport, options.save);
 
     process.exit(0);
   },
@@ -92,19 +98,24 @@ const commandModule: CommandModule<Options, Options> = {
 
 export default commandModule;
 
-async function runGasReport(): Promise<GasReport> {
+async function runGasReport(options: Options): Promise<GasReport> {
   console.log("Running gas report");
   const gasReport: GasReport = [];
 
   // Extract the logs from the child process
-  let stdout: string;
+  let logs: string;
   try {
-    // Run the generated file using forge
-    const child = execa("forge", ["test", "-vvv"], {
-      stdio: ["inherit", "pipe", "inherit"],
-      env: { GAS_REPORTER_ENABLED: "true" },
-    });
-    stdout = (await child).stdout;
+    if (options.stdin) {
+      // Read the logs from stdin
+      logs = await readStdIn();
+    } else {
+      // Run the default test command to capture the logs
+      const child = execa("forge", ["test", "-vvv"], {
+        stdio: ["inherit", "pipe", "inherit"],
+        env: { GAS_REPORTER_ENABLED: "true" },
+      });
+      logs = (await child).stdout;
+    }
   } catch (error: any) {
     console.log(error.stdout ?? error);
     console.log(chalk.red("\n-----------\nError while running the gas report (see above)"));
@@ -112,7 +123,7 @@ async function runGasReport(): Promise<GasReport> {
   }
 
   // Extract the gas reports from the logs
-  const lines = stdout.split("\n").map(stripAnsi);
+  const lines = logs.split("\n").map(stripAnsi);
   const gasReportPattern = /^\s*GAS REPORT: (\d+) (.*)$/;
   const testFunctionPattern = /^\[(?:PASS|FAIL).*\] (\w+)\(\)/;
   const testFilePattern = /^Running \d+ tests? for (.*):(.*)$/;
@@ -191,3 +202,17 @@ function saveGasReport(gasReport: GasReport, path: string) {
   console.log(chalk.bold(`Saving gas report to ${path}`));
   writeFileSync(path, `${JSON.stringify(gasReport, null, 2)}\n`);
 }
+
+function readStdIn(): Promise<string> {
+  return new Promise((resolve) => {
+    let data = "";
+
+    process.stdin.on("data", (chunk) => {
+      data += chunk;
+    });
+
+    process.stdin.on("end", () => {
+      resolve(data);
+    });
+  });
+}