From 0a3825e79eabb05ce2fa3bd71f4bb86f27e49558 Mon Sep 17 00:00:00 2001
From: Raul Jordan <raul@prysmaticlabs.com>
Date: Thu, 24 Oct 2019 21:32:04 +0800
Subject: [PATCH] Optimize Integer Square Root (#3785)

---
 shared/mathutil/math_helper.go      | 32 +++++++++++++++++--------
 shared/mathutil/math_helper_test.go | 36 ++++++++++++++++++++++++++++-
 2 files changed, 58 insertions(+), 10 deletions(-)

diff --git a/shared/mathutil/math_helper.go b/shared/mathutil/math_helper.go
index 5eabb991b..0f2ffb236 100644
--- a/shared/mathutil/math_helper.go
+++ b/shared/mathutil/math_helper.go
@@ -1,18 +1,32 @@
 package mathutil
 
-import "math"
+import (
+	"math"
+)
 
 // IntegerSquareRoot defines a function that returns the
-// largest possible integer root of a number.
+// largest possible integer root of a number using a divide and conquer
+// binary search approach:
+//
+// inspiration: https://www.geeksforgeeks.org/square-root-of-an-integer
 func IntegerSquareRoot(n uint64) uint64 {
-	x := n
-	y := (x + 1) / 2
-
-	for y < x {
-		x = y
-		y = (x + n/x) / 2
+	x := uint64(0)
+	y := uint64(1 << 32)
+	for {
+		if y <= 1+x {
+			return x
+		}
+		sqt := x + ((y - x) >> 1)
+		sq := sqt * sqt
+		if sq == n {
+			return sqt
+		}
+		if sq > n {
+			y = sqt
+		} else {
+			x = sqt
+		}
 	}
-	return x
 }
 
 // CeilDiv8 divides the input number by 8
diff --git a/shared/mathutil/math_helper_test.go b/shared/mathutil/math_helper_test.go
index 34c8bc2c7..146dc97e0 100644
--- a/shared/mathutil/math_helper_test.go
+++ b/shared/mathutil/math_helper_test.go
@@ -31,12 +31,46 @@ func TestIntegerSquareRoot(t *testing.T) {
 			number: 97282,
 			root:   311,
 		},
+		{
+			number: 1 << 32,
+			root:   1 << 16,
+		},
+		{
+			number: (1 << 32) + 1,
+			root:   1 << 16,
+		},
+		{
+			number: 1 << 33,
+			root:   92681,
+		},
+		{
+			number: 1 << 60,
+			root:   1 << 30,
+		},
+		{
+			number: 1 << 53,
+			root:   94906265,
+		},
+		{
+			number: 1 << 62,
+			root:   1 << 31,
+		},
 	}
 
 	for _, testVals := range tt {
 		root := mathutil.IntegerSquareRoot(testVals.number)
 		if testVals.root != root {
-			t.Fatalf("expected root and computed root are not equal %d, %d", testVals.root, root)
+			t.Errorf("For %d, expected root and computed root are not equal want %d, got %d", testVals.number, testVals.root, root)
+		}
+	}
+}
+
+func BenchmarkIntegerSquareRoot(b *testing.B) {
+	val := uint64(1 << 62)
+	for i := 0; i < b.N; i++ {
+		root := mathutil.IntegerSquareRoot(val)
+		if root != 1<<31 {
+			b.Fatalf("Expected root and computed root are not equal 1<<31, %d", root)
 		}
 	}
 }