More accurately compute the required enclave heap memory

gen_enclave_config.py

@@ -31,25 +31,31 @@ def get_heap_size(N, M, T, B, PRIVATE_ROUTE=True, PRO=1, PRI=1, PUO=1, PUI=1, nu
     # Base heap of 2 MB per thread
     heap_size = 2000000 * T
 
-    num_buffers = 5
     num_out_mult = PRO
-    if(not(PRIVATE_ROUTE)):
+    if not PRIVATE_ROUTE:
         num_out_mult = PUO
+    num_in_mult = PRI
+    if not PRIVATE_ROUTE:
+        num_in_mult = PUI
 
     # Storage and Ingestion data stored per_client = 52 bytes
     heap_size += clients_per_server * (B + 60)
 
-    # 5 Buffers of clients_per_server items of B size each for
-    # private routing
-    heap_size += (clients_per_server * B * num_buffers) * num_out_mult
+    # 2 Buffers of clients_per_server items of B size each, plus 1 of
+    # size (clients_per_server + M) items, for private routing
+    heap_size += (clients_per_server * B * 2) * num_out_mult
+    heap_size += ((clients_per_server + M) * B) * num_in_mult
     # Additional buffers for public routing
-    if(not(PRIVATE_ROUTE)):
-        heap_size += (clients_per_server * B * 3) + ((M-1)**2 * B * 4)
+    wn_size = max(clients_per_server, 2*(M-1)**2)
+    # Round up to a multiple of M
+    colsort_size = int((wn_size + M - 1) / M) * M
+
+    if not PRIVATE_ROUTE:
+        heap_size += (colsort_size * B * 3) + (2 * (M-1)**2 * B * 2)
 
     # num_WN_to_precompute times size of each WN
-    wn_size = max(clients_per_server, 2*(M-1)**2)
     heap_size += (num_WN_to_precompute * num_out_mult * \
-        (wn_size * math.ceil(math.log(wn_size,2)) * 8))
+        (wn_size * math.ceil(math.log(wn_size,2) + 1) * 8))
 
     heap_size_page_aligned = math.ceil(heap_size/4096) * 4096
     return heap_size_page_aligned
@@ -82,7 +88,7 @@ def generate_config(N, M, T, B, PRIVATE_ROUTE=True, PRO=1, PRI=1, PUO=1, PUI=1,
 
 if __name__ == "__main__":
 
-    if(len(sys.argv)!=5):
+    if len(sys.argv) != 5:
         print("Incorrect usage!\n")
         print("./gen_enclave_config.py expects 4 parameters.")
         print("Usage: ./gen_enclave_config.py <N = number of clients> <M = number of servers> <T = number of threads> <B = message_size>")