Account for inserted/cloned statements when mapping pre-transformed s…

…tatements

logic/decompiler_output.dl

@@ -319,20 +319,40 @@ TAC_Block_Head(block, irstmt) :-
 
 
 // unmodified statements
-TAC_Statement_OriginalStatementList(irstmt, [stmt, nil]),
-TAC_Statement_OriginalStatement(irstmt, stmt) :-
+
+// [Old]
+// YB: Does not account for cloned statements
+// TAC_Statement_OriginalStatementList(irstmt, [stmt, nil]),
+// TAC_Statement_OriginalStatement(irstmt, stmt) :-
+//   Statement_IRStatement(stmt, _, irstmt),
+//   preTrans.Statement_Block(stmt, _).
+
+TAC_Statement_OriginalStatementList(irstmt, [ogStmt, nil]),
+TAC_Statement_OriginalStatement(irstmt, ogStmt) :-
   Statement_IRStatement(stmt, _, irstmt),
-  preTrans.Statement_Block(stmt, _).
+  preTrans.Statement_Block(stmt, _),
+  PreTransStatement_OriginalStatement(stmt, ogStmt). //this predicate takes care of block cloning as well
+
 
 // Inserted statements that didn't exist before
-TAC_Statement_OriginalStatementList(irstmt, [stmt_nexthead, nil]),
-TAC_Statement_OriginalStatement(irstmt, stmt_nexthead) :-
+
+// [Old]
+// YB: Sometimes maps to a statement that is not present in the original bytecode
+// TAC_Statement_OriginalStatementList(irstmt, [stmt_nexthead, nil]),
+// TAC_Statement_OriginalStatement(irstmt, stmt_nexthead) :-
+//   Statement_IRStatement(stmt, _, irstmt),
+//   !preTrans.Statement_Block(stmt, _), // didn't exist
+//   IRStatement_Block(irstmt, irblock),
+//   LocalBlockEdge(irblock, irblock_next),
+//   Block_IRBlock(block_next, _, irblock_next),
+//   stmt_nexthead = as(block_next, Statement).
+
+// recover the insertion in the original bytecode (regardless of whether the original statement had been removed )
+TAC_Statement_OriginalStatementList(irstmt, [ogStmt, nil]),
+TAC_Statement_OriginalStatement(irstmt, ogStmt) :-
   Statement_IRStatement(stmt, _, irstmt),
   !preTrans.Statement_Block(stmt, _), // didn't exist
-  IRStatement_Block(irstmt, irblock),
-  LocalBlockEdge(irblock, irblock_next),
-  Block_IRBlock(block_next, _, irblock_next),
-  stmt_nexthead = as(block_next, Statement).
+  insertor.InsertedOpNewStatement(ogStmt, _, stmt).
 
 .output TAC_Statement_OriginalStatement, TAC_Statement_OriginalStatementList
 

logic/local.dl

@@ -779,6 +779,16 @@ ByteCodeHex(substr(bytecode, 2, strlen(bytecode))):-
  ***********
  */
 
+// Map each pre-tranformed statement to the corresponding statement in the original bytecode 
+// (accounting for all of the early block-cloning rounds)
+.decl PreTransStatement_OriginalStatement(stmt: Statement, ogStmt: Statement)
+DEBUG_OUTPUT(PreTransStatement_OriginalStatement)
+PreTransStatement_OriginalStatement(ogStmt, ogStmt) :-
+  factReader.Statement_Opcode(ogStmt, _).
+
+PreTransStatement_OriginalStatement(rvtClStmt, ogStmt) :-
+  revertCloner.StatementToClonedStatement(_, _, ogStmt, rvtClStmt).
+
 .init preTrans = PreTransLocalAnalysis
 
 .init revertCloner = RevertBlockCloner
@@ -791,6 +801,14 @@ COPY_OUTPUT(blockCloner, revertCloner)
 blockCloner.Prev_Block_OriginalBlock(block, originalBlock):- revertCloner.Block_OriginalBlock(block, originalBlock).
 COPY_OUTPUT(preTrans, blockCloner)
 
+PreTransStatement_OriginalStatement(heurClStmt, ogStmt) :-
+  blockCloner.StatementToClonedStatement(_, _, rvtClStmt, heurClStmt),
+  revertCloner.StatementToClonedStatement(_, _, ogStmt, rvtClStmt).
+
+PreTransStatement_OriginalStatement(heurClStmt, ogStmt) :-
+  blockCloner.StatementToClonedStatement(_, _, ogStmt, heurClStmt),
+  !revertCloner.StatementToClonedStatement(_, _, _, ogStmt).
+
 #else
 
 COPY_OUTPUT(preTrans, revertCloner)