Update gen_enclave_config to track the updated buffer sizes

gen_enclave_config.py

@@ -43,16 +43,19 @@ def get_heap_size(N, M, T, B, PRIVATE_ROUTE=True, PRO=1, PRI=1, PUO=1, PUI=1, nu
 
     # 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 ) * num_out_mult
-    heap_size += ((clients_per_server + (M-1)**2) * B ) * num_out_mult
+    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
-    wn_size = max(clients_per_server, 2*(M-1)**2)
+
+    # Round (M-1)^2 up to a multiple of M
+    round1b_size = (M-1)*M
+    wn_size = max(clients_per_server, 2*round1b_size)
     # 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)
+        heap_size += (colsort_size * B * 3) + (2 * round1b_size * B)
 
     # num_WN_to_precompute times size of each WN
     heap_size += (num_WN_to_precompute * num_out_mult * \