Implement more swigless interface

min-bindings/python/ckzg.c

@@ -2,8 +2,14 @@
 #include <Python.h>
 #include "c_kzg_4844.h"
 
-static void free_capsule(PyObject *c) {
-  free(PyCapsule_GetPointer(c, PyCapsule_GetName(c)));
+static void free_BLSFieldElement(PyObject *c) {
+  free(PyCapsule_GetPointer(c, "BLSFieldElement"));
+}
+
+static void free_PolynomialEvalForm(PyObject *c) {
+  PolynomialEvalForm *p = PyCapsule_GetPointer(c, "PolynomialEvalForm");
+  free_polynomial(p);
+  free(p);
 }
 
 static PyObject* bytes_to_bls_field_wrap(PyObject *self, PyObject *args) {
@@ -19,26 +25,74 @@ static PyObject* bytes_to_bls_field_wrap(PyObject *self, PyObject *args) {
 
   bytes_to_bls_field(out, (const uint8_t*)PyBytes_AsString((PyObject*)pybytes));
 
-  return PyCapsule_New(out, "BLSFieldElement", free_capsule);
+  return PyCapsule_New(out, "BLSFieldElement", free_BLSFieldElement);
 }
 
-static PyObject* uint64s_from_BLSFieldElement_wrap(PyObject *self, PyObject *args) {
+static PyObject* int_from_BLSFieldElement(PyObject *self, PyObject *args) {
   PyObject *c;
   if (!PyArg_UnpackTuple(args, "uint64s_from_BLSFieldElement", 1, 1, &c) ||
       !PyCapsule_IsValid(c, "BLSFieldElement"))
     return PyErr_Format(PyExc_ValueError, "expected a BLSFieldElement capsule");
-  uint64_t out[4];
-  uint64s_from_BLSFieldElement(out, PyCapsule_GetPointer(c, PyCapsule_GetName(c)));
-  return PyTuple_Pack(4,
-      PyLong_FromUnsignedLong(out[0]),
-      PyLong_FromUnsignedLong(out[1]),
-      PyLong_FromUnsignedLong(out[2]),
-      PyLong_FromUnsignedLong(out[3]));
+  uint64_t u[4];
+  uint64s_from_BLSFieldElement(u, PyCapsule_GetPointer(c, PyCapsule_GetName(c)));
+  PyObject *out = PyLong_FromUnsignedLong(0);
+  PyObject *mult = PyLong_FromUnsignedLong(1);
+  PyObject *two64 = PyNumber_Power(PyLong_FromUnsignedLong(2), PyLong_FromUnsignedLong(64), Py_None);
+  for (int i = 0; i < 4; i++) {
+    out = PyNumber_Add(out, PyNumber_Multiply(mult, PyLong_FromUnsignedLong(u[i])));
+    mult = PyNumber_Multiply(mult, two64);
+  }
+  return out;
+}
+
+static PyObject* alloc_polynomial_wrap(PyObject *self, PyObject *args) {
+  PyObject *a;
+  if (!PyArg_UnpackTuple(args, "alloc_polynomial_wrap", 1, 1, &a) ||
+      !PySequence_Check(a))
+    return PyErr_Format(PyExc_ValueError, "expected sequence");
+
+  PolynomialEvalForm *p = (PolynomialEvalForm*)malloc(sizeof(PolynomialEvalForm));
+
+  if (p == NULL) return PyErr_NoMemory();
+
+  Py_ssize_t n = PySequence_Length(a);
+  p->length = n;
+
+  if (alloc_polynomial(p, n) != C_KZG_OK)
+    return PyErr_Format(PyExc_RuntimeError, "error allocating polynomial");
+
+  PyObject *e;
+  for (Py_ssize_t i = 0; i < n; i++) {
+    e = PySequence_GetItem(a, i);
+    if (!PyCapsule_IsValid(e, "BLSFieldElement")) {
+      free_polynomial(p);
+      free(p);
+      return PyErr_Format(PyExc_ValueError, "expected BLSFieldElement capsules");
+    }
+    p->values[i] = *(BLSFieldElement*)PyCapsule_GetPointer(e, "BLSFieldElement");
+  }
+
+  return PyCapsule_New(p, "PolynomialEvalForm", free_PolynomialEvalForm);
+}
+
+static PyObject* bytes_from_G1_wrap(PyObject *self, PyObject *args) {
+  PyObject *c;
+  if (!PyArg_UnpackTuple(args, "bytes_from_G1", 1, 1, &c) ||
+      !PyCapsule_IsValid(c, "G1"))
+    return PyErr_Format(PyExc_ValueError, "expected G1 capsule");
+  uint8_t bytes[48];
+  bytes_from_G1(bytes, PyCapsule_GetPointer(c, "G1"));
+  return PyBytes_FromStringAndSize((char*)bytes, 48);
 }
 
 static PyMethodDef ckzgmethods[] = {
-  {"uint64s_from_BLSFieldElement", uint64s_from_BLSFieldElement_wrap, METH_VARARGS, "Convert a field element to a 4-tuple of uint64s"},
-  {"bytes_to_bls_field", bytes_to_bls_field_wrap, METH_VARARGS, "Convert 32 bytes to a field element"},
+  {"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"},
+  {"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"},
+  // {"compute_powers",           compute_powers_wrap,         METH_VARARGS, "Create a list of powers of a field element"},
   {NULL, NULL, 0, NULL}
 };
 

min-bindings/python/tests.py

@@ -1,25 +1,9 @@
-import atexit
 import ckzg
 import random
-import ssz
-
-def int_from_fr(fr):
-    digits = ckzg.uint64s_from_BLSFieldElement(fr)
-    res, mult = 0, 1
-    for x in digits:
-        res += mult * x
-        mult *= 2**64
-    return res
 
 # Simple test of bytes_to_bls_field
 
 bs = (329).to_bytes(32, "little")
-fr = ckzg.bytes_to_bls_field(bs)
-assert int_from_fr(fr) == 329
+assert 329 == ckzg.int_from_BLSFieldElement(ckzg.bytes_to_bls_field(bs))
 
 print('Tests passed')
-
-def cleanup():
-    pass
-
-atexit.register(cleanup)