diff --git a/openfeature-provider-local/src/main/java/com/spotify/confidence/SwapWasmResolverApi.java b/openfeature-provider-local/src/main/java/com/spotify/confidence/SwapWasmResolverApi.java
index a8afd538..f3807ff0 100644
--- a/openfeature-provider-local/src/main/java/com/spotify/confidence/SwapWasmResolverApi.java
+++ b/openfeature-provider-local/src/main/java/com/spotify/confidence/SwapWasmResolverApi.java
@@ -167,23 +167,21 @@ private ResolveWithStickyRequest handleMissingMaterializations(
                 } catch (InterruptedException | ExecutionException e) {
                   throw new RuntimeException(e);
                 }
-                materializationPerUnitMap.put(
+                materializationPerUnitMap.computeIfAbsent(
                     unit,
-                    MaterializationMap.newBuilder()
-                        .putAllInfoMap(
-                            loadedAssignments.entrySet().stream()
-                                .collect(
-                                    Collectors.toMap(
-                                        Map.Entry::getKey,
-                                        e -> {
-                                          final var info = e.getValue();
-                                          return com.spotify.confidence.flags.resolver.v1
-                                              .MaterializationInfo.newBuilder()
-                                              .setUnitInInfo(info.isUnitInMaterialization())
-                                              .putAllRuleToVariant(info.ruleToVariant())
-                                              .build();
-                                        })))
-                        .build());
+                    k -> MaterializationMap.newBuilder().putAllInfoMap(new HashMap<>()).build());
+                materializationPerUnitMap.computeIfPresent(
+                    unit,
+                    (k, v) -> {
+                      final Map<
+                              String, com.spotify.confidence.flags.resolver.v1.MaterializationInfo>
+                          map = new HashMap<>();
+                      loadedAssignments.forEach(
+                          (s, materializationInfo) -> {
+                            map.put(s, materializationInfo.toProto());
+                          });
+                      return v.toBuilder().putAllInfoMap(map).build();
+                    });
               });
         });