[GR-80729] Optimize creating dicts in JSON scanner

graalpython/com.oracle.graal.python.test/src/tests/test_json.py

@@ -60,6 +60,39 @@
 
 
 class JsonTest(unittest.TestCase):
+    def test_load_duplicate_nested_keys(self):
+        text = '{"a": {"old": 0}, "b": [{"x": 1}], "a": {"last": [{"y": 2}]}, "c": {}}'
+        result = json.loads(text)
+        assert list(result.items()) == [("a", {"last": [{"y": 2}]}), ("b", [{"x": 1}]), ("c", {})]
+        result["d"] = 3
+        del result["b"]
+        assert list(result) == ["a", "c", "d"]
+        assert result.popitem() == ("d", 3)
+        assert json.loads(json.dumps(result)) == result
+
+    def test_load_dict_index_sizes(self):
+        for size in (0, 1, 4, 8, 32, 130, 1024, 33000):
+            expected = {str(i): i for i in range(size)}
+            result = json.loads(json.dumps(expected))
+            assert list(result.items()) == list(expected.items())
+            result["extra"] = None
+            for i in range(0, size, 3):
+                assert result.pop(str(i)) == i
+            result.update(expected)
+            assert result == dict(expected, extra=None)
+
+    def test_load_duplicate_colliding_keys(self):
+        # These strings have the same Java-style string hash.
+        keys = [""]
+        for _ in range(8):
+            keys = [key + suffix for key in keys for suffix in ("Aa", "BB")]
+        pairs = [(key, i) for i, key in enumerate(keys)]
+        pairs += [(key, -i) for i, key in enumerate(reversed(keys))]
+        text = "{" + ",".join(json.dumps(key) + ":" + str(value) for key, value in pairs) + "}"
+        result = json.loads(text)
+        assert list(result.items()) == list(dict(pairs).items())
+        assert result == json.loads(text, object_pairs_hook=dict)
+
     def test_callable_object_hook(self):
         called = []
         class Hook:

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/modules/json/JSONScannerBuiltins.java

@@ -459,6 +459,7 @@ private Object scanOnceUnicode(VirtualFrame frame, BoundaryCallData boundaryCall
             ObjectSequenceStorage currentListStorage = null;
             ObjectSequenceStorage currentPairsStorage = null;
             EconomicMapStorage currentDictStorage = null;
+            ArrayBuilder<?> currentDictEntries = null;
             ScannerState state = ScannerState.initial;
             boolean commaSeen = false;
             int commaIndex = -1;
@@ -482,7 +483,7 @@ private Object scanOnceUnicode(VirtualFrame frame, BoundaryCallData boundaryCall
                         TruffleString parentKey = null;
                         final Object dict;
                         /*
-                         * If no hooks are present, the stack contains only the parent dicts or
+                         * If no hooks are present, the stack contains only the parent entry buffers or
                          * lists the current object is nested in. If a objectHook or objectPairsHook
                          * is present, the stack also stores the property keys of nested dicts. For
                          * example, when parsing the innermost dict of '{"a":{"b":{"c":"d"}}}', the
@@ -508,9 +509,10 @@ private Object scanOnceUnicode(VirtualFrame frame, BoundaryCallData boundaryCall
                             assert ((PDict) topOfStack).getDictStorage() == currentDictStorage;
                             dict = callObjectHook.executeObject(objectHook, topOfStack);
                         } else {
-                            assert topOfStack instanceof PDict;
-                            assert ((PDict) topOfStack).getDictStorage() == currentDictStorage;
-                            dict = topOfStack;
+                            // Adopt the entry buffer and build the hash index only once, at }.
+                            assert topOfStack == currentDictEntries;
+                            EconomicMapStorage storage = EconomicMapStorage.fromStringEntries(currentDictEntries.getInternalArray(), currentDictEntries.size() / 2, hashCodeNode, equalNode);
+                            dict = PFactory.createDict(language, storage);
                         }
                         if (stack.isEmpty()) {
                             return dict;
@@ -537,11 +539,13 @@ private Object scanOnceUnicode(VirtualFrame frame, BoundaryCallData boundaryCall
                                 }
                             } else {
                                 Object parent = stack.peek();
-                                if (parent instanceof PDict parentDict) {
-                                    currentDictStorage = (EconomicMapStorage) parentDict.getDictStorage();
+                                if (parent instanceof ArrayBuilder<?> entries) {
+                                    currentDictEntries = entries;
+                                    currentDictEntries.set(currentDictEntries.size() - 1, dict);
                                     nextState = ScannerState.dict;
                                 } else {
                                     currentListStorage = (ObjectSequenceStorage) ((PList) parent).getSequenceStorage();
+                                    currentListStorage.setObjectItemNormalized(currentListStorage.length() - 1, dict);
                                     nextState = ScannerState.list;
                                 }
                             }
@@ -593,9 +597,12 @@ private Object scanOnceUnicode(VirtualFrame frame, BoundaryCallData boundaryCall
                             } else if (hasPairsHook) {
                                 currentPairsStorage = (ObjectSequenceStorage) parent;
                                 nextState = ScannerState.dict;
-                            } else {
+                            } else if (hasObjectHook) {
                                 currentDictStorage = (EconomicMapStorage) ((PDict) parent).getDictStorage();
                                 nextState = ScannerState.dict;
+                            } else {
+                                currentDictEntries = (ArrayBuilder<?>) parent;
+                                nextState = ScannerState.dict;
                             }
                         }
                     }
@@ -612,8 +619,10 @@ private Object scanOnceUnicode(VirtualFrame frame, BoundaryCallData boundaryCall
                         idx++;
                         if (hasPairsHook) {
                             value = new ObjectSequenceStorage(4);
-                        } else {
+                        } else if (hasObjectHook) {
                             value = PFactory.createDict(language, EconomicMapStorage.create());
+                        } else {
+                            value = new ArrayBuilder<>(8);
                         }
                         nextState = ScannerState.dict;
                     } else if (c == '[') {
@@ -639,9 +648,12 @@ private Object scanOnceUnicode(VirtualFrame frame, BoundaryCallData boundaryCall
                         assert propertyKey != null;
                         if (hasPairsHook) {
                             currentPairsStorage.appendItem(PFactory.createTuple(language, new Object[]{propertyKey, value}));
-                        } else {
+                        } else if (hasObjectHook) {
                             HashingStorage newStorage = hashingStorageSetItem.execute(inliningTarget, currentDictStorage, propertyKey, value);
                             assert newStorage == currentDictStorage;
+                        } else {
+                            currentDictEntries.add(propertyKey);
+                            currentDictEntries.add(value);
                         }
                     } else if (state == ScannerState.list) {
                         assert propertyKey == null;
@@ -667,6 +679,9 @@ private Object scanOnceUnicode(VirtualFrame frame, BoundaryCallData boundaryCall
                                  */
                                 stack.add(propertyKey);
                             }
+                        } else if (!hasObjectHook) {
+                            assert nextState == ScannerState.dict;
+                            currentDictEntries = (ArrayBuilder<?>) value;
                         } else {
                             assert nextState == ScannerState.dict;
                             currentDictStorage = (EconomicMapStorage) ((PDict) value).getDictStorage();

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

@@ -83,6 +83,15 @@ private EconomicMapStorage(int initialCapacity) {
         super(initialCapacity);
     }
 
+    /** Takes ownership of the array; callers must not use it after this call. */
+    public static EconomicMapStorage fromStringEntries(Object[] entries, int entryCount, TruffleString.HashCodeNode hashCodeNode, TruffleString.EqualNode equalNode) {
+        return new EconomicMapStorage(entries, entryCount, hashCodeNode, equalNode);
+    }
+
+    private EconomicMapStorage(Object[] entries, int entryCount, TruffleString.HashCodeNode hashCodeNode, TruffleString.EqualNode equalNode) {
+        super(entries, entryCount, hashCodeNode, equalNode);
+    }
+
     private EconomicMapStorage() {
         this(4);
     }

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

@@ -40,10 +40,14 @@
  */
 package com.oracle.graal.python.builtins.objects.common;
 
+import static com.oracle.graal.python.util.PythonUtils.TS_ENCODING;
 import static com.oracle.truffle.api.CompilerDirectives.SLOWPATH_PROBABILITY;
 
+import java.util.Arrays;
+
 import com.oracle.graal.python.builtins.objects.common.ObjectHashMapFactory.PutNodeGen;
 import com.oracle.graal.python.builtins.objects.common.ObjectHashMapFactory.RemoveNodeGen;
+import com.oracle.graal.python.lib.PyObjectHashNode;
 import com.oracle.graal.python.lib.PyObjectRichCompareBool;
 import com.oracle.graal.python.runtime.sequence.storage.ObjectSequenceStorage;
 import com.oracle.graal.python.util.ArrayBuilder;
@@ -65,6 +69,7 @@
 import com.oracle.truffle.api.profiles.InlinedBranchProfile;
 import com.oracle.truffle.api.profiles.InlinedConditionProfile;
 import com.oracle.truffle.api.profiles.InlinedCountingConditionProfile;
+import com.oracle.truffle.api.strings.TruffleString;
 
 /**
  * Generic dictionary/set backing storage implementation.
@@ -227,6 +232,52 @@ public ObjectHashMap() {
         allocateData(INITIAL_INDICES_SIZE);
     }
 
+    /**
+     * Takes ownership of an alternating string key/value array and builds its hash index. Duplicate
+     * keys are compacted in place, retaining the first key and the last value.
+     */
+    protected ObjectHashMap(Object[] entries, int entryCount, TruffleString.HashCodeNode hashCodeNode, TruffleString.EqualNode equalNode) {
+        assert (entries.length & 1) == 0 && entries.length > 0;
+        assert entryCount >= 0 && entryCount <= entries.length / 2;
+        int entryCapacity = entries.length / 2;
+        int bucketsCount = getMinBucketsCount(entryCapacity);
+        ensureArraySizesFit(bucketsCount, entryCapacity);
+        keysAndValues = entries;
+        metadata = createMetadata(bucketsCount, entryCapacity);
+        int indexByteSize = getIndexByteSize(entryCapacity);
+        int indicesOffset = getIndicesOffset(entryCapacity);
+        int collisionMask = getPhysicalCollisionMaskForIndexByteSize(indexByteSize);
+        for (int i = 0; i < entryCount; i++) {
+            TruffleString key = (TruffleString) entries[i * 2];
+            Object value = entries[i * 2 + 1];
+            long hash = PyObjectHashNode.hash(key, hashCodeNode);
+            int bucket = getIndex(bucketsCount, hash);
+            long perturb = hash;
+            while (true) {
+                int index = getIndex(metadata, indicesOffset, indexByteSize, bucket);
+                if (index == EMPTY_INDEX) {
+                    setIndex(metadata, indicesOffset, indexByteSize, collisionMask, bucket, size + INDEX_OFFSET);
+                    entries[size * 2] = key;
+                    entries[size * 2 + 1] = value;
+                    setHash(metadata, size, hash);
+                    size++;
+                    usedHashes++;
+                    usedIndices++;
+                    break;
+                }
+                int existingIndex = unwrapIndex(index);
+                if (getHash(metadata, existingIndex) == hash && equalNode.execute((TruffleString) getKey(existingIndex, entries), key, TS_ENCODING)) {
+                    entries[existingIndex * 2 + 1] = value;
+                    break;
+                }
+                markCollision(metadata, indicesOffset, indexByteSize, collisionMask, bucket);
+                perturb >>>= PERTURB_SHIFT;
+                bucket = nextIndex(bucketsCount, bucket, perturb);
+            }
+        }
+        Arrays.fill(entries, size * 2, entryCount * 2, null);
+    }
+
     protected ObjectHashMap(ObjectHashMap original) {
         size = original.size;
         usedHashes = original.usedHashes;

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/util/ArrayBuilder.java

@@ -83,6 +83,11 @@ public T get(int index) {
         return (T) data[index];
     }
 
+    public void set(int index, Object item) {
+        assert index >= 0 && index < size;
+        data[index] = item;
+    }
+
     @SuppressWarnings("unchecked")
     public T pop() {
         assert size > 0;
@@ -129,6 +134,13 @@ public ObjectSequenceStorage toObjectSequenceStorage() {
         return new ObjectSequenceStorage(data, size);
     }
 
+    /**
+     * Returns the backing array, which may be overallocated.
+     */
+    public Object[] getInternalArray() {
+        return data;
+    }
+
     public int size() {
         return size;
     }