From 04c8d4bcec72bfb6545eb125a418382f4ab14b31 Mon Sep 17 00:00:00 2001
From: David Peter <mail@david-peter.de>
Date: Tue, 4 Jun 2024 20:04:47 +0200
Subject: [PATCH] Change quadratic_equation to return a list of solutions

---
 book/src/list-functions.md       |  5 +++--
 numbat/modules/extra/algebra.nbt | 20 +++++++++++++----
 numbat/tests/interpreter.rs      | 37 +++++++++++++++-----------------
 3 files changed, 36 insertions(+), 26 deletions(-)

diff --git a/book/src/list-functions.md b/book/src/list-functions.md
index c638685f..69f705fb 100644
--- a/book/src/list-functions.md
+++ b/book/src/list-functions.md
@@ -117,10 +117,11 @@ fn rand_lognormal(μ: Scalar, σ: Scalar) -> Scalar
 fn rand_pareto<T>(α: Scalar, min: T) -> T
 ```
 
-### Algebra (experimental)
+### Algebra
 
 ```nbt
-fn quadratic_equation<A2, B2>(a: A2, b: B2, c: B2²/A2) -> String
+# Returns the solutions of the equation a x² + b x + c = 0
+quadratic_equation<A: Dim, B: Dim>(a: A, b: B, c: B² / A) -> List<B / A>
 ```
 
 ## Date and time
diff --git a/numbat/modules/extra/algebra.nbt b/numbat/modules/extra/algebra.nbt
index 5a1f79be..54c7aaa3 100644
--- a/numbat/modules/extra/algebra.nbt
+++ b/numbat/modules/extra/algebra.nbt
@@ -1,9 +1,21 @@
+use core::error
 use math::functions
 
-fn _qe_solution<A: Dim, B: Dim>(a: A, b: B, c: B²/A, sign: Scalar) -> B/A =
+fn _qe_solution<A: Dim, B: Dim>(a: A, b: B, c: B² / A, sign: Scalar) -> B / A =
   (-b + sign × sqrt(b² - 4 a c)) / 2 a
 
-fn quadratic_equation<A: Dim, B: Dim>(a: A, b: B, c: B²/A) -> String =
+@name("Solve quadratic equations")
+@url("https://en.wikipedia.org/wiki/Quadratic_equation")
+@description("Returns the solutions of the equation a x² + b x + c = 0")
+fn quadratic_equation<A: Dim, B: Dim>(a: A, b: B, c: B² / A) -> List<B / A> =
   if a == 0
-    then if b == 0 then if c == 0 then "infinitely many solutions" else "no solution" else "x = {-c / b}"
-    else if b² < 4 a c then "no real-valued solution" else if b² == 4 a c then "x = {-b / 2 a}" else "x₁ = {_qe_solution(a, b, c, 1)}; x₂ = {_qe_solution(a, b, c, -1)}"
+    then if b == 0
+      then if c == 0
+        then error("infinitely many solutions")
+        else []
+      else [-c / b]
+    else if b² < 4 a c
+      then []
+      else if b² == 4 a c
+        then [-b / 2 a]
+        else [_qe_solution(a, b, c, 1), _qe_solution(a, b, c, -1)]
diff --git a/numbat/tests/interpreter.rs b/numbat/tests/interpreter.rs
index 8e2b7477..77557852 100644
--- a/numbat/tests/interpreter.rs
+++ b/numbat/tests/interpreter.rs
@@ -47,8 +47,7 @@ fn expect_output(code: &str, expected_output: impl AsRef<str>) {
 }
 
 #[track_caller]
-fn expect_failure(code: &str, msg_part: &str) {
-    let mut ctx = get_test_context();
+fn expect_failure_with_context(ctx: &mut Context, code: &str, msg_part: &str) {
     if let Err(e) = ctx.interpret(code, CodeSource::Internal) {
         let error_message = e.to_string();
         println!("{}", error_message);
@@ -61,6 +60,12 @@ fn expect_failure(code: &str, msg_part: &str) {
     }
 }
 
+#[track_caller]
+fn expect_failure(code: &str, msg_part: &str) {
+    let mut ctx = get_test_context();
+    expect_failure_with_context(&mut ctx, code, msg_part)
+}
+
 #[track_caller]
 fn get_error_message(code: &str) -> String {
     let mut ctx = get_test_context();
@@ -236,27 +241,19 @@ fn test_algebra() {
     let _ = ctx
         .interpret("use extra::algebra", CodeSource::Internal)
         .unwrap();
-    expect_output_with_context(
-        &mut ctx,
-        "quadratic_equation(1, 0, -1)",
-        "\"x₁ = 1; x₂ = -1\"",
-    );
-    expect_output_with_context(&mut ctx, "quadratic_equation(0, 9, 3)", "\"x = -0.333333\"");
-    expect_output_with_context(&mut ctx, "quadratic_equation(0, 0, 1)", "\"no solution\"");
-    expect_output_with_context(&mut ctx, "quadratic_equation(9, -126, 441)", "\"x = 7\"");
-    expect_output_with_context(&mut ctx, "quadratic_equation(1, -2, 1)", "\"x = 1\"");
-    expect_output_with_context(&mut ctx, "quadratic_equation(0, 1, 1)", "\"x = -1\"");
-    expect_output_with_context(&mut ctx, "quadratic_equation(1, 0, 0)", "\"x = 0\"");
-    expect_output_with_context(
+    expect_output_with_context(&mut ctx, "quadratic_equation(1, 0, -1)", "[1, -1]");
+    expect_output_with_context(&mut ctx, "quadratic_equation(0, 9, 3)", "[-0.333333]");
+    expect_output_with_context(&mut ctx, "quadratic_equation(0, 0, 1)", "[]");
+    expect_output_with_context(&mut ctx, "quadratic_equation(9, -126, 441)", "[7]");
+    expect_output_with_context(&mut ctx, "quadratic_equation(1, -2, 1)", "[1]");
+    expect_output_with_context(&mut ctx, "quadratic_equation(0, 1, 1)", "[-1]");
+    expect_output_with_context(&mut ctx, "quadratic_equation(1, 0, 0)", "[0]");
+    expect_failure_with_context(
         &mut ctx,
         "quadratic_equation(0, 0, 0)",
-        "\"infinitely many solutions\"",
-    );
-    expect_output_with_context(
-        &mut ctx,
-        "quadratic_equation(1, 1, 1)",
-        "\"no real-valued solution\"",
+        "infinitely many solutions",
     );
+    expect_output_with_context(&mut ctx, "quadratic_equation(1, 1, 1)", "[]");
 }
 
 #[test]