[GR-79939] Improve performance of pyperformance pickle benchmarks

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/modules/pickle/MemoTable.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2024, 2025, Oracle and/or its affiliates. All rights reserved.
+ * Copyright (c) 2024, 2026, Oracle and/or its affiliates. All rights reserved.
  * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
  *
  * The Universal Permissive License (UPL), Version 1.0
@@ -42,6 +42,8 @@
 
 import static com.oracle.graal.python.runtime.exception.PythonErrorType.PicklingError;
 
+import java.util.Arrays;
+
 import com.oracle.graal.python.nodes.ErrorMessages;
 import com.oracle.graal.python.nodes.PRaiseNode;
 import com.oracle.graal.python.util.PythonUtils;
@@ -50,50 +52,22 @@
 
 public final class MemoTable {
 
-    private static final int INITIAL_CAPACITY = 8;
+    static final int INITIAL_CAPACITY = 8;
     private static final int OCCUPANCY_EXPONENT = 1; // 2^X relation between capacity and size
     private static final int CAPACITY_INC_EXPONENT = 2; // 2^X increase in capacity when resizing
 
-    public static final class MemoIterator {
-
-        private int index;
-        private final Object[] keys;
-        private final int[] values;
-
-        public MemoIterator(MemoTable table) {
-            this.keys = table.keys;
-            this.values = table.values;
-            this.index = -1;
-        }
-
-        public boolean advance() {
-            while (true) {
-                index++;
-                if (index >= keys.length) {
-                    return false;
-                }
-                if (keys[index] != null) {
-                    return true;
-                }
-            }
-        }
-
-        public Object key() {
-            return keys[index];
-        }
-
-        public int value() {
-            return values[index];
-        }
-    }
-
     private Object[] keys;
     private int[] values;
     private int mask;
     private int size;
 
     public MemoTable() {
-        initArrays(INITIAL_CAPACITY);
+        this(INITIAL_CAPACITY);
+    }
+
+    public MemoTable(int initialCapacity) {
+        assert initialCapacity >= INITIAL_CAPACITY && Integer.bitCount(initialCapacity) == 1;
+        initArrays(initialCapacity);
     }
 
     private MemoTable(MemoTable map) {
@@ -112,7 +86,8 @@ public int size() {
     }
 
     public void clear() {
-        initArrays(INITIAL_CAPACITY);
+        Arrays.fill(keys, null);
+        Arrays.fill(values, 0);
         this.size = 0;
     }
 
@@ -176,11 +151,20 @@ private void resize() {
             throw PRaiseNode.raiseStatic(null, PicklingError, ErrorMessages.STRUCT_SIZE_TOO_LONG);
         }
 
-        MemoIterator iterator = iterator(); // captures the current contents
+        Object[] oldKeys = keys;
+        int[] oldValues = values;
         initArrays(newLength);
 
-        while (iterator.advance()) {
-            setInternal(iterator.key(), iterator.value());
+        for (int i = 0; i < oldKeys.length; i++) {
+            Object key = oldKeys[i];
+            if (key != null) {
+                int index = getIndex(key);
+                while (keys[index] != null) {
+                    index = (index + 1) & mask;
+                }
+                keys[index] = key;
+                values[index] = oldValues[i];
+            }
         }
     }
 
@@ -192,7 +176,15 @@ public void set(Object key, int value) {
         }
     }
 
-    public MemoIterator iterator() {
-        return new MemoIterator(this);
+    int capacity() {
+        return keys.length;
+    }
+
+    Object getKeyAt(int index) {
+        return keys[index];
+    }
+
+    int getValueAt(int index) {
+        return values[index];
     }
 }