From: Mark Wooding <mdw@distorted.org.uk>
Date: Mon, 25 Nov 2019 12:57:55 +0000 (+0000)
Subject: pgen.c: Add a binding for Catacomb's `primeiter'.
X-Git-Url: https://git.distorted.org.uk/~mdw/catacomb-python/commitdiff_plain/2b08f130f77549bba01e22a8f0a3b2a1905b8081

pgen.c: Add a binding for Catacomb's `primeiter'.
---

diff --git a/catacomb-python.h b/catacomb-python.h
index 67117fa..aa55cc5 100644
--- a/catacomb-python.h
+++ b/catacomb-python.h
@@ -81,6 +81,7 @@ PUBLIC_SYMBOLS;
 #include <catacomb/mp-fibonacci.h>
 
 #include <catacomb/pgen.h>
+#include <catacomb/primeiter.h>
 #include <catacomb/pfilt.h>
 #include <catacomb/strongprime.h>
 #include <catacomb/limlee.h>
diff --git a/pgen.c b/pgen.c
index f1a11f4..e08c6fe 100644
--- a/pgen.c
+++ b/pgen.c
@@ -921,6 +921,83 @@ static const PyTypeObject pgtest_pytype_skel = {
   0					/* @tp_is_gc@ */
 };
 
+/*----- Prime iteration ---------------------------------------------------*/
+
+static PyTypeObject *piter_pytype;
+
+typedef struct piter_pyobj {
+  PyObject_HEAD
+  primeiter i;
+} piter_pyobj;
+#define PITER_I(o) (&((piter_pyobj *)(o))->i)
+
+static PyObject *piter_pynew(PyTypeObject *ty, PyObject *arg, PyObject *kw)
+{
+  piter_pyobj *rc;
+  mp *n = 0;
+  static const char *const kwlist[] = { "start", 0 };
+
+  if (!PyArg_ParseTupleAndKeywords(arg, kw, "|O&:new", KWLIST, convmp, &n))
+    return (0);
+  rc = (piter_pyobj *)ty->tp_alloc(ty, 0);
+  primeiter_create(&rc->i, n);
+  return ((PyObject *)rc);
+}
+
+static void piter_pydealloc(PyObject *me)
+  { primeiter_destroy(PITER_I(me)); FREEOBJ(me); }
+
+static PyObject *piter_pynext(PyObject *me)
+  { return (mp_pywrap(primeiter_next(PITER_I(me), 0))); }
+
+static const PyTypeObject piter_pytype_skel = {
+  PyVarObject_HEAD_INIT(0, 0)		/* Header */
+  "PrimeIter",				/* @tp_name@ */
+  sizeof(piter_pyobj),			/* @tp_basicsize@ */
+  0,					/* @tp_itemsize@ */
+
+  piter_pydealloc,			/* @tp_dealloc@ */
+  0,					/* @tp_print@ */
+  0,					/* @tp_getattr@ */
+  0,					/* @tp_setattr@ */
+  0,					/* @tp_compare@ */
+  0,					/* @tp_repr@ */
+  0,					/* @tp_as_number@ */
+  0,					/* @tp_as_sequence@ */
+  0,					/* @tp_as_mapping@ */
+  0,					/* @tp_hash@ */
+  0,					/* @tp_call@ */
+  0,					/* @tp_str@ */
+  0,					/* @tp_getattro@ */
+  0,					/* @tp_setattro@ */
+  0,					/* @tp_as_buffer@ */
+  Py_TPFLAGS_DEFAULT |			/* @tp_flags@ */
+    Py_TPFLAGS_BASETYPE,
+
+  /* @tp_doc@ */
+  "PrimeIter([start = N]): Prime-number iterator.",
+
+  0,					/* @tp_traverse@ */
+  0,					/* @tp_clear@ */
+  0,					/* @tp_richcompare@ */
+  0,					/* @tp_weaklistoffset@ */
+  PyObject_SelfIter,			/* @tp_iter@ */
+  piter_pynext,				/* @tp_iternext@ */
+  0,					/* @tp_methods@ */
+  0,					/* @tp_members@ */
+  0,					/* @tp_getset@ */
+  0,					/* @tp_base@ */
+  0,					/* @tp_dict@ */
+  0,					/* @tp_descr_get@ */
+  0,					/* @tp_descr_set@ */
+  0,					/* @tp_dictoffset@ */
+  0,					/* @tp_init@ */
+  PyType_GenericAlloc,			/* @tp_alloc@ */
+  piter_pynew,				/* @tp_new@ */
+  0,					/* @tp_free@ */
+  0					/* @tp_is_gc@ */
+};
+
 /*----- Prime generation functions ----------------------------------------*/
 
 void pgenerr(struct excinfo *exc)
@@ -1096,6 +1173,7 @@ void pgen_pyinit(void)
   INITTYPE(pgstep, pgev);
   INITTYPE(pgjump, pgev);
   INITTYPE(pgtest, pgev);
+  INITTYPE(piter, root);
   addmethods(methods);
 }
 
@@ -1112,6 +1190,7 @@ void pgen_pyinsert(PyObject *mod)
   INSERT("pgen_stdev", pgev_stdev(pgen_ev));
   INSERT("pgen_spinev", pgev_stdev(pgen_evspin));
   INSERT("pgen_subev", pgev_stdev(pgen_subev));
+  INSERT("PrimeIter", piter_pytype);
   setconstants(mod, consts);
 }
 
diff --git a/t/t-pgen.py b/t/t-pgen.py
index dd66903..651ce70 100644
--- a/t/t-pgen.py
+++ b/t/t-pgen.py
@@ -223,6 +223,22 @@ class TestPGen (U.TestCase):
     me.assertEqual(p.nbits, 512)
     me.assertEqual(ev.events, "[p [h]:F17/P6/D][p:F60/P26/D]")
 
+###--------------------------------------------------------------------------
+class TestPrimeIter (U.TestCase):
+
+  def test(me):
+    me.assertEqual(list(I.islice(C.PrimeIter(0), 5)), [2, 3, 5, 7, 11])
+    me.assertEqual(list(I.islice(C.PrimeIter(1000), 5)),
+                   [1009, 1013, 1019, 1021, 1031])
+    me.assertEqual(list(I.islice(C.PrimeIter(1000000), 5)),
+                   [1000003, 1000033, 1000037, 1000039, 1000081])
+    me.assertEqual(list(I.islice(C.PrimeIter(1000000000000000000000000000000000000), 5)),
+                   [1000000000000000000000000000000000067,
+                    1000000000000000000000000000000000123,
+                    1000000000000000000000000000000000141,
+                    1000000000000000000000000000000000159,
+                    1000000000000000000000000000000000163])
+
 ###----- That's all, folks --------------------------------------------------
 
 if __name__ == "__main__": U.main()