Support all opcodes in a trace parser

4 files changed, 45 insertions, 55 deletions

diff --git a/packages/sol-cov/src/trace.ts b/packages/sol-cov/src/trace.ts
index 6bc28989d..febb1034e 100644
--- a/packages/sol-cov/src/trace.ts
+++ b/packages/sol-cov/src/trace.ts
@@ -1,5 +1,5 @@
-import { StructLog, TransactionTrace } from '@0xproject/types';
-import { BigNumber } from '@0xproject/utils';
+import { OpCode, StructLog, TransactionTrace } from '@0xproject/types';
+import { BigNumber, logUtils } from '@0xproject/utils';
 import { addHexPrefix, stripHexPrefix } from 'ethereumjs-util';
 import * as fs from 'fs';
 import * as _ from 'lodash';
@@ -26,68 +26,44 @@ export function getTracesByContractAddress(structLogs: StructLog[], startAddress
         // That means that we can always safely pop from it
         currentTraceSegment.push(structLog);
 
-        if (structLog.op === 'DELEGATECALL') {
+        if (
+            structLog.op === OpCode.CallCode ||
+            structLog.op === OpCode.StaticCall ||
+            structLog.op === OpCode.Call ||
+            structLog.op === OpCode.DelegateCall
+        ) {
             const currentAddress = _.last(callStack) as string;
-            const jumpAddressOffset = 4;
+            const jumpAddressOffset = structLog.op === OpCode.DelegateCall ? 4 : 2;
             const newAddress = padZeros(new BigNumber(addHexPrefix(structLog.stack[jumpAddressOffset])).toString(16));
             callStack.push(newAddress);
             traceByContractAddress[currentAddress] = (traceByContractAddress[currentAddress] || []).concat(
                 currentTraceSegment,
             );
             currentTraceSegment = [];
-        } else if (structLog.op === 'RETURN') {
+        } else if (
+            structLog.op === OpCode.Return ||
+            structLog.op === OpCode.Stop ||
+            structLog.op === OpCode.Revert ||
+            structLog.op === OpCode.Invalid ||
+            structLog.op === OpCode.SelfDestruct
+        ) {
             const currentAddress = callStack.pop() as string;
             traceByContractAddress[currentAddress] = (traceByContractAddress[currentAddress] || []).concat(
                 currentTraceSegment,
             );
             currentTraceSegment = [];
-        } else if (structLog.op === 'STOP') {
-            const currentAddress = callStack.pop() as string;
-            traceByContractAddress[currentAddress] = (traceByContractAddress[currentAddress] || []).concat(
-                currentTraceSegment,
-            );
-            currentTraceSegment = [];
-            if (i !== structLogs.length - 1) {
-                throw new Error('Malformed trace. STOP is not the last opcode executed');
+            if (structLog.op === OpCode.SelfDestruct) {
+                // TODO: Record contract bytecode before the selfdestruct and handle that scenario
+                logUtils.warn(
+                    "Detected a selfdestruct. Sol-cov currently doesn't support that scenario. We'll just skip the trace part for a destructed contract",
+                );
             }
-        } else if (structLog.op === 'CREATE') {
-            console.warn(
+        } else if (structLog.op === OpCode.Create) {
+            // TODO: Extract the new contract address from the stack and handle that scenario
+            logUtils.warn(
                 "Detected a contract created from within another contract. Sol-cov currently doesn't support that scenario. We'll just skip that trace",
             );
             return traceByContractAddress;
-        } else if (structLog.op === 'CALL') {
-            const currentAddress = _.last(callStack) as string;
-            const jumpAddressOffset = 2;
-            const newAddress = padZeros(new BigNumber(addHexPrefix(structLog.stack[jumpAddressOffset])).toString(16));
-            callStack.push(newAddress);
-            traceByContractAddress[currentAddress] = (traceByContractAddress[currentAddress] || []).concat(
-                currentTraceSegment,
-            );
-            currentTraceSegment = [];
-        } else if (structLog.op === 'CALLCODE') {
-            throw new Error('CALLCODE opcode unsupported by coverage');
-        } else if (structLog.op === 'STATICCALL') {
-            throw new Error('STATICCALL opcode unsupported by coverage');
-        } else if (structLog.op === 'REVERT') {
-            const currentAddress = callStack.pop() as string;
-            traceByContractAddress[currentAddress] = (traceByContractAddress[currentAddress] || []).concat(
-                currentTraceSegment,
-            );
-            currentTraceSegment = [];
-            if (i !== structLogs.length - 1) {
-                throw new Error('Malformed trace. REVERT is not the last opcode executed');
-            }
-        } else if (structLog.op === 'INVALID') {
-            const currentAddress = callStack.pop() as string;
-            traceByContractAddress[currentAddress] = (traceByContractAddress[currentAddress] || []).concat(
-                currentTraceSegment,
-            );
-            currentTraceSegment = [];
-            if (i !== structLogs.length - 1) {
-                throw new Error('Malformed trace. INVALID is not the last opcode executed');
-            }
-        } else if (structLog.op === 'SELFDESTRUCT') {
-            throw new Error('SELFDESTRUCT opcode unsupported by coverage');
         }
     }
     if (callStack.length !== 0) {
diff --git a/packages/sol-cov/test/trace_test.ts b/packages/sol-cov/test/trace_test.ts
index b9d846732..58b9203b0 100644
--- a/packages/sol-cov/test/trace_test.ts
+++ b/packages/sol-cov/test/trace_test.ts
@@ -1,4 +1,4 @@
-import { StructLog } from '@0xproject/types';
+import { OpCode, StructLog } from '@0xproject/types';
 import * as chai from 'chai';
 import * as fs from 'fs';
 import * as _ from 'lodash';
@@ -15,13 +15,13 @@ const DEFAULT_STRUCT_LOG: StructLog = {
     gas: 0,
     gasCost: 0,
     memory: [],
-    op: 'DEFAULT',
+    op: OpCode.Invalid,
     pc: 0,
     stack: [],
     storage: {},
 };
 
-function addDefaultStructLogFields(compactStructLog: Partial<StructLog> & { op: string; depth: number }): StructLog {
+function addDefaultStructLogFields(compactStructLog: Partial<StructLog> & { op: OpCode; depth: number }): StructLog {
     return { ...DEFAULT_STRUCT_LOG, ...compactStructLog };
 }
 
@@ -31,16 +31,16 @@ describe('Trace', () => {
             const delegateCallAddress = '0x0000000000000000000000000000000000000002';
             const trace = [
                 {
-                    op: 'DELEGATECALL',
+                    op: OpCode.DelegateCall,
                     stack: ['0x', '0x', delegateCallAddress],
                     depth: 0,
                 },
                 {
-                    op: 'RETURN',
+                    op: OpCode.Return,
                     depth: 1,
                 },
                 {
-                    op: 'RETURN',
+                    op: OpCode.Return,
                     depth: 0,
                 },
             ];
diff --git a/packages/types/src/index.ts b/packages/types/src/index.ts
index f28c2f9a3..055c47e0a 100644
--- a/packages/types/src/index.ts
+++ b/packages/types/src/index.ts
@@ -58,7 +58,18 @@ export interface DataItem {
     components?: DataItem[];
 }
 
-export type OpCode = string;
+export enum OpCode {
+    DelegateCall = 'DELEGATECALL',
+    Revert = 'REVERT',
+    Create = 'CREATE',
+    Stop = 'STOP',
+    Invalid = 'INVALID',
+    CallCode = 'CALLCODE',
+    StaticCall = 'STATICCALL',
+    Return = 'RETURN',
+    Call = 'CALL',
+    SelfDestruct = 'SELFDESTRUCT',
+}
 
 export interface StructLog {
     depth: number;
diff --git a/packages/utils/src/log_utils.ts b/packages/utils/src/log_utils.ts
index d0f0e34c9..87f8479b5 100644
--- a/packages/utils/src/log_utils.ts
+++ b/packages/utils/src/log_utils.ts
@@ -2,4 +2,7 @@ export const logUtils = {
     log(...args: any[]): void {
         console.log(...args); // tslint:disable-line:no-console
     },
+    warn(...args: any[]): void {
+        console.warn(...args); // tslint:disable-line:no-console
+    },
 };