Stop making an assumption that contract code is immutable

3 files changed, 73 insertions, 52 deletions

diff --git a/packages/sol-cov/src/coverage_manager.ts b/packages/sol-cov/src/coverage_manager.ts
index aced9208f..5d11800f9 100644
--- a/packages/sol-cov/src/coverage_manager.ts
+++ b/packages/sol-cov/src/coverage_manager.ts
@@ -23,6 +23,8 @@ import {
     StatementDescription,
     StatementMap,
     TraceInfo,
+    TraceInfoExistingContract,
+    TraceInfoNewContract,
 } from './types';
 import { utils } from './utils';
 
@@ -72,9 +74,8 @@ function getSingleFileCoverageForTrace(
 }
 
 export class CoverageManager {
-    private _traceInfoByAddress: { [address: string]: TraceInfo[] } = {};
+    private _traceInfos: TraceInfo[] = [];
     private _contractsData: ContractData[] = [];
-    private _txDataByHash: { [txHash: string]: string } = {};
     private _getContractCodeAsync: (address: string) => Promise<string>;
     constructor(
         artifactsPath: string,
@@ -85,14 +86,8 @@ export class CoverageManager {
         this._getContractCodeAsync = getContractCodeAsync;
         this._contractsData = collectContractsData(artifactsPath, sourcesPath, networkId);
     }
-    public setTxDataByHash(txHash: string, data: string): void {
-        this._txDataByHash[txHash] = data;
-    }
-    public appendTraceInfo(address: string, traceInfo: TraceInfo): void {
-        if (_.isUndefined(this._traceInfoByAddress[address])) {
-            this._traceInfoByAddress[address] = [];
-        }
-        this._traceInfoByAddress[address].push(traceInfo);
+    public appendTraceInfo(traceInfo: TraceInfo): void {
+        this._traceInfos.push(traceInfo);
     }
     public async writeCoverageAsync(): Promise<void> {
         const finalCoverage = await this._computeCoverageAsync();
@@ -103,13 +98,13 @@ export class CoverageManager {
     }
     private async _computeCoverageAsync(): Promise<Coverage> {
         const collector = new Collector();
-        for (const address of _.keys(this._traceInfoByAddress)) {
-            if (address !== constants.NEW_CONTRACT) {
+        for (const traceInfo of this._traceInfos) {
+            if (traceInfo.address !== constants.NEW_CONTRACT) {
                 // Runtime transaction
-                const runtimeBytecode = await this._getContractCodeAsync(address);
+                const runtimeBytecode = (traceInfo as TraceInfoExistingContract).runtimeBytecode;
                 const contractData = _.find(this._contractsData, { runtimeBytecode }) as ContractData;
                 if (_.isUndefined(contractData)) {
-                    throw new Error(`Transaction to an unknown address: ${address}`);
+                    throw new Error(`Transaction to an unknown address: ${traceInfo.address}`);
                 }
                 const bytecodeHex = contractData.runtimeBytecode.slice(2);
                 const sourceMap = contractData.sourceMapRuntime;
@@ -120,44 +115,40 @@ export class CoverageManager {
                     contractData.sources,
                 );
                 for (let fileIndex = 0; fileIndex < contractData.sources.length; fileIndex++) {
-                    _.forEach(this._traceInfoByAddress[address], (traceInfo: TraceInfo) => {
-                        const singleFileCoverageForTrace = getSingleFileCoverageForTrace(
-                            contractData,
-                            traceInfo.coveredPcs,
-                            pcToSourceRange,
-                            fileIndex,
-                        );
-                        collector.add(singleFileCoverageForTrace);
-                    });
+                    const singleFileCoverageForTrace = getSingleFileCoverageForTrace(
+                        contractData,
+                        traceInfo.coveredPcs,
+                        pcToSourceRange,
+                        fileIndex,
+                    );
+                    collector.add(singleFileCoverageForTrace);
                 }
             } else {
                 // Contract creation transaction
-                _.forEach(this._traceInfoByAddress[address], (traceInfo: TraceInfo) => {
-                    const bytecode = this._txDataByHash[traceInfo.txHash];
-                    const contractData = _.find(this._contractsData, contractDataCandidate =>
-                        bytecode.startsWith(contractDataCandidate.bytecode),
-                    ) as ContractData;
-                    if (_.isUndefined(contractData)) {
-                        throw new Error(`Unknown contract creation transaction`);
-                    }
-                    const bytecodeHex = contractData.bytecode.slice(2);
-                    const sourceMap = contractData.sourceMap;
-                    const pcToSourceRange = parseSourceMap(
-                        contractData.sourceCodes,
-                        sourceMap,
-                        bytecodeHex,
-                        contractData.sources,
+                const bytecode = (traceInfo as TraceInfoNewContract).bytecode;
+                const contractData = _.find(this._contractsData, contractDataCandidate =>
+                    bytecode.startsWith(contractDataCandidate.bytecode),
+                ) as ContractData;
+                if (_.isUndefined(contractData)) {
+                    throw new Error(`Unknown contract creation transaction`);
+                }
+                const bytecodeHex = contractData.bytecode.slice(2);
+                const sourceMap = contractData.sourceMap;
+                const pcToSourceRange = parseSourceMap(
+                    contractData.sourceCodes,
+                    sourceMap,
+                    bytecodeHex,
+                    contractData.sources,
+                );
+                for (let fileIndex = 0; fileIndex < contractData.sources.length; fileIndex++) {
+                    const singleFileCoverageForTrace = getSingleFileCoverageForTrace(
+                        contractData,
+                        traceInfo.coveredPcs,
+                        pcToSourceRange,
+                        fileIndex,
                     );
-                    for (let fileIndex = 0; fileIndex < contractData.sources.length; fileIndex++) {
-                        const singleFileCoverageForTrace = getSingleFileCoverageForTrace(
-                            contractData,
-                            traceInfo.coveredPcs,
-                            pcToSourceRange,
-                            fileIndex,
-                        );
-                        collector.add(singleFileCoverageForTrace);
-                    }
-                });
+                    collector.add(singleFileCoverageForTrace);
+                }
             }
         }
         // TODO: Submit a PR to DT
diff --git a/packages/sol-cov/src/coverage_subprovider.ts b/packages/sol-cov/src/coverage_subprovider.ts
index ef425ee81..d3783abb2 100644
--- a/packages/sol-cov/src/coverage_subprovider.ts
+++ b/packages/sol-cov/src/coverage_subprovider.ts
@@ -5,6 +5,7 @@ import * as Web3 from 'web3';
 
 import { constants } from './constants';
 import { CoverageManager } from './coverage_manager';
+import { TraceInfoExistingContract, TraceInfoNewContract } from './types';
 
 /*
  * This class implements the web3-provider-engine subprovider interface and collects traces of all transactions that were sent and all calls that were executed.
@@ -84,15 +85,32 @@ export class CoverageSubprovider extends Subprovider {
         cb();
     }
     private async _recordTxTraceAsync(address: string, data: string | undefined, txHash: string): Promise<void> {
-        this._coverageManager.setTxDataByHash(txHash, data || '');
-        const payload = {
+        let payload = {
             method: 'debug_traceTransaction',
             params: [txHash, { disableMemory: true, disableStack: true, disableStorage: true }], // TODO For now testrpc just ignores those parameters https://github.com/trufflesuite/ganache-cli/issues/489
         };
         const jsonRPCResponsePayload = await this.emitPayloadAsync(payload);
         const trace: Web3.TransactionTrace = jsonRPCResponsePayload.result;
         const coveredPcs = _.map(trace.structLogs, log => log.pc);
-        this._coverageManager.appendTraceInfo(address, { coveredPcs, txHash });
+        if (address === constants.NEW_CONTRACT) {
+            const traceInfo: TraceInfoNewContract = {
+                coveredPcs,
+                txHash,
+                address,
+                bytecode: data as string,
+            };
+            this._coverageManager.appendTraceInfo(traceInfo);
+        } else {
+            payload = { method: 'eth_getCode', params: [address, 'latest'] };
+            const runtimeBytecode = (await this.emitPayloadAsync(payload)).result;
+            const traceInfo: TraceInfoExistingContract = {
+                coveredPcs,
+                txHash,
+                address,
+                runtimeBytecode,
+            };
+            this._coverageManager.appendTraceInfo(traceInfo);
+        }
     }
     private async _recordCallTraceAsync(callData: Partial<Web3.CallData>, blockNumber: Web3.BlockParam): Promise<void> {
         const snapshotId = Number((await this.emitPayloadAsync({ method: 'evm_snapshot' })).result);
diff --git a/packages/sol-cov/src/types.ts b/packages/sol-cov/src/types.ts
index 5d07cd01b..d6491100b 100644
--- a/packages/sol-cov/src/types.ts
+++ b/packages/sol-cov/src/types.ts
@@ -83,7 +83,19 @@ export interface ContractData {
     sources: string[];
 }
 
-export interface TraceInfo {
+export interface TraceInfoBase {
     coveredPcs: number[];
     txHash: string;
 }
+
+export interface TraceInfoNewContract extends TraceInfoBase {
+    address: 'NEW_CONTRACT';
+    bytecode: string;
+}
+
+export interface TraceInfoExistingContract extends TraceInfoBase {
+    address: string;
+    runtimeBytecode: string;
+}
+
+export type TraceInfo = TraceInfoNewContract | TraceInfoExistingContract;