From f9d8176e94cf355dc5c4c35fbdd7e125eb29bcef Mon Sep 17 00:00:00 2001
From: expikr <77922942+expikr@users.noreply.github.com>
Date: Mon, 15 Jan 2024 10:04:30 +0800
Subject: [PATCH] Update atan2.zig (#17840)

Co-authored-by: castholm <carl@astholm.se>
---
 lib/std/math/atan2.zig        | 49 ++++++++++++++++++++---------------
 lib/std/math/complex/arg.zig  |  5 ++--
 lib/std/math/complex/atan.zig |  4 +--
 3 files changed, 32 insertions(+), 26 deletions(-)

diff --git a/lib/std/math/atan2.zig b/lib/std/math/atan2.zig
index 93b1b88f32..a60fa73e8c 100644
--- a/lib/std/math/atan2.zig
+++ b/lib/std/math/atan2.zig
@@ -10,25 +10,28 @@ const expect = std.testing.expect;
 
 /// Returns the arc-tangent of y/x.
 ///
-/// Special Cases:
-///  - atan2(y, nan)     = nan
-///  - atan2(nan, x)     = nan
-///  - atan2(+0, x>=0)   = +0
-///  - atan2(-0, x>=0)   = -0
-///  - atan2(+0, x<=-0)  = +pi
-///  - atan2(-0, x<=-0)  = -pi
-///  - atan2(y>0, 0)     = +pi/2
-///  - atan2(y<0, 0)     = -pi/2
-///  - atan2(+inf, +inf) = +pi/4
-///  - atan2(-inf, +inf) = -pi/4
-///  - atan2(+inf, -inf) = 3pi/4
-///  - atan2(-inf, -inf) = -3pi/4
-///  - atan2(y, +inf)    = 0
-///  - atan2(y>0, -inf)  = +pi
-///  - atan2(y<0, -inf)  = -pi
-///  - atan2(+inf, x)    = +pi/2
-///  - atan2(-inf, x)    = -pi/2
-pub fn atan2(comptime T: type, y: T, x: T) T {
+///      Special Cases:
+/// |   y   |   x   | radians |
+/// |-------|-------|---------|
+/// |  fin  |  nan  |   nan   |
+/// |  nan  |  fin  |   nan   |
+/// |  +0   | >=+0  |   +0    |
+/// |  -0   | >=+0  |   -0    |
+/// |  +0   | <=-0  |   pi    |
+/// |  -0   | <=-0  |  -pi    |
+/// |  pos  |   0   |  +pi/2  |
+/// |  neg  |   0   |  -pi/2  |
+/// | +inf  | +inf  |  +pi/4  |
+/// | -inf  | +inf  |  -pi/4  |
+/// | +inf  | -inf  |  3pi/4  |
+/// | -inf  | -inf  | -3pi/4  |
+/// |  fin  | +inf  |    0    |
+/// |  pos  | -inf  |  +pi    |
+/// |  neg  | -inf  |  -pi    |
+/// | +inf  |  fin  |  +pi/2  |
+/// | -inf  |  fin  |  -pi/2  |
+pub fn atan2(y: anytype, x: anytype) @TypeOf(x, y) {
+    const T = @TypeOf(x, y);
     return switch (T) {
         f32 => atan2_32(y, x),
         f64 => atan2_64(y, x),
@@ -212,8 +215,12 @@ fn atan2_64(y: f64, x: f64) f64 {
 }
 
 test "math.atan2" {
-    try expect(atan2(f32, 0.2, 0.21) == atan2_32(0.2, 0.21));
-    try expect(atan2(f64, 0.2, 0.21) == atan2_64(0.2, 0.21));
+    const y32: f32 = 0.2;
+    const x32: f32 = 0.21;
+    const y64: f64 = 0.2;
+    const x64: f64 = 0.21;
+    try expect(atan2(y32, x32) == atan2_32(0.2, 0.21));
+    try expect(atan2(y64, x64) == atan2_64(0.2, 0.21));
 }
 
 test "math.atan2_32" {
diff --git a/lib/std/math/complex/arg.zig b/lib/std/math/complex/arg.zig
index b2e960ae3e..bf5e0eff42 100644
--- a/lib/std/math/complex/arg.zig
+++ b/lib/std/math/complex/arg.zig
@@ -5,9 +5,8 @@ const cmath = math.complex;
 const Complex = cmath.Complex;
 
 /// Returns the angular component (in radians) of z.
-pub fn arg(z: anytype) @TypeOf(z.re) {
-    const T = @TypeOf(z.re);
-    return math.atan2(T, z.im, z.re);
+pub fn arg(z: anytype) @TypeOf(z.re, z.im) {
+    return math.atan2(z.im, z.re);
 }
 
 const epsilon = 0.0001;
diff --git a/lib/std/math/complex/atan.zig b/lib/std/math/complex/atan.zig
index b932b6d624..028a9f746c 100644
--- a/lib/std/math/complex/atan.zig
+++ b/lib/std/math/complex/atan.zig
@@ -54,7 +54,7 @@ fn atan32(z: Complex(f32)) Complex(f32) {
         return Complex(f32).init(maxnum, maxnum);
     }
 
-    var t = 0.5 * math.atan2(f32, 2.0 * x, a);
+    var t = 0.5 * math.atan2(2.0 * x, a);
     const w = redupif32(t);
 
     t = y - 1.0;
@@ -103,7 +103,7 @@ fn atan64(z: Complex(f64)) Complex(f64) {
         return Complex(f64).init(maxnum, maxnum);
     }
 
-    var t = 0.5 * math.atan2(f64, 2.0 * x, a);
+    var t = 0.5 * math.atan2(2.0 * x, a);
     const w = redupif64(t);
 
     t = y - 1.0;