primeiter: New functions for iterating over small primes.
[u/mdw/catacomb] / primeiter.c
diff --git a/primeiter.c b/primeiter.c
new file mode 100644 (file)
index 0000000..d7126ba
--- /dev/null
@@ -0,0 +1,254 @@
+/* -*-c-*-
+ *
+ * Iterate over small primes efficiently
+ *
+ * (c) 2007 Straylight/Edgeware
+ */
+
+/*----- Licensing notice --------------------------------------------------*
+ *
+ * This file is part of Catacomb.
+ *
+ * Catacomb is free software; you can redistribute it and/or modify
+ * it under the terms of the GNU Library General Public License as
+ * published by the Free Software Foundation; either version 2 of the
+ * License, or (at your option) any later version.
+ *
+ * Catacomb is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU Library General Public License for more details.
+ *
+ * You should have received a copy of the GNU Library General Public
+ * License along with Catacomb; if not, write to the Free
+ * Software Foundation, Inc., 59 Temple Place - Suite 330, Boston,
+ * MA 02111-1307, USA.
+ */
+
+/*----- Header files ------------------------------------------------------*/
+
+#include "fibrand.h"
+#include "mp.h"
+#include "pgen.h"
+#include "primeiter.h"
+#include "primetab.h"
+#include "wheel.h"
+
+/*----- Theory ------------------------------------------------------------*
+ *
+ * For small primes, we can just pluck them out of the small primes table.
+ * For larger primes, we can test them individually, or build a sieve or
+ * something, but since we don't know when to stop, that could be tricky.
+ *
+ * We've built a `wheel', as follows.  Let %$m$% be the product of the first
+ * %$n$% primes.  There are %$\phi(m)$% integers %$n_i$%, with %$0 < n_i <
+ * m$% coprime to %$m$%, and any integer %$j > n$% must be congruent to some
+ * %$n_i$% modulo %$m$%.  The wheel itself doesn't list the %$n_i$%, but
+ * rather the differences %$\delta_i = n_i - n_{i-1}$% (wrapping around
+ * appropriately at the ends), so you can just add simple offsets to step
+ * onwards.  The wheel assumes you start at 1 and move on round.
+ */
+
+/*----- Main code ---------------------------------------------------------*/
+
+/* --- @wheelsync@ --- *
+ *
+ * Arguments:  @primeiter *pi@ = iterator to synchronize
+ *             @mp *where@ = value to synchronize
+ *
+ * Returns:    ---
+ *
+ * Use:                Sets up the wheel index to match the given integer.  After
+ *             this, we can step along the wheel to find candidate primes.
+ */
+
+static void wheelsync(primeiter *pi, mp *where)
+{
+  mpw w;
+  mp t;
+  mpw rr;
+  mp *r = MP_NEW;
+  unsigned i, n;
+
+  w = WHEELMOD;
+  mp_build(&t, &w, &w + 1);
+  mp_div(0, &r, where, &t);
+  rr = MP_ZEROP(r) ? 0 : r->v[0];
+
+  for (i = 0, n = 1; rr > n; n += wheel[i], i++);
+  w = n - rr;
+  pi->p = mp_add(MP_NEW, where, &t);
+  pi->i = i;
+  pi->r = fibrand_create(0);
+  MP_DROP(r);
+}
+
+/* --- @primeiter_create@ --- *
+ *
+ * Arguments:  @primeiter *pi@ = pointer to an iterator structure
+ *             @mp *start@ = where to start
+ *
+ * Returns:    ---
+ *
+ * Use:                Initializes a prime iterator.  The first output will be the
+ *             smallest prime not less than @start@.
+ */
+
+void primeiter_create(primeiter *pi, mp *start)
+{
+  mpw n;
+  unsigned l, h, m;
+
+  if (!start || MP_CMP(start, <=, MP_TWO))
+    start = MP_TWO;
+
+  if (MP_LEN(start) <= 1) {
+    n = start->v[0];
+    if (n <= MAXPRIME) {
+      l = 0;
+      h = NPRIME;
+      for (;;) {
+       m = l + (h - l)/2;
+       if (primetab[m] == n) break;
+       else if (m == l) { m++; break; }
+       else if (primetab[m] < n) l = m;
+       else h = m;
+      }
+      pi->i = m;
+      pi->mode = PIMODE_PTAB;
+      mp_build(&pi->pp, &pi->w, &pi->w + 1);
+      pi->p = &pi->pp;
+      return;
+    }
+  }
+
+  wheelsync(pi, start);
+  pi->mode = PIMODE_STALL;
+}
+
+/* --- @primeiter_destroy@ --- *
+ *
+ * Arguments:  @primeiter *pi@ = pointer to iterator structure
+ *
+ * Returns:    ---
+ *
+ * Use:                Frees up an iterator structure when it's no longer wanted.
+ */
+
+void primeiter_destroy(primeiter *pi)
+{
+  switch (pi->mode) {
+    case PIMODE_PTAB:
+      break;
+    case PIMODE_STALL:
+    case PIMODE_WHEEL:
+      MP_DROP(pi->p);
+      GR_DESTROY(pi->r);
+      break;
+    default:
+      abort();
+  }
+}
+
+/* --- @primeiter_next@ --- *
+ *
+ * Arguments:  @primeiter *pi@ = pointer to an iterator structure
+ *             @mp *d@ = fake destination
+ *
+ * Returns:    The next prime number from the iterator.
+ *
+ * Use:                Returns a new prime number.
+ */
+
+mp *primeiter_next(primeiter *pi, mp *d)
+{
+  mp *p;
+
+  switch (pi->mode) {
+    case PIMODE_PTAB:
+      pi->w = primetab[pi->i++];
+      if (pi->i >= NPRIME) {
+       wheelsync(pi, pi->p);
+       pi->mode = PIMODE_WHEEL;
+      }
+      p = MP_COPY(pi->p);
+      MP_SPLIT(p);
+      break;
+    case PIMODE_STALL:
+      pi->mode = PIMODE_WHEEL;
+      goto loop;
+    case PIMODE_WHEEL:
+      do {
+       MP_DEST(pi->p, MP_LEN(pi->p) + 1, pi->p->f);
+       MPX_UADDN(pi->p->v, pi->p->vl, wheel[pi->i++]);
+       MP_SHRINK(pi->p);
+       if (pi->i >= WHEELN) pi->i = 0;
+      loop:;
+      } while (!pgen_primep(pi->p, pi->r));
+      p = MP_COPY(pi->p);
+      break;
+    default:
+      abort();
+  }
+  if (d) MP_DROP(d);
+  return (p);
+}
+
+/*----- Test rig ----------------------------------------------------------*/
+
+#ifdef TEST_RIG
+
+#include <mLib/macros.h>
+#include <mLib/testrig.h>
+
+static int test(dstr *v)
+{
+  mp *start = *(mp **)v[0].buf;
+  mp *pp[5], *ret[5];
+  int i;
+  primeiter pi;
+  int ok = 1;
+
+  for (i = 0; i < N(pp); i++)
+    pp[i] = *(mp **)v[i + 1].buf;
+  primeiter_create(&pi, start);
+  for (i = 0; i < N(pp); i++) {
+    ret[i] = primeiter_next(&pi, MP_NEW);
+    if (!MP_EQ(ret[i], pp[i])) ok = 0;
+  }
+  if (!ok) {
+    fprintf(stderr, "\n*** primeiter test failure:\n***   start = ");
+    mp_writefile(start, stderr, 10);
+    for (i = 0; i < N(pp); i++) {
+      fprintf(stderr, "\n***   p[%d] = ", i);
+      mp_writefile(ret[i], stderr, 10);
+      fprintf(stderr, " %s ", MP_EQ(ret[i], pp[i]) ? "==" : "!=");
+      mp_writefile(pp[i], stderr, 10);
+    }
+    fputc('\n', stderr);
+  }
+  for (i = 0; i < N(pp); i++) {
+    MP_DROP(pp[i]);
+    MP_DROP(ret[i]);
+  }
+  primeiter_destroy(&pi);
+  MP_DROP(start);
+  assert(mparena_count(MPARENA_GLOBAL) == 0);
+  return (ok);
+}
+
+static test_chunk tests[] = {
+  { "primeiter", test,
+    { &type_mp,  &type_mp, &type_mp, &type_mp, &type_mp, &type_mp, } },
+  { 0 }
+};
+
+int main(int argc, char *argv[])
+{
+  test_run(argc, argv, tests, SRCDIR "/tests/pgen");
+  return (0);
+}
+
+#endif
+
+/*----- That's all, folks -------------------------------------------------*/