[GR-65227] Ensure writing same value for all keys in a HashingStorage handles storage generalization

Author: timfel

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/common/HashingCollectionNodes.java

@@ -62,6 +62,7 @@
 import com.oracle.graal.python.nodes.PGuards;
 import com.oracle.graal.python.nodes.PNodeWithContext;
 import com.oracle.graal.python.nodes.util.CastToTruffleStringNode;
+import com.oracle.truffle.api.CompilerDirectives;
 import com.oracle.truffle.api.HostCompilerDirectives.InliningCutoff;
 import com.oracle.truffle.api.dsl.Bind;
 import com.oracle.truffle.api.dsl.Cached;
@@ -122,8 +123,8 @@ abstract static class SetValueHashingStorageNode extends PNodeWithContext {
 
         @Specialization
         static HashingStorage doEconomicStorage(VirtualFrame frame, Node inliningTarget, EconomicMapStorage map, Object value,
-                        @Cached PutUnsafeNode putNode,
-                        @Cached InlinedLoopConditionProfile loopProfile) {
+                        @Shared("putNode") @Cached PutUnsafeNode putNode,
+                        @Shared("loopProfile") @Cached InlinedLoopConditionProfile loopProfile) {
             // We want to avoid calling __hash__() during map.put
             map.setValueForAllKeys(frame, inliningTarget, value, putNode, loopProfile);
             return map;
@@ -135,14 +136,25 @@ static HashingStorage doGeneric(VirtualFrame frame, Node inliningTarget, Hashing
                         @Cached HashingStorageSetItem setItem,
                         @Cached HashingStorageGetIterator getIterator,
                         @Cached HashingStorageIteratorNext itNext,
-                        @Cached HashingStorageIteratorKey itKey) {
+                        @Cached HashingStorageIteratorKey itKey,
+                        @Shared("putNode") @Cached PutUnsafeNode putNode,
+                        @Shared("loopProfile") @Cached InlinedLoopConditionProfile loopProfile) {
             HashingStorageIterator it = getIterator.execute(inliningTarget, map);
-            HashingStorage storage = map;
             while (itNext.execute(inliningTarget, map, it)) {
-                Object key = itKey.execute(inliningTarget, storage, it);
-                storage = setItem.execute(frame, inliningTarget, storage, key, value);
+                Object key = itKey.execute(inliningTarget, map, it);
+                HashingStorage newStorage = setItem.execute(frame, inliningTarget, map, key, value);
+                if (newStorage != map) {
+                    // when the storage changes, the iterator state is not a reliable cursor
+                    // anymore and we need to restart.
+                    if (newStorage instanceof EconomicMapStorage mapStorage) {
+                        mapStorage.setValueForAllKeys(frame, inliningTarget, value, putNode, loopProfile);
+                        return mapStorage;
+                    } else {
+                        throw CompilerDirectives.shouldNotReachHere("We only generalize to EconomicMapStorage");
+                    }
+                }
             }
-            return storage;
+            return map;
         }
 
         protected static boolean isEconomicMapStorage(Object o) {