Add range tracking to server accumulator

PiperOrigin-RevId: 853836142

Author: schoppmp

MODULE.bazel

@@ -107,5 +107,9 @@ crate.spec(
     package = "sha3",
     version = "0.10.8",
 )
+crate.spec(
+    package = "rangemap",
+    version = "1.7.1",
+)
 crate.from_specs()
 use_repo(crate, crate_index = "crates")

willow/proto/willow/server_accumulator.proto

@@ -24,10 +24,22 @@ option java_outer_classname = "ServerAccumulatorProto";
 
 message ServerAccumulatorState {
   ServerStateProto server_state = 1;
-  VerifierStateProto verifier_state = 2;
   AggregationConfigProto aggregation_config = 3;
+  // We have one verifier state per range of nonces processed by this
+  // accumulator. States gat merged when adjacent ranges are processed.
+  repeated VerifierStateProto verifier_states = 2;
+  // The ranges of nonces processed by this accumulator. In the same order as
+  // the corresponding verifier states.
+  repeated NonceRange processed_nonce_ranges = 4;
 }
 
-message ClientMessageList {
+message ClientMessageRange {
   repeated ClientMessage client_messages = 1;
+  NonceRange nonce_range = 2;
+}
+
+// A range of nonces. Used to group client messages by nonce.
+message NonceRange {
+  bytes start = 1;  // Inclusive.
+  bytes end = 2;    // Exclusive.
 }

willow/src/api/BUILD

@@ -69,13 +69,17 @@ cc_test(
     name = "server_accumulator_test",
     srcs = ["server_accumulator_test.cc"],
     deps = [
+        ":client_cc",
         ":server_accumulator",
         ":server_accumulator_cxx",  # fixdeps: keep
         "@googletest//:gtest_main",
         "@abseil-cpp//absl/status",
         "@abseil-cpp//absl/status:statusor",
+        "//shell_wrapper:status_matchers",
         "//willow/proto/willow:aggregation_config_cc_proto",
         "//willow/proto/willow:server_accumulator_cc_proto",
+        "//willow/src/input_encoding:codec",
+        "//willow/src/testing_utils:shell_testing_decryptor_cc",
     ],
 )
 
@@ -94,6 +98,7 @@ rust_library(
         ":aggregation_config",
         "@protobuf//rust:protobuf",
         "@cxx.rs//:cxx",
+        "@crate_index//:rangemap",
         "//shell_wrapper:status",
         "//willow/proto/willow:aggregation_config_rust_proto",
         "//willow/proto/willow:server_accumulator_rust_proto",

willow/src/api/server_accumulator.cc

@@ -59,11 +59,17 @@ WillowShellServerAccumulator::CreateFromSerializedState(
 }
 
 absl::Status WillowShellServerAccumulator::ProcessClientMessages(
-    willow::ClientMessageList client_messages) {
+    willow::ClientMessageRange client_messages) {
+  auto serialized_client_messages = client_messages.SerializeAsString();
   client_messages.Clear();
+  return ProcessClientMessages(std::move(serialized_client_messages));
+}
+
+absl::Status WillowShellServerAccumulator::ProcessClientMessages(
+    std::string serialized_client_messages) {
   std::unique_ptr<std::string> status_message;
   int status_code = accumulator_->ProcessClientMessages(
-      std::make_unique<std::string>(client_messages.SerializeAsString()),
+      std::make_unique<std::string>(std::move(serialized_client_messages)),
       &status_message);
   if (status_code != 0) {
     return absl::Status(absl::StatusCode(status_code), *status_message);

willow/src/api/server_accumulator.h

@@ -48,7 +48,12 @@ class WillowShellServerAccumulator {
 
   // Processes a list of client messages. If an invalid message is encountered,
   // an error is logged and processing continues.
-  absl::Status ProcessClientMessages(willow::ClientMessageList client_messages);
+  absl::Status ProcessClientMessages(
+      willow::ClientMessageRange client_messages);
+
+  // Processes a list of client messages, given as a serialized
+  // ClientMessageList proto.
+  absl::Status ProcessClientMessages(std::string serialized_client_messages);
 
   // Merges the state of `other` into the current accumulator.
   absl::Status Merge(std::unique_ptr<WillowShellServerAccumulator> other);