state-diff configs & kv functions (#15903)

* state-diff configs

* state diff kv functions

* potuz's comments

* Update config.go

* fix merge conflicts

* apply bazel's suggestion and fix some bugs

* preston's feedback

beacon-chain/db/kv/BUILD.bazel

@@ -27,6 +27,9 @@ go_library(
         "p2p.go",
         "schema.go",
         "state.go",
+        "state_diff.go",
+        "state_diff_cache.go",
+        "state_diff_helpers.go",
         "state_summary.go",
         "state_summary_cache.go",
         "utils.go",
@@ -41,10 +44,12 @@ go_library(
         "//beacon-chain/db/iface:go_default_library",
         "//beacon-chain/state:go_default_library",
         "//beacon-chain/state/state-native:go_default_library",
+        "//cmd/beacon-chain/flags:go_default_library",
         "//config/features:go_default_library",
         "//config/fieldparams:go_default_library",
         "//config/params:go_default_library",
         "//consensus-types/blocks:go_default_library",
+        "//consensus-types/hdiff:go_default_library",
         "//consensus-types/interfaces:go_default_library",
         "//consensus-types/light-client:go_default_library",
         "//consensus-types/primitives:go_default_library",
@@ -53,6 +58,7 @@ go_library(
         "//encoding/ssz/detect:go_default_library",
         "//genesis:go_default_library",
         "//io/file:go_default_library",
+        "//math:go_default_library",
         "//monitoring/progress:go_default_library",
         "//monitoring/tracing:go_default_library",
         "//monitoring/tracing/trace:go_default_library",
@@ -98,6 +104,7 @@ go_test(
         "migration_block_slot_index_test.go",
         "migration_state_validators_test.go",
         "p2p_test.go",
+        "state_diff_test.go",
         "state_summary_test.go",
         "state_test.go",
         "utils_test.go",
@@ -111,6 +118,7 @@ go_test(
         "//beacon-chain/db/iface:go_default_library",
         "//beacon-chain/state:go_default_library",
         "//beacon-chain/state/state-native:go_default_library",
+        "//cmd/beacon-chain/flags:go_default_library",
         "//config/features:go_default_library",
         "//config/fieldparams:go_default_library",
         "//config/params:go_default_library",
@@ -120,6 +128,7 @@ go_test(
         "//consensus-types/primitives:go_default_library",
         "//encoding/bytesutil:go_default_library",
         "//genesis:go_default_library",
+        "//math:go_default_library",
         "//proto/dbval:go_default_library",
         "//proto/engine/v1:go_default_library",
         "//proto/prysm/v1alpha1:go_default_library",

beacon-chain/db/kv/key.go

@@ -2,6 +2,13 @@ package kv
 
 import "bytes"
 
+func hasPhase0Key(enc []byte) bool {
+	if len(phase0Key) >= len(enc) {
+		return false
+	}
+	return bytes.Equal(enc[:len(phase0Key)], phase0Key)
+}
+
 // In order for an encoding to be Altair compatible, it must be prefixed with altair key.
 func hasAltairKey(enc []byte) bool {
 	if len(altairKey) >= len(enc) {

beacon-chain/db/kv/kv.go

@@ -91,6 +91,7 @@ type Store struct {
 	blockCache          *ristretto.Cache[string, interfaces.ReadOnlySignedBeaconBlock]
 	validatorEntryCache *ristretto.Cache[[]byte, *ethpb.Validator]
 	stateSummaryCache   *stateSummaryCache
+	stateDiffCache      *stateDiffCache
 	ctx                 context.Context
 }
 
@@ -112,6 +113,7 @@ var Buckets = [][]byte{
 	lightClientUpdatesBucket,
 	lightClientBootstrapBucket,
 	lightClientSyncCommitteeBucket,
+	stateDiffBucket,
 	// Indices buckets.
 	blockSlotIndicesBucket,
 	stateSlotIndicesBucket,
@@ -201,6 +203,14 @@ func NewKVStore(ctx context.Context, dirPath string, opts ...KVStoreOption) (*St
 		return nil, err
 	}
 
+	if features.Get().EnableStateDiff {
+		sdCache, err := newStateDiffCache(kv)
+		if err != nil {
+			return nil, err
+		}
+		kv.stateDiffCache = sdCache
+	}
+
 	return kv, nil
 }
 

beacon-chain/db/kv/schema.go

@@ -16,6 +16,7 @@ var (
 	stateValidatorsBucket = []byte("state-validators")
 	feeRecipientBucket    = []byte("fee-recipient")
 	registrationBucket    = []byte("registration")
+	stateDiffBucket       = []byte("state-diff")
 
 	// Light Client Updates Bucket
 	lightClientUpdatesBucket       = []byte("light-client-updates")
@@ -46,6 +47,7 @@ var (
 
 	// Below keys are used to identify objects are to be fork compatible.
 	// Objects that are only compatible with specific forks should be prefixed with such keys.
+	phase0Key                  = []byte("phase0")
 	altairKey                  = []byte("altair")
 	bellatrixKey               = []byte("merge")
 	bellatrixBlindKey          = []byte("blind-bellatrix")

beacon-chain/db/kv/state_diff.go

@@ -0,0 +1,232 @@
+package kv
+
+import (
+	"context"
+
+	"github.com/OffchainLabs/prysm/v7/beacon-chain/state"
+	"github.com/OffchainLabs/prysm/v7/cmd/beacon-chain/flags"
+	"github.com/OffchainLabs/prysm/v7/consensus-types/hdiff"
+	"github.com/OffchainLabs/prysm/v7/consensus-types/primitives"
+	"github.com/OffchainLabs/prysm/v7/monitoring/tracing/trace"
+	"github.com/pkg/errors"
+	bolt "go.etcd.io/bbolt"
+)
+
+const (
+	stateSuffix     = "_s"
+	validatorSuffix = "_v"
+	balancesSuffix  = "_b"
+)
+
+/*
+	We use a level-based approach to save state diffs. Each level corresponds to an exponent of 2 (exponents[lvl]).
+	The data at level 0 is saved every 2**exponent[0] slots and always contains a full state snapshot that is used as a base for the delta saved at other levels.
+*/
+
+// saveStateByDiff takes a state and decides between saving a full state snapshot or a diff.
+func (s *Store) saveStateByDiff(ctx context.Context, st state.ReadOnlyBeaconState) error {
+	_, span := trace.StartSpan(ctx, "BeaconDB.saveStateByDiff")
+	defer span.End()
+
+	if st == nil {
+		return errors.New("state is nil")
+	}
+
+	slot := st.Slot()
+	offset := s.getOffset()
+	if uint64(slot) < offset {
+		return ErrSlotBeforeOffset
+	}
+
+	// Find the level to save the state.
+	lvl := computeLevel(offset, slot)
+	if lvl == -1 {
+		return nil
+	}
+
+	// Save full state if level is 0.
+	if lvl == 0 {
+		return s.saveFullSnapshot(st)
+	}
+
+	// Get anchor state to compute the diff from.
+	anchorState, err := s.getAnchorState(offset, lvl, slot)
+	if err != nil {
+		return err
+	}
+
+	return s.saveHdiff(lvl, anchorState, st)
+}
+
+// stateByDiff retrieves the full state for a given slot.
+func (s *Store) stateByDiff(ctx context.Context, slot primitives.Slot) (state.BeaconState, error) {
+	offset := s.getOffset()
+	if uint64(slot) < offset {
+		return nil, ErrSlotBeforeOffset
+	}
+
+	snapshot, diffChain, err := s.getBaseAndDiffChain(offset, slot)
+	if err != nil {
+		return nil, err
+	}
+
+	for _, diff := range diffChain {
+		if err := ctx.Err(); err != nil {
+			return nil, err
+		}
+
+		snapshot, err = hdiff.ApplyDiff(ctx, snapshot, diff)
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	return snapshot, nil
+}
+
+// saveHdiff computes the diff between the anchor state and the current state and saves it to the database.
+// This function needs to be called only with the latest finalized state, and in a strictly increasing slot order.
+func (s *Store) saveHdiff(lvl int, anchor, st state.ReadOnlyBeaconState) error {
+	slot := uint64(st.Slot())
+	key := makeKeyForStateDiffTree(lvl, slot)
+
+	diff, err := hdiff.Diff(anchor, st)
+	if err != nil {
+		return err
+	}
+
+	err = s.db.Update(func(tx *bolt.Tx) error {
+		bucket := tx.Bucket(stateDiffBucket)
+		if bucket == nil {
+			return bolt.ErrBucketNotFound
+		}
+		buf := append(key, stateSuffix...)
+		if err := bucket.Put(buf, diff.StateDiff); err != nil {
+			return err
+		}
+		buf = append(key, validatorSuffix...)
+		if err := bucket.Put(buf, diff.ValidatorDiffs); err != nil {
+			return err
+		}
+		buf = append(key, balancesSuffix...)
+		if err := bucket.Put(buf, diff.BalancesDiff); err != nil {
+			return err
+		}
+		return nil
+	})
+	if err != nil {
+		return err
+	}
+
+	// Save the full state to the cache (if not the last level).
+	if lvl != len(flags.Get().StateDiffExponents)-1 {
+		err = s.stateDiffCache.setAnchor(lvl, st)
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+// SaveFullSnapshot saves the full level 0 state snapshot to the database.
+func (s *Store) saveFullSnapshot(st state.ReadOnlyBeaconState) error {
+	slot := uint64(st.Slot())
+	key := makeKeyForStateDiffTree(0, slot)
+	stateBytes, err := st.MarshalSSZ()
+	if err != nil {
+		return err
+	}
+	// add version key to value
+	enc, err := addKey(st.Version(), stateBytes)
+	if err != nil {
+		return err
+	}
+
+	err = s.db.Update(func(tx *bolt.Tx) error {
+		bucket := tx.Bucket(stateDiffBucket)
+		if bucket == nil {
+			return bolt.ErrBucketNotFound
+		}
+
+		if err := bucket.Put(key, enc); err != nil {
+			return err
+		}
+
+		return nil
+	})
+	if err != nil {
+		return err
+	}
+	// Save the full state to the cache, and invalidate other levels.
+	s.stateDiffCache.clearAnchors()
+	err = s.stateDiffCache.setAnchor(0, st)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func (s *Store) getDiff(lvl int, slot uint64) (hdiff.HdiffBytes, error) {
+	key := makeKeyForStateDiffTree(lvl, slot)
+	var stateDiff []byte
+	var validatorDiff []byte
+	var balancesDiff []byte
+
+	err := s.db.View(func(tx *bolt.Tx) error {
+		bucket := tx.Bucket(stateDiffBucket)
+		if bucket == nil {
+			return bolt.ErrBucketNotFound
+		}
+		buf := append(key, stateSuffix...)
+		stateDiff = bucket.Get(buf)
+		if stateDiff == nil {
+			return errors.New("state diff not found")
+		}
+		buf = append(key, validatorSuffix...)
+		validatorDiff = bucket.Get(buf)
+		if validatorDiff == nil {
+			return errors.New("validator diff not found")
+		}
+		buf = append(key, balancesSuffix...)
+		balancesDiff = bucket.Get(buf)
+		if balancesDiff == nil {
+			return errors.New("balances diff not found")
+		}
+		return nil
+	})
+
+	if err != nil {
+		return hdiff.HdiffBytes{}, err
+	}
+
+	return hdiff.HdiffBytes{
+		StateDiff:      stateDiff,
+		ValidatorDiffs: validatorDiff,
+		BalancesDiff:   balancesDiff,
+	}, nil
+}
+
+func (s *Store) getFullSnapshot(slot uint64) (state.BeaconState, error) {
+	key := makeKeyForStateDiffTree(0, slot)
+	var enc []byte
+
+	err := s.db.View(func(tx *bolt.Tx) error {
+		bucket := tx.Bucket(stateDiffBucket)
+		if bucket == nil {
+			return bolt.ErrBucketNotFound
+		}
+		enc = bucket.Get(key)
+		if enc == nil {
+			return errors.New("state not found")
+		}
+		return nil
+	})
+
+	if err != nil {
+		return nil, err
+	}
+
+	return decodeStateSnapshot(enc)
+}

beacon-chain/db/kv/state_diff_cache.go

@@ -0,0 +1,77 @@
+package kv
+
+import (
+	"encoding/binary"
+	"errors"
+	"sync"
+
+	"github.com/OffchainLabs/prysm/v7/beacon-chain/state"
+	"github.com/OffchainLabs/prysm/v7/cmd/beacon-chain/flags"
+	"go.etcd.io/bbolt"
+)
+
+type stateDiffCache struct {
+	sync.RWMutex
+	anchors []state.ReadOnlyBeaconState
+	offset  uint64
+}
+
+func newStateDiffCache(s *Store) (*stateDiffCache, error) {
+	var offset uint64
+
+	err := s.db.View(func(tx *bbolt.Tx) error {
+		bucket := tx.Bucket(stateDiffBucket)
+		if bucket == nil {
+			return bbolt.ErrBucketNotFound
+		}
+
+		offsetBytes := bucket.Get([]byte("offset"))
+		if offsetBytes == nil {
+			return errors.New("state diff cache: offset not found")
+		}
+		offset = binary.LittleEndian.Uint64(offsetBytes)
+		return nil
+	})
+	if err != nil {
+		return nil, err
+	}
+
+	return &stateDiffCache{
+		anchors: make([]state.ReadOnlyBeaconState, len(flags.Get().StateDiffExponents)-1), // -1 because last level doesn't need to be cached
+		offset:  offset,
+	}, nil
+}
+
+func (c *stateDiffCache) getAnchor(level int) state.ReadOnlyBeaconState {
+	c.RLock()
+	defer c.RUnlock()
+	return c.anchors[level]
+}
+
+func (c *stateDiffCache) setAnchor(level int, anchor state.ReadOnlyBeaconState) error {
+	c.Lock()
+	defer c.Unlock()
+	if level >= len(c.anchors) || level < 0 {
+		return errors.New("state diff cache: anchor level out of range")
+	}
+	c.anchors[level] = anchor
+	return nil
+}
+
+func (c *stateDiffCache) getOffset() uint64 {
+	c.RLock()
+	defer c.RUnlock()
+	return c.offset
+}
+
+func (c *stateDiffCache) setOffset(offset uint64) {
+	c.Lock()
+	defer c.Unlock()
+	c.offset = offset
+}
+
+func (c *stateDiffCache) clearAnchors() {
+	c.Lock()
+	defer c.Unlock()
+	c.anchors = make([]state.ReadOnlyBeaconState, len(flags.Get().StateDiffExponents)-1) // -1 because last level doesn't need to be cached
+}

beacon-chain/db/kv/state_diff_helpers.go

@@ -0,0 +1,248 @@
+package kv
+
+import (
+	"context"
+	"encoding/binary"
+	"errors"
+	"fmt"
+
+	"github.com/OffchainLabs/prysm/v7/beacon-chain/state"
+	statenative "github.com/OffchainLabs/prysm/v7/beacon-chain/state/state-native"
+	"github.com/OffchainLabs/prysm/v7/cmd/beacon-chain/flags"
+	"github.com/OffchainLabs/prysm/v7/consensus-types/hdiff"
+	"github.com/OffchainLabs/prysm/v7/consensus-types/primitives"
+	"github.com/OffchainLabs/prysm/v7/math"
+	ethpb "github.com/OffchainLabs/prysm/v7/proto/prysm/v1alpha1"
+	"github.com/OffchainLabs/prysm/v7/runtime/version"
+	"go.etcd.io/bbolt"
+)
+
+var (
+	offsetKey           = []byte("offset")
+	ErrSlotBeforeOffset = errors.New("slot is before root offset")
+)
+
+func makeKeyForStateDiffTree(level int, slot uint64) []byte {
+	buf := make([]byte, 16)
+	buf[0] = byte(level)
+	binary.LittleEndian.PutUint64(buf[1:], slot)
+	return buf
+}
+
+func (s *Store) getAnchorState(offset uint64, lvl int, slot primitives.Slot) (anchor state.ReadOnlyBeaconState, err error) {
+	if lvl <= 0 || lvl > len(flags.Get().StateDiffExponents) {
+		return nil, errors.New("invalid value for level")
+	}
+
+	if uint64(slot) < offset {
+		return nil, ErrSlotBeforeOffset
+	}
+	relSlot := uint64(slot) - offset
+	prevExp := flags.Get().StateDiffExponents[lvl-1]
+	if prevExp < 2 || prevExp >= 64 {
+		return nil, fmt.Errorf("state diff exponent %d out of range for uint64", prevExp)
+	}
+	span := math.PowerOf2(uint64(prevExp))
+	anchorSlot := primitives.Slot(uint64(slot) - relSlot%span)
+
+	// anchorLvl can be [0, lvl-1]
+	anchorLvl := computeLevel(offset, anchorSlot)
+	if anchorLvl == -1 {
+		return nil, errors.New("could not compute anchor level")
+	}
+
+	// Check if we have the anchor in cache.
+	anchor = s.stateDiffCache.getAnchor(anchorLvl)
+	if anchor != nil {
+		return anchor, nil
+	}
+
+	// If not, load it from the database.
+	anchor, err = s.stateByDiff(context.Background(), anchorSlot)
+	if err != nil {
+		return nil, err
+	}
+
+	// Save it in the cache.
+	err = s.stateDiffCache.setAnchor(anchorLvl, anchor)
+	if err != nil {
+		return nil, err
+	}
+	return anchor, nil
+}
+
+// computeLevel computes the level in the diff tree. Returns -1 in case slot should not be in tree.
+func computeLevel(offset uint64, slot primitives.Slot) int {
+	rel := uint64(slot) - offset
+	for i, exp := range flags.Get().StateDiffExponents {
+		if exp < 2 || exp >= 64 {
+			return -1
+		}
+		span := math.PowerOf2(uint64(exp))
+		if rel%span == 0 {
+			return i
+		}
+	}
+	// If rel isn’t on any of the boundaries, we should ignore saving it.
+	return -1
+}
+
+func (s *Store) setOffset(slot primitives.Slot) error {
+	err := s.db.Update(func(tx *bbolt.Tx) error {
+		bucket := tx.Bucket(stateDiffBucket)
+		if bucket == nil {
+			return bbolt.ErrBucketNotFound
+		}
+
+		offsetBytes := bucket.Get(offsetKey)
+		if offsetBytes != nil {
+			return fmt.Errorf("offset already set to %d", binary.LittleEndian.Uint64(offsetBytes))
+		}
+
+		offsetBytes = make([]byte, 8)
+		binary.LittleEndian.PutUint64(offsetBytes, uint64(slot))
+		if err := bucket.Put(offsetKey, offsetBytes); err != nil {
+			return err
+		}
+		return nil
+	})
+	if err != nil {
+		return err
+	}
+
+	// Save the offset in the cache.
+	s.stateDiffCache.setOffset(uint64(slot))
+	return nil
+}
+
+func (s *Store) getOffset() uint64 {
+	return s.stateDiffCache.getOffset()
+}
+
+func keyForSnapshot(v int) ([]byte, error) {
+	switch v {
+	case version.Fulu:
+		return fuluKey, nil
+	case version.Electra:
+		return ElectraKey, nil
+	case version.Deneb:
+		return denebKey, nil
+	case version.Capella:
+		return capellaKey, nil
+	case version.Bellatrix:
+		return bellatrixKey, nil
+	case version.Altair:
+		return altairKey, nil
+	case version.Phase0:
+		return phase0Key, nil
+	default:
+		return nil, errors.New("unsupported fork")
+	}
+}
+
+func addKey(v int, bytes []byte) ([]byte, error) {
+	key, err := keyForSnapshot(v)
+	if err != nil {
+		return nil, err
+	}
+	enc := make([]byte, len(key)+len(bytes))
+	copy(enc, key)
+	copy(enc[len(key):], bytes)
+	return enc, nil
+}
+
+func decodeStateSnapshot(enc []byte) (state.BeaconState, error) {
+	switch {
+	case hasFuluKey(enc):
+		var fuluState ethpb.BeaconStateFulu
+		if err := fuluState.UnmarshalSSZ(enc[len(fuluKey):]); err != nil {
+			return nil, err
+		}
+		return statenative.InitializeFromProtoUnsafeFulu(&fuluState)
+	case HasElectraKey(enc):
+		var electraState ethpb.BeaconStateElectra
+		if err := electraState.UnmarshalSSZ(enc[len(ElectraKey):]); err != nil {
+			return nil, err
+		}
+		return statenative.InitializeFromProtoUnsafeElectra(&electraState)
+	case hasDenebKey(enc):
+		var denebState ethpb.BeaconStateDeneb
+		if err := denebState.UnmarshalSSZ(enc[len(denebKey):]); err != nil {
+			return nil, err
+		}
+		return statenative.InitializeFromProtoUnsafeDeneb(&denebState)
+	case hasCapellaKey(enc):
+		var capellaState ethpb.BeaconStateCapella
+		if err := capellaState.UnmarshalSSZ(enc[len(capellaKey):]); err != nil {
+			return nil, err
+		}
+		return statenative.InitializeFromProtoUnsafeCapella(&capellaState)
+	case hasBellatrixKey(enc):
+		var bellatrixState ethpb.BeaconStateBellatrix
+		if err := bellatrixState.UnmarshalSSZ(enc[len(bellatrixKey):]); err != nil {
+			return nil, err
+		}
+		return statenative.InitializeFromProtoUnsafeBellatrix(&bellatrixState)
+	case hasAltairKey(enc):
+		var altairState ethpb.BeaconStateAltair
+		if err := altairState.UnmarshalSSZ(enc[len(altairKey):]); err != nil {
+			return nil, err
+		}
+		return statenative.InitializeFromProtoUnsafeAltair(&altairState)
+	case hasPhase0Key(enc):
+		var phase0State ethpb.BeaconState
+		if err := phase0State.UnmarshalSSZ(enc[len(phase0Key):]); err != nil {
+			return nil, err
+		}
+		return statenative.InitializeFromProtoUnsafePhase0(&phase0State)
+	default:
+		return nil, errors.New("unsupported fork")
+	}
+}
+
+func (s *Store) getBaseAndDiffChain(offset uint64, slot primitives.Slot) (state.BeaconState, []hdiff.HdiffBytes, error) {
+	if uint64(slot) < offset {
+		return nil, nil, ErrSlotBeforeOffset
+	}
+	rel := uint64(slot) - offset
+	lvl := computeLevel(offset, slot)
+	if lvl == -1 {
+		return nil, nil, errors.New("slot not in tree")
+	}
+
+	exponents := flags.Get().StateDiffExponents
+
+	baseSpan := math.PowerOf2(uint64(exponents[0]))
+	baseAnchorSlot := uint64(slot) - rel%baseSpan
+
+	type diffItem struct {
+		level int
+		slot  uint64
+	}
+
+	var diffChainItems []diffItem
+	for i, exp := range exponents[1 : lvl+1] {
+		span := math.PowerOf2(uint64(exp))
+		diffSlot := rel / span * span
+		if diffSlot == baseAnchorSlot {
+			continue
+		}
+		diffChainItems = append(diffChainItems, diffItem{level: i + 1, slot: diffSlot + offset})
+	}
+
+	baseSnapshot, err := s.getFullSnapshot(baseAnchorSlot)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	diffChain := make([]hdiff.HdiffBytes, 0, len(diffChainItems))
+	for _, item := range diffChainItems {
+		diff, err := s.getDiff(item.level, item.slot)
+		if err != nil {
+			return nil, nil, err
+		}
+		diffChain = append(diffChain, diff)
+	}
+
+	return baseSnapshot, diffChain, nil
+}

beacon-chain/db/kv/state_diff_test.go

@@ -0,0 +1,622 @@
+package kv
+
+import (
+	"context"
+	"encoding/binary"
+	"fmt"
+	"math/rand"
+	"testing"
+
+	"github.com/OffchainLabs/prysm/v7/beacon-chain/state"
+	"github.com/OffchainLabs/prysm/v7/cmd/beacon-chain/flags"
+	"github.com/OffchainLabs/prysm/v7/config/params"
+	"github.com/OffchainLabs/prysm/v7/consensus-types/primitives"
+	"github.com/OffchainLabs/prysm/v7/math"
+	ethpb "github.com/OffchainLabs/prysm/v7/proto/prysm/v1alpha1"
+	"github.com/OffchainLabs/prysm/v7/runtime/version"
+	"github.com/OffchainLabs/prysm/v7/testing/require"
+	"github.com/OffchainLabs/prysm/v7/testing/util"
+	"go.etcd.io/bbolt"
+)
+
+func TestStateDiff_LoadOrInitOffset(t *testing.T) {
+	setDefaultExponents()
+
+	db := setupDB(t)
+	err := setOffsetInDB(db, 10)
+	require.NoError(t, err)
+	offset := db.getOffset()
+	require.Equal(t, uint64(10), offset)
+
+	err = db.setOffset(10)
+	require.ErrorContains(t, "offset already set", err)
+	offset = db.getOffset()
+	require.Equal(t, uint64(10), offset)
+}
+
+func TestStateDiff_ComputeLevel(t *testing.T) {
+	db := setupDB(t)
+	setDefaultExponents()
+
+	err := setOffsetInDB(db, 0)
+	require.NoError(t, err)
+
+	offset := db.getOffset()
+
+	// 2 ** 21
+	lvl := computeLevel(offset, primitives.Slot(math.PowerOf2(21)))
+	require.Equal(t, 0, lvl)
+
+	// 2 ** 21 * 3
+	lvl = computeLevel(offset, primitives.Slot(math.PowerOf2(21)*3))
+	require.Equal(t, 0, lvl)
+
+	// 2 ** 18
+	lvl = computeLevel(offset, primitives.Slot(math.PowerOf2(18)))
+	require.Equal(t, 1, lvl)
+
+	// 2 ** 18 * 3
+	lvl = computeLevel(offset, primitives.Slot(math.PowerOf2(18)*3))
+	require.Equal(t, 1, lvl)
+
+	// 2 ** 16
+	lvl = computeLevel(offset, primitives.Slot(math.PowerOf2(16)))
+	require.Equal(t, 2, lvl)
+
+	// 2 ** 16 * 3
+	lvl = computeLevel(offset, primitives.Slot(math.PowerOf2(16)*3))
+	require.Equal(t, 2, lvl)
+
+	// 2 ** 13
+	lvl = computeLevel(offset, primitives.Slot(math.PowerOf2(13)))
+	require.Equal(t, 3, lvl)
+
+	// 2 ** 13 * 3
+	lvl = computeLevel(offset, primitives.Slot(math.PowerOf2(13)*3))
+	require.Equal(t, 3, lvl)
+
+	// 2 ** 11
+	lvl = computeLevel(offset, primitives.Slot(math.PowerOf2(11)))
+	require.Equal(t, 4, lvl)
+
+	// 2 ** 11 * 3
+	lvl = computeLevel(offset, primitives.Slot(math.PowerOf2(11)*3))
+	require.Equal(t, 4, lvl)
+
+	// 2 ** 9
+	lvl = computeLevel(offset, primitives.Slot(math.PowerOf2(9)))
+	require.Equal(t, 5, lvl)
+
+	// 2 ** 9 * 3
+	lvl = computeLevel(offset, primitives.Slot(math.PowerOf2(9)*3))
+	require.Equal(t, 5, lvl)
+
+	// 2 ** 5
+	lvl = computeLevel(offset, primitives.Slot(math.PowerOf2(5)))
+	require.Equal(t, 6, lvl)
+
+	// 2 ** 5 * 3
+	lvl = computeLevel(offset, primitives.Slot(math.PowerOf2(5)*3))
+	require.Equal(t, 6, lvl)
+
+	// 2 ** 7
+	lvl = computeLevel(offset, primitives.Slot(math.PowerOf2(7)))
+	require.Equal(t, 6, lvl)
+
+	// 2 ** 5 + 1
+	lvl = computeLevel(offset, primitives.Slot(math.PowerOf2(5)+1))
+	require.Equal(t, -1, lvl)
+
+	// 2 ** 5 + 16
+	lvl = computeLevel(offset, primitives.Slot(math.PowerOf2(5)+16))
+	require.Equal(t, -1, lvl)
+
+	// 2 ** 5 + 32
+	lvl = computeLevel(offset, primitives.Slot(math.PowerOf2(5)+32))
+	require.Equal(t, 6, lvl)
+
+}
+
+func TestStateDiff_SaveFullSnapshot(t *testing.T) {
+	setDefaultExponents()
+
+	for v := range version.All() {
+		t.Run(version.String(v), func(t *testing.T) {
+			db := setupDB(t)
+
+			// Create state with slot 0
+			st, enc := createState(t, 0, v)
+
+			err := setOffsetInDB(db, 0)
+			require.NoError(t, err)
+
+			err = db.saveStateByDiff(context.Background(), st)
+			require.NoError(t, err)
+
+			err = db.db.View(func(tx *bbolt.Tx) error {
+				bucket := tx.Bucket(stateDiffBucket)
+				if bucket == nil {
+					return bbolt.ErrBucketNotFound
+				}
+				s := bucket.Get(makeKeyForStateDiffTree(0, uint64(0)))
+				if s == nil {
+					return bbolt.ErrIncompatibleValue
+				}
+				require.DeepSSZEqual(t, enc, s)
+				return nil
+			})
+			require.NoError(t, err)
+		})
+	}
+}
+
+func TestStateDiff_SaveAndReadFullSnapshot(t *testing.T) {
+	setDefaultExponents()
+
+	for v := range version.All() {
+		t.Run(version.String(v), func(t *testing.T) {
+			db := setupDB(t)
+
+			st, _ := createState(t, 0, v)
+
+			err := setOffsetInDB(db, 0)
+			require.NoError(t, err)
+
+			err = db.saveStateByDiff(context.Background(), st)
+			require.NoError(t, err)
+
+			readSt, err := db.stateByDiff(context.Background(), 0)
+			require.NoError(t, err)
+			require.NotNil(t, readSt)
+
+			stSSZ, err := st.MarshalSSZ()
+			require.NoError(t, err)
+			readStSSZ, err := readSt.MarshalSSZ()
+			require.NoError(t, err)
+			require.DeepSSZEqual(t, stSSZ, readStSSZ)
+		})
+	}
+}
+
+func TestStateDiff_SaveDiff(t *testing.T) {
+	setDefaultExponents()
+
+	for v := range version.All() {
+		t.Run(version.String(v), func(t *testing.T) {
+			db := setupDB(t)
+
+			// Create state with slot 2**21
+			slot := primitives.Slot(math.PowerOf2(21))
+			st, enc := createState(t, slot, v)
+
+			err := setOffsetInDB(db, uint64(slot))
+			require.NoError(t, err)
+
+			err = db.saveStateByDiff(context.Background(), st)
+			require.NoError(t, err)
+
+			err = db.db.View(func(tx *bbolt.Tx) error {
+				bucket := tx.Bucket(stateDiffBucket)
+				if bucket == nil {
+					return bbolt.ErrBucketNotFound
+				}
+				s := bucket.Get(makeKeyForStateDiffTree(0, uint64(slot)))
+				if s == nil {
+					return bbolt.ErrIncompatibleValue
+				}
+				require.DeepSSZEqual(t, enc, s)
+				return nil
+			})
+			require.NoError(t, err)
+
+			// create state with slot 2**18 (+2**21)
+			slot = primitives.Slot(math.PowerOf2(18) + math.PowerOf2(21))
+			st, _ = createState(t, slot, v)
+
+			err = db.saveStateByDiff(context.Background(), st)
+			require.NoError(t, err)
+
+			key := makeKeyForStateDiffTree(1, uint64(slot))
+			err = db.db.View(func(tx *bbolt.Tx) error {
+				bucket := tx.Bucket(stateDiffBucket)
+				if bucket == nil {
+					return bbolt.ErrBucketNotFound
+				}
+				buf := append(key, "_s"...)
+				s := bucket.Get(buf)
+				if s == nil {
+					return bbolt.ErrIncompatibleValue
+				}
+				buf = append(key, "_v"...)
+				v := bucket.Get(buf)
+				if v == nil {
+					return bbolt.ErrIncompatibleValue
+				}
+				buf = append(key, "_b"...)
+				b := bucket.Get(buf)
+				if b == nil {
+					return bbolt.ErrIncompatibleValue
+				}
+				return nil
+			})
+			require.NoError(t, err)
+		})
+	}
+}
+
+func TestStateDiff_SaveAndReadDiff(t *testing.T) {
+	setDefaultExponents()
+
+	for v := range version.All() {
+		t.Run(version.String(v), func(t *testing.T) {
+			db := setupDB(t)
+
+			st, _ := createState(t, 0, v)
+
+			err := setOffsetInDB(db, 0)
+			require.NoError(t, err)
+
+			err = db.saveStateByDiff(context.Background(), st)
+			require.NoError(t, err)
+
+			slot := primitives.Slot(math.PowerOf2(5))
+			st, _ = createState(t, slot, v)
+
+			err = db.saveStateByDiff(context.Background(), st)
+			require.NoError(t, err)
+
+			readSt, err := db.stateByDiff(context.Background(), slot)
+			require.NoError(t, err)
+			require.NotNil(t, readSt)
+
+			stSSZ, err := st.MarshalSSZ()
+			require.NoError(t, err)
+			readStSSZ, err := readSt.MarshalSSZ()
+			require.NoError(t, err)
+			require.DeepSSZEqual(t, stSSZ, readStSSZ)
+		})
+	}
+}
+
+func TestStateDiff_SaveAndReadDiff_MultipleLevels(t *testing.T) {
+	setDefaultExponents()
+
+	for v := range version.All() {
+		t.Run(version.String(v), func(t *testing.T) {
+			db := setupDB(t)
+
+			st, _ := createState(t, 0, v)
+
+			err := setOffsetInDB(db, 0)
+			require.NoError(t, err)
+
+			err = db.saveStateByDiff(context.Background(), st)
+			require.NoError(t, err)
+
+			slot := primitives.Slot(math.PowerOf2(11))
+			st, _ = createState(t, slot, v)
+
+			err = db.saveStateByDiff(context.Background(), st)
+			require.NoError(t, err)
+
+			readSt, err := db.stateByDiff(context.Background(), slot)
+			require.NoError(t, err)
+			require.NotNil(t, readSt)
+
+			stSSZ, err := st.MarshalSSZ()
+			require.NoError(t, err)
+			readStSSZ, err := readSt.MarshalSSZ()
+			require.NoError(t, err)
+			require.DeepSSZEqual(t, stSSZ, readStSSZ)
+
+			slot = primitives.Slot(math.PowerOf2(11) + math.PowerOf2(9))
+			st, _ = createState(t, slot, v)
+
+			err = db.saveStateByDiff(context.Background(), st)
+			require.NoError(t, err)
+
+			readSt, err = db.stateByDiff(context.Background(), slot)
+			require.NoError(t, err)
+			require.NotNil(t, readSt)
+
+			stSSZ, err = st.MarshalSSZ()
+			require.NoError(t, err)
+			readStSSZ, err = readSt.MarshalSSZ()
+			require.NoError(t, err)
+			require.DeepSSZEqual(t, stSSZ, readStSSZ)
+
+			slot = primitives.Slot(math.PowerOf2(11) + math.PowerOf2(9) + math.PowerOf2(5))
+			st, _ = createState(t, slot, v)
+
+			err = db.saveStateByDiff(context.Background(), st)
+			require.NoError(t, err)
+
+			readSt, err = db.stateByDiff(context.Background(), slot)
+			require.NoError(t, err)
+			require.NotNil(t, readSt)
+
+			stSSZ, err = st.MarshalSSZ()
+			require.NoError(t, err)
+			readStSSZ, err = readSt.MarshalSSZ()
+			require.NoError(t, err)
+			require.DeepSSZEqual(t, stSSZ, readStSSZ)
+		})
+	}
+}
+
+func TestStateDiff_SaveAndReadDiffForkTransition(t *testing.T) {
+	setDefaultExponents()
+
+	for v := range version.All()[:len(version.All())-1] {
+		t.Run(version.String(v), func(t *testing.T) {
+			db := setupDB(t)
+
+			st, _ := createState(t, 0, v)
+
+			err := setOffsetInDB(db, 0)
+			require.NoError(t, err)
+
+			err = db.saveStateByDiff(context.Background(), st)
+			require.NoError(t, err)
+
+			slot := primitives.Slot(math.PowerOf2(5))
+			st, _ = createState(t, slot, v+1)
+
+			err = db.saveStateByDiff(context.Background(), st)
+			require.NoError(t, err)
+
+			readSt, err := db.stateByDiff(context.Background(), slot)
+			require.NoError(t, err)
+			require.NotNil(t, readSt)
+
+			stSSZ, err := st.MarshalSSZ()
+			require.NoError(t, err)
+			readStSSZ, err := readSt.MarshalSSZ()
+			require.NoError(t, err)
+			require.DeepSSZEqual(t, stSSZ, readStSSZ)
+		})
+	}
+}
+
+func TestStateDiff_OffsetCache(t *testing.T) {
+	setDefaultExponents()
+
+	// test for slot numbers 0 and 1 for every version
+	for slotNum := range 2 {
+		for v := range version.All() {
+			t.Run(fmt.Sprintf("slotNum=%d,%s", slotNum, version.String(v)), func(t *testing.T) {
+				db := setupDB(t)
+
+				slot := primitives.Slot(slotNum)
+				err := setOffsetInDB(db, uint64(slot))
+				require.NoError(t, err)
+				st, _ := createState(t, slot, v)
+				err = db.saveStateByDiff(context.Background(), st)
+				require.NoError(t, err)
+
+				offset := db.stateDiffCache.getOffset()
+				require.Equal(t, uint64(slotNum), offset)
+
+				slot2 := primitives.Slot(uint64(slotNum) + math.PowerOf2(uint64(flags.Get().StateDiffExponents[0])))
+				st2, _ := createState(t, slot2, v)
+				err = db.saveStateByDiff(context.Background(), st2)
+				require.NoError(t, err)
+
+				offset = db.stateDiffCache.getOffset()
+				require.Equal(t, uint64(slot), offset)
+			})
+		}
+	}
+}
+
+func TestStateDiff_AnchorCache(t *testing.T) {
+	setDefaultExponents()
+
+	for v := range version.All() {
+		t.Run(version.String(v), func(t *testing.T) {
+			exponents := flags.Get().StateDiffExponents
+			localCache := make([]state.ReadOnlyBeaconState, len(exponents)-1)
+			db := setupDB(t)
+			err := setOffsetInDB(db, 0) // lvl 0
+			require.NoError(t, err)
+
+			// at first the cache should be empty
+			for i := 0; i < len(flags.Get().StateDiffExponents)-1; i++ {
+				anchor := db.stateDiffCache.getAnchor(i)
+				require.IsNil(t, anchor)
+			}
+
+			// add level 0
+			slot := primitives.Slot(0) // offset 0 is already set
+			st, _ := createState(t, slot, v)
+			err = db.saveStateByDiff(context.Background(), st)
+			require.NoError(t, err)
+			localCache[0] = st
+
+			// level 0 should be the same
+			require.DeepEqual(t, localCache[0], db.stateDiffCache.getAnchor(0))
+
+			// rest of the cache should be nil
+			for i := 1; i < len(exponents)-1; i++ {
+				require.IsNil(t, db.stateDiffCache.getAnchor(i))
+			}
+
+			// skip last level as it does not get cached
+			for i := len(exponents) - 2; i > 0; i-- {
+				slot = primitives.Slot(math.PowerOf2(uint64(exponents[i])))
+				st, _ := createState(t, slot, v)
+				err = db.saveStateByDiff(context.Background(), st)
+				require.NoError(t, err)
+				localCache[i] = st
+
+				// anchor cache must match local cache
+				for i := 0; i < len(exponents)-1; i++ {
+					if localCache[i] == nil {
+						require.IsNil(t, db.stateDiffCache.getAnchor(i))
+						continue
+					}
+					localSSZ, err := localCache[i].MarshalSSZ()
+					require.NoError(t, err)
+					anchorSSZ, err := db.stateDiffCache.getAnchor(i).MarshalSSZ()
+					require.NoError(t, err)
+					require.DeepSSZEqual(t, localSSZ, anchorSSZ)
+				}
+			}
+
+			// moving to a new tree should invalidate the cache except for level 0
+			twoTo21 := math.PowerOf2(21)
+			slot = primitives.Slot(twoTo21)
+			st, _ = createState(t, slot, v)
+			err = db.saveStateByDiff(context.Background(), st)
+			require.NoError(t, err)
+			localCache = make([]state.ReadOnlyBeaconState, len(exponents)-1)
+			localCache[0] = st
+
+			// level 0 should be the same
+			require.DeepEqual(t, localCache[0], db.stateDiffCache.getAnchor(0))
+
+			// rest of the cache should be nil
+			for i := 1; i < len(exponents)-1; i++ {
+				require.IsNil(t, db.stateDiffCache.getAnchor(i))
+			}
+		})
+	}
+}
+
+func TestStateDiff_EncodingAndDecoding(t *testing.T) {
+	for v := range version.All() {
+		t.Run(version.String(v), func(t *testing.T) {
+			st, enc := createState(t, 0, v) // this has addKey called inside
+			stDecoded, err := decodeStateSnapshot(enc)
+			require.NoError(t, err)
+			st1ssz, err := st.MarshalSSZ()
+			require.NoError(t, err)
+			st2ssz, err := stDecoded.MarshalSSZ()
+			require.NoError(t, err)
+			require.DeepSSZEqual(t, st1ssz, st2ssz)
+		})
+	}
+}
+
+func createState(t *testing.T, slot primitives.Slot, v int) (state.ReadOnlyBeaconState, []byte) {
+	p := params.BeaconConfig()
+	var st state.BeaconState
+	var err error
+	switch v {
+	case version.Phase0:
+		st, err = util.NewBeaconState()
+		require.NoError(t, err)
+		err = st.SetFork(&ethpb.Fork{
+			PreviousVersion: p.GenesisForkVersion,
+			CurrentVersion:  p.GenesisForkVersion,
+			Epoch:           0,
+		})
+		require.NoError(t, err)
+	case version.Altair:
+		st, err = util.NewBeaconStateAltair()
+		require.NoError(t, err)
+		err = st.SetFork(&ethpb.Fork{
+			PreviousVersion: p.GenesisForkVersion,
+			CurrentVersion:  p.AltairForkVersion,
+			Epoch:           p.AltairForkEpoch,
+		})
+		require.NoError(t, err)
+	case version.Bellatrix:
+		st, err = util.NewBeaconStateBellatrix()
+		require.NoError(t, err)
+		err = st.SetFork(&ethpb.Fork{
+			PreviousVersion: p.AltairForkVersion,
+			CurrentVersion:  p.BellatrixForkVersion,
+			Epoch:           p.BellatrixForkEpoch,
+		})
+		require.NoError(t, err)
+	case version.Capella:
+		st, err = util.NewBeaconStateCapella()
+		require.NoError(t, err)
+		err = st.SetFork(&ethpb.Fork{
+			PreviousVersion: p.BellatrixForkVersion,
+			CurrentVersion:  p.CapellaForkVersion,
+			Epoch:           p.CapellaForkEpoch,
+		})
+		require.NoError(t, err)
+	case version.Deneb:
+		st, err = util.NewBeaconStateDeneb()
+		require.NoError(t, err)
+		err = st.SetFork(&ethpb.Fork{
+			PreviousVersion: p.CapellaForkVersion,
+			CurrentVersion:  p.DenebForkVersion,
+			Epoch:           p.DenebForkEpoch,
+		})
+		require.NoError(t, err)
+	case version.Electra:
+		st, err = util.NewBeaconStateElectra()
+		require.NoError(t, err)
+		err = st.SetFork(&ethpb.Fork{
+			PreviousVersion: p.DenebForkVersion,
+			CurrentVersion:  p.ElectraForkVersion,
+			Epoch:           p.ElectraForkEpoch,
+		})
+		require.NoError(t, err)
+	case version.Fulu:
+		st, err = util.NewBeaconStateFulu()
+		require.NoError(t, err)
+		err = st.SetFork(&ethpb.Fork{
+			PreviousVersion: p.ElectraForkVersion,
+			CurrentVersion:  p.FuluForkVersion,
+			Epoch:           p.FuluForkEpoch,
+		})
+		require.NoError(t, err)
+	default:
+		t.Fatalf("unsupported version: %d", v)
+	}
+
+	err = st.SetSlot(slot)
+	require.NoError(t, err)
+	slashings := make([]uint64, 8192)
+	slashings[0] = uint64(rand.Intn(10))
+	err = st.SetSlashings(slashings)
+	require.NoError(t, err)
+	stssz, err := st.MarshalSSZ()
+	require.NoError(t, err)
+	enc, err := addKey(v, stssz)
+	require.NoError(t, err)
+	return st, enc
+}
+
+func setOffsetInDB(s *Store, offset uint64) error {
+	err := s.db.Update(func(tx *bbolt.Tx) error {
+		bucket := tx.Bucket(stateDiffBucket)
+		if bucket == nil {
+			return bbolt.ErrBucketNotFound
+		}
+
+		offsetBytes := bucket.Get(offsetKey)
+		if offsetBytes != nil {
+			return fmt.Errorf("offset already set to %d", binary.LittleEndian.Uint64(offsetBytes))
+		}
+
+		offsetBytes = make([]byte, 8)
+		binary.LittleEndian.PutUint64(offsetBytes, offset)
+		if err := bucket.Put(offsetKey, offsetBytes); err != nil {
+			return err
+		}
+		return nil
+	})
+	if err != nil {
+		return err
+	}
+
+	sdCache, err := newStateDiffCache(s)
+	if err != nil {
+		return err
+	}
+	s.stateDiffCache = sdCache
+	return nil
+}
+
+func setDefaultExponents() {
+	globalFlags := flags.GlobalFlags{
+		StateDiffExponents: []int{21, 18, 16, 13, 11, 9, 5},
+	}
+	flags.Init(&globalFlags)
+}

beacon-chain/node/node.go

@@ -280,7 +280,10 @@ func configureBeacon(cliCtx *cli.Context) error {
 		return errors.Wrap(err, "could not configure beacon chain")
 	}
 
-	flags.ConfigureGlobalFlags(cliCtx)
+	err := flags.ConfigureGlobalFlags(cliCtx)
+	if err != nil {
+		return errors.Wrap(err, "could not configure global flags")
+	}
 
 	if err := configureChainConfig(cliCtx); err != nil {
 		return errors.Wrap(err, "could not configure chain config")