Add more benchmark functionality

Programs/Source/benchmark_priority_queue.mpc

@@ -1,7 +1,8 @@
-from Compiler import library as lib, oram
+from Compiler import library as lib, oram, path_oram
 from Compiler.dijkstra import HeapQ
 from Compiler.path_oblivious_heap import (
     PathObliviousHeap,
+    POHToHeapQAdapter,
     POHVariant,
     path_oblivious_sort,
 )
@@ -9,6 +10,10 @@ from Compiler.types import Array, sint
 
 DEBUG = True
 
+from Compiler import path_oblivious_heap
+
+path_oblivious_heap.DEBUG = DEBUG
+
 
 def noop(*args, **kwargs):
     pass
@@ -18,20 +23,31 @@ def noop(*args, **kwargs):
 dprint = lib.print_ln if DEBUG else noop
 
 # Benchmark types
-INSERT = False
+INSERT = True
+EXTRACT = True
 SORTING = True
 
+INSERT = INSERT or EXTRACT  # Always insert if we are going to extract
+
 # Benchmark parameters
-## Insert
-RANGE = [2**i for i in range(2, 11)]
-OPERATIONS_PER_STEP = 1
+## Insert / ExtractMin
+RANGE = [2**i for i in range(1, 14)]  # TODO: Test this range
+OPERATIONS_PER_STEP = 3
+TIME_INIT = True
+TREE = False
+TREE_PATH = False
+OPTIMAL_TREE = True
+OPTIMAL_PATH = True
+POH_PATH = True
+
 ## Sorting
-LENGTHS = [2, 4, 6, 8, 10]
+LENGTHS = [5, 10, 15, 20, 25, 30]
+
+# Timing with consecutive ids
 
 timer_offset = 1000  # Hopefully run timers in an unused range
 
 
-# Timing with consecutive ids
 def start_fancy_timer(id: int | None = None) -> int:
     global timer_offset
     _id = id if id is not None else timer_offset
@@ -48,36 +64,86 @@ def stop_fancy_timer(id):
 # BENCHMARK
 
 if INSERT:
-    for capacity in RANGE:
-        # Benchmark binary heap built on ORAM
-        dprint(f"[ORAM Heap] Initializing empty structure with capacity {capacity}")
-        id = start_fancy_timer()
-        bin_heap = HeapQ(capacity, oram_type=oram.RecursiveORAM)
-        stop_fancy_timer(id)
-        dprint(f"[ORAM Heap] Running updates for capacity {capacity}")
-        id = start_fancy_timer()
-        for i in range(OPERATIONS_PER_STEP):
-            stop_fancy_timer(id)
-            dprint(f"\n[ORAM Heap] Update {i} for capacity {capacity}")
-            start_fancy_timer(id)
-            bin_heap.update(0, i)
-        stop_fancy_timer(id)
 
-        # Benchmark Path Oblivious Heap
-        dprint(f"[POH] Initializing empty structure with capacity {capacity}")
-        id = start_fancy_timer()
-        poh = PathObliviousHeap(capacity, bucket_size=2, variant=POHVariant.PATH)
-        stop_fancy_timer(id)
-        dprint(f"[POH] Running updates for capacity {capacity}")
-        id = start_fancy_timer()
-        for i in range(OPERATIONS_PER_STEP):
+    def operation_round(q, apply_op, capacity, tag=""):
+        global timer_offset
+        dprint(f"\n[{tag}] Initializing empty structure with capacity {capacity}")
+        id = None
+        if TIME_INIT:
+            id = start_fancy_timer()
+        if TIME_INIT:
             stop_fancy_timer(id)
-            dprint(f"[POH] Update {i} for capacity {capacity}")
+        dprint(f"\n[{tag}] Running updates for capacity {capacity}")
+        id = timer_offset  # Handle id manually to avoid incrementing for each operation
+        for i in range(OPERATIONS_PER_STEP):
+            dprint(f"\n[{tag}] Update {i} for capacity {capacity}")
             start_fancy_timer(id)
-            poh.insert(0, i)
-        stop_fancy_timer(id)
+            apply_op(q, i)
+            stop_fancy_timer(id)
+        timer_offset += 1
+
+    def benchmark_operations(q_init, capacity, *args, tag="", **kwargs):
+        apply_insert = lambda q, i: q.update(0, i)
+        apply_extract = lambda q, _: q.pop()
+        q = q_init(capacity, *args, **kwargs)
+        if INSERT:
+            operation_round(q, apply_insert, capacity, tag=tag + " insert")
+            if EXTRACT:
+                operation_round(q, apply_extract, capacity, tag=tag + " extract_min")
+
+    dprint(f"\n\nBENCHMARKING INSERT{'AND EXTRACT ' if EXTRACT else ''} TIME")
+    for capacity in RANGE:
+
+        dprint(f"\nCAPACITY {capacity}")
+
+        if TREE:
+            # Benchmark binary heap built on ORAM (Tree ORAM variant)
+            benchmark_operations(
+                HeapQ,
+                capacity,
+                oram_type=oram.RecursiveORAM,
+                tag="ORAM Heap (Tree)",
+            )
+
+        if TREE_PATH:
+            # Benchmark binary heap built on ORAM (Path ORAM variant)
+            benchmark_operations(
+                HeapQ,
+                capacity,
+                oram_type=path_oram.RecursivePathORAM,
+                tag="ORAM Heap (Path)",
+            )
+
+        if OPTIMAL_TREE:
+            # Benchmark binary heap built on ORAM (OptimalORAM variant)
+            benchmark_operations(
+                HeapQ,
+                capacity,
+                oram_type=oram.OptimalORAM,
+                tag="ORAM Heap (Optimal Tree)",
+            )
+
+        if OPTIMAL_PATH:
+            # Benchmark binary heap built on ORAM (OptimalORAM Path variant)
+            benchmark_operations(
+                HeapQ,
+                capacity,
+                oram_type=path_oram.OptimalORAM,
+                tag="ORAM Heap (Optimal Path)",
+            )
+
+        if POH_PATH:
+            # Benchmark Path Oblivious Heap (Path variant)
+            benchmark_operations(
+                POHToHeapQAdapter,
+                capacity,
+                bucket_size=2,
+                variant=POHVariant.PATH,
+                tag="POH (Path)",
+            )
 
 if SORTING:
+    dprint("\n\nBENCHMARKING SORTING TIME")
     for n in LENGTHS:
         a = Array(n, sint)
         b = Array(n, sint)