From de2f2407ced08aba944f6949c58bb930b1003906 Mon Sep 17 00:00:00 2001
From: Francesco Guardiani <me@slinkydeveloper.com>
Date: Tue, 4 Jun 2024 19:44:19 +0200
Subject: [PATCH] Remove stream, seems unneeded. (#345)

---
 .../sdk/core/InvocationStateMachine.java       |  9 +--------
 .../dev/restate/sdk/core/UserStateStore.java   | 18 ++++++++----------
 2 files changed, 9 insertions(+), 18 deletions(-)

diff --git a/sdk-core/src/main/java/dev/restate/sdk/core/InvocationStateMachine.java b/sdk-core/src/main/java/dev/restate/sdk/core/InvocationStateMachine.java
index 24979b41..70c86b30 100644
--- a/sdk-core/src/main/java/dev/restate/sdk/core/InvocationStateMachine.java
+++ b/sdk-core/src/main/java/dev/restate/sdk/core/InvocationStateMachine.java
@@ -183,14 +183,7 @@ void onStartMessage(MessageLite msg) {
     this.entriesToReplay = startMessage.getKnownEntries();
 
     // Set up the state cache
-    this.userStateStore =
-        new UserStateStore(
-            startMessage.getPartialState(),
-            startMessage.getStateMapList().stream()
-                .collect(
-                    Collectors.toMap(
-                        Protocol.StartMessage.StateEntry::getKey,
-                        Protocol.StartMessage.StateEntry::getValue)));
+    this.userStateStore = new UserStateStore(startMessage);
 
     // Tracing and logging setup
     this.loggingContextSetter.set(
diff --git a/sdk-core/src/main/java/dev/restate/sdk/core/UserStateStore.java b/sdk-core/src/main/java/dev/restate/sdk/core/UserStateStore.java
index e66136cb..95408633 100644
--- a/sdk-core/src/main/java/dev/restate/sdk/core/UserStateStore.java
+++ b/sdk-core/src/main/java/dev/restate/sdk/core/UserStateStore.java
@@ -9,11 +9,10 @@
 package dev.restate.sdk.core;
 
 import com.google.protobuf.ByteString;
+import dev.restate.generated.service.protocol.Protocol;
 import java.nio.ByteBuffer;
 import java.util.HashMap;
-import java.util.Map;
 import java.util.Set;
-import java.util.stream.Collectors;
 
 final class UserStateStore {
 
@@ -46,14 +45,13 @@ public ByteBuffer getValue() {
   private boolean isPartial;
   private final HashMap<ByteString, State> map;
 
-  UserStateStore(boolean isPartial, Map<ByteString, ByteString> map) {
-    this.isPartial = isPartial;
-    this.map =
-        new HashMap<>(
-            map.entrySet().stream()
-                .collect(
-                    Collectors.toMap(
-                        Map.Entry::getKey, e -> new Value(e.getValue().asReadOnlyByteBuffer()))));
+  UserStateStore(Protocol.StartMessage startMessage) {
+    this.isPartial = startMessage.getPartialState();
+    this.map = new HashMap<>(startMessage.getStateMapCount());
+    for (int i = 0; i < startMessage.getStateMapCount(); i++) {
+      Protocol.StartMessage.StateEntry entry = startMessage.getStateMap(i);
+      this.map.put(entry.getKey(), new Value(entry.getValue().asReadOnlyByteBuffer()));
+    }
   }
 
   public State get(ByteString key) {