Merge pull request #1014 from tskovlund/master

Add Path Oblivious Heap

Compiler/dijkstra.py

@@ -47,9 +47,11 @@ class HeapEntry(object):
         print_ln('empty %s, prio %s, value %s', *(reveal(x) for x in self))
 
 class HeapORAM(object):
-    def __init__(self, size, oram_type, init_rounds, int_type):
+    def __init__(self, size, oram_type, init_rounds, int_type, entry_size=None):
+        if entry_size is None:
+            entry_size = (32,log2(size))
         self.int_type = int_type
-        self.oram = oram_type(size, entry_size=(32,log2(size)), \
+        self.oram = oram_type(size, entry_size=entry_size, \
                                   init_rounds=init_rounds, \
                                   value_type=int_type.basic_type)
     def __getitem__(self, index):
@@ -74,13 +76,15 @@ class HeapORAM(object):
         return len(self.oram)
 
 class HeapQ(object):
-    def __init__(self, max_size, oram_type=ORAM, init_rounds=-1, int_type=sint):
+    def __init__(self, max_size, oram_type=ORAM, init_rounds=-1, int_type=sint, entry_size=None):
+        if entry_size is None:
+            entry_size = (32, log2(max_size))
         basic_type = int_type.basic_type
         self.max_size = max_size
         self.levels = log2(max_size)
         self.depth = self.levels - 1
-        self.heap = HeapORAM(2**self.levels, oram_type, init_rounds, int_type)
-        self.value_index = oram_type(max_size, entry_size=log2(max_size), \
+        self.heap = HeapORAM(2**self.levels, oram_type, init_rounds, int_type, entry_size=entry_size)
+        self.value_index = oram_type(max_size, entry_size=entry_size[1], \
                                          init_rounds=init_rounds, \
                                          value_type=basic_type)
         self.size = MemValue(int_type(0))

Programs/Source/benchmark_priority_queue.mpc

@@ -0,0 +1,313 @@
+from Compiler import library as lib, oram, path_oram
+from Compiler.dijkstra import HeapQ
+from Compiler.path_oblivious_heap import (
+    POHToHeapQAdapter,
+    POHVariant,
+    UniquePOHToHeapQAdapter,
+    path_oblivious_sort,
+)
+from Compiler.types import Array, sint
+from Compiler.util import log2
+
+DEBUG = True
+
+from Compiler import path_oblivious_heap
+
+path_oblivious_heap.DEBUG = False
+
+
+def noop(*args, **kwargs):
+    pass
+
+
+# Only print if DEBUG is enabled
+dprint = lib.print_ln if DEBUG else noop
+
+# Benchmark types
+INSERT = True
+EXTRACT = True
+SORTING = False
+
+INSERT = INSERT or EXTRACT  # Always insert if we are going to extract
+
+# Benchmark parameters
+## General
+KEY_SIZE = lambda n: 32
+VALUE_SIZE = lambda n: log2(n)
+N_THREADS = 1
+N_PARALLEL = 1
+
+## Insert / ExtractMin
+RANGE = [2**i for i in range(1, 18)]
+# Important: there must be space to insert this amount of entries in the queue
+OPERATIONS_PER_STEP = 10
+TIME_INIT = True
+TREE_HEAP = False
+TREE_PATH_HEAP = False
+LINEAR_HEAP = False
+OPTIMAL_TREE_HEAP = False
+OPTIMAL_PATH_HEAP = False
+POH_PATH = True
+POH_PATH_CONSTANT_STASH = True
+UNIQUE_POH_PATH_LINEAR = False
+UNIQUE_POH_PATH_PATH = False
+UNIQUE_POH_PATH_CONSTANT_STASH_LINEAR = False
+UNIQUE_POH_PATH_CONSTANT_STASH_PATH = False
+
+## Sorting
+LENGTHS = [2**i for i in range(1, 10)]
+SORTING_BITS = 32
+RADIX_SORT = True
+POS = True
+POS_CONSTANT_STASH = True
+
+# Set module variables based on parameters
+oram.n_threads = N_THREADS
+oram.n_parallel = N_PARALLEL
+oram.n_threads_for_tree = N_THREADS
+
+# Timing with consecutive ids
+timer_offset = 1000  # Hopefully run timers in an unused range
+
+
+def start_fancy_timer(id: int | None = None) -> int:
+    global timer_offset
+    _id = id if id is not None else timer_offset
+    lib.start_timer(_id)
+    if id is None:
+        timer_offset += 1
+    return _id
+
+
+def stop_fancy_timer(id):
+    lib.stop_timer(id)
+
+
+# BENCHMARK
+
+if INSERT:
+
+    def operation_round(i, id, q, apply_op, capacity, tag=""):
+        dprint(f"\n[{tag}] Update %s for capacity {capacity}", i.reveal())
+        start_fancy_timer(id)
+        apply_op(q, i)
+        stop_fancy_timer(id)
+
+    def benchmark_operations(q_init, capacity, *args, tag="", **kwargs):
+        global timer_offset
+        apply_insert = lambda q, i: q.update(0, i)
+        apply_extract = lambda q, _: q.pop()
+        init_id = timer_offset
+        insert_id = init_id + 1
+        extract_id = insert_id + 1
+        dprint(
+            f"\n[{tag}] Running {OPERATIONS_PER_STEP} update{'s' if OPERATIONS_PER_STEP > 1 else ''} for capacity {capacity}"
+        )
+        @lib.for_range(OPERATIONS_PER_STEP)
+        def _(i):
+            dprint(f"\n[{tag}] Initializing empty structure with capacity {capacity}")
+            if TIME_INIT:
+                start_fancy_timer(init_id)
+            q = q_init(capacity, *args, **kwargs)
+            if TIME_INIT:
+                stop_fancy_timer(init_id)
+            if INSERT:
+                operation_round(
+                    i, insert_id, q, apply_insert, capacity, tag=tag + " insert"
+                )
+                if EXTRACT:
+                    operation_round(
+                        i,
+                        extract_id,
+                        q,
+                        apply_extract,
+                        capacity,
+                        tag=tag + " extract_min",
+                    )
+        timer_offset += 3
+
+    dprint(f"\n\nBENCHMARKING INSERT {'AND EXTRACT ' if EXTRACT else ''}TIME")
+    for capacity in RANGE:
+        entry_size = (KEY_SIZE(capacity), VALUE_SIZE(capacity))
+
+        dprint(f"\nCAPACITY {capacity}")
+
+        if TREE_HEAP:
+            # Benchmark binary heap built on ORAM (Tree ORAM variant)
+            benchmark_operations(
+                HeapQ,
+                capacity,
+                oram_type=oram.RecursiveORAM,
+                entry_size=entry_size,
+                tag="ORAM Heap (Tree)",
+            )
+
+        if TREE_PATH_HEAP:
+            # Benchmark binary heap built on ORAM (Path ORAM variant)
+            benchmark_operations(
+                HeapQ,
+                capacity,
+                oram_type=path_oram.RecursivePathORAM,
+                entry_size=entry_size,
+                tag="ORAM Heap (Path)",
+            )
+
+        if LINEAR_HEAP:
+            # Benchmark binary heap built on ORAM (Linear ORAM variant)
+            benchmark_operations(
+                HeapQ,
+                capacity,
+                oram_type=oram.LinearORAM,
+                entry_size=entry_size,
+                tag="ORAM Heap (Linear)",
+            )
+
+        if OPTIMAL_TREE_HEAP:
+            # Benchmark binary heap built on ORAM (OptimalORAM variant)
+            benchmark_operations(
+                HeapQ,
+                capacity,
+                oram_type=oram.OptimalORAM,
+                entry_size=entry_size,
+                tag="ORAM Heap (Optimal Tree)",
+            )
+
+        if OPTIMAL_PATH_HEAP:
+            # Benchmark binary heap built on ORAM (OptimalORAM Path variant)
+            benchmark_operations(
+                HeapQ,
+                capacity,
+                oram_type=path_oram.OptimalORAM,
+                entry_size=entry_size,
+                tag="ORAM Heap (Optimal Path)",
+            )
+
+        if POH_PATH:
+            # Benchmark Path Oblivious Heap (Path variant)
+            benchmark_operations(
+                POHToHeapQAdapter,
+                capacity,
+                bucket_size=2,
+                stash_size=log2(capacity) ** 2,
+                variant=POHVariant.PATH,
+                entry_size=entry_size,
+                tag="POH (Path (superlogarithmic stash size))",
+            )
+
+        if POH_PATH_CONSTANT_STASH:
+            # Benchmark Path Oblivious Heap (Path variant with constant stash size)
+            benchmark_operations(
+                POHToHeapQAdapter,
+                capacity,
+                bucket_size=2,
+                stash_size=20,  # based on empirical analysis by Keller and Scholl
+                variant=POHVariant.PATH,
+                entry_size=entry_size,
+                tag="POH (Path (constant stash size))",
+            )
+
+        if UNIQUE_POH_PATH_LINEAR:
+            # Benchmark Unique Path Oblivious Heap (Path variant with constant stash size and linear ORAM)
+            benchmark_operations(
+                UniquePOHToHeapQAdapter,
+                capacity,
+                bucket_size=2,
+                stash_size=log2(capacity) ** 2,
+                variant=POHVariant.PATH,
+                oram_type=oram.LinearORAM,
+                entry_size=entry_size,
+                tag="Unique POH (Path (superlogarithmic stash size)) (Linear ORAM)",
+            )
+
+        if UNIQUE_POH_PATH_PATH:
+            # Benchmark Unique Path Oblivious Heap (Path variant with constant stash size and linear ORAM)
+            benchmark_operations(
+                UniquePOHToHeapQAdapter,
+                capacity,
+                bucket_size=2,
+                stash_size=log2(capacity) ** 2,
+                variant=POHVariant.PATH,
+                oram_type=path_oram.RecursivePathORAM,
+                entry_size=entry_size,
+                tag="Unique POH (Path (superlogarithmic stash size)) (Path ORAM)",
+            )
+
+        if UNIQUE_POH_PATH_CONSTANT_STASH_LINEAR:
+            benchmark_operations(
+                UniquePOHToHeapQAdapter,
+                capacity,
+                bucket_size=2,
+                stash_size=20,
+                variant=POHVariant.PATH,
+                oram_type=oram.LinearORAM,
+                entry_size=entry_size,
+                tag="Unique POH (Path (constant stash size)) (Linear ORAM)",
+            )
+
+        if UNIQUE_POH_PATH_CONSTANT_STASH_PATH:
+            benchmark_operations(
+                UniquePOHToHeapQAdapter,
+                capacity,
+                bucket_size=2,
+                stash_size=20,
+                variant=POHVariant.PATH,
+                oram_type=path_oram.RecursivePathORAM,
+                entry_size=entry_size,
+                tag="Unique POH (Path (constant stash size)) (Path ORAM)",
+            )
+
+if SORTING:
+    dprint("\n\nBENCHMARKING SORTING TIME")
+    for n in LENGTHS:
+        dprint(f"\nLENGTH {n}")
+        a = Array(n, sint)
+
+        @lib.for_range(n)
+        def _(i):
+            a[i] = sint.get_random_int(SORTING_BITS)
+
+        if RADIX_SORT:
+            a_ = Array(n, sint).assign(a)
+            lib.print_ln(
+                "\n[Sorting (Radix)] Unsorted array of length %s: %s", n, a_.reveal()
+            )
+            lib.print_ln("[Sorting (Radix)] Sorting array...")
+            id = start_fancy_timer()
+            a_.sort()
+            stop_fancy_timer(id)
+            lib.print_ln("[Sorting (Radix)] Sorted array: %s", a_.reveal())
+
+        if POS:
+            a_ = Array(n, sint).assign(a)
+            lib.print_ln(
+                "\n[Sorting (POH) superlogarithmic stash size] Unsorted array of length %s: %s",
+                n,
+                a_.reveal(),
+            )
+            lib.print_ln("[Sorting (POH) superlogarithmic stash size] Sorting array...")
+            id = start_fancy_timer()
+            path_oblivious_sort(
+                a_, a_, SORTING_BITS, stash_size=log2(n) ** 2, variant=POHVariant.PATH
+            )
+            stop_fancy_timer(id)
+            lib.print_ln(
+                "[Sorting (POH) superlogarithmic stash size] Sorted array: %s",
+                a_.reveal(),
+            )
+
+        if POS_CONSTANT_STASH:
+            a_ = Array(n, sint).assign(a)
+            lib.print_ln(
+                "\n[Sorting (POH) constant stash size] Unsorted array of length %s: %s",
+                n,
+                a_.reveal(),
+            )
+            lib.print_ln("[Sorting (POH) constant stash size] Sorting array...")
+            id = start_fancy_timer()
+            path_oblivious_sort(
+                a_, a_, SORTING_BITS, stash_size=20, variant=POHVariant.PATH
+            )
+            stop_fancy_timer(id)
+            lib.print_ln(
+                "[Sorting (POH) constant stash size] Sorted array: %s", a_.reveal()
+            )

Programs/Source/test_path_oblivious_heap.mpc

@@ -0,0 +1,169 @@
+from Compiler import library as lib, path_oblivious_heap as poh, util
+from Compiler.path_oram import OptimalORAM, RecursivePathORAM
+from Compiler import dijkstra
+from Compiler.types import sint
+
+### SETUP ###
+
+# Which tests are we running?
+SCRATCHPAD = False
+UNIT = True
+DIJKSTRA = True
+
+# Module settings
+from Compiler import oram, path_oblivious_heap
+
+oram.crash_on_overflow = False
+poh.CRASH_ON_EMPTY = True
+poh.TRACE = False
+poh.DEBUG = True
+
+### UTILITY ###
+
+# Timing with consecutive ids
+
+timer_offset = 1000  # Hopefully run timers in an unused range
+
+
+def start_fancy_timer(id: int | None = None) -> int:
+    global timer_offset
+    _id = id if id is not None else timer_offset
+    lib.start_timer(_id)
+    if id is None:
+        timer_offset += 1
+    return _id
+
+
+def stop_fancy_timer(id):
+    lib.stop_timer(id)
+
+
+if SCRATCHPAD:
+    q = poh.UniquePOHToHeapQAdapter(3)
+    q.insert(0, 10)
+    q.insert(1, 1)
+    q.insert(2, 1)
+
+    lib.print_ln("%s", q.extract_min().reveal())
+    lib.print_ln("%s", q.extract_min().reveal())
+    lib.print_ln("%s", q.extract_min().reveal())
+
+if UNIT:
+    # POH
+    lib.print_ln("Testing PathObliviousHeap.__init__...")
+    q = poh.PathObliviousHeap(10)
+    lib.print_ln("Initialization was successful!")
+
+    lib.print_ln("Testing PathObliviousHeap.insert of two identical entries...")
+    lib.print_ln("Inserting {value: 5, priority: 10}...")
+    q.insert(5, 10)
+    lib.print_ln("Inserting {value: 5, priority: 10}...")
+    q.insert(5, 10)
+    lib.print_ln("Inserts were successful!")
+
+    lib.print_ln("Testing PathObliviousHeap.extract_min of two identical entries...")
+    lib.print_ln("Extracted values should be 5, 5.")
+    v = q.extract_min()
+    lib.print_ln("Extracted value %s", v.reveal())
+    v = q.extract_min()
+    lib.print_ln("Extracted value %s", v.reveal())
+
+    # UniquePOH
+    lib.print_ln("Testing UniquePathObliviousHeap.__init__...")
+    q = poh.UniquePathObliviousHeap(4)
+    lib.print_ln("Initialization was successful!")
+    lib.print_ln("Testing UniquePathObliviousHeap.insert...")
+    lib.print_ln("Inserting {value: 5, priority: 10}...")
+    q.insert(5, 10)
+    lib.print_ln("Inserting {value: 5, priority: 10}...")
+    q.insert(5, 10)
+    lib.print_ln("Inserting {value: 6, priority: 11}...")
+    q.insert(6, 11)
+    lib.print_ln("Inserts were successful!")
+
+    lib.print_ln("Testing UniquePathObliviousHeap.extract_min twice...")
+    lib.print_ln("Extracted values should be 5, 6.")
+    v = q.extract_min()
+    lib.print_ln("Extracted value %s", v.reveal())
+    v = q.extract_min()
+    lib.print_ln("Extracted value %s", v.reveal())
+
+    lib.print_ln("Testing UniquePathObliviousHeap.insert and update...")
+    lib.print_ln("Inserting {value: 7, priority: 14}...")
+    q.insert(7, 14)
+    lib.print_ln("Updating {value: 7, priority: 1}...")
+    q.update(7, 1)
+    lib.print_ln("Inserting {value: 2, priority: 10}...")
+    q.update(7, 1)
+    lib.print_ln("Inserts and updates were successful!")
+
+    lib.print_ln("Testing UniquePathObliviousHeap.find_min...")
+    v = q.find_min()
+    lib.print_ln("Found value should be 7.")
+    lib.print_ln("Found value %s", v.reveal())
+
+    # SubtreeMinEntry
+    lib.print_ln("Testing SubtreeMinEntry comparisons...")
+    id = start_fancy_timer()
+    poh.test_SubtreeMinEntry_cmp()
+    stop_fancy_timer(id)
+
+if DIJKSTRA:
+    dijkstra.HeapQ = poh.POHToHeapQAdapter
+
+    # example code for graph with vertices 0,1,2 and with following weights
+    # 0 -> 1: 5
+    # 0 -> 2: 20
+    # 1 -> 2: 10
+
+    lib.print_ln("Testing dijkstra on Path Oblivious Heap...")
+    lib.print_ln("Output should be the following:")
+    lib.print_ln("from 0 to 0 at cost 0 via vertex 0")
+    lib.print_ln("from 0 to 1 at cost 5 via vertex 0")
+    lib.print_ln("from 0 to 2 at cost 15 via vertex 1")
+
+    # structure for edges
+    # contains tuples of form (neighbor, cost, last neighbor bit)
+    edges = OptimalORAM(
+        4,  # number of edges
+        entry_size=(
+            2,  # enough bits for vertices
+            5,  # enough bits for costs
+            1,
+        ),  # always one
+    )
+
+    # first edge from vertex 0
+    edges[0] = (1, 5, 0)
+    # second and last edge from vertex 0
+    edges[1] = (2, 20, 1)
+    # edge from vertex 1
+    edges[2] = (2, 10, 1)
+    # dummy edge from vertex 2 to itself
+    edges[3] = (2, 0, 1)
+
+    # structure assigning edge list indices to vertices
+    e_index = OptimalORAM(
+        3, entry_size=2  # number vertices
+    )  # enough bits for edge indices
+
+    # edges from 0 start at 0
+    e_index[0] = 0
+    # edges from 1 start at 2
+    e_index[1] = 2
+    # edges from 2 start at 3
+    e_index[2] = 3
+
+    source = sint(0)
+
+    res = dijkstra.dijkstra(source, edges, e_index, OptimalORAM)
+
+    @lib.for_range(res.size)
+    def _(i):
+        lib.print_ln(
+            "from %s to %s at cost %s via vertex %s",
+            source.reveal(),
+            i,
+            res[i][0].reveal(),
+            res[i][1].reveal(),
+        )