perf: reduce the overhead of recording branchesnedbat/faster-branches

3 files changed, 47 insertions, 8 deletions

diff --git a/CHANGES.rst b/CHANGES.rst
index f2e2f1bc..81e02732 100644
--- a/CHANGES.rst
+++ b/CHANGES.rst
@@ -27,6 +27,9 @@ Unreleased
   file10.py will appear after file9.py.  This applies to file names, module
   names, environment variables, and test contexts.
 
+- Branch coverage measurement is faster, though you might only notice on
+  code that is executed many times, such as long-running loops.
+
 
 .. _changes_602:
 
diff --git a/coverage/collector.py b/coverage/collector.py
index 733b6f32..a72129e2 100644
--- a/coverage/collector.py
+++ b/coverage/collector.py
@@ -157,9 +157,11 @@ class Collector:
         if self._trace_class is CTracer:
             self.file_disposition_class = CFileDisposition
             self.supports_plugins = True
+            self.packed_arcs = True
         else:
             self.file_disposition_class = FileDisposition
             self.supports_plugins = False
+            self.packed_arcs = False
 
     def __repr__(self):
         return f"<Collector at 0x{id(self):x}: {self.tracer_name()}>"
@@ -437,7 +439,25 @@ class Collector:
             return False
 
         if self.branch:
-            self.covdata.add_arcs(self.mapped_file_dict(self.data))
+            if self.packed_arcs:
+                # Unpack the line number pairs packed into integers.  See
+                # tracer.c:CTracer_record_pair for the C code that creates
+                # these packed ints.
+                data = {}
+                for fname, packeds in self.data.items():
+                    tuples = []
+                    for packed in packeds:
+                        l1 = packed & 0xFFFFF
+                        l2 = (packed & (0xFFFFF << 20)) >> 20
+                        if packed & (1 << 40):
+                            l1 *= -1
+                        if packed & (1 << 41):
+                            l2 *= -1
+                        tuples.append((l1, l2))
+                    data[fname] = tuples
+            else:
+                data = self.data
+            self.covdata.add_arcs(self.mapped_file_dict(data))
         else:
             self.covdata.add_lines(self.mapped_file_dict(self.data))
 
diff --git a/coverage/ctracer/tracer.c b/coverage/ctracer/tracer.c
index 00d9f106..f508a85d 100644
--- a/coverage/ctracer/tracer.c
+++ b/coverage/ctracer/tracer.c
@@ -174,22 +174,38 @@ static int
 CTracer_record_pair(CTracer *self, int l1, int l2)
 {
     int ret = RET_ERROR;
-
-    PyObject * t = NULL;
-
-    t = Py_BuildValue("(ii)", l1, l2);
-    if (t == NULL) {
+    PyObject * packed_obj = NULL;
+    long packed = 0;
+
+    // Conceptually, data is a set of tuples (l1, l2), but that literally
+    // making a set of tuples would require us to construct a tuple just to
+    // see if we'd already recorded an arc.  On many-times-executed code,
+    // that would mean we construct a tuple, find the tuple is already in the
+    // set, then discard the tuple.  We can avoid that overhead by packing
+    // the two line numbers into one integer instead.
+    // See collector.py:flush_data for the Python code that unpacks this.
+    if (l1 < 0) {
+        packed |= (1L << 40);
+        l1 = -l1;
+    }
+    if (l2 < 0) {
+        packed |= (1L << 41);
+        l2 = -l2;
+    }
+    packed |= (((long)l2) << 20) + (long)l1;
+    packed_obj = PyLong_FromLong(packed);
+    if (packed_obj == NULL) {
         goto error;
     }
 
-    if (PySet_Add(self->pcur_entry->file_data, t) < 0) {
+    if (PySet_Add(self->pcur_entry->file_data, packed_obj) < 0) {
         goto error;
     }
 
     ret = RET_OK;
 
 error:
-    Py_XDECREF(t);
+    Py_XDECREF(packed_obj);
 
     return ret;
 }