feat(sqrt): add sqrt support

Rubilmax · Feb 19, 2024 · b3c8023 · b3c8023
1 parent 6ef4b4e
commit b3c8023
Show file tree

Hide file tree

Showing 6 changed files with 64 additions and 3 deletions.
diff --git a/src/index.ts b/src/index.ts
@@ -9,6 +9,7 @@ import {
   mulDivDown,
   approxEqAbs,
   abs,
+  sqrt,
 } from "./utils";
 import { format, toFloat } from "./format";
 import { compDiv, compMul } from "./comp";
@@ -31,6 +32,7 @@ import {
   percentToWad,
   toPercentFloat,
   formatPercent,
+  percentSqrt,
 } from "./percent";
 import {
   wadAdd,
@@ -51,6 +53,7 @@ import {
   wadToPercent,
   toWadFloat,
   formatWad,
+  wadSqrt,
 } from "./wad";
 import {
   rayAdd,
@@ -71,6 +74,7 @@ import {
   rayToPercent,
   toRayFloat,
   formatRay,
+  raySqrt,
 } from "./ray";
 
 declare global {
@@ -87,6 +91,8 @@ declare global {
     mulDivUp: (other: bigint, scale: bigint) => bigint;
     mulDivDown: (other: bigint, scale: bigint) => bigint;
 
+    sqrt: () => bigint;
+
     compMul: (other: bigint) => bigint;
     compDiv: (other: bigint) => bigint;
 
@@ -103,6 +109,7 @@ declare global {
     percentPowUp: (exponent: bigint) => bigint;
     percentPowDown: (exponent: bigint) => bigint;
     percentExpN: (exponent: bigint) => bigint;
+    percentSqrt: () => bigint;
     percentToDecimals: (decimals: number) => bigint;
     percentToWad: () => bigint;
     percentToRay: () => bigint;
@@ -122,6 +129,7 @@ declare global {
     wadPowUp: (exponent: bigint) => bigint;
     wadPowDown: (exponent: bigint) => bigint;
     wadExpN: (exponent: bigint) => bigint;
+    wadSqrt: () => bigint;
     wadToDecimals: (decimals: number) => bigint;
     wadToPercent: () => bigint;
     wadToRay: () => bigint;
@@ -141,6 +149,7 @@ declare global {
     rayPowUp: (exponent: bigint) => bigint;
     rayPowDown: (exponent: bigint) => bigint;
     rayExpN: (exponent: bigint) => bigint;
+    raySqrt: () => bigint;
     rayToDecimals: (decimals: number) => bigint;
     rayToPercent: () => bigint;
     rayToWad: () => bigint;
@@ -196,6 +205,10 @@ BigInt.prototype.mulDivDown = function (other: bigint, scale: bigint) {
   return mulDivDown(this as bigint, other, scale);
 };
 
+BigInt.prototype.sqrt = function () {
+  return sqrt(this as bigint);
+};
+
 BigInt.prototype.compMul = function (other: bigint) {
   return compMul(this as bigint, other);
 };
@@ -242,6 +255,9 @@ BigInt.prototype.percentPowDown = function (exponent: bigint) {
 BigInt.prototype.percentExpN = function (N: bigint) {
   return percentExpN(this as bigint, N);
 };
+BigInt.prototype.percentSqrt = function () {
+  return percentSqrt(this as bigint);
+};
 BigInt.prototype.percentToDecimals = function (decimals: number) {
   return percentToDecimals(this as bigint, decimals);
 };
@@ -297,6 +313,9 @@ BigInt.prototype.wadPowDown = function (exponent: bigint) {
 BigInt.prototype.wadExpN = function (N: bigint) {
   return wadExpN(this as bigint, N);
 };
+BigInt.prototype.wadSqrt = function () {
+  return wadSqrt(this as bigint);
+};
 BigInt.prototype.wadToDecimals = function (decimals: number) {
   return wadToDecimals(this as bigint, decimals);
 };
@@ -352,6 +371,9 @@ BigInt.prototype.rayPowDown = function (exponent: bigint) {
 BigInt.prototype.rayExpN = function (N: bigint) {
   return rayExpN(this as bigint, N);
 };
+BigInt.prototype.raySqrt = function () {
+  return raySqrt(this as bigint);
+};
 BigInt.prototype.rayToDecimals = function (decimals: number) {
   return rayToDecimals(this as bigint, decimals);
 };

diff --git a/src/percent.ts b/src/percent.ts
@@ -1,5 +1,5 @@
 import { PERCENT } from "./constants";
-import { avgHalfUp, expN, mulDivDown, mulDivHalfUp, mulDivUp, pow } from "./utils";
+import { avgHalfUp, expN, mulDivDown, mulDivHalfUp, mulDivUp, pow, sqrt } from "./utils";
 import { format, toDecimals, toFloat } from "./format";
 
 export const percentAdd = (x: bigint, percent: bigint) => {
@@ -54,6 +54,10 @@ export const percentExpN = (x: bigint, N: bigint) => {
   return expN(x, N, PERCENT, mulDivDown);
 };
 
+export const percentSqrt = (x: bigint) => {
+  return sqrt(x, PERCENT);
+};
+
 export const percentToDecimals = (x: bigint, decimals: number) => {
   return toDecimals(x, decimals, 4);
 };

diff --git a/src/ray.ts b/src/ray.ts
@@ -1,5 +1,5 @@
 import { RAY } from "./constants";
-import { avgHalfUp, expN, mulDivDown, mulDivHalfUp, mulDivUp, pow } from "./utils";
+import { avgHalfUp, expN, mulDivDown, mulDivHalfUp, mulDivUp, pow, sqrt } from "./utils";
 import { format, toDecimals, toFloat } from "./format";
 
 export const rayAdd = (x: bigint, ray: bigint) => {
@@ -54,6 +54,10 @@ export const rayExpN = (x: bigint, N: bigint) => {
   return expN(x, N, RAY, mulDivDown);
 };
 
+export const raySqrt = (x: bigint) => {
+  return sqrt(x, RAY);
+};
+
 export const rayToDecimals = (x: bigint, decimals: number) => {
   return toDecimals(x, decimals, 27);
 };

diff --git a/src/utils.ts b/src/utils.ts
@@ -68,6 +68,22 @@ export const pow = (
   return mulDiv(x, pow(xSquared, (exponent - 1n) / 2n, scale, mulDiv), scale);
 };
 
+export const sqrt = (x: bigint, scale = 1n, mulDiv: MulDiv = mulDivHalfUp) => {
+  if (x === 0n) return 0n;
+
+  const scale2 = 2n * scale;
+
+  let x0 = mulDiv(x, scale, scale2);
+  let x1 = mulDiv(x0 + mulDivDown(x, scale, x0), scale, scale2);
+
+  while (x1 < x0) {
+    x0 = x1;
+    x1 = mulDiv(x0 + mulDivDown(x, scale, x0), scale, scale2);
+  }
+
+  return x0;
+};
+
 export const expN = (x: bigint, N: bigint, scale: bigint, mulDiv: MulDiv = mulDivDown) => {
   let res = scale;
   let monomial = scale;

diff --git a/src/wad.ts b/src/wad.ts
@@ -1,5 +1,5 @@
 import { WAD } from "./constants";
-import { avgHalfUp, expN, mulDivDown, mulDivHalfUp, mulDivUp, pow } from "./utils";
+import { avgHalfUp, expN, mulDivDown, mulDivHalfUp, mulDivUp, pow, sqrt } from "./utils";
 import { format, toDecimals, toFloat } from "./format";
 
 export const wadAdd = (x: bigint, wad: bigint) => {
@@ -54,6 +54,10 @@ export const wadExpN = (x: bigint, N: bigint) => {
   return expN(x, N, WAD, mulDivDown);
 };
 
+export const wadSqrt = (x: bigint) => {
+  return sqrt(x, WAD);
+};
+
 export const wadToDecimals = (x: bigint, decimals: number) => {
   return toDecimals(x, decimals, 18);
 };

diff --git a/test/index.spec.ts b/test/index.spec.ts
@@ -148,4 +148,15 @@ describe("evm-maths", () => {
   it("should format with 0 digit", async () => {
     expect(BigInt.WAD.formatWad(0)).toEqual("1");
   });
+
+  it("should calculate sqrt", async () => {
+    expect(4n.sqrt().toString()).toEqual("2");
+    expect(5n.sqrt().toString()).toEqual("2");
+    expect((BigInt.PERCENT * 4n).percentSqrt().formatPercent()).toEqual("2.0000");
+    expect((BigInt.PERCENT * 5n).percentSqrt().formatPercent()).toEqual("2.2361");
+    expect((BigInt.WAD * 4n).wadSqrt().formatWad()).toEqual("2.000000000000000000");
+    expect((BigInt.WAD * 5n).wadSqrt().formatWad()).toEqual("2.236067977499789696");
+    expect((BigInt.RAY * 4n).raySqrt().formatRay()).toEqual("2.000000000000000000000000000");
+    expect((BigInt.RAY * 5n).raySqrt().formatRay()).toEqual("2.236067977499789696409173669");
+  });
 });