Add in Object Mapping For Types (#9381)

Co-authored-by: Raul Jordan <raul@prysmaticlabs.com>
Co-authored-by: prylabs-bulldozer[bot] <58059840+prylabs-bulldozer[bot]@users.noreply.github.com>

beacon-chain/p2p/types/BUILD.bazel

@@ -3,6 +3,7 @@ load("@prysm//tools/go:def.bzl", "go_library", "go_test")
 go_library(
     name = "go_default_library",
     srcs = [
+        "object_mapping.go",
         "rpc_errors.go",
         "rpc_goodbye_codes.go",
         "types.go",
@@ -15,18 +16,28 @@ go_library(
         "//validator/client:__pkg__",
     ],
     deps = [
+        "//proto/prysm/v1alpha1:go_default_library",
+        "//proto/prysm/v1alpha1/block:go_default_library",
+        "//proto/prysm/v1alpha1/metadata:go_default_library",
+        "//proto/prysm/v1alpha1/wrapper:go_default_library",
+        "//shared/bytesutil:go_default_library",
         "//shared/params:go_default_library",
         "@com_github_ferranbt_fastssz//:go_default_library",
         "@com_github_pkg_errors//:go_default_library",
         "@com_github_prysmaticlabs_eth2_types//:go_default_library",
+        "@org_golang_google_protobuf//proto:go_default_library",
     ],
 )
 
 go_test(
     name = "go_default_test",
-    srcs = ["types_test.go"],
+    srcs = [
+        "object_mapping_test.go",
+        "types_test.go",
+    ],
     embed = [":go_default_library"],
     deps = [
+        "//shared/bytesutil:go_default_library",
         "//shared/params:go_default_library",
         "//shared/testutil/assert:go_default_library",
         "//shared/testutil/require:go_default_library",

beacon-chain/p2p/types/object_mapping.go

@@ -0,0 +1,64 @@
+package types
+
+import (
+	eth "github.com/prysmaticlabs/prysm/proto/prysm/v1alpha1"
+	ethpb "github.com/prysmaticlabs/prysm/proto/prysm/v1alpha1"
+	"github.com/prysmaticlabs/prysm/proto/prysm/v1alpha1/block"
+	"github.com/prysmaticlabs/prysm/proto/prysm/v1alpha1/metadata"
+	wrapperv1 "github.com/prysmaticlabs/prysm/proto/prysm/v1alpha1/wrapper"
+	wrapperv2 "github.com/prysmaticlabs/prysm/proto/prysm/v1alpha1/wrapper"
+	"github.com/prysmaticlabs/prysm/shared/bytesutil"
+	"github.com/prysmaticlabs/prysm/shared/params"
+	"google.golang.org/protobuf/proto"
+)
+
+func init() {
+	// Initialize data maps.
+	InitializeDataMaps()
+}
+
+// This file provides a mapping of fork version to the respective data type. This is
+// to allow any service to appropriately use the correct data type with a provided
+// fork-version.
+
+var (
+	// BlockMap maps the fork-version to the underlying data type for that
+	// particular fork period.
+	BlockMap map[[4]byte]func() (block.SignedBeaconBlock, error)
+	// StateMap maps the fork-version to the underlying data type for that
+	// particular fork period.
+	StateMap map[[4]byte]proto.Message
+	// MetaDataMap maps the fork-version to the underlying data type for that
+	// particular fork period.
+	MetaDataMap map[[4]byte]func() metadata.Metadata
+)
+
+// InitializeDataMaps initializes all the relevant object maps. This function is called to
+// reset maps and reinitialize them.
+func InitializeDataMaps() {
+	// Reset our block map.
+	BlockMap = map[[4]byte]func() (block.SignedBeaconBlock, error){
+		bytesutil.ToBytes4(params.BeaconConfig().GenesisForkVersion): func() (block.SignedBeaconBlock, error) {
+			return wrapperv1.WrappedPhase0SignedBeaconBlock(&eth.SignedBeaconBlock{}), nil
+		},
+		bytesutil.ToBytes4(params.BeaconConfig().AltairForkVersion): func() (block.SignedBeaconBlock, error) {
+			return wrapperv2.WrappedAltairSignedBeaconBlock(&ethpb.SignedBeaconBlockAltair{Block: &ethpb.BeaconBlockAltair{}})
+		},
+	}
+
+	// Reset our state map.
+	StateMap = map[[4]byte]proto.Message{
+		bytesutil.ToBytes4(params.BeaconConfig().GenesisForkVersion): &ethpb.BeaconState{},
+		bytesutil.ToBytes4(params.BeaconConfig().AltairForkVersion):  &ethpb.BeaconStateAltair{},
+	}
+
+	// Reset our metadata map.
+	MetaDataMap = map[[4]byte]func() metadata.Metadata{
+		bytesutil.ToBytes4(params.BeaconConfig().GenesisForkVersion): func() metadata.Metadata {
+			return wrapperv2.WrappedMetadataV0(&ethpb.MetaDataV0{})
+		},
+		bytesutil.ToBytes4(params.BeaconConfig().AltairForkVersion): func() metadata.Metadata {
+			return wrapperv2.WrappedMetadataV1(&ethpb.MetaDataV1{})
+		},
+	}
+}

beacon-chain/p2p/types/object_mapping_test.go

@@ -0,0 +1,51 @@
+package types
+
+import (
+	"testing"
+
+	"github.com/prysmaticlabs/prysm/shared/bytesutil"
+	"github.com/prysmaticlabs/prysm/shared/params"
+	"github.com/prysmaticlabs/prysm/shared/testutil/assert"
+)
+
+func TestInitializeDataMaps(t *testing.T) {
+	params.SetupTestConfigCleanup(t)
+	tests := []struct {
+		name   string
+		action func()
+		exists bool
+	}{
+		{
+			name: "no change",
+			action: func() {
+			},
+			exists: true,
+		},
+		{
+			name: "fork version changes",
+			action: func() {
+				cfg := params.BeaconConfig()
+				cfg.GenesisForkVersion = []byte{0x01, 0x02, 0x00, 0x00}
+				params.OverrideBeaconConfig(cfg)
+			},
+			exists: false,
+		},
+		{
+			name: "fork version changes with reset",
+			action: func() {
+				cfg := params.BeaconConfig()
+				cfg.GenesisForkVersion = []byte{0x01, 0x02, 0x00, 0x00}
+				params.OverrideBeaconConfig(cfg)
+				InitializeDataMaps()
+			},
+			exists: true,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			tt.action()
+			_, ok := BlockMap[bytesutil.ToBytes4(params.BeaconConfig().GenesisForkVersion)]
+			assert.Equal(t, tt.exists, ok)
+		})
+	}
+}