primes: avoid reinventing the gmemoize wheel

Technologicat · Technologicat · commit bd0d392686a3 · 2019-02-25T16:51:06.000+02:00
diff --git a/unpythonic/mathseq.py b/unpythonic/mathseq.py
@@ -25,7 +25,6 @@
            "cauchyprod", "diagonal_reduce",
            "fibonacci", "primes"]
 
-from threading import RLock
 from itertools import repeat, chain, count, takewhile
 from operator import add as primitive_add, mul as primitive_mul, \
                      pow as primitive_pow, mod as primitive_mod, \
@@ -34,7 +33,7 @@
                      neg as primitive_neg, pos as primitive_pos
 
 from .it import take, rev
-from .gmemo import imemoize
+from .gmemo import imemoize, gmemoize
 
 # stuff to support float, mpf and SymPy expressions transparently
 #
@@ -718,37 +717,23 @@ def diagonal():
 
 def fibonacci():
     """Return the Fibonacci numbers 1, 1, 2, 3, 5, 8, ... as a lazy sequence."""
-    a, b = 1, 1
-    while True:
-        yield a
-        a, b = b, a + b
-
-# Manual memoization for speed. One global memo to save memory when several
-# instances of primes() are created in the same session.
-_primesmemo = [2, 3]
-_primeslock = RLock()
+    def fibos():
+        a, b = 1, 1
+        while True:
+            yield a
+            a, b = b, a + b
+    return m(fibos())
+
+@gmemoize
+def _primes():  # at the top level; we want only one global memo shared across all instances
+    yield 2
+    for n in chain([3, 5, 7], (d + k for d in count(10, step=10)
+                                     for k in [1, 3, 7, 9])):
+        if not any(n % p == 0 for p in takewhile(lambda x: x*x <= n, _primes())):
+            yield n
 def primes():
     """Return the prime numbers 2, 3, 5, 7, 11, 13, ... as a lazy sequence.
 
-    Implementation is based on an FP version of the sieve of Eratosthenes,
-    with internal memoization.
+    FP sieve of Eratosthenes with memoization.
     """
-    # TODO: utilize the memo also when already in the loop.
-    def memoprimes():
-        # to avoid a race condition, we must know the last p actually yielded by us
-        # (since some other thread may insert a new entry to _primesmemo between
-        #  our last yield and when we read _primesmemo[-1] for the divmod check)
-        for p in _primesmemo:
-            yield p
-        # find the next candidate
-        decade, mod = divmod(p, 10)
-        inits = range(mod + 2, 10, 2) if mod < 9 else []
-        inits = (10*decade + k for k in inits)
-        for n in chain(inits, (d + k for d in count(10*(decade + 1), step=10)
-                                     for k in [1, 3, 7, 9])):
-            if not any(n % p == 0 for p in takewhile(lambda x: x*x <= n, _primesmemo)):
-                with _primeslock:  # lock to make test-and-update atomic.
-                    if _primesmemo[-1] < n:
-                        _primesmemo.append(n)
-                yield n
-    return m(memoprimes())
+    return m(_primes())
diff --git a/unpythonic/test/test_mathseq.py b/unpythonic/test/test_mathseq.py
@@ -219,9 +219,6 @@ def ulp(x):  # Unit in the Last Place
     assert tuple(take(10, primes())) == (2, 3, 5, 7, 11, 13, 17, 19, 23, 29)
     assert tuple(take(10, fibonacci())) == (1, 1, 2, 3, 5, 8, 13, 21, 34, 55)
 
-    assert tuple(take(20, primes())) == tuple(take(20, primes()))  # caching
-    assert last(take(21, primes())) == 73
-
     print("All tests PASSED")
 
 if __name__ == '__main__':
