Deduplicate Copy Functions In State Packages (#9437)

* remove duplicated copy functions

* add all unit tests

Co-authored-by: Nishant Das <nishdas93@gmail.com>

shared/bytesutil/bytes.go

@@ -4,6 +4,7 @@ package bytesutil
 import (
 	"encoding/binary"
 	"errors"
+	"fmt"
 	"math/bits"
 	"regexp"
 
@@ -166,6 +167,20 @@ func ToLowInt64(x []byte) int64 {
 	return int64(binary.LittleEndian.Uint64(x))
 }
 
+// SafeCopyRootAtIndex takes a copy of an 32-byte slice in a slice of byte slices. Returns error if index out of range.
+func SafeCopyRootAtIndex(input [][]byte, idx uint64) ([]byte, error) {
+	if input == nil {
+		return nil, nil
+	}
+
+	if uint64(len(input)) <= idx {
+		return nil, fmt.Errorf("index %d out of range", idx)
+	}
+	item := make([]byte, 32)
+	copy(item, input[idx])
+	return item, nil
+}
+
 // SafeCopyBytes will copy and return a non-nil byte array, otherwise it returns nil.
 func SafeCopyBytes(cp []byte) []byte {
 	if cp != nil {
@@ -176,8 +191,8 @@ func SafeCopyBytes(cp []byte) []byte {
 	return nil
 }
 
-// Copy2dBytes will copy and return a non-nil 2d byte array, otherwise it returns nil.
-func Copy2dBytes(ary [][]byte) [][]byte {
+// SafeCopy2dBytes will copy and return a non-nil 2d byte array, otherwise it returns nil.
+func SafeCopy2dBytes(ary [][]byte) [][]byte {
 	if ary != nil {
 		copied := make([][]byte, len(ary))
 		for i, a := range ary {

shared/bytesutil/bytes_test.go

@@ -1,6 +1,7 @@
 package bytesutil_test
 
 import (
+	"reflect"
 	"testing"
 
 	"github.com/prysmaticlabs/prysm/shared/bytesutil"
@@ -397,3 +398,77 @@ func TestIsHex(t *testing.T) {
 		assert.Equal(t, tt.b, isHex)
 	}
 }
+
+func TestSafeCopyRootAtIndex(t *testing.T) {
+	tests := []struct {
+		name    string
+		input   [][]byte
+		idx     uint64
+		want    []byte
+		wantErr bool
+	}{
+		{
+			name:    "index out of range in non-empty slice",
+			input:   [][]byte{{0x1}, {0x2}},
+			idx:     2,
+			wantErr: true,
+		},
+		{
+			name:    "index out of range in empty slice",
+			input:   [][]byte{},
+			idx:     0,
+			wantErr: true,
+		},
+		{
+			name:  "nil input",
+			input: nil,
+			idx:   3,
+			want:  nil,
+		},
+		{
+			name:  "correct copy",
+			input: [][]byte{{0x1}, {0x2}},
+			idx:   1,
+			want:  bytesutil.PadTo([]byte{0x2}, 32),
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, err := bytesutil.SafeCopyRootAtIndex(tt.input, tt.idx)
+			if (err != nil) != tt.wantErr {
+				t.Errorf("SafeCopyRootAtIndex() error = %v, wantErr %v", err, tt.wantErr)
+				return
+			}
+			if !reflect.DeepEqual(got, tt.want) {
+				t.Errorf("SafeCopyRootAtIndex() got = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}
+
+func TestSafeCopy2dBytes(t *testing.T) {
+	tests := []struct {
+		name  string
+		input [][]byte
+	}{
+		{
+			name:  "nil input",
+			input: nil,
+		},
+		{
+			name:  "correct copy",
+			input: [][]byte{{0x1}, {0x2}},
+		},
+		{
+			name:  "empty",
+			input: [][]byte{},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if got := bytesutil.SafeCopy2dBytes(tt.input); !reflect.DeepEqual(got, tt.input) {
+				t.Errorf("SafeCopy2dBytes() = %v, want %v", got, tt.input)
+			}
+		})
+	}
+}