Basic toy GSW FHE implementation

fhe.py

import numpy, random

dim = 2
qbits = 8
msgbits = 2
ctlength = 6
base_error = 1

modulus = 2 ** qbits
msg_modulus = 2 ** msgbits

# Generates fully random data (good for secret keys)
def random(shape):
    return numpy.random.randint(0, modulus, shape, dtype=numpy.uint64)

# Generates low-norm random data (good for "pads" and noise)
def random_low_norm(shape):
    return numpy.random.randint(
        0, base_error * 2 + 1, shape, dtype=numpy.uint64
    ) - base_error

def zeros(shape):
    return numpy.zeros(shape, dtype=numpy.uint64)

def matmul(A, B):
    return numpy.matmul(A, B) % modulus

# Generates a privkey and a pubkey
def keygen() -> ("vec[dim]", "mat[dim, ctlength+1]"):
    secret = random(dim)
    PUBKEY_RANDOM = random((dim, ctlength))
    PUBKEY_BOTTOM_ROW = numpy.reshape(
        matmul(secret, PUBKEY_RANDOM) + random_low_norm(ctlength),
        (1, ctlength)
    )
    return secret, numpy.concat(
        (PUBKEY_RANDOM,
        PUBKEY_BOTTOM_ROW),
    )

# Generates a matrix that looks like:
#
# 1 2 4 0 0 0 0 0 0
# 0 0 0 1 2 4 0 0 0
# 0 0 0 0 0 0 1 2 4

def gadget(dim):
    x = zeros((dim, dim * qbits))
    for i in range(dim):
        for j in range(qbits):
            x[i, i*qbits+j] = 2**j
    return x

G = gadget(dim+1)

# Converts
#
#           1  0
#           0  1
# 1 2 --->  0  0
# 3 4      [1] 0
#          [1] 0
#          [0] 1
#
# Notice how the highlighted bits are the LSB-first binary representation of 3
#
# Invariant: gadget * gadget_inv(x) = x

def gadget_inv(inp: "mat[w,h]") -> "mat[w*log(q),h]":
    o = zeros((inp.shape[0] * qbits, inp.shape[1]))
    for i in range(qbits):
        o[i::qbits, :] = (inp >> i) % 2
    return o

# s -> -s ++ [1]
def minus_secret_append_one(secret: "vec[n]") -> "vec[n+1]":
    return numpy.concat((-secret, numpy.ones(1, dtype=numpy.uint64)))

# Converts a plaintext to a ciphertext
def encrypt_single(message: "int", PUBKEY: "mat[dim+1,l]") -> "mat[dim+1,l*qbits]":
    R = random_low_norm((ctlength, (dim + 1) * qbits))
    return matmul(PUBKEY, R) + message * G

# Converts a ciphertext to a plaintext
def decrypt_single(ct: "mat[dim+1,l*qbits]", secret: "vec[dim]") -> "int":
    appended_secret = minus_secret_append_one(secret)
    gadgeted_out = matmul(appended_secret, ct)
    # The bottom row is the "gadget version" of the message (NB: the
    # "multiplied by G" version, NOT a bit-decomposition). Take the
    # rightmost entry that's still multiplied not quite far enough to
    # overflow out the message bits we want, adjust it to account for
    # symmetric error, left-shift it and there's our answer.
    half_msg_space = (modulus // msg_modulus // 2)
    val = (
        ((gadgeted_out[-msgbits] + half_msg_space) % modulus)
        >>
        (qbits - msgbits)
    )
    return val

# Adding two ciphertexts is trivial, as in basically any FHE scheme
def add_cts(ct1, ct2):
    return (ct1 + ct2) % modulus

# ct1 * gadget_inv(ct2) encodes the ciphertext of m1 * m2
# Multiplying over the bit-representation (low-norm) ensures that
# the error only blows up by a constant factor
def mul_cts(ct1, ct2):
    return matmul(ct1, gadget_inv(ct2))

def test():
    secret, PUBKEY = keygen()
    m1 = random(()) >> (qbits - msgbits)
    print('m1', m1)
    ct1 = encrypt_single(m1, PUBKEY)
    print('ct', ct1)
    o = decrypt_single(ct1, secret)
    print('decrypted', o)
    m2 = random(()) >> (qbits - msgbits)
    print('m2', m2)
    ct2 = encrypt_single(m2, PUBKEY)
    osum = decrypt_single(add_cts(ct1, ct2), secret)
    print(f'decrypted sum: {osum}, compare {m1}+{m2} = {(m1+m2)%2**msgbits}')
    oprod = decrypt_single(mul_cts(ct1, ct2), secret)
    print(f'decrypted product: {oprod}, compare {m1}*{m2} = {(m1*m2)%2**msgbits}')

if __name__ == '__main__':
    test()