From cc643ac4cc0ec45da73ee5183245e60df6933888 Mon Sep 17 00:00:00 2001
From: Preston Van Loon <preston@prysmaticlabs.com>
Date: Wed, 11 May 2022 09:25:11 -0500
Subject: [PATCH] native-state: Simplify MarshalSSZ (#10677)

Co-authored-by: prylabs-bulldozer[bot] <58059840+prylabs-bulldozer[bot]@users.noreply.github.com>
---
 beacon-chain/state/state-native/BUILD.bazel |  1 +
 beacon-chain/state/state-native/ssz.go      | 30 +++++----------------
 2 files changed, 7 insertions(+), 24 deletions(-)

diff --git a/beacon-chain/state/state-native/BUILD.bazel b/beacon-chain/state/state-native/BUILD.bazel
index 70508e05a..0b2f870ca 100644
--- a/beacon-chain/state/state-native/BUILD.bazel
+++ b/beacon-chain/state/state-native/BUILD.bazel
@@ -68,6 +68,7 @@ go_library(
         "//encoding/ssz:go_default_library",
         "//proto/prysm/v1alpha1:go_default_library",
         "//runtime/version:go_default_library",
+        "@com_github_ferranbt_fastssz//:go_default_library",
         "@com_github_pkg_errors//:go_default_library",
         "@com_github_prysmaticlabs_go_bitfield//:go_default_library",
         "@io_opencensus_go//trace:go_default_library",
diff --git a/beacon-chain/state/state-native/ssz.go b/beacon-chain/state/state-native/ssz.go
index 1c3ab7ea4..f7d998823 100644
--- a/beacon-chain/state/state-native/ssz.go
+++ b/beacon-chain/state/state-native/ssz.go
@@ -1,36 +1,18 @@
 package state_native
 
 import (
+	ssz "github.com/ferranbt/fastssz"
 	"github.com/pkg/errors"
-	ethpb "github.com/prysmaticlabs/prysm/proto/prysm/v1alpha1"
-	"github.com/prysmaticlabs/prysm/runtime/version"
 )
 
 var errAssertionFailed = errors.New("failed to convert interface to proto state")
-var errUnsupportedVersion = errors.New("unsupported beacon state version")
 
 func (b *BeaconState) MarshalSSZ() ([]byte, error) {
 	proto := b.ToProto()
-	switch b.Version() {
-	case version.Phase0:
-		s, ok := proto.(*ethpb.BeaconState)
-		if !ok {
-			return nil, errAssertionFailed
-		}
-		return s.MarshalSSZ()
-	case version.Altair:
-		s, ok := proto.(*ethpb.BeaconStateAltair)
-		if !ok {
-			return nil, errAssertionFailed
-		}
-		return s.MarshalSSZ()
-	case version.Bellatrix:
-		s, ok := proto.(*ethpb.BeaconStateBellatrix)
-		if !ok {
-			return nil, errAssertionFailed
-		}
-		return s.MarshalSSZ()
-	default:
-		return nil, errUnsupportedVersion
+
+	s, ok := proto.(ssz.Marshaler)
+	if !ok {
+		return nil, errAssertionFailed
 	}
+	return s.MarshalSSZ()
 }