From c6c0e85ea4f233579be8b5397d77e8b20f5d3ff8 Mon Sep 17 00:00:00 2001
From: weiihann <weihan774237@gmail.com>
Date: Tue, 31 Dec 2024 10:27:16 +0800
Subject: [PATCH] add SetBit()

---
 core/trie/bitarray.go      | 12 ++++++++++++
 core/trie/bitarray_test.go | 34 ++++++++++++++++++++++++++++++++++
 2 files changed, 46 insertions(+)

diff --git a/core/trie/bitarray.go b/core/trie/bitarray.go
index 1e5de6408..e466a4089 100644
--- a/core/trie/bitarray.go
+++ b/core/trie/bitarray.go
@@ -479,6 +479,18 @@ func (b *BitArray) SetUint64(length uint8, data uint64) *BitArray {
 	return b
 }
 
+// Sets the bit array to a single bit.
+func (b *BitArray) SetBit(bit bool) *BitArray {
+	b.len = 1
+	if bit {
+		b.words[0] = 1
+	} else {
+		b.words[0] = 0
+	}
+	b.truncateToLength()
+	return b
+}
+
 // Returns the length of the encoded bit array in bytes.
 func (b *BitArray) EncodedLen() uint {
 	return b.byteCount() + 1
diff --git a/core/trie/bitarray_test.go b/core/trie/bitarray_test.go
index 7cb59493a..933b6fdd4 100644
--- a/core/trie/bitarray_test.go
+++ b/core/trie/bitarray_test.go
@@ -1713,3 +1713,37 @@ func TestSetFeltValidation(t *testing.T) {
 		})
 	}
 }
+
+func TestSetBit(t *testing.T) {
+	tests := []struct {
+		name string
+		bit  bool
+		want BitArray
+	}{
+		{
+			name: "set bit false",
+			bit:  false,
+			want: BitArray{
+				len:   1,
+				words: [4]uint64{0, 0, 0, 0},
+			},
+		},
+		{
+			name: "set bit true",
+			bit:  true,
+			want: BitArray{
+				len:   1,
+				words: [4]uint64{1, 0, 0, 0},
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got := new(BitArray).SetBit(tt.bit)
+			if !got.Equal(&tt.want) {
+				t.Errorf("SetBit(%v) = %v, want %v", tt.bit, got, tt.want)
+			}
+		})
+	}
+}