Fix primality test for very small numbers

.github/workflows/test.yaml

@@ -45,6 +45,8 @@ jobs:
         ./2weso_test
         echo "Running prover_test"
         ./prover_test
+        echo "Running primes_test"
+        ./primes_test
 
     - name: Test vdf-client
       if: matrix.config != 'optimized=1'

.gitignore

@@ -25,6 +25,7 @@ src/vdf_client
 src/1weso_test
 src/2weso_test
 src/prover_test
+src/primes_test
 /verifier
 /verifier_test
 /build/*

src/1weso_test.cpp

@@ -1,9 +1,9 @@
 #include "vdf.h"
+#include "vdf_assert.h"
 #include "create_discriminant.h"
 #include "verifier.h"
 
 #include <atomic>
-#include <cassert>
 
 int segments = 7;
 int thread_count = 3;
@@ -60,7 +60,7 @@ int main(int argc, char const* argv[]) try
     form proof_form = DeserializeForm(D, proof.proof.data(), proof.proof.size());
     VerifyWesolowskiProof(D, x_init, y, proof_form, iter, is_valid);
     std::cout << "Verify result: " << is_valid << "\n";
-    assert(is_valid);
+    VDF_ASSERT(is_valid);
 }
 catch (std::exception const& e) {
     std::cerr << "Exception " << e.what() << '\n';

src/Makefile.vdf-client

@@ -27,15 +27,17 @@ endif
 
 .PHONY: all clean
 
-all: vdf_client prover_test 1weso_test 2weso_test vdf_bench
+ALL = vdf_client primes_test prover_test 1weso_test 2weso_test vdf_bench
+
+all: $(ALL)
 
 clean:
-	rm -f *.o vdf_client prover_test 1weso_test 2weso_test compile_asm vdf_bench
+	rm -f *.o $(ALL) compile_asm avx512_test
 
-vdf_client vdf_bench prover_test 1weso_test 2weso_test avx512_test: %: %.o lzcnt.o asm_compiled.o avx2_asm_compiled.o avx512_asm_compiled.o
+$(ALL) avx512_test: %: %.o lzcnt.o asm_compiled.o avx2_asm_compiled.o avx512_asm_compiled.o
 	$(CXX) $(LDFLAGS) -o $@ $^ $(LDLIBS)
 
-vdf_client.o vdf_bench.o prover_test.o 1weso_test.o 2weso_test.o avx512_test.o: CXXFLAGS += $(OPT_CFLAGS)
+$(addsuffix .o,$(ALL)) avx512_test.o: CXXFLAGS += $(OPT_CFLAGS)
 
 lzcnt.o: refcode/lzcnt.c
 	$(CC) -c refcode/lzcnt.c

src/pprods.h

@@ -386,4 +386,6 @@ static const uint64_t pprods[] = {
     0xaf157a9999de1f,    // 14887*14891*14897*14923
     0xb1110edec9a5e7,    // 14929*14939*14947*14951
 };
+
+static const uint32_t pprods_max_prime = 14951;
 #endif // PPRODS_H

src/primes_test.cpp

@@ -0,0 +1,41 @@
+#include "vdf_assert.h"
+#include "vdf.h"
+
+int main(int argc, char **argv)
+{
+    char small_primes[] = {2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61};
+    int i;
+    uint64_t small_primes_bitmap = 0;
+    mpz_t n;
+
+    for (i = 0; i < sizeof(small_primes); i++) {
+        small_primes_bitmap |= 1UL << small_primes[i];
+    }
+
+    mpz_init(n);
+    for (i = 0; i < 64; i++) {
+        /* Test all numbers from 0 to 63. */
+        mpz_set_si(n, i);
+        VDF_ASSERT(!!is_prime_bpsw(n) == !!(small_primes_bitmap & 1UL << i));
+    }
+
+    /* Test rejection of negative numbers. */
+    mpz_set_si(n, -11);
+    VDF_ASSERT(!is_prime_bpsw(n));
+
+    /* Known prime. */
+    mpz_set_str(n, "0xc8fd905ec4869c162aa4fde5f2630fd51a44784a380aaff309d9df647b2060fc5e63b9bceb19113e8f3f2291397239ed415c9cd4e5e600e574204eec5a39884cfc9fe40f1f6c68b533b5c5721512c8924cb8792668a2e5eab823c27c5dc8509b37455a314514fa25a82584bc764500518cf9e43931151c41d824a2ba1f3e0897", 0);
+    VDF_ASSERT(is_prime_bpsw(n));
+
+    /* Make n composite. */
+    mpz_mul_ui(n, n, 5);
+    VDF_ASSERT(!is_prime_bpsw(n));
+
+    /* Known composite. */
+    mpz_set_str(n, "0xc01bcf1c6456e52f8945057d1f7d3d48a7baa6577504b4aa22c45b00cce31640f2d095d0b2b2d55d5e87ea21629d8c873905fc7f19f965f0466537c4898cd2604b82be985545bb745df1fad6fc9fa0ef268f85d8a5dae77e4ba9e4aa2ea2f484b3e469e0572a7df8630f74cdaa1ae271ae1333feaaceeca87916662556c311c8bd229d1628668141f0a5e8cbeacff8b1b4c9736da5567f6ab65961ae833cfc93049f9d16b9d4ab6b302c83ece81e0338baeaebf8cf685a325c2b5c23071c92e6168c124fd3868336c1ea3510f27a7c8390b32cd785518eab0aceff64cb0636e449354ddc8d7cc3d095c22ef1617a7e217fb27d532c6a66e46081d142f50d2ae3", 0);
+    VDF_ASSERT(!is_prime_bpsw(n));
+
+    fprintf(stderr, "Test of primes successful!\n");
+    mpz_clear(n);
+    return 0;
+}

src/primetest.h

@@ -141,9 +141,10 @@ static int is_prime_bpsw(const mpz_t n)
     int min_pprods = 5, n_pprods = sizeof(pprods) / sizeof(*pprods);
     mpz_t b, d;
 
-    /* Discard even numbers. */
-    if (!mpz_tstbit(n, 0) && mpz_cmp_ui(n, 2))
-        return 0;
+    /* Discard even and non-positive numbers. */
+    if (!mpz_tstbit(n, 0) || n->_mp_size < 1) {
+        return mpz_cmp_ui(n, 2) ? 0 : 2;
+    }
 
     /* Adjust the number of GCDs to compute with products of small primes
      * if bit length of n is less than 1024. As the computational cost of
@@ -155,8 +156,16 @@ static int is_prime_bpsw(const mpz_t n)
     }
 
     for (i = 0; i < n_pprods; i++) {
-        if (mpn_gcd_1(n->_mp_d, n->_mp_size, pprods[i]) != 1)
-            return 0;
+        if (mpn_gcd_1(n->_mp_d, n->_mp_size, pprods[i]) != 1) {
+            if (n->_mp_size > 1 || n->_mp_d[0] > pprods_max_prime) {
+                return 0;
+            } else if (n->_mp_d[0] < 8) {
+                /* Special case for primes 3, 5, 7. */
+                return (1 << n->_mp_d[0]) & 0xac ? 2 : 0;
+            } else {
+                break;
+            }
+        }
     }
 
     mpz_init2(b, n->_mp_size * sizeof(n->_mp_d[0]) * 8);

src/vdf_assert.h

@@ -0,0 +1,19 @@
+#ifndef VDF_ASSERT_H
+#define VDF_ASSERT_H
+
+#include <string>
+#include <stdexcept>
+
+static void vdf_assert_fail(const char *file, int line)
+{
+    throw std::runtime_error("Assertion failure at " + std::string(file) +
+            ": " + std::to_string(line));
+}
+
+#define VDF_ASSERT(expr) do { \
+    if (!(expr)) { \
+        vdf_assert_fail(__FILE__, __LINE__); \
+    } \
+} while (0)
+
+#endif // VDF_ASSERT_H

tools/gen_pprods.py

@@ -73,6 +73,8 @@ def print_pprods_h(prods, factors):
 static const uint64_t pprods[] = {
 %s
 };
+
+static const uint32_t pprods_max_prime = %d;
 #endif // PPRODS_H
 """
     tmpl = tmpl.strip()
@@ -81,7 +83,7 @@ def print_pprods_h(prods, factors):
         factors_str = "*".join(map(str, factors[i]))
         extra_spaces = " " * (16 - (pr.bit_length() + 3) // 4)
         prods_str += tab + "0x%x,  %s// %s\n" % (pr, extra_spaces, factors_str)
-    print(tmpl % (prods_str.rstrip(),))
+    print(tmpl % (prods_str.rstrip(), factors[-1][-1]))
 
 
 n = int(sys.argv[1])