Refactor sol-cov to avoid keeping traceInfo in memory

5 files changed, 76 insertions, 92 deletions

diff --git a/packages/sol-cov/src/coverage_manager.ts b/packages/sol-cov/src/coverage_manager.ts
index 3ab363b52..a55967243 100644
--- a/packages/sol-cov/src/coverage_manager.ts
+++ b/packages/sol-cov/src/coverage_manager.ts
@@ -36,7 +36,8 @@ const mkdirpAsync = promisify<undefined>(mkdirp);
 export class CoverageManager {
     private _artifactAdapter: AbstractArtifactAdapter;
     private _logger: Logger;
-    private _traceInfos: TraceInfo[] = [];
+    private _contractsData!: ContractData[];
+    private _collector = new Collector();
     /**
      * Computed partial coverage for a single file & subtrace
      * @param contractData      Contract metadata (source, srcMap, bytecode)
@@ -130,49 +131,39 @@ export class CoverageManager {
         this._logger = getLogger('sol-cov');
         this._logger.setLevel(isVerbose ? levels.TRACE : levels.ERROR);
     }
-    public appendTraceInfo(traceInfo: TraceInfo): void {
-        this._traceInfos.push(traceInfo);
-    }
     public async writeCoverageAsync(): Promise<void> {
-        const finalCoverage = await this._computeCoverageAsync();
+        const finalCoverage = this._collector.getFinalCoverage();
         const stringifiedCoverage = JSON.stringify(finalCoverage, null, '\t');
         await mkdirpAsync('coverage');
         fs.writeFileSync('coverage/coverage.json', stringifiedCoverage);
     }
-    private async _computeCoverageAsync(): Promise<Coverage> {
-        const contractsData = await this._artifactAdapter.collectContractsDataAsync();
-        const collector = new Collector();
-        for (const traceInfo of this._traceInfos) {
-            const isContractCreation = traceInfo.address === constants.NEW_CONTRACT;
-            const bytecode = isContractCreation
-                ? (traceInfo as TraceInfoNewContract).bytecode
-                : (traceInfo as TraceInfoExistingContract).runtimeBytecode;
-            const contractData = utils.getContractDataIfExists(contractsData, bytecode);
-            if (_.isUndefined(contractData)) {
-                const errMsg = isContractCreation
-                    ? `Unknown contract creation transaction`
-                    : `Transaction to an unknown address: ${traceInfo.address}`;
-                this._logger.warn(errMsg);
-                continue;
-            }
-            const bytecodeHex = stripHexPrefix(bytecode);
-            const sourceMap = isContractCreation ? contractData.sourceMap : contractData.sourceMapRuntime;
-            const pcToSourceRange = parseSourceMap(
-                contractData.sourceCodes,
-                sourceMap,
-                bytecodeHex,
-                contractData.sources,
+    public async computeCoverageAsync(traceInfo: TraceInfo): Promise<void> {
+        if (_.isUndefined(this._contractsData)) {
+            this._contractsData = await this._artifactAdapter.collectContractsDataAsync();
+        }
+        const isContractCreation = traceInfo.address === constants.NEW_CONTRACT;
+        const bytecode = isContractCreation
+            ? (traceInfo as TraceInfoNewContract).bytecode
+            : (traceInfo as TraceInfoExistingContract).runtimeBytecode;
+        const contractData = utils.getContractDataIfExists(this._contractsData, bytecode);
+        if (_.isUndefined(contractData)) {
+            const errMsg = isContractCreation
+                ? `Unknown contract creation transaction`
+                : `Transaction to an unknown address: ${traceInfo.address}`;
+            this._logger.warn(errMsg);
+            return;
+        }
+        const bytecodeHex = stripHexPrefix(bytecode);
+        const sourceMap = isContractCreation ? contractData.sourceMap : contractData.sourceMapRuntime;
+        const pcToSourceRange = parseSourceMap(contractData.sourceCodes, sourceMap, bytecodeHex, contractData.sources);
+        for (let fileIndex = 0; fileIndex < contractData.sources.length; fileIndex++) {
+            const singleFileCoverageForTrace = CoverageManager._getSingleFileCoverageForSubtrace(
+                contractData,
+                traceInfo.subtrace,
+                pcToSourceRange,
+                fileIndex,
             );
-            for (let fileIndex = 0; fileIndex < contractData.sources.length; fileIndex++) {
-                const singleFileCoverageForTrace = CoverageManager._getSingleFileCoverageForSubtrace(
-                    contractData,
-                    traceInfo.subtrace,
-                    pcToSourceRange,
-                    fileIndex,
-                );
-                collector.add(singleFileCoverageForTrace);
-            }
+            this._collector.add(singleFileCoverageForTrace);
         }
-        return collector.getFinalCoverage();
     }
 }
diff --git a/packages/sol-cov/src/coverage_subprovider.ts b/packages/sol-cov/src/coverage_subprovider.ts
index 174b7c6ac..1dcbfcbfe 100644
--- a/packages/sol-cov/src/coverage_subprovider.ts
+++ b/packages/sol-cov/src/coverage_subprovider.ts
@@ -3,6 +3,7 @@ import * as _ from 'lodash';
 import { AbstractArtifactAdapter } from './artifact_adapters/abstract_artifact_adapter';
 import { CoverageManager } from './coverage_manager';
 import { TraceCollectionSubprovider } from './trace_collection_subprovider';
+import { TraceInfo } from './types';
 
 /**
  * This class implements the [web3-provider-engine](https://github.com/MetaMask/provider-engine) subprovider interface.
@@ -25,12 +26,13 @@ export class CoverageSubprovider extends TraceCollectionSubprovider {
         super(defaultFromAddress, traceCollectionSubproviderConfig);
         this._coverageManager = new CoverageManager(artifactAdapter, isVerbose);
     }
+    public async handleTraceInfoAsync(traceInfo: TraceInfo): Promise<void> {
+        return this._coverageManager.computeCoverageAsync(traceInfo);
+    }
     /**
      * Write the test coverage results to a file in Istanbul format.
      */
     public async writeCoverageAsync(): Promise<void> {
-        const traceInfos = this.getCollectedTraceInfos();
-        _.forEach(traceInfos, traceInfo => this._coverageManager.appendTraceInfo(traceInfo));
-        await this._coverageManager.writeCoverageAsync();
+        return this._coverageManager.writeCoverageAsync();
     }
 }
diff --git a/packages/sol-cov/src/profiler_manager.ts b/packages/sol-cov/src/profiler_manager.ts
index 0ab0ea544..bec92f424 100644
--- a/packages/sol-cov/src/profiler_manager.ts
+++ b/packages/sol-cov/src/profiler_manager.ts
@@ -31,7 +31,8 @@ const mkdirpAsync = promisify<undefined>(mkdirp);
 export class ProfilerManager {
     private _artifactAdapter: AbstractArtifactAdapter;
     private _logger: Logger;
-    private _traceInfos: TraceInfo[] = [];
+    private _contractsData!: ContractData[];
+    private _collector = new Collector();
     /**
      * Computed partial coverage for a single file & subtrace
      * @param contractData      Contract metadata (source, srcMap, bytecode)
@@ -86,49 +87,39 @@ export class ProfilerManager {
         this._logger = getLogger('sol-cov');
         this._logger.setLevel(isVerbose ? levels.TRACE : levels.ERROR);
     }
-    public appendTraceInfo(traceInfo: TraceInfo): void {
-        this._traceInfos.push(traceInfo);
-    }
     public async writeProfilerOutputAsync(): Promise<void> {
-        const finalCoverage = await this._computeCoverageAsync();
+        const finalCoverage = this._collector.getFinalCoverage();
         const stringifiedCoverage = JSON.stringify(finalCoverage, null, '\t');
         await mkdirpAsync('coverage');
         fs.writeFileSync('coverage/coverage.json', stringifiedCoverage);
     }
-    private async _computeCoverageAsync(): Promise<Coverage> {
-        const contractsData = await this._artifactAdapter.collectContractsDataAsync();
-        const collector = new Collector();
-        for (const traceInfo of this._traceInfos) {
-            const isContractCreation = traceInfo.address === constants.NEW_CONTRACT;
-            const bytecode = isContractCreation
-                ? (traceInfo as TraceInfoNewContract).bytecode
-                : (traceInfo as TraceInfoExistingContract).runtimeBytecode;
-            const contractData = utils.getContractDataIfExists(contractsData, bytecode);
-            if (_.isUndefined(contractData)) {
-                const errMsg = isContractCreation
-                    ? `Unknown contract creation transaction`
-                    : `Transaction to an unknown address: ${traceInfo.address}`;
-                this._logger.warn(errMsg);
-                continue;
-            }
-            const bytecodeHex = stripHexPrefix(bytecode);
-            const sourceMap = isContractCreation ? contractData.sourceMap : contractData.sourceMapRuntime;
-            const pcToSourceRange = parseSourceMap(
-                contractData.sourceCodes,
-                sourceMap,
-                bytecodeHex,
-                contractData.sources,
+    public async computeCoverageAsync(traceInfo: TraceInfo): Promise<void> {
+        if (_.isUndefined(this._contractsData)) {
+            this._contractsData = await this._artifactAdapter.collectContractsDataAsync();
+        }
+        const isContractCreation = traceInfo.address === constants.NEW_CONTRACT;
+        const bytecode = isContractCreation
+            ? (traceInfo as TraceInfoNewContract).bytecode
+            : (traceInfo as TraceInfoExistingContract).runtimeBytecode;
+        const contractData = utils.getContractDataIfExists(this._contractsData, bytecode);
+        if (_.isUndefined(contractData)) {
+            const errMsg = isContractCreation
+                ? `Unknown contract creation transaction`
+                : `Transaction to an unknown address: ${traceInfo.address}`;
+            this._logger.warn(errMsg);
+            return;
+        }
+        const bytecodeHex = stripHexPrefix(bytecode);
+        const sourceMap = isContractCreation ? contractData.sourceMap : contractData.sourceMapRuntime;
+        const pcToSourceRange = parseSourceMap(contractData.sourceCodes, sourceMap, bytecodeHex, contractData.sources);
+        for (let fileIndex = 0; fileIndex < contractData.sources.length; fileIndex++) {
+            const singleFileCoverageForTrace = ProfilerManager._getSingleFileCoverageForSubtrace(
+                contractData,
+                traceInfo.subtrace,
+                pcToSourceRange,
+                fileIndex,
             );
-            for (let fileIndex = 0; fileIndex < contractData.sources.length; fileIndex++) {
-                const singleFileCoverageForTrace = ProfilerManager._getSingleFileCoverageForSubtrace(
-                    contractData,
-                    traceInfo.subtrace,
-                    pcToSourceRange,
-                    fileIndex,
-                );
-                collector.add(singleFileCoverageForTrace);
-            }
+            this._collector.add(singleFileCoverageForTrace);
         }
-        return collector.getFinalCoverage();
     }
 }
diff --git a/packages/sol-cov/src/profiler_subprovider.ts b/packages/sol-cov/src/profiler_subprovider.ts
index ac878c070..c66cd1dec 100644
--- a/packages/sol-cov/src/profiler_subprovider.ts
+++ b/packages/sol-cov/src/profiler_subprovider.ts
@@ -3,6 +3,7 @@ import * as _ from 'lodash';
 import { AbstractArtifactAdapter } from './artifact_adapters/abstract_artifact_adapter';
 import { ProfilerManager } from './profiler_manager';
 import { TraceCollectionSubprovider } from './trace_collection_subprovider';
+import { TraceInfo } from './types';
 
 /**
  * This class implements the [web3-provider-engine](https://github.com/MetaMask/provider-engine) subprovider interface.
@@ -25,12 +26,13 @@ export class ProfilerSubprovider extends TraceCollectionSubprovider {
         super(defaultFromAddress, traceCollectionSubproviderConfig);
         this._profilerManager = new ProfilerManager(artifactAdapter, isVerbose);
     }
+    public async handleTraceInfoAsync(traceInfo: TraceInfo): Promise<void> {
+        return this._profilerManager.computeCoverageAsync(traceInfo);
+    }
     /**
      * Write the test profiler results to a file in Istanbul format.
      */
     public async writeProfilerOutputAsync(): Promise<void> {
-        const traceInfos = this.getCollectedTraceInfos();
-        _.forEach(traceInfos, traceInfo => this._profilerManager.appendTraceInfo(traceInfo));
-        await this._profilerManager.writeProfilerOutputAsync();
+        return this._profilerManager.writeProfilerOutputAsync();
     }
 }
diff --git a/packages/sol-cov/src/trace_collection_subprovider.ts b/packages/sol-cov/src/trace_collection_subprovider.ts
index fe73546e8..5e90e3907 100644
--- a/packages/sol-cov/src/trace_collection_subprovider.ts
+++ b/packages/sol-cov/src/trace_collection_subprovider.ts
@@ -29,12 +29,11 @@ export interface TraceCollectionSubproviderConfig {
  * This class implements the [web3-provider-engine](https://github.com/MetaMask/provider-engine) subprovider interface.
  * It collects traces of all transactions that were sent and all calls that were executed through JSON RPC.
  */
-export class TraceCollectionSubprovider extends Subprovider {
+export abstract class TraceCollectionSubprovider extends Subprovider {
     // Lock is used to not accept normal transactions while doing call/snapshot magic because they'll be reverted later otherwise
     private _lock = new Lock();
     private _defaultFromAddress: string;
     private _web3Wrapper!: Web3Wrapper;
-    private _traceInfos: TraceInfo[] = [];
     private _isEnabled = true;
     private _config: TraceCollectionSubproviderConfig;
     /**
@@ -47,12 +46,6 @@ export class TraceCollectionSubprovider extends Subprovider {
         this._config = config;
     }
     /**
-     * Returns all trace infos collected by the subprovider so far
-     */
-    public getCollectedTraceInfos(): TraceInfo[] {
-        return this._traceInfos;
-    }
-    /**
      * Starts trace collection
      */
     public start(): void {
@@ -65,6 +58,11 @@ export class TraceCollectionSubprovider extends Subprovider {
         this._isEnabled = false;
     }
     /**
+     * Called for each subtrace.
+     * @param traceInfo Trace info for this subtrace.
+     */
+    public abstract handleTraceInfoAsync(traceInfo: TraceInfo): Promise<void>;
+    /**
      * This method conforms to the web3-provider-engine interface.
      * It is called internally by the ProviderEngine when it is this subproviders
      * turn to handle a JSON RPC request.
@@ -192,7 +190,7 @@ export class TraceCollectionSubprovider extends Subprovider {
                         runtimeBytecode,
                     };
                 }
-                this._traceInfos.push(traceInfo);
+                this.handleTraceInfoAsync(traceInfo);
             }
         } else {
             for (const subcallAddress of subcallAddresses) {
@@ -204,7 +202,7 @@ export class TraceCollectionSubprovider extends Subprovider {
                     address: subcallAddress,
                     runtimeBytecode,
                 };
-                this._traceInfos.push(traceInfo);
+                this.handleTraceInfoAsync(traceInfo);
             }
         }
     }