mp.c: Accept arbitrary iterables as inputs to `MPCRT'.

[catacomb-python] / mp.c

diff --git a/mp.c b/mp.c
index 265a5c8..3549ccb 100644 (file)
--- a/mp.c
+++ b/mp.c
@@ -1784,7 +1784,7 @@ typedef struct mpcrt_pyobj {
 static PyObject *mcmeth_solve(PyObject *me, PyObject *arg)
 {
   mpcrt *c = MPCRT_PY(me);
-  PyObject *q = 0, *x, *z = 0;
+  PyObject *q = 0, *it, *x, *z = 0;
   mp *xx;
   mp **v = 0;
   Py_ssize_t i = 0, n = c->k;
@@ -1793,15 +1793,20 @@ static PyObject *mcmeth_solve(PyObject *me, PyObject *arg)
     q = arg;
   else if (!PyArg_ParseTuple(arg, "O:solve", &q))
     goto end;
-  if (!PySequence_Check(q)) TYERR("want a sequence of residues");
-  i = PySequence_Size(q); if (i < 0) goto end;
-  if (i != n) VALERR("residue count mismatch");
-  v = xmalloc(n * sizeof(*v));
-  for (i = 0; i < n; i++) {
-    if ((x = PySequence_GetItem(q, i)) == 0) goto end;
+  v = xmalloc(n*sizeof(*v));
+  it = PyObject_GetIter(q); if (!it) goto end;
+  while (i < n) {
+    x = PyIter_Next(it);
+    if (!x) {
+      if (PyErr_Occurred()) goto end;
+      VALERR("residue count mismatch");
+    }
     xx = getmp(x); Py_DECREF(x); if (!xx) goto end;
-    v[i] = xx;
+    v[i++] = xx; xx = 0;
   }
+  x = PyIter_Next(it);
+  if (x) { Py_DECREF(x); VALERR("residue count mismatch"); }
+  else if (PyErr_Occurred()) goto end;
   z = mp_pywrap(mpcrt_solve(c, MP_NEW, v));
 end:
   if (v) {
@@ -1823,7 +1828,7 @@ static PyObject *mpcrt_pynew(PyTypeObject *ty, PyObject *arg, PyObject *kw)
   mpcrt_mod *v = 0;
   Py_ssize_t n, i = 0, j;
   static const char *const kwlist[] = { "mv", 0 };
-  PyObject *q = 0, *x;
+  PyObject *q, *it = 0, *x;
   mp *xx = MP_NEW, *y = MP_NEW, *g = MP_NEW;
   mpmul mm;
   mpcrt_pyobj *c = 0;
@@ -1832,16 +1837,26 @@ static PyObject *mpcrt_pynew(PyTypeObject *ty, PyObject *arg, PyObject *kw)
     q = arg;
   else if (!PyArg_ParseTupleAndKeywords(arg, kw, "O:new", KWLIST, &q))
     goto end;
-  if (!PySequence_Check(q)) TYERR("want a sequence of moduli");
-  n = PySequence_Size(q); if (n < 0) goto end;
-  if (!n) VALERR("want at least one modulus");
-  v = xmalloc(n * sizeof(*v));
-  for (i = 0; i < n; i++) {
-    if ((x = PySequence_GetItem(q, i)) == 0) goto end;
+
+  if (!PySequence_Check(q))
+    n = 16;
+  else {
+    n = PySequence_Size(arg);
+    if (n == (size_t)-1 && PyErr_Occurred()) goto end;
+  }
+
+  v = xmalloc(n*sizeof(*v));
+  it = PyObject_GetIter(q); if (!it) goto end;
+  for (;;) {
+    x = PyIter_Next(it); if (!x) break;
     xx = getmp(x); Py_DECREF(x); if (!xx) goto end;
     if (MP_CMP(xx, <=, MP_ZERO)) VALERR("moduli must be positive");
-    v[i].m = xx; v[i].n = 0; v[i].ni = 0; v[i].nni = 0; xx = MP_NEW;
+    if (i >= n) { n *= 2; v = xrealloc(v, n*sizeof(*v), i*sizeof(*v)); }
+    v[i].m = xx; v[i].n = 0; v[i].ni = 0; v[i].nni = 0; i++; xx = MP_NEW;
   }
+  if (PyErr_Occurred()) goto end;
+  Py_DECREF(it); it = 0;
+
   mpmul_init(&mm);
   for (j = 0; j < i; j++) mpmul_add(&mm, v[j].m);
   xx = mpmul_done(&mm);
@@ -1852,7 +1867,7 @@ static PyObject *mpcrt_pynew(PyTypeObject *ty, PyObject *arg, PyObject *kw)
   }
 
   c = (mpcrt_pyobj *)ty->tp_alloc(ty, 0);
-  mpcrt_create(&c->c, v, n, 0);
+  mpcrt_create(&c->c, v, i, 0);
   mp_drop(xx); mp_drop(y); mp_drop(g);
   return ((PyObject *)c);
 
@@ -1862,6 +1877,7 @@ end:
     xfree(v);
   }
   mp_drop(xx); mp_drop(y); mp_drop(g);
+  Py_XDECREF(it);
   return (0);
 }