Compute aggregated_poly in swigless python

min-bindings/python/ckzg.c

@@ -6,12 +6,22 @@ static void free_BLSFieldElement(PyObject *c) {
   free(PyCapsule_GetPointer(c, "BLSFieldElement"));
 }
 
+static void free_G1(PyObject *c) {
+  free(PyCapsule_GetPointer(c, "G1"));
+}
+
 static void free_PolynomialEvalForm(PyObject *c) {
   PolynomialEvalForm *p = PyCapsule_GetPointer(c, "PolynomialEvalForm");
   free_polynomial(p);
   free(p);
 }
 
+static void free_KZGSettings(PyObject *c) {
+  KZGSettings *s = PyCapsule_GetPointer(c, "KZGSettings");
+  free_trusted_setup(s);
+  free(s);
+}
+
 static PyObject* bytes_to_bls_field_wrap(PyObject *self, PyObject *args) {
   PyBytesObject *pybytes;
 
@@ -28,7 +38,7 @@ static PyObject* bytes_to_bls_field_wrap(PyObject *self, PyObject *args) {
   return PyCapsule_New(out, "BLSFieldElement", free_BLSFieldElement);
 }
 
-static PyObject* int_from_BLSFieldElement(PyObject *self, PyObject *args) {
+static PyObject* int_from_bls_field(PyObject *self, PyObject *args) {
   PyObject *c;
 
   if (!PyArg_UnpackTuple(args, "uint64s_from_BLSFieldElement", 1, 1, &c) ||
@@ -108,7 +118,7 @@ static PyObject* compute_powers_wrap(PyObject *self, PyObject *args) {
 
   if (out == NULL) return PyErr_NoMemory();
 
-  BLSFieldElement *a = (BLSFieldElement*)malloc(sizeof(BLSFieldElement) * z);
+  BLSFieldElement *a = (BLSFieldElement*)calloc(z, sizeof(BLSFieldElement));
 
   if (a == NULL) return PyErr_NoMemory();
 
@@ -131,14 +141,168 @@ static PyObject* compute_powers_wrap(PyObject *self, PyObject *args) {
   return out;
 }
 
+static PyObject* load_trusted_setup_wrap(PyObject *self, PyObject *args) {
+  PyObject *f;
+
+  if (!PyArg_ParseTuple(args, "U", &f))
+    return PyErr_Format(PyExc_ValueError, "expected a string");
+
+  KZGSettings *s = (KZGSettings*)malloc(sizeof(KZGSettings));
+
+  if (s == NULL) return PyErr_NoMemory();
+
+  if (load_trusted_setup(s, fopen(PyUnicode_AsUTF8(f), "r")) != C_KZG_OK)
+    return PyErr_Format(PyExc_RuntimeError, "error loading trusted setup");
+
+  return PyCapsule_New(s, "KZGSettings", free_KZGSettings);
+}
+
+static PyObject* blob_to_kzg_commitment_wrap(PyObject *self, PyObject *args) {
+  PyObject *a;
+  PyObject *c;
+
+  if (!PyArg_UnpackTuple(args, "alloc_polynomial_wrap", 2, 2, &a, &c) ||
+      !PySequence_Check(a) ||
+      !PyCapsule_IsValid(c, "KZGSettings"))
+    return PyErr_Format(PyExc_ValueError, "expected sequence and trusted setup");
+
+  Py_ssize_t n = PySequence_Length(a);
+
+  BLSFieldElement *blob = (BLSFieldElement*)calloc(n, sizeof(BLSFieldElement));
+
+  if (blob == NULL) return PyErr_NoMemory();
+
+  PyObject *e;
+  for (Py_ssize_t i = 0; i < n; i++) {
+    e = PySequence_GetItem(a, i);
+    if (!PyCapsule_IsValid(e, "BLSFieldElement")) {
+      free(blob);
+      return PyErr_Format(PyExc_ValueError, "expected BLSFieldElement capsules");
+    }
+    // TODO: could avoid copying if blob_to_kzg_commitment expected pointers instead of an array
+    blob[i] = *(BLSFieldElement*)PyCapsule_GetPointer(e, "BLSFieldElement");
+  }
+
+  KZGCommitment *k = (KZGCommitment*)malloc(sizeof(KZGCommitment));
+
+  if (k == NULL) return PyErr_NoMemory();
+
+  blob_to_kzg_commitment(k, blob, PyCapsule_GetPointer(c, "KZGSettings"));
+
+  free(blob);
+
+  return PyCapsule_New(k, "G1", free_G1);
+}
+
+static PyObject* vector_lincomb_wrap(PyObject *self, PyObject *args) {
+  PyObject *vs;
+  PyObject *fs;
+
+  if (!PyArg_UnpackTuple(args, "vector_lincomb", 2, 2, &vs, &fs) ||
+      !PySequence_Check(vs) ||
+      !PySequence_Check(fs))
+    return PyErr_Format(PyExc_ValueError, "expected two sequences");
+
+  Py_ssize_t n = PySequence_Length(vs);
+  if (PySequence_Length(fs) != n)
+    return PyErr_Format(PyExc_ValueError, "expected same-length sequences");
+
+  if (n == 0) { return fs; }
+
+  if (!PySequence_Check(PySequence_GetItem(vs, 0)))
+    return PyErr_Format(PyExc_ValueError, "expected sequence of sequences");
+
+  Py_ssize_t i, j, m = PySequence_Length(PySequence_GetItem(vs, 0));
+
+  const BLSFieldElement* *vectors = (const BLSFieldElement**)calloc(n * m, sizeof(BLSFieldElement*));
+
+  if (vectors == NULL) return PyErr_NoMemory();
+
+  PyObject *tmp, *out;
+
+  for (i = 0; i < n; i++) {
+    if (!PySequence_Check(PySequence_GetItem(vs, i))) {
+      free(vectors);
+      return PyErr_Format(PyExc_ValueError, "expected sequence of sequences");
+    }
+    tmp = PySequence_GetItem(vs, i);
+    if (PySequence_Length(tmp) != m) {
+      free(vectors);
+      return PyErr_Format(PyExc_ValueError, "expected vectors of same length");
+    }
+    for (j = 0; j < m; j++) {
+      out = PySequence_GetItem(tmp, j);
+      if (!PyCapsule_IsValid(out, "BLSFieldElement")) {
+        free(vectors);
+        return PyErr_Format(PyExc_ValueError, "expected vectors of BLSFieldElement capsules");
+      }
+      vectors[i * m + j] = (BLSFieldElement*)PyCapsule_GetPointer(out, "BLSFieldElement");
+    }
+  }
+
+  const BLSFieldElement* *scalars = (const BLSFieldElement**)calloc(n, sizeof(BLSFieldElement*));
+
+  if (scalars == NULL) {
+    free(vectors);
+    return PyErr_NoMemory();
+  }
+
+  for (i = 0; i < n; i++) {
+    tmp = PySequence_GetItem(fs, i);
+    if (!PyCapsule_IsValid(tmp, "BLSFieldElement")) {
+      free(scalars);
+      free(vectors);
+      return PyErr_Format(PyExc_ValueError, "expected a BLSFieldElement capsule");
+    }
+    scalars[i] = (BLSFieldElement*)PyCapsule_GetPointer(tmp, "BLSFieldElement");
+  }
+
+  BLSFieldElement *r = (BLSFieldElement*)calloc(m, sizeof(BLSFieldElement));
+
+  if (r == NULL) {
+    free(scalars);
+    free(vectors);
+    return PyErr_NoMemory();
+  }
+
+  vector_lincomb(r, vectors, scalars, n, m);
+
+  free(scalars);
+  free(vectors);
+
+  out = PyList_New(m);
+
+  if (out == NULL) {
+    free(r);
+    return PyErr_NoMemory();
+  }
+
+  BLSFieldElement *f;
+
+  for (j = 0; j < m; j++) {
+    f = (BLSFieldElement*)malloc(sizeof(BLSFieldElement));
+    if (f == NULL) {
+      free(r);
+      return PyErr_NoMemory();
+    }
+    *f = r[j];
+    PyList_SetItem(out, j, PyCapsule_New(f, "BLSFieldElement", free_BLSFieldElement));
+  }
+
+  free(r);
+
+  return out;
+}
+
 static PyMethodDef ckzgmethods[] = {
   {"bytes_from_G1",            bytes_from_G1_wrap,          METH_VARARGS, "Convert a group element to 48 bytes"},
-  {"int_from_BLSFieldElement", int_from_BLSFieldElement,    METH_VARARGS, "Convert a field element to a 256-bit int"},
+  {"int_from_bls_field",       int_from_bls_field,          METH_VARARGS, "Convert a field element to a 256-bit int"},
   {"bytes_to_bls_field",       bytes_to_bls_field_wrap,     METH_VARARGS, "Convert 32 bytes to a field element"},
   {"alloc_polynomial",         alloc_polynomial_wrap,       METH_VARARGS, "Create a PolynomialEvalForm from a sequence of field elements"},
-  // {"load_trusted_setup",       load_trusted_setup_wrap,     METH_VARARGS, "Load trusted setup from file path"},
-  // {"blob_to_kzg_commitment",   blob_to_kzg_commitment_wrap, METH_VARARGS, "Create a commitment from a sequence of field elements"},
+  {"load_trusted_setup",       load_trusted_setup_wrap,     METH_VARARGS, "Load trusted setup from file path"},
+  {"blob_to_kzg_commitment",   blob_to_kzg_commitment_wrap, METH_VARARGS, "Create a commitment from a sequence of field elements"},
   {"compute_powers",           compute_powers_wrap,         METH_VARARGS, "Create a list of powers of a field element"},
+  {"vector_lincomb",           vector_lincomb_wrap,         METH_VARARGS, "Multiply a matrix of field elements with a vector"},
   {NULL, NULL, 0, NULL}
 };
 

min-bindings/python/tests.py

@@ -1,10 +1,11 @@
 import ckzg
 import random
+import ssz
 
 # Simple test of bytes_to_bls_field
 
 bs = (329).to_bytes(32, "little")
-assert 329 == ckzg.int_from_BLSFieldElement(ckzg.bytes_to_bls_field(bs))
+assert 329 == ckzg.int_from_bls_field(ckzg.bytes_to_bls_field(bs))
 
 # Simple test of compute_powers
 
@@ -15,8 +16,36 @@ powers = ckzg.compute_powers(ckzg.bytes_to_bls_field(x.to_bytes(32, "little")),
 
 p_check = 1
 for p in powers:
-    assert p_check == ckzg.int_from_BLSFieldElement(p)
+    assert p_check == ckzg.int_from_bls_field(p)
     p_check *= x
     p_check %= 2**256
 
+# Commit to a few random blobs
+
+BLOB_SIZE = 4096
+MAX_BLOBS_PER_BLOCK = 16
+
+blobs_sedes = ssz.List(ssz.Vector(ssz.uint256, BLOB_SIZE), MAX_BLOBS_PER_BLOCK)
+kzg_commitments_sedes = ssz.List(ssz.bytes48, MAX_BLOBS_PER_BLOCK)
+
+blobs = [[ckzg.bytes_to_bls_field(random.randbytes(32)) for _ in range(BLOB_SIZE)] for _ in range(3)]
+
+ts = ckzg.load_trusted_setup("../../src/trusted_setup.txt")
+
+kzg_commitments = [ckzg.blob_to_kzg_commitment(blob, ts) for blob in blobs]
+
+# Compute polynomial commitments for these blobs
+# We don't follow the spec exactly to get the hash, but it shouldn't matter since it's random data
+
+encoded_blobs = ssz.encode([[ckzg.int_from_bls_field(fr) for fr in blob] for blob in blobs], blobs_sedes)
+encoded_commitments = ssz.encode([ckzg.bytes_from_G1(c) for c in kzg_commitments], kzg_commitments_sedes)
+hashed = ssz.hash.hashlib.sha256(encoded_blobs + encoded_commitments).digest()
+
+r = ckzg.bytes_to_bls_field(hashed)
+r_powers = ckzg.compute_powers(r, len(blobs))
+
+values = ckzg.vector_lincomb(blobs, r_powers)
+
+aggregated_poly = ckzg.alloc_polynomial(values)
+
 print('Tests passed')