Add CachedState wrapper to improve process_slots logger parameter design by guha-rahul · Pull Request #610 · blockblaz/zeam

guha-rahul · 2026-02-27T11:12:28Z

Summary

Add CachedState struct wrapping BeamState with logger and allocator, removing awkward logger: anytype parameter from callers
Refactor apply_raw_block and apply_transition to use CachedState

Fixes #274

Copilot

Pull request overview

This PR introduces a CachedState wrapper around BeamState to improve the design of state transition functions by removing the awkward logger: anytype parameter and adding justifications caching to avoid expensive flatten/unflatten operations on every block. The CachedState struct encapsulates the state, allocator, logger, and a lazy-loaded justifications cache that is only flushed back to SSZ fields when needed for hash tree root computations.

Changes:

Added CachedState struct with lazy-loaded justifications caching to eliminate repeated SSZ serialization overhead
Refactored apply_raw_block and apply_transition to use CachedState instead of passing logger as a parameter
Changed visibility of BeamState.process_block from public to private and BeamState.shiftJustifiedSlots from private to public

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 2 comments.

File	Description
pkgs/types/src/cached_state.zig	New file implementing CachedState wrapper with justifications caching and delegated state transition methods
pkgs/types/src/state.zig	Changed visibility: shiftJustifiedSlots made public (needed by CachedState), process_block made private (replaced by CachedState.process_block)
pkgs/types/src/lib.zig	Exports CachedState for public API consumption
pkgs/state-transition/src/transition.zig	Updated apply_raw_block and apply_transition to use CachedState with explicit flush calls before hashTreeRoot

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Copilot · 2026-02-27T11:22:37Z

pkgs/types/src/state.zig

    }

-    pub fn process_block(self: *Self, allocator: Allocator, staged_block: BeamBlock, logger: zeam_utils.ModuleLogger, cache: ?*utils.RootToSlotCache) !void {
+    fn process_block(self: *Self, allocator: Allocator, staged_block: BeamBlock, logger: zeam_utils.ModuleLogger, cache: ?*utils.RootToSlotCache) !void {


Changing process_block visibility from pub fn to fn breaks existing tests in state.zig. Multiple test functions (lines 814, 899, 918, 937, 967, 986, 995, 1000, 1082, 1108, 1123, 1133) directly call state.process_block() which will no longer be accessible after this visibility change. These tests need to be updated to use CachedState.process_block instead, or process_block should remain public.

Suggested change

fn process_block(self: *Self, allocator: Allocator, staged_block: BeamBlock, logger: zeam_utils.ModuleLogger, cache: ?*utils.RootToSlotCache) !void {

pub fn process_block(self: *Self, allocator: Allocator, staged_block: BeamBlock, logger: zeam_utils.ModuleLogger, cache: ?*utils.RootToSlotCache) !void {

Copilot · 2026-02-27T11:22:38Z

pkgs/types/src/cached_state.zig

+    fn process_attestations(self: *Self, attestations: AggregatedAttestations, cache: ?*RootToSlotCache) !void {
+        const allocator = self.allocator;
+        const logger = self.logger;
+
+        const attestations_timer = zeam_metrics.lean_state_transition_attestations_processing_time_seconds.start();
+        defer _ = attestations_timer.observe();
+
+        if (comptime !zeam_metrics.isZKVM()) {
+            const attestation_count: u64 = @intCast(attestations.constSlice().len);
+            zeam_metrics.metrics.lean_state_transition_attestations_processed_total.incrBy(attestation_count);
+        }
+
+        logger.debug("process attestations slot={d} \n prestate:historical hashes={d} justified slots={d} attestations={d}, ", .{ self.state.slot, self.state.historical_block_hashes.len(), self.state.justified_slots.len(), attestations.constSlice().len });
+        const justified_str = try self.state.latest_justified.toJsonString(allocator);
+        defer allocator.free(justified_str);
+        const finalized_str = try self.state.latest_finalized.toJsonString(allocator);
+        defer allocator.free(finalized_str);
+
+        logger.debug("prestate justified={s} finalized={s}", .{ justified_str, finalized_str });
+
+        const justifications = try self.ensureJustificationsLoaded();
+
+        var finalized_slot: Slot = self.state.latest_finalized.slot;
+
+        // Use the global cache directly if provided, otherwise build a local cache.
+        var owned_cache: ?RootToSlotCache = if (cache == null) RootToSlotCache.init(allocator) else null;
+        defer if (owned_cache) |*oc| oc.deinit();
+        const block_cache = cache orelse &(owned_cache.?);
+        if (owned_cache != null) {
+            try self.state.initRootToSlotCache(block_cache);
+        }
+
+        const num_validators: usize = @intCast(self.state.validatorCount());
+        for (attestations.constSlice()) |aggregated_attestation| {
+            var validator_indices = try attestation.aggregationBitsToValidatorIndices(&aggregated_attestation.aggregation_bits, allocator);
+            defer validator_indices.deinit(allocator);
+
+            if (validator_indices.items.len == 0) {
+                continue;
+            }
+
+            const attestation_data = aggregated_attestation.data;
+            const source_slot: Slot = attestation_data.source.slot;
+            const target_slot: Slot = attestation_data.target.slot;
+            const attestation_str = try attestation_data.toJsonString(allocator);
+            defer allocator.free(attestation_str);
+
+            logger.debug("processing attestation={s} validators_count={d}\n", .{ attestation_str, validator_indices.items.len });
+
+            const historical_len: Slot = @intCast(self.state.historical_block_hashes.len());
+            if (source_slot >= historical_len) {
+                return StateTransitionError.InvalidSlotIndex;
+            }
+            if (target_slot >= historical_len) {
+                return StateTransitionError.InvalidSlotIndex;
+            }
+
+            const is_source_justified = try utils.isSlotJustified(finalized_slot, &self.state.justified_slots, source_slot);
+            const is_target_already_justified = try utils.isSlotJustified(finalized_slot, &self.state.justified_slots, target_slot);
+            const stored_source_root = try self.state.historical_block_hashes.get(@intCast(source_slot));
+            const stored_target_root = try self.state.historical_block_hashes.get(@intCast(target_slot));
+            const is_zero_source = std.mem.eql(u8, &attestation_data.source.root, &utils.ZERO_HASH);
+            const is_zero_target = std.mem.eql(u8, &attestation_data.target.root, &utils.ZERO_HASH);
+            if (is_zero_source or is_zero_target) {
+                logger.debug("skipping the attestation as not viable: source_zero_root={} target_zero_root={}", .{
+                    is_zero_source,
+                    is_zero_target,
+                });
+                continue;
+            }
+            const has_correct_source_root = std.mem.eql(u8, &attestation_data.source.root, &stored_source_root);
+            const has_correct_target_root = std.mem.eql(u8, &attestation_data.target.root, &stored_target_root);
+            const has_known_root = has_correct_source_root and has_correct_target_root;
+
+            const target_not_ahead = target_slot <= source_slot;
+            const is_target_justifiable = try utils.IsJustifiableSlot(self.state.latest_finalized.slot, target_slot);
+
+            if (!is_source_justified or
+                is_target_already_justified or
+                !has_known_root or
+                target_not_ahead or
+                !is_target_justifiable)
+            {
+                logger.debug("skipping the attestation as not viable: !(source_justified={}) or target_already_justified={} !(known_root={}) or target_not_ahead={} or !(target_justifiable={})", .{
+                    is_source_justified,
+                    is_target_already_justified,
+                    has_known_root,
+                    target_not_ahead,
+                    is_target_justifiable,
+                });
+                continue;
+            }
+
+            var target_justifications = justifications.get(attestation_data.target.root) orelse targetjustifications: {
+                const targetjustifications = try allocator.alloc(u8, num_validators);
+                @memset(targetjustifications, 0);
+                try justifications.put(allocator, attestation_data.target.root, targetjustifications);
+                break :targetjustifications targetjustifications;
+            };
+
+            for (validator_indices.items) |validator_index| {
+                if (validator_index >= num_validators) {
+                    return StateTransitionError.InvalidValidatorId;
+                }
+                target_justifications[validator_index] = 1;
+            }
+            try justifications.put(allocator, attestation_data.target.root, target_justifications);
+            var target_justifications_count: usize = 0;
+            for (target_justifications) |justified| {
+                if (justified == 1) {
+                    target_justifications_count += 1;
+                }
+            }
+            logger.debug("target jcount={d} target_root=0x{x} justifications_len={d}\n", .{ target_justifications_count, &attestation_data.target.root, target_justifications.len });
+
+            if (3 * target_justifications_count >= 2 * num_validators) {
+                self.state.latest_justified = attestation_data.target;
+                try utils.setSlotJustified(finalized_slot, &self.state.justified_slots, target_slot, true);
+                if (justifications.fetchRemove(attestation_data.target.root)) |kv| {
+                    allocator.free(kv.value);
+                }
+                logger.debug("justified root=0x{x} slot={d}", .{ &self.state.latest_justified.root, self.state.latest_justified.slot });
+
+                var can_target_finalize = true;
+                const start_slot_usize: usize = @intCast(source_slot + 1);
+                const end_slot_usize: usize = @intCast(target_slot);
+                for (start_slot_usize..end_slot_usize) |slot_usize| {
+                    const slot: Slot = @intCast(slot_usize);
+                    if (try utils.IsJustifiableSlot(self.state.latest_finalized.slot, slot)) {
+                        can_target_finalize = false;
+                        break;
+                    }
+                }
+                logger.debug("----------------can_target_finalize ({d})={any}----------\n\n", .{ source_slot, can_target_finalize });
+                if (can_target_finalize == true) {
+                    const old_finalized_slot = finalized_slot;
+                    self.state.latest_finalized = attestation_data.source;
+                    finalized_slot = self.state.latest_finalized.slot;
+
+                    const delta: Slot = finalized_slot - old_finalized_slot;
+                    if (delta > 0) {
+                        try self.state.shiftJustifiedSlots(delta, allocator);
+
+                        var roots_to_remove: std.ArrayList(Root) = .empty;
+                        defer roots_to_remove.deinit(allocator);
+                        var iter = justifications.iterator();
+                        while (iter.next()) |entry| {
+                            const root = entry.key_ptr.*;
+                            const slot = block_cache.get(root) orelse return StateTransitionError.InvalidJustificationRoot;
+                            if (slot <= finalized_slot) {
+                                try roots_to_remove.append(allocator, root);
+                            }
+                        }
+                        for (roots_to_remove.items) |root| {
+                            if (justifications.fetchRemove(root)) |kv| {
+                                allocator.free(kv.value);
+                            }
+                        }
+                    }
+                    const finalized_str_new = try self.state.latest_finalized.toJsonString(allocator);
+                    defer allocator.free(finalized_str_new);
+
+                    logger.debug("finalized={s}", .{finalized_str_new});
+                }
+            }
+        }
+
+        logger.debug("poststate:historical hashes={d} justified slots={d}\n justifications_roots:{d}\n justifications_validators={d}\n", .{ self.state.historical_block_hashes.len(), self.state.justified_slots.len(), self.state.justifications_roots.len(), self.state.justifications_validators.len() });
+        const justified_str_final = try self.state.latest_justified.toJsonString(allocator);
+        defer allocator.free(justified_str_final);
+        const finalized_str_final = try self.state.latest_finalized.toJsonString(allocator);
+        defer allocator.free(finalized_str_final);
+
+        logger.debug("poststate: justified={s} finalized={s}", .{ justified_str_final, finalized_str_final });
+    }


The process_attestations function is duplicated between BeamState (state.zig line 347-546) and CachedState (cached_state.zig line 88-262). This creates a maintenance burden where bug fixes or feature changes need to be applied in two places. Consider refactoring to have a single implementation that both can use, perhaps by extracting the core logic into a shared helper function that takes the justifications map as a parameter.

Copilot

Pull request overview

Copilot reviewed 4 out of 4 changed files in this pull request and generated 4 comments.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Copilot · 2026-03-06T09:55:37Z

pkgs/types/src/state.zig


 const block = @import("./block.zig");
 const attestation = @import("./attestation.zig");
+const cached_state_mod = @import("./cached_state.zig");


state.zig now imports ./cached_state.zig, but cached_state.zig also imports ./state.zig, creating an import cycle and (in non-test builds) leaving an unused top-level cached_state_mod decl. To avoid build failures, remove this import from state.zig and move the updated tests into a separate test file/module that imports both, or refactor cached_state.zig to not import state.zig (e.g. make CachedState generic over the state type).

Suggested change

const cached_state_mod = @import("./cached_state.zig");

Copilot · 2026-03-06T09:55:38Z

pkgs/types/src/state.zig

    pub fn genGenesisBlock(self: *const Self, allocator: Allocator, genesis_block: *block.BeamBlock) !void {
        var state_root: [32]u8 = undefined;


This PR removes pub fn process_block / process_attestations from BeamState. Since BeamState is part of the public @zeam/types API, this is a breaking change for downstream users. Consider keeping thin deprecated wrappers on BeamState (delegating to CachedState) or documenting the migration and bumping the relevant package version accordingly.

Copilot · 2026-03-06T09:55:38Z

pkgs/types/src/cached_state.zig

+        if (self.justifications) |*justifications| {
+            var iterator = justifications.iterator();
+            while (iterator.next()) |entry| {
+                self.allocator.free(entry.value_ptr.*);
+            }
+            justifications.deinit(self.allocator);
+            self.justifications = null;


CachedState.deinit() currently frees the cached justifications without flushing them back to BeamState. If process_attestations/process_block were called and the caller forgets to call flushJustifications(), the underlying BeamState.justifications_* SSZ fields will remain stale/inconsistent. Consider either flushing in deinit() (and making deinit fallible), providing a separate finish()/commit() API that must be called, or adding a debug assertion when dropping dirty cached state without flushing.

Suggested change

if (self.justifications) |*justifications| {

var iterator = justifications.iterator();

while (iterator.next()) |entry| {

self.allocator.free(entry.value_ptr.*);

}

justifications.deinit(self.allocator);

self.justifications = null;

// Ensure any loaded justifications are flushed back to BeamState

// before dropping the cache. This prevents silently losing

// in-memory updates if the caller forgot to call flushJustifications().

if (self.justifications != null) {

self.flushJustifications() catch |err| {

std.debug.panic(

"CachedState.deinit: failed to flush justifications: {s}",

.{@errorName(err)},

);

};

Copilot · 2026-03-06T09:55:38Z

pkgs/types/src/cached_state.zig

+            }
+        }
+
+        logger.debug("poststate:historical hashes={d} justified slots={d}\n justifications_roots:{d}\n justifications_validators={d}\n", .{ self.state.historical_block_hashes.len(), self.state.justified_slots.len(), self.state.justifications_roots.len(), self.state.justifications_validators.len() });


The poststate debug log reads self.state.justifications_roots.len() / justifications_validators.len(), but those SSZ fields are not updated until flushJustifications() runs. While the cache is loaded, this log will report stale lengths. Consider logging the cached justifications map size instead (or flush before logging) so debug output reflects the current in-memory state.

Suggested change

logger.debug("poststate:historical hashes={d} justified slots={d}\n justifications_roots:{d}\n justifications_validators={d}\n", .{ self.state.historical_block_hashes.len(), self.state.justified_slots.len(), self.state.justifications_roots.len(), self.state.justifications_validators.len() });

logger.debug("poststate:historical hashes={d} justified slots={d}\n justifications_roots:{d}\n justifications_validators={d}\n", .{

self.state.historical_block_hashes.len(),

self.state.justified_slots.len(),

justifications.count(),

justifications.count(),

});

guha-rahul added 4 commits February 27, 2026 16:31

make shiftJustifiedSlots public for CachedState access

c00a910

add CachedState wrapper with lazy justifications cache

1689367

use CachedState in apply_raw_block and apply_transition

410f499

make BeamState.process_block private

5d6529b

Copilot AI review requested due to automatic review settings February 27, 2026 11:12

Copilot started reviewing on behalf of guha-rahul February 27, 2026 11:12 View session

Copilot AI reviewed Feb 27, 2026

View reviewed changes

guha-rahul and others added 2 commits March 6, 2026 15:00

Merge branch 'blockblaz:main' into cached-state-refactor

51b726c

move process_block and process_attestations to CachedState

567cedc

Copilot AI review requested due to automatic review settings March 6, 2026 09:48

Copilot started reviewing on behalf of guha-rahul March 6, 2026 09:48 View session

Copilot AI reviewed Mar 6, 2026

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add CachedState wrapper to improve process_slots logger parameter design #610

Add CachedState wrapper to improve process_slots logger parameter design #610
guha-rahul wants to merge 6 commits intoblockblaz:mainfrom
guha-rahul:cached-state-refactor

guha-rahul commented Feb 27, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Copilot AI Feb 27, 2026

Uh oh!

Copilot AI Feb 27, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Copilot AI Mar 6, 2026

Uh oh!

Copilot AI Mar 6, 2026

Uh oh!

Copilot AI Mar 6, 2026

Uh oh!

Copilot AI Mar 6, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

	fn process_block(self: Self, allocator: Allocator, staged_block: BeamBlock, logger: zeam_utils.ModuleLogger, cache: ?utils.RootToSlotCache) !void {
	pub fn process_block(self: Self, allocator: Allocator, staged_block: BeamBlock, logger: zeam_utils.ModuleLogger, cache: ?utils.RootToSlotCache) !void {

		pub fn genGenesisBlock(self: const Self, allocator: Allocator, genesis_block: block.BeamBlock) !void {
		var state_root: [32]u8 = undefined;

-        if (self.justifications) |*justifications| {
-            var iterator = justifications.iterator();
-            while (iterator.next()) |entry| {
-                self.allocator.free(entry.value_ptr.*);
-            }
-            justifications.deinit(self.allocator);
-            self.justifications = null;
+        // Ensure any loaded justifications are flushed back to BeamState
+        // before dropping the cache. This prevents silently losing
+        // in-memory updates if the caller forgot to call flushJustifications().
+        if (self.justifications != null) {
+            self.flushJustifications() catch |err| {
+                std.debug.panic(
+                    "CachedState.deinit: failed to flush justifications: {s}",
+                    .{@errorName(err)},
+                );
+            };

Conversation

guha-rahul commented Feb 27, 2026

Summary

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Copilot AI Feb 27, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Feb 27, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Copilot AI Mar 6, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Mar 6, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Mar 6, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Mar 6, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants