Adding relations to non edged bytecode chunks that appear in the debug trace

Author: fergarrui

src/main/java/net/nandgr/debugger/cfg/CFGCreatorDefault.java

@@ -4,6 +4,7 @@
 import net.nandgr.debugger.cfg.beans.BytecodeSection;
 import net.nandgr.debugger.cfg.beans.ContractBytecode;
 import net.nandgr.debugger.cfg.beans.OpcodeSource;
+import net.nandgr.debugger.node.response.json.DebugTraceTransactionLog;
 import net.nandgr.eth.Opcode;
 import net.nandgr.eth.Opcodes;
 import java.math.BigInteger;
@@ -20,10 +21,11 @@ public class CFGCreatorDefault {
             Opcodes.JUMP,
             Opcodes.STOP,
             Opcodes.REVERT,
-            Opcodes.RETURN
+            Opcodes.RETURN,
+            Opcodes.INVALID
     ).collect(Collectors.toList());
 
-    public ContractBytecode createContractBytecode(List<OpcodeSource> contractOpcodes) {
+    public ContractBytecode createContractBytecode(List<OpcodeSource> contractOpcodes, Map<Integer, DebugTraceTransactionLog> trace) {
         BigInteger codeOffset = BigInteger.valueOf(0);
         BytecodeSection constructorSection = null;
         boolean constructorFound = false;
@@ -52,11 +54,51 @@ public ContractBytecode createContractBytecode(List<OpcodeSource> contractOpcode
 
         Map<Integer, BytecodeChunk> functionsChunks = splitInChunks(functionsOpcodes);
         createRelations(functionsChunks);
+        checkChunksWithTrace(functionsChunks, trace);
         BytecodeSection functionsSection = new BytecodeSection(functionsChunks);
 
         return new ContractBytecode(constructorSection, functionsSection);
     }
 
+    private void checkChunksWithTrace(Map<Integer, BytecodeChunk> functionsChunks, Map<Integer, DebugTraceTransactionLog> trace) {
+        for (Map.Entry<Integer, BytecodeChunk> entry : functionsChunks.entrySet()) {
+            BytecodeChunk chunk = entry.getValue();
+            OpcodeSource lastOpcode = chunk.getOpcodes().get(chunk.getOpcodes().size() - 1);
+
+            if (trace.containsKey(lastOpcode.getOffset()) && chunk.hasEmptyRelations() && isJumpOpcode(lastOpcode)) {
+                DebugTraceTransactionLog jumpTrace = trace.get(lastOpcode.getOffset());
+                Integer jumpDestination = Integer.valueOf(jumpTrace.getStack().get(jumpTrace.getStack().size() - 1), 16);
+                BytecodeChunk destChunk;
+                if (functionsChunks.containsKey(jumpDestination)) {
+                    destChunk = functionsChunks.get(jumpDestination);
+                } else {
+                    destChunk = searchDestinationChunk(functionsChunks, jumpDestination);
+                }
+                chunk.setBranchA(destChunk);
+            }
+        }
+    }
+
+    private boolean isJumpOpcode(OpcodeSource lastOpcode) {
+        return isJump(lastOpcode) || lastOpcode.getOpcode().equals(Opcodes.JUMPI);
+    }
+
+    private boolean isJump(OpcodeSource lastOpcode) {
+        return lastOpcode.getOpcode().equals(Opcodes.JUMP);
+    }
+
+    private BytecodeChunk searchDestinationChunk(Map<Integer, BytecodeChunk> functionsChunks, Integer jumpDestination) {
+        for (Map.Entry<Integer, BytecodeChunk> entry : functionsChunks.entrySet()) {
+            BytecodeChunk chunk = entry.getValue();
+            for (OpcodeSource opcodeSource : chunk.getOpcodes()) {
+                if (opcodeSource.getOffset() == jumpDestination) {
+                    return chunk;
+                }
+            }
+        }
+        return null;
+    }
+
     private static void adjustOffsets(List<OpcodeSource> functionsOpcodes, int offset) {
         for (Opcode opcode : functionsOpcodes) {
             opcode.setOffset(opcode.getOffset() - offset);

src/main/java/net/nandgr/debugger/transformers/SolidityTransformer.java

@@ -124,7 +124,7 @@ public List<ContractObject> loadContracts(String sourceCodeFile) throws Transfor
 
             CFGCreatorDefault cfgCreatorDefault = new CFGCreatorDefault();
 
-            ContractBytecode contractBytecode = cfgCreatorDefault.createContractBytecode(opcodeSources);
+            ContractBytecode contractBytecode = cfgCreatorDefault.createContractBytecode(opcodeSources, traceDataContract);
             Map<Integer, BytecodeChunk> runtimeChunks = contractBytecode.getRuntime().getChunks();
 
             GraphVizCreator graphVizCreator = new GraphVizCreator(runtimeChunks, traceDataContract, cName);