From 10cc0d2af0b3e68071b233549e1892ef63e22ee8 Mon Sep 17 00:00:00 2001
From: Lumi Pakkanen <lumi.pakkanen@gmail.com>
Date: Sun, 28 Apr 2024 14:33:42 +0300
Subject: [PATCH] Implement utilities for obtaining the nth prime or a range of
 primes

ref #32
---
 src/__tests__/primes.spec.ts | 64 +++++++++++++++++++++++++++-
 src/primes.ts                | 81 ++++++++++++++++++++++++++++++++++++
 2 files changed, 144 insertions(+), 1 deletion(-)

diff --git a/src/__tests__/primes.spec.ts b/src/__tests__/primes.spec.ts
index 771196e..79b9e95 100644
--- a/src/__tests__/primes.spec.ts
+++ b/src/__tests__/primes.spec.ts
@@ -1,5 +1,5 @@
 import {describe, it, expect} from 'vitest';
-import {PRIMES, isPrime, primes} from '../primes';
+import {PRIMES, isPrime, nthPrime, primeRange, primes} from '../primes';
 
 describe('Array of prime numbers', () => {
   it('has no gaps', () => {
@@ -84,3 +84,65 @@ describe('Lists of primes', () => {
     expect(primes(7920, 8094)).toEqual(LARGER_PRIMES);
   });
 });
+
+describe('Nth odd prime generator', () => {
+  it('knowns 3 is the 1st odd prime', () => {
+    expect(nthPrime(1)).toBe(3);
+  });
+
+  it('knowns 2 is the prime at index 0', () => {
+    expect(nthPrime(0)).toBe(2);
+  });
+
+  it('returns undefined for a negative index', () => {
+    expect(nthPrime(-1)).toBe(undefined);
+  });
+
+  it('returns undefined for a fractional index', () => {
+    expect(nthPrime(0.5)).toBe(undefined);
+  });
+
+  it('knows that 7919 is the 999th odd prime', () => {
+    expect(nthPrime(999)).toBe(7919);
+  });
+
+  it('knows that 7927 is the 1000th odd prime', () => {
+    expect(nthPrime(1000)).toBe(7927);
+  });
+
+  it('knows the about larger primes', () => {
+    for (let i = 0; i < LARGER_PRIMES.length; ++i) {
+      expect(nthPrime(1000 + i)).toBe(LARGER_PRIMES[i]);
+    }
+  });
+});
+
+describe('Prime range generator', () => {
+  it('produces an empty range for start > end', () => {
+    expect(primeRange(99999, 9999)).toEqual([]);
+  });
+
+  it('produces an empty range for start === end', () => {
+    expect(primeRange(9999, 9999)).toEqual([]);
+  });
+
+  it('produces the first 4 primes', () => {
+    expect(primeRange(0, 4)).toEqual([2, 3, 5, 7]);
+  });
+
+  it('produces 5 primes starting from the 999th odd prime', () => {
+    expect(primeRange(999, 999 + 5)).toEqual([7919, 7927, 7933, 7937, 7949]);
+  });
+
+  it('produces the larger primes', () => {
+    expect(primeRange(1000, 1000 + LARGER_PRIMES.length)).toEqual(
+      LARGER_PRIMES
+    );
+  });
+
+  it('produces the larger primes but one', () => {
+    expect(primeRange(1001, 1000 + LARGER_PRIMES.length)).toEqual(
+      LARGER_PRIMES.slice(1)
+    );
+  });
+});
diff --git a/src/primes.ts b/src/primes.ts
index cb3767e..c82b4f3 100644
--- a/src/primes.ts
+++ b/src/primes.ts
@@ -128,6 +128,30 @@ export function isPrime(n: number) {
   return true;
 }
 
+/**
+ * Obtain the nth odd prime.
+ * @param n 1-based ordinal of the nth odd prime, or zero to obtain prime two.
+ * @returns The nth odd prime number or `2` if `n === 0`.
+ */
+export function nthPrime(n: number) {
+  if (n < PRIMES.length) {
+    return PRIMES[n];
+  }
+  n -= PRIMES.length;
+  let prime = 7927;
+  while (n) {
+    // Skip over known composites mod 6.
+    prime += 4;
+    n -= Number(isPrime(prime));
+    if (!n) {
+      return prime;
+    }
+    prime += 2;
+    n -= Number(isPrime(prime));
+  }
+  return prime;
+}
+
 /**
  * Generate an array of prime numbers.
  * @param start Smallest prime number to include (or the next smallest prime).
@@ -180,3 +204,60 @@ export function primes(start: number, end?: number) {
   }
   return result;
 }
+
+/**
+ * Obtain a range of odd primes starting at ordinal the given ordinal.
+ * @param start 1-based ordinal of the nth odd prime to start from, or zero to include prime two.
+ * @param end Range end. `end - start` elements are returned.
+ * @returns The primes in the range.
+ */
+export function primeRange(start: number, end: number) {
+  start = Math.round(Math.max(0, start));
+  end = Math.round(end);
+  if (end <= start) {
+    return [];
+  }
+  if (end <= PRIMES.length) {
+    return PRIMES.slice(start, end);
+  }
+  const result = PRIMES.slice(start);
+  start = Math.max(0, start - PRIMES.length);
+  end -= PRIMES.length;
+
+  // We skip over known composites mod 6.
+  let prime = 7927;
+  let delta: number;
+  let a = 4;
+  let b = 2;
+  while (start) {
+    prime += 4;
+    delta = Number(isPrime(prime));
+    start -= delta;
+    end -= delta;
+    if (!start) {
+      a = 2;
+      b = 4;
+      break;
+    }
+    prime += 2;
+    delta = Number(isPrime(prime));
+    start -= delta;
+    end -= delta;
+  }
+  while (end) {
+    if (isPrime(prime)) {
+      result.push(prime);
+      end--;
+      if (!end) {
+        return result;
+      }
+    }
+    prime += a;
+    if (isPrime(prime)) {
+      result.push(prime);
+      end--;
+    }
+    prime += b;
+  }
+  return result;
+}