Add primality testing and next_prime()/prev_prime()

6 files changed, 69 insertions, 4 deletions

diff --git a/builtins/integers.c b/builtins/integers.c
index 0bf7dc22..87d1f0ad 100644
--- a/builtins/integers.c
+++ b/builtins/integers.c
@@ -330,6 +330,33 @@ public Int_t Int$from_text(CORD text) {
     return Int$from_mpz(i);
 }
 
+public bool Int$is_prime(Int_t x, Int_t reps)
+{
+    mpz_t p;
+    mpz_init_set_int(p, x);
+    if (Int$compare(&reps, (Int_t[1]){I_small(9999)}, &$Int) > 0)
+        fail("Number of prime-test repetitions should not be above 9999");
+    int reps_int = Int_to_Int32(reps, false);
+    return (mpz_probab_prime_p(p, reps_int) != 0);
+}
+
+public Int_t Int$next_prime(Int_t x)
+{
+    mpz_t p;
+    mpz_init_set_int(p, x);
+    mpz_nextprime(p, p);
+    return Int$from_mpz(p);
+}
+
+public Int_t Int$prev_prime(Int_t x)
+{
+    mpz_t p;
+    mpz_init_set_int(p, x);
+    if (mpz_prevprime(p, p) == 0)
+        fail("There is no prime number before %r", Int$as_text(&x, false, &$Int));
+    return Int$from_mpz(p);
+}
+
 public const TypeInfo $Int = {
     .size=sizeof(Int_t),
     .align=__alignof__(Int_t),
diff --git a/builtins/integers.h b/builtins/integers.h
index edfd07dd..ed3e293c 100644
--- a/builtins/integers.h
+++ b/builtins/integers.h
@@ -117,6 +117,9 @@ Int_t Int$slow_bit_xor(Int_t x, Int_t y);
 Int_t Int$slow_negative(Int_t x);
 Int_t Int$slow_negated(Int_t x);
 Int_t Int$abs(Int_t x);
+bool Int$is_prime(Int_t x, Int_t reps);
+Int_t Int$next_prime(Int_t x);
+Int_t Int$prev_prime(Int_t x);
 
 extern const TypeInfo $Int;
 
diff --git a/compile.c b/compile.c
index c0f52c4f..dffa145e 100644
--- a/compile.c
+++ b/compile.c
@@ -200,7 +200,7 @@ static CORD compile_lvalue(env_t *env, ast_t *ast)
             ast_t *subject = ast->tag == Index ? Match(ast, Index)->indexed : Match(ast, FieldAccess)->fielded;
             code_err(subject, "This is an immutable value, you can't assign to it");
         } else {
-            code_err(ast, "This is a value of type %T and can't be assigned to", get_type(env, ast));
+            code_err(ast, "This is a value of type %T and can't be used as an assignment target", get_type(env, ast));
         }
     }
 
@@ -209,7 +209,7 @@ static CORD compile_lvalue(env_t *env, ast_t *ast)
         type_t *container_t = get_type(env, index->indexed);
         if (!index->index && container_t->tag == PointerType) {
             if (Match(container_t, PointerType)->is_optional)
-                code_err(index->indexed, "This pointer might be null, so it can't be safely assigned to");
+                code_err(index->indexed, "This pointer might be null, so it can't be safely used as an assignment target");
             return compile(env, ast);
         }
         container_t = value_type(container_t);
diff --git a/environment.c b/environment.c
index 7f7d7e7f..e13c9c13 100644
--- a/environment.c
+++ b/environment.c
@@ -110,6 +110,9 @@ env_t *new_compilation_unit(CORD *libname)
             {"abs", "Int$abs", "func(x:Int)->Int"},
             {"sqrt", "Int$sqrt", "func(x:Int)->Int"},
             {"power", "Int$power", "func(base:Int,exponent:Int)->Int"},
+            {"is_prime", "Int$is_prime", "func(x:Int,reps=50)->Bool"},
+            {"next_prime", "Int$next_prime", "func(x:Int)->Int"},
+            {"prev_prime", "Int$prev_prime", "func(x:Int)->Int"},
         )},
         {"Int64", Type(IntType, .bits=TYPE_IBITS64), "Int64_t", "$Int64", TypedArray(ns_entry_t,
             {"format", "Int64$format", "func(i:Int64, digits=0)->Text"},
diff --git a/parse.c b/parse.c
index 740791b5..3b5974aa 100644
--- a/parse.c
+++ b/parse.c
@@ -930,10 +930,13 @@ ast_t *parse_comprehension_suffix(parse_ctx_t *ctx, ast_t *expr) {
 
     expect_str(ctx, start, &pos, "in", "I expected an 'in' for this 'for'");
     ast_t *iter = expect(ctx, start, &pos, parse_expr, "I expected an iterable value for this 'for'");
-    whitespace(&pos);
+    const char *next_pos = pos;
+    whitespace(&next_pos);
     ast_t *filter = NULL;
-    if (match_word(&pos, "if"))
+    if (match_word(&next_pos, "if")) {
+        pos = next_pos;
         filter = expect(ctx, pos-2, &pos, parse_expr, "I expected a condition for this 'if'");
+    }
     return NewAST(ctx->file, start, pos, Comprehension, .expr=expr, .vars=vars, .iter=iter, .filter=filter);
 }
 
diff --git a/test/integers.tm b/test/integers.tm
index 11603b22..5743cbb2 100644
--- a/test/integers.tm
+++ b/test/integers.tm
@@ -82,3 +82,32 @@ func main():
 			//! n=$n, d=$d:
 			>> (n/d)*d + (n mod d) == n
 			= yes
+
+	>> 0:next_prime()
+	= 2
+	>> 7:next_prime()
+	= 11
+	>> 11:prev_prime()
+	= 7
+	>> (and) p:is_prime() for p in [
+		2, 3, 5, 7,
+		137372146048179869781170214707,
+		811418847921670560768224995279,
+		292590241572454328697048860273,
+		754893741683930091960170890717,
+		319651808258437169510475301537,
+		323890224935694708770556249787,
+		507626552342376235511933571091,
+		548605069630614185274710840981,
+		121475876690852432982324195553,
+		771958616175795150904761471637,
+	]
+	= yes
+
+	>> (or) p:is_prime() for p in [
+		-1, 0, 1, 4, 6,
+		137372146048179869781170214707*2,
+		811418847921670560768224995279*3,
+		292590241572454328697048860273*754893741683930091960170890717,
+	]
+	= no