Refactor CompressedPrimeSieve: add docs, clean up, add tests

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: williamfiset

src/main/java/com/williamfiset/algorithms/math/CompressedPrimeSieve.java

@@ -1,76 +1,70 @@
 /**
- * Generate a compressed prime sieve using bit manipulation. The idea is that each bit represents a
- * boolean value indicating whether a number is prime or not. This saves a lot of room when creating
- * the sieve. In this implementation I store all odd numbers in individual longs meaning that for
- * each long I use I can represent a range of 128 numbers (even numbers are omitted because they are
- * not prime, with the exception of 2 which is handled as a special case).
+ * Generates a compressed prime sieve using bit manipulation. Each bit represents whether an odd
+ * number is prime or not. Even numbers are omitted (except 2, handled as a special case), so each
+ * long covers a range of 128 numbers.
  *
- * <p>Time Complexity: ~O(nloglogn)
+ * <p>Time: ~O(n log(log(n)))
  *
- * <p>Compile: javac -d src/main/java
- * src/main/java/com/williamfiset/algorithms/math/CompressedPrimeSieve.java
- *
- * <p>Run: java -cp src/main/java com/williamfiset/algorithms/math/CompressedPrimeSieve
+ * <p>Space: O(n / 128) longs
  *
  * @author William Fiset, william.alexandre.fiset@gmail.com
  */
 package com.williamfiset.algorithms.math;
 
 public class CompressedPrimeSieve {
+
   private static final double NUM_BITS = 128.0;
   private static final int NUM_BITS_SHIFT = 7; // 2^7 = 128
 
-  // Sets the bit representing n to 1 indicating this number is not prime
+  // Marks n as not prime by setting its bit to 1.
   private static void setBit(long[] arr, int n) {
-    if ((n & 1) == 0) return; // n is even
+    if ((n & 1) == 0)
+      return;
     arr[n >> NUM_BITS_SHIFT] |= 1L << ((n - 1) >> 1);
   }
 
-  // Returns true if the bit for n is off (meaning n is a prime).
-  // Note: do use this method to access numbers outside your prime sieve range!
+  // Returns true if n's bit is unset (meaning n is prime).
   private static boolean isNotSet(long[] arr, int n) {
-    if (n < 2) return false; // n is not prime
-    if (n == 2) return true; // two is prime
-    if ((n & 1) == 0) return false; // n is even
+    if (n < 2)
+      return false;
+    if (n == 2)
+      return true;
+    if ((n & 1) == 0)
+      return false;
     long chunk = arr[n >> NUM_BITS_SHIFT];
     long mask = 1L << ((n - 1) >> 1);
     return (chunk & mask) != mask;
   }
 
-  // Returns true/false depending on whether n is prime.
+  /** Returns true if n is prime according to the given sieve. */
   public static boolean isPrime(long[] sieve, int n) {
     return isNotSet(sieve, n);
   }
 
-  // Returns an array of longs with each bit indicating whether a number
-  // is prime or not. Use the isNotSet and setBit methods to toggle to bits for each number.
+  /**
+   * Builds a compressed prime sieve for all numbers up to {@code limit}.
+   *
+   * @param limit the upper bound (inclusive) for the sieve.
+   * @return a bit-packed array where each bit indicates whether an odd number is composite.
+   */
   public static long[] primeSieve(int limit) {
-    final int numChunks = (int) Math.ceil(limit / NUM_BITS);
-    final int sqrtLimit = (int) Math.sqrt(limit);
-    // if (limit < 2) return 0; // uncomment for primeCount purposes
-    // int primeCount = (int) Math.ceil(limit / 2.0); // Counts number of primes <= limit
+    int numChunks = (int) Math.ceil(limit / NUM_BITS);
+    int sqrtLimit = (int) Math.sqrt(limit);
     long[] chunks = new long[numChunks];
-    chunks[0] = 1; // 1 as not prime
+    chunks[0] = 1; // Mark 1 as not prime.
     for (int i = 3; i <= sqrtLimit; i += 2)
       if (isNotSet(chunks, i))
         for (int j = i * i; j <= limit; j += i)
-          if (isNotSet(chunks, j)) {
+          if (isNotSet(chunks, j))
             setBit(chunks, j);
-            // primeCount--;
-          }
     return chunks;
   }
 
-  /* Example usage. */
-
   public static void main(String[] args) {
-    final int limit = 200;
-    long[] sieve = CompressedPrimeSieve.primeSieve(limit);
-
-    for (int i = 0; i <= limit; i++) {
-      if (CompressedPrimeSieve.isPrime(sieve, i)) {
+    int limit = 200;
+    long[] sieve = primeSieve(limit);
+    for (int i = 0; i <= limit; i++)
+      if (isPrime(sieve, i))
         System.out.printf("%d is prime!\n", i);
-      }
-    }
   }
 }

src/test/java/com/williamfiset/algorithms/math/BUILD

@@ -0,0 +1,27 @@
+load("@rules_java//java:defs.bzl", "java_test")
+
+# Common dependencies for JUnit 5 tests
+JUNIT5_DEPS = [
+    "@maven//:org_junit_jupiter_junit_jupiter_api",
+    "@maven//:org_junit_jupiter_junit_jupiter_engine",
+]
+
+JUNIT5_RUNTIME_DEPS = [
+    "@maven//:org_junit_platform_junit_platform_console",
+]
+
+TEST_DEPS = [
+    "//src/main/java/com/williamfiset/algorithms/math:math",
+    "@maven//:com_google_truth_truth",
+] + JUNIT5_DEPS
+
+# bazel test //src/test/java/com/williamfiset/algorithms/math:CompressedPrimeSieveTest
+java_test(
+    name = "CompressedPrimeSieveTest",
+    srcs = ["CompressedPrimeSieveTest.java"],
+    main_class = "org.junit.platform.console.ConsoleLauncher",
+    use_testrunner = False,
+    args = ["--select-class=com.williamfiset.algorithms.math.CompressedPrimeSieveTest"],
+    runtime_deps = JUNIT5_RUNTIME_DEPS,
+    deps = TEST_DEPS,
+)

src/test/java/com/williamfiset/algorithms/math/CompressedPrimeSieveTest.java

@@ -0,0 +1,77 @@
+package com.williamfiset.algorithms.math;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import org.junit.jupiter.api.*;
+
+public class CompressedPrimeSieveTest {
+
+  // First 25 primes (all primes <= 100).
+  private static final int[] PRIMES_UP_TO_100 = {
+    2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61, 67, 71, 73, 79, 83, 89,
+    97
+  };
+
+  @Test
+  public void testPrimesUpTo100() {
+    long[] sieve = CompressedPrimeSieve.primeSieve(100);
+    java.util.Set<Integer> expected = new java.util.HashSet<>();
+    for (int p : PRIMES_UP_TO_100)
+      expected.add(p);
+
+    for (int i = 0; i <= 100; i++) {
+      if (expected.contains(i))
+        assertThat(CompressedPrimeSieve.isPrime(sieve, i)).isTrue();
+      else
+        assertThat(CompressedPrimeSieve.isPrime(sieve, i)).isFalse();
+    }
+  }
+
+  @Test
+  public void testZeroAndOneAreNotPrime() {
+    long[] sieve = CompressedPrimeSieve.primeSieve(10);
+    assertThat(CompressedPrimeSieve.isPrime(sieve, 0)).isFalse();
+    assertThat(CompressedPrimeSieve.isPrime(sieve, 1)).isFalse();
+  }
+
+  @Test
+  public void testTwoIsPrime() {
+    long[] sieve = CompressedPrimeSieve.primeSieve(10);
+    assertThat(CompressedPrimeSieve.isPrime(sieve, 2)).isTrue();
+  }
+
+  @Test
+  public void testEvenNumbersAreNotPrime() {
+    long[] sieve = CompressedPrimeSieve.primeSieve(200);
+    for (int i = 4; i <= 200; i += 2)
+      assertThat(CompressedPrimeSieve.isPrime(sieve, i)).isFalse();
+  }
+
+  @Test
+  public void testPrimeCount() {
+    // There are 168 primes <= 1000.
+    long[] sieve = CompressedPrimeSieve.primeSieve(1000);
+    int count = 0;
+    for (int i = 0; i <= 1000; i++)
+      if (CompressedPrimeSieve.isPrime(sieve, i))
+        count++;
+    assertThat(count).isEqualTo(168);
+  }
+
+  @Test
+  public void testLargePrimes() {
+    long[] sieve = CompressedPrimeSieve.primeSieve(10000);
+    // Some known primes near the upper range.
+    assertThat(CompressedPrimeSieve.isPrime(sieve, 9973)).isTrue();
+    assertThat(CompressedPrimeSieve.isPrime(sieve, 9967)).isTrue();
+    // Some known composites near the upper range.
+    assertThat(CompressedPrimeSieve.isPrime(sieve, 9999)).isFalse();
+    assertThat(CompressedPrimeSieve.isPrime(sieve, 10000)).isFalse();
+  }
+
+  @Test
+  public void testSmallLimit() {
+    long[] sieve = CompressedPrimeSieve.primeSieve(2);
+    assertThat(CompressedPrimeSieve.isPrime(sieve, 2)).isTrue();
+  }
+}