test_template.txt

import sys
import itertools
import json
import numpy as np
from time import time


class ArrayPreprocessing:
    def apply(self, arr):
        raise NotImplemented


class PermuteDims(ArrayPreprocessing):
    def __init__(self, dim_i, dim_j):
        self.dim_i = dim_i
        self.dim_j = dim_j

    def apply(self, arr):
        slices = []
        for ind in range(arr.shape[self.dim_i]):
            i_slice = np.take(arr, ind, axis=self.dim_i)
            roll_dim = self.dim_j if self.dim_j < self.dim_i else self.dim_j - 1
            perm_i_slice = np.roll(i_slice, -ind, axis=roll_dim)
            slices.append(perm_i_slice)

        return np.stack(slices, axis=self.dim_i)

    def __str__(self):
        return "Permute({},{})".format(self.dim_i, self.dim_j)


class VectorDimContent:
    def __init__(self):
        pass

    def offset_left(self):
        raise NotImplemented

    def size(self):
        raise NotImplemented


class FilledDim(VectorDimContent):
    def __init__(self, dim, extent, stride=1, oob_left=0, oob_right=0, pad_left=0, pad_right=0):
        super().__init__()
        self.dim = dim
        self.extent = extent
        self.stride = stride
        self.oob_left = oob_left
        self.oob_right = oob_right
        self.pad_left = pad_left
        self.pad_right = pad_right

    def offset_left(self):
        return self.oob_left + self.pad_left

    def size(self):
        return self.pad_left + self.oob_left + self.extent + self.oob_right + self.pad_right


class EmptyDim(VectorDimContent):
    def __init__(self, extent, pad_left=0, pad_right=0, oob_right=0):
        super().__init__()
        self.extent = extent
        self.pad_left = pad_left
        self.pad_right = pad_right
        self.oob_right = oob_right

    def offset_left(self):
        return self.pad_left

    def size(self):
        return self.pad_left + self.extent + self.oob_right + self.pad_right


class AbstractArray:
    def __init__(self, size: int, extents: list[int]):
        if len(extents) == 0:
            self.single = CiphertextVector(size)

        else:
            self.map = {}
            coords = itertools.product(*map(lambda e: tuple(range(e)), extents))
            for coord in coords:
                self.map[coord] = CiphertextVector(size)

    def set(self, coord: list[int], val):
        if len(coord) == 0:
            self.single = val

        else:
            self.map[tuple(coord)] = val

    def get(self, coord=[]):
        if len(coord) == 0:
            return self.single

        else:
            return self.map[tuple(coord)]

    def create_vector(self, size):
        pass

    def show(self):
        if self.single is not None:
            print(self.single)

        else:
            for coord, val in self.map.items():
                print("{} => {}", coord, val)


class CiphertextArray(AbstractArray):
    def __init__(self, size: int, extents: list[int]):
        super().__init__(size, extents)

    def create_vector(self, size: int):
        return CiphertextVector(size)


class PlaintextArray(AbstractArray):
    def __init__(self, size: int, extents: list[int]):
        super().__init__(size, extents)

    def create_vector(self, size: int):
        return PlaintextVector(size)


class NativeArray(AbstractArray):
    def __init__(self, size: int, extents: list[int]):
        super().__init__(size, extents)

    def create_vector(self, size: int):
        return NativeVector(size)


class AbstractVector:
    def __init__(self, size: int, array=None):
        self.size = size
        if array is None:
            self.array = np.zeros((size,))

        else:
            self.array = array

    def load(self, x):
        n = len(x)
        assert(self.size >= n)
        self.array[:n] = np.array(x)

    def get(self, n):
        return self.array[:n]

    def add_inplace(self, y):
        assert(self.validate_operand(y))
        self.array = self.array + y.array

    def sub_inplace(self, y):
        assert(self.validate_operand(y))
        self.array = self.array - y.array

    def multiply_inplace(self, y):
        assert(self.validate_operand(y))
        self.array = self.array * y.array

    def rotate_inplace(self, n):
        self.array = self.array[[(i-n) % self.size for i in range(self.size)]]

    def __add__(self, x):
        assert(self.validate_operand(x))
        return self.create(self.size, self.array + x.array)

    def __sub__(self, x):
        assert(self.validate_operand(x))
        return self.create(self.size, self.array - x.array)
        
    def __mul__(self, x):
        assert(self.validate_operand(x))
        return self.create(self.size, self.array * x.array)

    def rotate(self, n):
        return self.create(self.size, self.array[[(i-n) % self.size for i in range(self.size)]])

    def __str__(self):
        return str(self.array)

    def create(self, size, array):
        pass

    def validate_operand(self, y):
        pass

class CiphertextVector(AbstractVector):
    def __init__(self, size, array=None):
        super().__init__(size, array)

    def create(self, size, array):
        return CiphertextVector(size, array)

    def validate_operand(self, y):
        return isinstance(y, CiphertextVector) or isinstance(y, PlaintextVector)


class PlaintextVector(AbstractVector):
    def __init__(self, size, array):
        super().__init__(size, array)

    def create(self, size, array):
        return CiphertextVector(size, array)

    def validate_operand(self, y):
        return isinstance(y, CiphertextVector) or isinstance(y, PlaintextVector)


class NativeVector(AbstractVector):
    def __init__(self, size, array):
        super().__init__(size, array)

    def create(self, size, array):
        return CiphertextVector(size, array)

    def validate_operand(self, y):
        return isinstance(y, NativeVector)


class TestWrapper:
    def __init__(self, size, client_inputs={}, server_inputs={}):
        self.size = size
        self.client_inputs = client_inputs
        self.server_inputs = server_inputs
        self.client_arrays = {}
        self.server_arrays = {}
        self.client_buffer = {}
        self.server_buffer = {}
        self.party = None

    def set_party(self, party):
        self.party = party

    def client_input(self, name):
        assert(self.party == "client")
        self.client_arrays[(name, "")] = self.client_inputs[name]

    def server_input(self, name):
        assert(self.party == "server")
        self.server_arrays[(name, "")] = self.server_inputs[name]

    def client_output(self, arr):
        assert(self.party == "client")
        print("client output:")
        arr.show()

    def client_send(self, name, arr):
        assert(self.party == "client")
        vec = arr.get([])
        assert(isinstance(vec, NativeVector))

        # "encrypt" whatever the client sends
        self.client_buffer[name] = CiphertextVector(self.size, vec.array)

    def server_recv(self, name):
        assert(self.party == "server")
        return self.vec_to_array(self.client_buffer[name])

    def server_send(self, name, arr):
        assert(self.party == "server")
        assert(isinstance(arr, CiphertextArray))
        self.server_buffer[name] = arr

    def client_recv(self, name):
        assert(self.party == "client")
        return self.server_buffer[name]

    def vec_to_array(self, vec):
        arr = None
        if isinstance(vec, NativeVector):
            arr = NativeArray(self.size, [])

        elif isinstance(vec, CiphertextVector):
            arr = CiphertextArray(self.size, [])

        elif isinstance(vec, PlaintextVector):
            arr = PlaintextArray(self.size, [])

        else:
            raise Exception("unknown type")

        arr.set([], vec)
        return arr

    def get_array(self, name, preprocess=None):
        preprocess_str = str(preprocess) if preprocess is not None else ""

        arrays = None
        if self.party == "server":
            arrays = self.server_arrays

        elif self.party == "client":
            arrays = self.client_arrays

        else:
            raise Exception("party not set")

        if not (name, preprocess_str) in arrays:
            arr = arrays[(name, "")]
            parr = preprocess.apply(arr)
            arrays[(name, preprocess_str)] = parr
            return parr

        else:
            return arrays[(name, preprocess_str)]

    def native_array(self, extents: list[list[int]]):
        return NativeArray(self.size, extents)

    def ciphertext_array(self, extents):
        return CiphertextArray(self.size, extents)

    def plaintext_array(self, extents):
        return PlaintextArray(self.size, extents)

    def build_vector(self, name: str, preprocess: ArrayPreprocessing, src_offset: list[int], dims: list[VectorDimContent]) -> NativeArray:
        array = self.get_array(name, preprocess)
        if len(dims) == 0:
            npvec = np.zeros((self.size,))
            npvec[0] = array[tuple(src_offset)]
            vec = NativeVector(self.size, npvec)
            return self.vec_to_array(vec)

        else:
            dst_offset = []
            dst_shape = []
            stride_map = {}
            coords = []
            dst_size = 1
            for i, dim in enumerate(dims):
                dst_offset.append(dim.offset_left())
                dst_shape.append(dim.size())
                dst_size = dim.size()
                coords.append(range(dim.extent))

                if isinstance(dim, FilledDim):
                    stride_map[i] = (dim.dim, dim.stride)

            dst = np.zeros(dst_shape, dtype=int)
            for coord in itertools.product(*coords):
                src_coord = src_offset[:]
                dst_coord = dst_offset[:]
                for i, coord in enumerate(coord):
                    dst_coord[i] += coord
                    if i in stride_map:
                        (src_dim, stride) = stride_map[i]
                        src_coord[src_dim] += coord * stride

                dst_coord_tup = tuple(dst_coord)
                src_coord_tup = tuple(src_coord)
                src_in_bounds = all(map(lambda t: t[0] > t[1], zip(array.shape, src_coord_tup)))
                if src_in_bounds:
                    dst[dst_coord_tup] = array[tuple(src_coord)]

                else:
                    dst[dst_coord_tup] = 0

            dst_flat = dst.reshape(dst_size)

            # should we actually repeat? or just pad with zeros?
            repeats = int(self.size / dst_size)
            if self.size % dst_size != 0:
                repeats += 1

            dst_vec = np.tile(dst_flat, repeats)[:self.size]
            return self.vec_to_array(NativeVector(self.size, dst_vec))


    def const(self, const: int):
        return self.vec_to_array(NativeVector(self.size, np.array(self.size * [const])))

    def mask(self, mask: list[(int, int, int)]):
        raise NotImplemented

    def set(self, arr, coord, val):
        arr.set(coord, val)

    def encode(self, arr, coord):
        vec = arr.get(coord)
        assert(isinstance(vec, NativeVector))
        arr.set(coord, PlaintextVector(vec.size, vec.array))

    def encrypt(self, x):
        return x

    def add(self, x, y):
        return x + y

    def add_plain(self, x, y):
        return x + y

    def add_inplace(self, x, y):
        x.add_inplace(y)

    def add_plain_inplace(self, x, y):
        x.add_inplace(y)

    def multiply(self, x, y):
        return x * y

    def multiply_plain(self, x, y):
        return x * y

    def multiply_inplace(self, x, y):
        x.multiply_inplace(y)

    def multiply_plain_inplace(self, x, y):
        x.multiply_inplace(y)

    def rotate_rows(self, x, amt):
        return x.rotate(amt)

    def rotate_rows_inplace(self, x, amt):
        x.rotate_inplace(amt)

    def relinearize_inplace(self, x):
        pass

    def invariant_noise_budget(self, x):
        return 0

### START GENERATED CODE
{{{ program }}}
### END GENERATED CODE


input_str = None
client_inputs = {}
server_inputs = {}
if len(sys.argv) >= 2:
    with open(sys.argv[1]) as f:
        input_str = f.read()

else:
    input_str = sys.stdin.read()

inputs = json.loads(input_str)
for key, val in inputs["client"].items():
    client_inputs[key] = np.array(val)

for key, val in inputs["server"].items():
    server_inputs[key] = np.array(val)

vec_size = {{{ size }}}
seal = TestWrapper(vec_size, client_inputs, server_inputs)

seal.set_party("client")
client_pre(seal)

exec_start_time = time()
seal.set_party("server")
server(seal)
exec_end_time = time()

seal.set_party("client")
client_post(seal)

print("exec duration: {}ms".format((exec_end_time - exec_start_time)*1000))