From ff3d5733d80fa9f9186d197f19fe4dcc35d5ff3c Mon Sep 17 00:00:00 2001
From: Cody Littley <56973212+cody-littley@users.noreply.github.com>
Date: Mon, 15 Jul 2024 15:52:37 -0500
Subject: [PATCH] Fix fft_fr.IsPowerOfTwo(). (#638)

Signed-off-by: Cody Littley <cody@eigenlabs.org>
---
 encoding/fft/fft_fr.go      |  7 ++++---
 encoding/fft/fft_fr_test.go | 21 +++++++++++++++++++++
 2 files changed, 25 insertions(+), 3 deletions(-)

diff --git a/encoding/fft/fft_fr.go b/encoding/fft/fft_fr.go
index 2fae4af69d..8592f7774d 100644
--- a/encoding/fft/fft_fr.go
+++ b/encoding/fft/fft_fr.go
@@ -109,9 +109,9 @@ func (fs *FFTSettings) InplaceFFT(vals []fr.Element, out []fr.Element, inv bool)
 	}
 	if inv {
 		var invLen fr.Element
-		
+
 		invLen.SetInt64(int64(n))
-		
+
 		invLen.Inverse(&invLen)
 		rootz := fs.ReverseRootsOfUnity[:fs.MaxWidth]
 		stride := fs.MaxWidth / n
@@ -132,6 +132,7 @@ func (fs *FFTSettings) InplaceFFT(vals []fr.Element, out []fr.Element, inv bool)
 	}
 }
 
+// IsPowerOfTwo returns true if the provided integer v is a power of 2.
 func IsPowerOfTwo(v uint64) bool {
-	return v&(v-1) == 0
+	return (v&(v-1) == 0) && (v != 0)
 }
diff --git a/encoding/fft/fft_fr_test.go b/encoding/fft/fft_fr_test.go
index a9c60b1dd7..f5f29b263e 100644
--- a/encoding/fft/fft_fr_test.go
+++ b/encoding/fft/fft_fr_test.go
@@ -24,6 +24,7 @@
 package fft
 
 import (
+	"math"
 	"testing"
 
 	"github.com/consensys/gnark-crypto/ecc/bn254/fr"
@@ -100,3 +101,23 @@ func TestInvFFT(t *testing.T) {
 		assert.True(t, res[i].Equal(&expected[i]))
 	}
 }
+
+func TestIsPowerOfTwo(t *testing.T) {
+	var i uint64
+	for i = 0; i <= 1024; i++ {
+		result := IsPowerOfTwo(i)
+
+		var expectedResult bool
+		if i == 0 {
+			// Special case: math.Log2() is undefined for 0
+			expectedResult = false
+		} else {
+			// If a number is not a power of two then the log base 2 of that number will not be a whole integer.
+			logBase2 := math.Log2(float64(i))
+			truncatedLogBase2 := float64(uint64(logBase2))
+			expectedResult = logBase2 == truncatedLogBase2
+		}
+
+		assert.Equal(t, expectedResult, result, "IsPowerOfTwo(%d) returned unexpected result '%t'.", i, result)
+	}
+}