Rename traces_diff -> generate_similarity_matrix

Rewrite the diff generation to generate a similarity
matrix instead

asm2vec/utils.py

@@ -26,78 +26,71 @@ class TraceData:
             trace. The tokens dictionary should be shared across all traces to
             ensure the same representation matches.
 
-        ID:
-            A representative identifier for the trace of a switch case. An ID
-            may point to more than one switch case, if the traces of those
-            switch cases are textually identical (ignoring constants).
+        fn_idx:
+            Index of the Function. This is necessary to index into the
+            training model embeddings. Different switch cases may share
+            the same fn_idx because they may be textually identical but
+            have different constants.
 
-        idx:
-            Index of the ID/Function. This is necessary to index into the
-            training model embeddings.
-
-        Opcode set:
+        Opcode Set/opcode_set:
             A set of opcodes that a single switch case covers.
 
         Pointer Opcode/ptr_opcode:
             A single opcode that represents the opcode set.
 
-        Constants vector:
+        Constants vector/constants_vector:
             Since the training process removes constants from the input text,
             the constants vector is a "signature" generated by looking at
-            constants used in the trace. These are correlated with ptr_opcodes
-            and not IDs since constants can differ among cases represented by
-            the same ID.
+            constants used in the trace.
     """
 
     def __init__(self, tokens):
         self.tokens = tokens
 
-        self.traces = dict()
+        self.__traces = dict()
         """
-        maps trace => ID.
+        Maps trace to trace idx
+        """
+
+        self.opcodes = dict()
+        """
+        maps ptr_opcode => { fn_idx, constants_vector, [opcodes...] }.
         See class docstring for more details.
         """
 
-        self.ids = dict()
-        """
-        maps ID => { idx, [ptr_opcodes...] }.
-        See class docstring for more details.
-        """
-
-        self.constants_vectors = dict()
-        """
-        maps ptr_opcode => constants_vector.
-        See class docstriing for more details.
-        """
-
-        self.opcode_sets = dict()
+        self.__opcode_sets = dict()
         """
         maps ptr_opcode => [opcodes...]
-        See class docstriing for more details.
         """
 
+    @property
+    def traces(self):
+        """
+        A list of parsed traces in the TraceData.
+
+        See class docstring for more details.
+        """
+        return list(self.__traces.keys())
+
     def __process_trace(self, ptr_opcode, text):
         fn = Function.load(text)
-        if fn in self.traces:
-            id = self.traces[fn]
-            self.ids[id]["ptr_opcodes"].append(ptr_opcode)
+        if fn in self.__traces:
+            fn_idx = self.__traces[fn]
         else:
-            # Use ptr_opcode as an ID
-            id = ptr_opcode
-
-            self.traces[fn] = id
+            fn_idx = len(self.__traces)
+            self.__traces[fn] = fn_idx
             self.tokens.add(fn.tokens())
-            self.ids[id] = {
-                "idx": len(self.ids),
-                "ptr_opcodes": [ptr_opcode],
-            }
 
-        self.constants_vectors[ptr_opcode] = self.__get_constants_vector(text)
+        self.opcodes[ptr_opcode] = {
+            "fn_idx": fn_idx,
+            "constants_vector": self.__get_constants_vector(text),
+            "opcodes": self.__opcode_sets[ptr_opcode],
+        }
 
     def __process_opcode_sets(self, opcode_sets_file):
         with open(opcode_sets_file) as f:
             data = json.load(f)
-            self.opcode_sets = {int(op): ops for op, ops in data.items()}
+            self.__opcode_sets = {int(op): ops for op, ops in data.items()}
 
     @staticmethod
     def __read_trace_from_file(f):
@@ -144,6 +137,9 @@ class TraceData:
                 filenames += [Path(path)]
 
         trace_data = TraceData(tokens)
+        # Process opcode_sets.json first, save traces for later
+
+        trace_files = dict()
         for filepath in filenames:
             filename = os.path.basename(filepath)
             file_split = os.path.splitext(filename)
@@ -152,9 +148,12 @@ class TraceData:
                 trace_data.__process_opcode_sets(filepath)
             elif file_ext == ".asm":
                 ptr_opcode = int(file_split[0], base=16)
-                with open(filepath) as f:
-                    text = trace_data.__read_trace_from_file(f)
-                    trace_data.__process_trace(ptr_opcode, text)
+                trace_files[ptr_opcode] = filepath
+
+        for ptr_opcode, trace_file in trace_files.items():
+            with open(trace_file) as f:
+                text = trace_data.__read_trace_from_file(f)
+                trace_data.__process_trace(ptr_opcode, text)
 
         return trace_data
 
@@ -200,7 +199,7 @@ def train(
     learning_rate=0.02,
 ):
     """Trains the model on the provided trace data."""
-    functions = trace_data.traces.keys()
+    functions = trace_data.traces
     tokens = trace_data.tokens
 
     if mode == "train":