fix(types): encode a beacon state without holding it twice

crates/blockchain/state_transition/src/beacon/helpers/test_state.rs

@@ -652,3 +652,42 @@ fn gloas_state(count: usize) -> gloas::BeaconState {
             .expect("the vector is built at its exact length"),
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    /// [`BeaconState::ssz_append_capacity`] for every fork. The test lives here
+    /// rather than beside that method in `ethlambda-types` because this module
+    /// is the one place a state of every fork is built.
+    ///
+    /// A small registry is the hard case. libssz's container encoder writes
+    /// each fixed-size field past the end of the buffer before moving it into
+    /// place, so a field needs room beyond the encoding's length unless enough
+    /// variable-size data still follows it, and a small state has little. A
+    /// fork that adds a fixed-size field larger than any the capacity counts
+    /// fails here rather than quietly regrowing a buffer the size of the state.
+    #[test]
+    fn ssz_append_never_regrows_a_buffer_with_its_capacity_to_spare() {
+        // A lean genesis with no validators: its fixed-size fields come before
+        // all of its variable-size ones, which hold almost nothing here.
+        let lean = BeaconState::Lean(ethlambda_types::state::State::from_genesis(7, vec![]));
+        let beacon = ForkName::ALL.map(|fork| with_validators_at(fork, 3));
+        for state in beacon.into_iter().chain([lean]) {
+            let fork = state.fork_name();
+            let capacity = state.ssz_append_capacity();
+
+            let encoded = state.to_ssz();
+            assert_eq!(encoded.capacity(), capacity, "{fork}: to_ssz regrew");
+            let decoded = BeaconState::from_ssz(fork, &encoded).expect("the encoding decodes");
+            assert_eq!(decoded, state, "{fork}");
+
+            // Appending after a prefix reserves past the prefix, not from the
+            // start of the buffer.
+            let mut framed = vec![0xAB];
+            state.ssz_append(&mut framed);
+            assert_eq!(framed.capacity(), 1 + capacity, "{fork}: ssz_append regrew");
+            assert_eq!(framed[1..], encoded, "{fork}");
+        }
+    }
+}

crates/common/types/src/beacon/containers/mod.rs

@@ -366,8 +366,74 @@ impl BeaconState {
     /// unlike the accessors that read a beacon field: this is the inverse of
     /// [`BeaconState::from_ssz`], which builds that variant, so refusing here
     /// would make decoding a state and re-encoding it panic.
+    ///
+    /// Goes through [`BeaconState::ssz_append`] rather than libssz's own
+    /// `to_ssz`, which reserves only the encoding's length and so holds a fulu
+    /// or gloas state twice while encoding it (see
+    /// [`BeaconState::ssz_append_capacity`]).
     pub fn to_ssz(&self) -> Vec<u8> {
-        dispatch_state_including_lean!(self, |state| state.to_ssz())
+        let mut buf = Vec::new();
+        self.ssz_append(&mut buf);
+        buf
+    }
+
+    /// How much [`BeaconState::ssz_append`] reserves past the end of the
+    /// buffer it appends to, so the append never regrows it: the encoding's
+    /// own length plus the largest fixed-size field of this variant's state.
+    /// Public for tests that check that bound.
+    ///
+    /// The headroom is for libssz's container encoder, which writes each
+    /// fixed-size field past the end of the buffer and then moves it into
+    /// place. A fixed-size field written after every variable-size one
+    /// (fulu's `proposer_lookahead`, gloas's `ptc_window`) therefore needs
+    /// room beyond the encoding's length, and without it a buffer the size of
+    /// the whole state is regrown to twice that, copying it, for the last few
+    /// hundred bytes.
+    pub fn ssz_append_capacity(&self) -> usize {
+        dispatch_state_including_lean!(self, |state| state.encoded_len())
+            + self.largest_fixed_field()
+    }
+
+    /// The largest fixed-size field in this variant's state: how far past the
+    /// encoding's end libssz's encoder can write (see
+    /// [`BeaconState::ssz_append_capacity`]).
+    ///
+    /// The beacon forks share one list. Every one of them has `randao_mixes`,
+    /// the largest on mainnet, and the only field that beats it on any preset
+    /// is gloas's `ptc_window` on minimal, by about a kilobyte, so a list per
+    /// fork would save nothing worth keeping eight lists for. A lean state is
+    /// a different container with fields a few hundred bytes wide, and
+    /// reserving a beacon state's megabytes for it would often outweigh the
+    /// state itself.
+    fn largest_fixed_field(&self) -> usize {
+        use libssz::SszEncode;
+        let largest = |sizes: &[usize]| sizes.iter().copied().max().unwrap_or(0);
+        match self {
+            BeaconState::Lean(_) => largest(&[
+                <crate::state::StateConfig as SszEncode>::fixed_size(),
+                <crate::block::BlockHeader as SszEncode>::fixed_size(),
+                <crate::checkpoint::Checkpoint as SszEncode>::fixed_size(),
+            ]),
+            _ => largest(&[
+                <shared::RandaoMixes as SszEncode>::fixed_size(),
+                <shared::BlockRoots as SszEncode>::fixed_size(),
+                <altair::SyncCommittee as SszEncode>::fixed_size(),
+                <gloas::PayloadTimelinessCommitteeWindow as SszEncode>::fixed_size(),
+            ]),
+        }
+    }
+
+    /// Appends [`BeaconState::to_ssz`]'s output to `buf`, after whatever it
+    /// already holds.
+    ///
+    /// For a caller that frames the encoding (the storage layer's fork tag):
+    /// a mainnet state encodes to hundreds of megabytes, and the append holds
+    /// one copy of it, never two. It first reserves
+    /// [`BeaconState::ssz_append_capacity`] past `buf`'s length, so a caller
+    /// need not size `buf` itself.
+    pub fn ssz_append(&self, buf: &mut Vec<u8>) {
+        buf.reserve_exact(self.ssz_append_capacity());
+        dispatch_state_including_lean!(self, |state| state.ssz_append(buf))
     }
 
     /// The state's merkle root, which a block's `state_root` must equal.

crates/storage/src/api/traits.rs

@@ -97,6 +97,16 @@ pub trait StorageWriteBatch: Send {
     /// Put multiple key-value pairs into a table.
     fn put_batch(&mut self, table: Table, batch: Vec<(Vec<u8>, Vec<u8>)>) -> Result<(), Error>;
 
+    /// Put one key-value pair into a table, from borrowed bytes.
+    ///
+    /// For a value the caller keeps after the write, so it need not clone one
+    /// only to hand it over: a beacon snapshot is hundreds of megabytes, and
+    /// RocksDB copies the value into its batch either way. The default goes
+    /// through [`put_batch`](Self::put_batch).
+    fn put(&mut self, table: Table, key: &[u8], value: &[u8]) -> Result<(), Error> {
+        self.put_batch(table, vec![(key.to_vec(), value.to_vec())])
+    }
+
     /// Delete multiple keys from a table.
     fn delete_batch(&mut self, table: Table, keys: Vec<Vec<u8>>) -> Result<(), Error>;
 

crates/storage/src/backend/rocksdb.rs

@@ -219,6 +219,16 @@ impl StorageWriteBatch for RocksDBWriteBatch {
         Ok(())
     }
 
+    fn put(&mut self, table: Table, key: &[u8], value: &[u8]) -> Result<(), Error> {
+        let cf = self
+            .db
+            .cf_handle(cf_name(table))
+            .ok_or_else(|| format!("Column family {} not found", cf_name(table)))?;
+        // The batch copies the bytes, so the caller's buffer is never cloned.
+        self.batch.put_cf(&cf, key, value);
+        Ok(())
+    }
+
     fn delete_batch(&mut self, table: Table, keys: Vec<Vec<u8>>) -> Result<(), Error> {
         let cf = self
             .db

crates/storage/src/backend/tests.rs

@@ -14,6 +14,7 @@ use crate::api::{StorageBackend, Table};
 /// Run the full test suite against a backend.
 pub fn run_backend_tests(backend: &dyn StorageBackend) {
     test_put_and_get(backend);
+    test_put_from_borrowed_bytes(backend);
     test_delete(backend);
     test_prefix_iterator(backend);
     test_nonexistent_key(backend);
@@ -47,6 +48,23 @@ fn test_put_and_get(backend: &dyn StorageBackend) {
     }
 }
 
+fn test_put_from_borrowed_bytes(backend: &dyn StorageBackend) {
+    let value = b"borrowed_value".to_vec();
+    {
+        let mut batch = backend.begin_write().unwrap();
+        batch
+            .put(Table::BlockHeaders, b"test_put_borrowed_key", &value)
+            .unwrap();
+        batch.commit().unwrap();
+    }
+
+    let view = backend.begin_read().unwrap();
+    let stored = view
+        .get(Table::BlockHeaders, b"test_put_borrowed_key")
+        .unwrap();
+    assert_eq!(stored, Some(value), "the caller's buffer is still its own");
+}
+
 fn test_contains(backend: &dyn StorageBackend) {
     {
         let mut batch = backend.begin_write().unwrap();

crates/storage/src/state_codec.rs

@@ -24,9 +24,11 @@ use ethlambda_types::{
 /// [`ForkName::Lean`]'s selector is not a variant index, so the byte must go
 /// back through [`ForkName::from_selector`] rather than being cast.
 pub(crate) fn encode_state_value(state: &BeaconState) -> Vec<u8> {
-    let mut bytes = Vec::new();
-    bytes.push(state.fork_name().selector());
-    bytes.extend_from_slice(&state.to_ssz());
+    // Encoded straight into the tagged buffer: the state writer runs this on
+    // every import, and a separate `Vec` would hold two copies of a mainnet
+    // state at once. `ssz_append` reserves its own room past the tag.
+    let mut bytes = vec![state.fork_name().selector()];
+    state.ssz_append(&mut bytes);
     bytes
 }
 
@@ -66,6 +68,22 @@ mod tests {
         assert_eq!(decode_state_value(&bytes), state);
     }
 
+    #[test]
+    fn a_beacon_state_value_is_its_tag_then_its_ssz_in_one_buffer() {
+        let electra = crate::beacon_state_delta::tests::electra_state_with_validators(64);
+        let state = BeaconState::Electra(electra);
+        let bytes = encode_state_value(&state);
+
+        assert_eq!(bytes[0], ForkName::Electra.selector());
+        assert_eq!(bytes[1..], state.to_ssz());
+        assert_eq!(
+            bytes.capacity(),
+            1 + state.ssz_append_capacity(),
+            "the buffer was regrown"
+        );
+        assert_eq!(decode_state_value(&bytes), state);
+    }
+
     #[test]
     fn the_selector_is_not_a_dense_index() {
         // ForkName::Lean is 255 so that beacon forks after fulu keep taking the

crates/storage/src/state_writer.rs

@@ -81,12 +81,16 @@ pub(crate) struct StateWrite {
     /// `Table::StateDiffs` value. `Some` for every lean state, and for a
     /// beacon state that is not an anchor.
     pub diff: Option<Vec<u8>>,
-    /// Beacon only: the target's [`encode_state_value`] bytes, for the
-    /// writer's snapshot memo. `Some` exactly when the write is a snapshot:
-    /// only a snapshot is ever a diff base, so a diff's own bytes are never
-    /// kept. `None` on the lean arm, which diffs in the field domain and never
-    /// reads a base's bytes.
-    pub encoded: Option<Vec<u8>>,
+    /// Beacon only: `snapshot` also becomes the writer's snapshot memo, the
+    /// base the next diffs are taken against. True exactly when a beacon state
+    /// is written as a snapshot: only a snapshot is ever a diff base, so a
+    /// diff's own bytes are never kept. False on the lean arm, which diffs in
+    /// the field domain and never reads a base's bytes.
+    ///
+    /// A flag rather than a second copy of the bytes: the writer puts the
+    /// snapshot into its batch from a borrow and then moves that same buffer
+    /// into the memo, so a mainnet state is never held twice for it.
+    pub memoize_snapshot: bool,
 }
 
 /// States handed to the writer but not yet committed to the backend.
@@ -504,7 +508,7 @@ pub(crate) fn plan_lean_state_write(state: State, parent: &State) -> StateWrite
     StateWrite {
         snapshot,
         diff: Some(diff),
-        encoded: None,
+        memoize_snapshot: false,
     }
 }
 
@@ -531,9 +535,9 @@ pub(crate) fn plan_beacon_state_write(
     let target = encode_state_value(state);
     match base {
         None => StateWrite {
-            snapshot: Some(target.clone()),
+            snapshot: Some(target),
             diff: None,
-            encoded: Some(target),
+            memoize_snapshot: true,
         },
         Some((snapshot_root, base)) => {
             let delta = beacon_state_delta::encode(&target, base);
@@ -552,7 +556,7 @@ pub(crate) fn plan_beacon_state_write(
             StateWrite {
                 snapshot: None,
                 diff: Some(framed),
-                encoded: None,
+                memoize_snapshot: false,
             }
         }
     }
@@ -693,17 +697,29 @@ impl StateWriter {
                 .put_batch(Table::StateDiffs, vec![(key.clone(), diff)])
                 .expect("put state diff");
         }
-        if let Some(snapshot) = write.snapshot {
-            batch
-                .put_batch(Table::States, vec![(key, snapshot)])
-                .expect("put state snapshot");
-        }
+        // A snapshot the memo keeps is put from a borrow, so the same buffer
+        // can then move into the memo; any other value is handed over.
+        let memo = match write.snapshot {
+            Some(snapshot) if write.memoize_snapshot => {
+                batch
+                    .put(Table::States, &key, &snapshot)
+                    .expect("put state snapshot");
+                Some(snapshot)
+            }
+            Some(snapshot) => {
+                batch
+                    .put_batch(Table::States, vec![(key, snapshot)])
+                    .expect("put state snapshot");
+                None
+            }
+            None => None,
+        };
         batch.commit().expect("commit");
 
         if matches!(&*request.state, BeaconState::Lean(_)) {
             return;
         }
-        match write.encoded {
+        match memo {
             // A snapshot is its own base, and the next diffs are taken against it.
             Some(encoded) => {
                 self.snapshot_memo = Some((root, encoded));
@@ -962,7 +978,7 @@ mod tests {
             Some(expected.as_slice()),
             "the snapshot holds the anchored state's own bytes, not its parent's"
         );
-        assert!(plan.encoded.is_none(), "the lean arm keeps no encoded memo");
+        assert!(!plan.memoize_snapshot, "the lean arm keeps no encoded memo");
 
         let mut inside = base_state();
         inside.slot = interval;
@@ -986,10 +1002,9 @@ mod tests {
             "no base means nothing to diff against"
         );
         assert!(plan.diff.is_none());
-        assert_eq!(
-            plan.encoded.as_deref(),
-            plan.snapshot.as_deref(),
-            "the memo is the very bytes that were snapshotted"
+        assert!(
+            plan.memoize_snapshot,
+            "the snapshotted bytes become the memo"
         );
     }
 
@@ -1014,7 +1029,7 @@ mod tests {
             "a base means a diff, not a snapshot"
         );
         assert!(
-            plan.encoded.is_none(),
+            !plan.memoize_snapshot,
             "only a snapshot is a diff base, so a diff keeps no memo bytes"
         );
         let framed = plan.diff.expect("a non-anchor records a diff");