Wrote test for sqrt , calcR in complex.t.sol and sqrt fuzz test for prbmath

src/ComplexHuff/Complex.huff

@@ -344,17 +344,28 @@
    swap1                  // [x,y]
 }
 
-
 ///@notice e^(a+bi) calculation
-
+//@dev made EXP_Z more accurate and reduced the case of overflow of integers by making the exponent value stick to 18 decimals
 #define macro EXP_Z() = takes(2) returns(2) {
  //INPUT STACK => [Re(A),Im(A)]
-  [e]                  // [e,Re(A),Im(A)]
-  exp                  // [e**Re(A),Im(A)]
+  [e]                  // [e,Re(A),Im(A)]                            
+  0x12                 // [18,e,Re(A),Im(A)]                         
+  dup3                 // [Re(A),18,e,Re(A),Im(A)]                                  
+  mul                  // [18*Re(A),e,Re(A),Im(A)]
+  0x12                 // [18,18*Re(A),e,Re(A),Im(A)]
+  swap1            
+  sub                  // [18*Re(A)-18,e,Re(A),Im(A)]
+  swap2          // [Re(A),e,18*Re(A)-18,Im(A)]
+  swap1             //[e , Re(A), ...]
+  exp                  // [e**Re(A),18*Re(A) - 18 , Im(A)]
+  swap1                // [18*Re(A) - 18 , e**Re(A) , Im(A)]
+  0x0a                // [10,18*Re(A) - 18 , e**Re(A) , Im(A)]
+  exp                 // [10**(18*Re(A) - 18) , e**Re(A) , Im(A)]
+  swap1               // [e**Re(A) , 10**(18*Re(A) - 18) , Im(A)]
+  sdiv               // [e**Re(A) / 10**(18*Re(A) - 18) , Im(A)]
   swap1
   FROM_POLAR()
 }
-
 ///@notice power of complex numbers
 
 #define macro POW() = takes(3) returns(2) {
@@ -393,4 +404,4 @@
     swap2 
     sdiv                //[r**n*cos(x)/1e18,r**n*sin(x)/1e18]
 
-}
+}

src/ComplexHuff/PRBMathWrapper.huff

@@ -7,7 +7,7 @@
 
 #define macro SQRT_LOCAL()=takes(0) returns(0){
     0x04 calldataload
-    // SQRT_PRB()
+     SQRT_PRB()
     0x00 mstore
     0x20 0x00 return
 }
@@ -41,4 +41,4 @@
 
     log_2:
     LOG_2_LOCAL()
-}
+}

test/Complex.t.sol

@@ -11,6 +11,7 @@ contract ComplexTest is Test {
     WRAPPER public complex;
     int256 scale = 1e18;
     int256 scale2 = 1e19;
+    uint256 constant pi = 3141592653589793238;
 
     function setUp() public {
         complex = WRAPPER(HuffDeployer.deploy("ComplexHuff/WRAPPER"));
@@ -254,34 +255,187 @@ contract ComplexTest is Test {
         assertEq(Ri, Rhi);
     }
 
-    // function testCalc_RFuzz(int256 ar, int256 ai, int256 br, int256 bi, int256 k) public {
-    //     ai = bound(ai, -1e9, 1e9);
-    //     bi = bound(bi, -1e9, 1e9);
-    //     ar = bound(ar, -1e9, 1e9);
-    //     br = bound(br, -1e9, 1e9);
-    //     k = bound(k, -1e9, 1e9);
-
-    //     (int256 mag1r,) = (complex.mulz(-ai, ai, ar, ar)); // ar^2 + ai^2
-    //     int256 mag1 = PRBMathSD59x18.sqrt(mag1r); // (ar^2+ai^2)^0.5
-    //     uint256 R1 = (complex.calcR(ai, ar)); // magnitude(A)
-    //     uint256 R2 = (complex.calcR(bi, br)); // magnitude(B)
-    //     (int256 A_Br, int256 A_Bi) = (complex.addz(bi, ai, br, ai)); // A+B
-    //     uint256 R3 = (complex.calcR(A_Bi, A_Br)); // magnitude(A+B)
-    //     uint256 R4 = (complex.calcR(k * ai, k * ar)); // magnitude(k*A)
-    //     uint256 magk = (complex.calcR(0, k));
-    //     uint256 _R1 = (complex.calcR(-ai, ar));
-
-    //     // Test by comparing
-    //     assertEq(uint256(mag1) / 1e9, R1);
-
-    //     // Test by property
-    //     // mag(A+B)<=mag(A)+mag(B)
-    //     assert(R3 <= (R1 + R2));
-    //     // mag(kA)=kmag(A)
-    //     assertEq(R4,(magk)*R1);
-    //     // mag(A)=mag(A')
-    //     assertEq(R1,_R1);
-    // }
+    function testCalc_RFuzz(int256 ar, int256 ai, int256 br, int256 bi, int256 k) public {
+        vm.assume(k != 0);
+        ai = bound(ai, -1e9, 1e9);
+        bi = bound(bi, -1e9, 1e9);
+        ar = bound(ar, -1e9, 1e9);
+        br = bound(br, -1e9, 1e9);
+        k = bound(k, -1e8, 1e8);
+
+        //ar^2+ai^2 = magnitude(A)**2
+        (int256 mag1r,) = (complex.mulz(-ai * scale, ai * scale, ar * scale, ar * scale)); // ar^2 + ai^2
+        int256 R1 = (int256((complex.calcR(ai * scale, ar * scale))) ** 2); // magnitude(A)
+        assertApproxEqAbs(mag1r / scale, R1 / scale, 1e10);
+
+        //mag(A) = mag(-A)
+        mag1r = int256(complex.calcR(-ai * scale, ar * scale) ** 2);
+        assertEq(mag1r / scale, R1 / scale);
+
+        //(mag(a))**2 = (k*mag(A))**2
+        (mag1r,) = (complex.mulz(-ai * scale, ai * scale, ar * scale, ar * scale)); // kA
+        int256 R3 = (int256(complex.calcR((k * ai) * scale, (k * ar) * scale)) ** 2); // magnitude(A)
+        assertApproxEqAbs(((k ** 2) * (mag1r / scale)), R3 / scale, 5e17);
+
+        //mag(z1z2) = mag(z1)*mag(z2)
+        (int256 ir, int256 mag1i) = complex.mulz(bi * scale, ai * scale, br * scale, ar * scale);
+        R1 = (int256((complex.calcR((ai * scale), (ar * scale)))));
+        int256 R2 = (int256(complex.calcR(bi * scale, br * scale))); // magnitude(B)
+        R3 = (int256(complex.calcR((mag1i / scale), (ir / scale))));
+        assertApproxEqAbs(R3, (R1 * R2) / scale, 1e10);
+
+        //magnitude(A+B) <= magnitude(A) + magnitude(B)
+        R3 = (int256(complex.calcR((bi + ai) * scale, (br + ar) * scale)));
+        assert(R1 + R2 >= R3);
+    }
+
+    function testSqrtfuzz(int256 ar, int256 ai) public {
+        ai = bound(ai, -1e10, 1e10);
+        ar = bound(ar, -1e10, 1e10);
+
+        (int256 sqrt_ar, int256 sqrt_ai) = complex.sqrt(ai * scale, ar * scale);
+        (int256 r, int256 i) = complex.mulz(sqrt_ai, sqrt_ai, sqrt_ar, sqrt_ar);
+        r >= 0 ? r = r : r = -r;
+        i >= 0 ? i = i : i = -i;
+        ar >= 0 ? ar = ar : ar = -ar;
+        ai >= 0 ? ai = ai : ai = -ai;
+        assertApproxEqAbs((r / (scale)), ar * scale, 5e10);
+        assertApproxEqAbs((i / (scale)), ai * scale, 5e10);
+    }
+
+    //tan(tan-1x) = x
+
+    function testAtan1to1Fuzz(int256 ar, int256 br) public {
+        ar = bound(ar, 1e16, 1e18);
+        br = bound(br, 1e16, 1e18);
+
+        int256 r = complex.atan1to1(ar);
+        assert(-157 * 1e16 <= r && r <= 157 * 1e16);
+
+        int256 r1 = (Trigonometry.sin(uint256(r)) * 1e18) / Trigonometry.cos(uint256(r));
+        assertApproxEqAbs(ar, r1, 5e15);
+
+        ar = bound(ar, 1e15, 1e16);
+        r = complex.atan1to1(ar);
+        r1 = (Trigonometry.sin(uint256(r)) * 1e18) / Trigonometry.cos(uint256(r));
+        assertApproxEqAbs(ar, r1, 5e14);
+
+        ar = bound(ar, 1e14, 1e15);
+        r = complex.atan1to1(ar);
+        r1 = (Trigonometry.sin(uint256(r)) * 1e18) / Trigonometry.cos(uint256(r));
+        assertApproxEqAbs(ar, r1, 5e13);
+    }
+
+    /*@dev testing using the following formula
+     ze^(itheta)/ze^(itheta2) = e^(i(theta-theta2))
+     z1*z2 = e^((i(theta-theta2)))
+     */
+
+    function testFromPolar_Fuzz(int256 r, int256 i, int256 i2) public {
+        r = bound(r, 1e18, 1e23);
+        i = bound(i, 1e18, 1e23);
+        i2 = bound(i, 1e18, 1e23);
+
+        (int256 r1, int256 i1) = complex.fromPolar(r, i);
+        (int256 mag1r) = int256(complex.calcR(((r1)), ((i1))));
+        assertApproxEqAbs(mag1r, r, 5e17);
+
+        (r1, i1) = complex.fromPolar(r, i);
+        (int256 r_1, int256 i_1) = complex.fromPolar(r, i2);
+        (int256 r3, int256 i3) = complex.fromPolar(1 * scale, i - i2);
+        (r1, i1) = (complex.divz(i_1, i1, r_1, r1));
+
+        assertApproxEqAbs(i1, i3, 0);
+
+        r = bound(r, 1e16, 1e19);
+        i = bound(i, 1e16, 1e19);
+        i2 = bound(i, 1e16, 1e19);
+        (r1, i1) = complex.fromPolar(r, i);
+        (mag1r) = int256(complex.calcR(((r1)), ((i1))));
+        assertApproxEqAbs(mag1r, r, 1e15);
+
+        (r1, i1) = complex.fromPolar(r, i);
+        (r_1, i_1) = complex.fromPolar(r, i2);
+        (r3, i3) = complex.fromPolar(1 * scale, i - i2);
+        (r1, i1) = (complex.divz(i_1, i1, r_1, r1));
+
+        assertApproxEqAbs(i1, i3, 0);
+
+        r = bound(r, 1e11, 1e15);
+        i = bound(i, 1e11, 1e15);
+        i2 = bound(i, 1e11, 1e15);
+
+        (r1, i1) = complex.fromPolar(r, i);
+        (mag1r) = int256(complex.calcR(((r1)), ((i1))));
+        assertApproxEqAbs(mag1r, r, 5e10);
+
+        (r1, i1) = complex.fromPolar(r, i);
+        (r_1, i_1) = complex.fromPolar(r, i2);
+        (r3, i3) = complex.fromPolar(1 * scale, i - i2);
+        (r1, i1) = (complex.divz(i_1, i1, r_1, r1));
+
+        assertApproxEqAbs(i1, i3, 0);
+
+        i = bound(i, -1e20, -1e16);
+        i2 = bound(i, -1e20, -1e16);
+
+        (r1, i1) = complex.fromPolar(r, i);
+        (mag1r) = int256(complex.calcR(((r1)), ((i1))));
+        assertApproxEqAbs(mag1r, r, 1e15);
+
+        (r1, i1) = complex.fromPolar(r, i);
+        (r_1, i_1) = complex.fromPolar(r, i2);
+        (r3, i3) = complex.fromPolar(1 * scale, i - i2);
+        (r1, i1) = (complex.divz(i_1, i1, r_1, r1));
+
+        assertApproxEqAbs(i1, i3, 0);
+
+        i = bound(i, -1e15, -1e11);
+        i2 = bound(i, -1e15, -1e11);
+
+        (r1, i1) = complex.fromPolar(r, i);
+        (mag1r) = int256(complex.calcR(((r1)), ((i1))));
+        assertApproxEqAbs(mag1r, r, 5e10);
+
+        (r1, i1) = complex.fromPolar(r, i);
+        (r_1, i_1) = complex.fromPolar(r, i2);
+        (r3, i3) = complex.fromPolar(1 * scale, i - i2);
+        (r1, i1) = (complex.divz(i_1, i1, r_1, r1));
+
+        assertApproxEqAbs(i1, i3, 0);
+    }
+
+    function testExpZ_Fuzz(int256 r, int256 i) public {
+        r = bound(r, 1, 1e2);
+        i = bound(i, 1e16, 1e18);
+
+        //periodicity check
+        (int256 r1, int256 i1) = complex.expZ(i, r);
+        (int256 r_1, int256 i_1) = complex.expZ(i + int256(2 * pi), r);
+        assertApproxEqAbs(r1, r_1, 0);
+        assertApproxEqAbs(i1, i_1, 0);
+
+        //conjugateof(e^z) = e^conjugateof(z)
+        (r1, i1) = complex.expZ(i, r);
+        (r_1, i_1) = complex.expZ(-i, r);
+        assertApproxEqAbs(r1, r_1, 1e15);
+        assertApproxEqAbs(-i1, i_1, 1e15);
+
+        r = bound(r, 1, 1e2);
+        i = bound(i, 1e12, 1e15);
+
+        //periodicity check
+        (r1, i1) = complex.expZ(i, r);
+        (r_1, i_1) = complex.expZ(i + int256(2 * pi), r);
+        assertApproxEqAbs(r1, r_1, 0);
+        assertApproxEqAbs(i1, i_1, 0);
+
+        //conjugateof(e^z) = e^conjugateof(z)
+        (r1, i1) = complex.expZ(i, r);
+        (r_1, i_1) = complex.expZ(-i, r);
+        assertApproxEqAbs(r1, r_1, 1e11);
+        assertApproxEqAbs(-i1, i_1, 1e11);
+    }
 }
 
 interface WRAPPER {

test/PRBMath.t.sol

@@ -14,10 +14,13 @@ contract PRBMathtest is Test {
         prb_huff = Wrapper(HuffDeployer.deploy("ComplexHuff/PRBMathWrapper"));
     }
 
-    // function test_sqrt(int256 num) public {
-    //     vm.assume(num > 0);
-    //     assertApproxEqAbs(PRBMathSD59x18.log2(num), prb_huff.log_2(num), 0);
-    // }
+    function test_sqrt(int256 num) public {
+        num = bound(num, 1, 1e38);
+        int256 res = num ** 2;
+        res = prb_huff.sqrt(res);
+        assertEq(res, num);
+        assertApproxEqAbs(PRBMathSD59x18.sqrt(num) / 1e9, prb_huff.sqrt(num), 0);
+    }
 
     function testFuzz_ln(int256 num) public {
         vm.assume(num > 0);