diff --git a/src/main/java/com/google/devtools/build/lib/packages/StarlarkProvider.java b/src/main/java/com/google/devtools/build/lib/packages/StarlarkProvider.java
index 24b83a04ba1ac9..57d102af4e7bea 100644
--- a/src/main/java/com/google/devtools/build/lib/packages/StarlarkProvider.java
+++ b/src/main/java/com/google/devtools/build/lib/packages/StarlarkProvider.java
@@ -484,7 +484,8 @@ Object optimizeField(int index, Object value) {
         return depset.getSet();
       }
       Class<?> elementClass = depset.getElementClass();
-      if (depsetTypePredictor.compareAndExchange(index, null, elementClass) == elementClass) {
+      Class<?> witness = depsetTypePredictor.compareAndExchange(index, null, elementClass);
+      if (witness == elementClass || witness == null) {
         return depset.getSet();
       }
     }
@@ -502,7 +503,7 @@ Object retrieveOptimizedField(int index, Object value) {
         // This matches empty depsets created in Starlark with `depset()`.
         return Depset.of(Object.class, nestedSet);
       }
-      @SuppressWarnings("unchecked") // can't parametrize Class literal by a non-raw type
+      @SuppressWarnings("unchecked") // can't parameterize Class literal by a non-raw type
       Depset depset = Depset.of((Class<Object>) depsetTypePredictor.get(index), nestedSet);
       return depset;
     }
diff --git a/src/test/java/com/google/devtools/build/lib/packages/StarlarkProviderTest.java b/src/test/java/com/google/devtools/build/lib/packages/StarlarkProviderTest.java
index 89e07f15482b9b..ef0e1d0d24d839 100644
--- a/src/test/java/com/google/devtools/build/lib/packages/StarlarkProviderTest.java
+++ b/src/test/java/com/google/devtools/build/lib/packages/StarlarkProviderTest.java
@@ -449,6 +449,29 @@ public void schemafulProvider_withDepset() throws Exception {
     assertThat(instance6.getValue("field")).isNull();
   }
 
+  @Test
+  public void schemafulProvider_optimizeField() {
+    StarlarkProvider provider =
+        StarlarkProvider.builder(Location.BUILTIN)
+            .setSchema(ImmutableList.of("field"))
+            .buildWithIdentityToken(generator.generate());
+
+    // The first set is unwrapped and the type String is registered in the predictor.
+    Depset depset1 = Depset.of(String.class, NestedSetBuilder.create(STABLE_ORDER, "a", "b", "c"));
+    assertThat(provider.optimizeField(0, depset1)).isSameInstanceAs(depset1.getSet());
+
+    // A set with Integer type does not match and cannot be optimized.
+    Depset depset2 =
+        Depset.of(
+            StarlarkInt.class,
+            NestedSetBuilder.create(STABLE_ORDER, StarlarkInt.of(1), StarlarkInt.of(2)));
+    assertThat(provider.optimizeField(0, depset2)).isSameInstanceAs(depset2);
+
+    // A third, matching Depset is unwrapped.
+    Depset depset3 = Depset.of(String.class, NestedSetBuilder.create(STABLE_ORDER, "d", "e"));
+    assertThat(provider.optimizeField(0, depset3)).isSameInstanceAs(depset3.getSet());
+  }
+
   @Test
   public void schemafulProvider_mutable() throws Exception {
     StarlarkProvider.Key key =