Create New SSZ Factory for Arrays of Roots

spectests/BUILD.bazel

@@ -28,6 +28,7 @@ go_test(
     tags = ["spectest"],
     deps = [
         "//:go_default_library",
+        "//types:go_default_library",
         "@com_github_ghodss_yaml//:go_default_library",
         "@io_bazel_rules_go//go/tools/bazel:go_default_library",
     ],

spectests/ssz_spec_test.go

@@ -12,8 +12,13 @@ import (
 	"github.com/bazelbuild/rules_go/go/tools/bazel"
 	"github.com/ghodss/yaml"
 	"github.com/prysmaticlabs/go-ssz"
+	"github.com/prysmaticlabs/go-ssz/types"
 )
 
+func init() {
+	types.ToggleCache(false)
+}
+
 // sszComparisonConfig is used to specify the value to marshal, unmarshal into,
 // as well as the expected results from the spec test YAML files.
 type sszComparisonConfig struct {

ssz.go

@@ -133,7 +133,7 @@ func HashTreeRoot(val interface{}) ([32]byte, error) {
 	if err != nil {
 		return [32]byte{}, errors.Wrapf(err, "could not generate tree hasher for type: %v", rval.Type())
 	}
-	return factory.Root(rval, rval.Type(), 0)
+	return factory.Root(rval, rval.Type(), "", 0)
 }
 
 // HashTreeRootBitlist determines the root hash of a bitfield.Bitlist type using SSZ's Merkleization.
@@ -162,7 +162,7 @@ func HashTreeRootWithCapacity(val interface{}, maxCapacity uint64) ([32]byte, er
 	if err != nil {
 		return [32]byte{}, errors.Wrapf(err, "could not generate tree hasher for type: %v", rval.Type())
 	}
-	return factory.Root(rval, rval.Type(), maxCapacity)
+	return factory.Root(rval, rval.Type(), "", maxCapacity)
 }
 
 // SigningRoot truncates the last property of the struct passed in

ssz_test.go

@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"encoding/hex"
 	"reflect"
+	"strconv"
 	"sync"
 	"testing"
 
@@ -12,6 +13,10 @@ import (
 	"github.com/prysmaticlabs/go-ssz/types"
 )
 
+type beaconState struct {
+	BlockRoots [][]byte `ssz-size:"65536,32"`
+}
+
 type fork struct {
 	PreviousVersion [4]byte
 	CurrentVersion  [4]byte
@@ -694,6 +699,74 @@ func TestSigningRoot_ConcurrentAccess(t *testing.T) {
 	wg.Wait()
 }
 
+func BenchmarkSSZ_NoCache(b *testing.B) {
+	b.StopTimer()
+	bs := &beaconState{
+		BlockRoots: make([][]byte, 65536),
+	}
+	for i := 0; i < len(bs.BlockRoots); i++ {
+		newItem := [32]byte{1, 2, 3}
+		bs.BlockRoots[i] = newItem[:]
+	}
+	b.StartTimer()
+	for i := 0; i < b.N; i++ {
+		if _, err := HashTreeRoot(bs); err != nil {
+			b.Fatal(err)
+		}
+	}
+	types.ToggleCache(true)
+}
+
+func BenchmarkSSZ_WithCache(b *testing.B) {
+	b.StopTimer()
+	types.ToggleCache(true)
+	bs := &beaconState{
+		BlockRoots: make([][]byte, 65536),
+	}
+	for i := 0; i < len(bs.BlockRoots); i++ {
+		newItem := [32]byte{1, 2, 3}
+		bs.BlockRoots[i] = newItem[:]
+	}
+	b.StartTimer()
+	for i := 0; i < b.N; i++ {
+		if _, err := HashTreeRoot(bs); err != nil {
+			b.Fatal(err)
+		}
+	}
+	types.ToggleCache(false)
+}
+
+func BenchmarkSSZ_SingleElementChanged(b *testing.B) {
+	b.StopTimer()
+	types.ToggleCache(true)
+	bs := &beaconState{
+		BlockRoots: make([][]byte, 65536),
+	}
+	for i := 0; i < len(bs.BlockRoots); i++ {
+		newItem := [32]byte{1, 2, 3}
+		bs.BlockRoots[i] = newItem[:]
+	}
+	if _, err := HashTreeRoot(bs); err != nil {
+		b.Fatal(err)
+	}
+	b.StartTimer()
+	for i := 0; i < b.N; i++ {
+		newItem := []byte(strconv.Itoa(i))
+		newRoot := toBytes32(newItem)
+		bs.BlockRoots[i%len(bs.BlockRoots)] = newRoot[:]
+		if _, err := HashTreeRoot(bs); err != nil {
+			b.Fatal(err)
+		}
+	}
+	types.ToggleCache(false)
+}
+
+func toBytes32(x []byte) [32]byte {
+	var y [32]byte
+	copy(y[:], x)
+	return y
+}
+
 func hexDecodeOrDie(t *testing.T, s string) []byte {
 	res, err := hex.DecodeString(s)
 	if err != nil {

types/BUILD.bazel

@@ -5,6 +5,7 @@ go_library(
     srcs = [
         "array_basic.go",
         "array_composite.go",
+        "array_roots.go",
         "basic.go",
         "determine_size.go",
         "factory.go",
@@ -30,7 +31,7 @@ go_library(
 go_test(
     name = "go_default_test",
     srcs = [
-        "array_basic_test.go",
+        "array_roots_test.go",
         "helpers_test.go",
         "struct_test.go",
     ],

types/array_basic.go

@@ -25,31 +25,18 @@ func newBasicArraySSZ() *basicArraySSZ {
 	}
 }
 
-func (b *basicArraySSZ) Root(val reflect.Value, typ reflect.Type, maxCapacity uint64) ([32]byte, error) {
+func (b *basicArraySSZ) Root(val reflect.Value, typ reflect.Type, fieldName string, maxCapacity uint64) ([32]byte, error) {
 	numItems := val.Len()
 	hashKeyElements := make([]byte, BytesPerChunk*numItems)
 	emptyKey := highwayhash.Sum(hashKeyElements, fastSumHashKey[:])
 	leaves := make([][]byte, numItems)
-	elemKind := typ.Elem().Kind()
 	offset := 0
-	var factory SSZAble
-	var err error
-	if numItems > 0 {
-		factory, err = SSZFactory(val.Index(0), typ.Elem())
-		if err != nil {
-			return [32]byte{}, err
-		}
+	factory, err := SSZFactory(val.Index(0), typ.Elem())
+	if err != nil {
+		return [32]byte{}, err
 	}
 	for i := 0; i < numItems; i++ {
-		// If we are marshaling an byte array of length 32, we shortcut the computations and
-		// simply return it as an identity root.
-		if elemKind == reflect.Array && typ.Elem().Elem().Kind() == reflect.Uint8 && val.Index(i).Len() == 32 {
-			leaves[i] = val.Index(i).Bytes()
-			copy(hashKeyElements[offset:offset+32], leaves[i])
-			offset += 32
-			continue
-		}
-		r, err := factory.Root(val.Index(i), typ.Elem(), 0)
+		r, err := factory.Root(val.Index(i), typ.Elem(), "", 0)
 		if err != nil {
 			return [32]byte{}, err
 		}
@@ -58,7 +45,7 @@ func (b *basicArraySSZ) Root(val reflect.Value, typ reflect.Type, maxCapacity ui
 		offset += 32
 	}
 	hashKey := highwayhash.Sum(hashKeyElements, fastSumHashKey[:])
-	if hashKey != emptyKey {
+	if enableCache && hashKey != emptyKey {
 		b.lock.Lock()
 		res := b.hashCache.Get(string(hashKey[:]))
 		b.lock.Unlock()
@@ -74,7 +61,7 @@ func (b *basicArraySSZ) Root(val reflect.Value, typ reflect.Type, maxCapacity ui
 	if err != nil {
 		return [32]byte{}, err
 	}
-	if hashKey != emptyKey {
+	if enableCache && hashKey != emptyKey {
 		b.lock.Lock()
 		b.hashCache.Set(string(hashKey[:]), root, time.Hour)
 		b.lock.Unlock()

types/array_composite.go

@@ -11,7 +11,7 @@ func newCompositeArraySSZ() *compositeArraySSZ {
 	return &compositeArraySSZ{}
 }
 
-func (b *compositeArraySSZ) Root(val reflect.Value, typ reflect.Type, maxCapacity uint64) ([32]byte, error) {
+func (b *compositeArraySSZ) Root(val reflect.Value, typ reflect.Type, fieldName string, maxCapacity uint64) ([32]byte, error) {
 	var factory SSZAble
 	var err error
 	numItems := val.Len()
@@ -30,7 +30,7 @@ func (b *compositeArraySSZ) Root(val reflect.Value, typ reflect.Type, maxCapacit
 	}
 	limit := (uint64(val.Len())*elemSize + 31) / 32
 	for i := 0; i < val.Len(); i++ {
-		r, err := factory.Root(val.Index(i), typ.Elem(), 0)
+		r, err := factory.Root(val.Index(i), typ.Elem(), "", 0)
 		if err != nil {
 			return [32]byte{}, err
 		}