use a simple PE-friendly stack implementation

Author: Nicolas Laurent

src/main/java/org/truffleruby/collections/SimpleStack.java

@@ -0,0 +1,54 @@
+/*
+ * Copyright (c) 2020 Oracle and/or its affiliates. All rights reserved. This
+ * code is released under a tri EPL/GPL/LGPL license. You can use it,
+ * redistribute it and/or modify it under the terms of the:
+ *
+ * Eclipse Public License version 2.0, or
+ * GNU General Public License version 2, or
+ * GNU Lesser General Public License version 2.1.
+ */
+package org.truffleruby.collections;
+
+import org.truffleruby.core.array.ArrayUtils;
+
+/** Simplistic array stack implementation that will partial-evaluate nicely, unlike {@link java.util.ArrayDeque}. */
+@SuppressWarnings("unchecked")
+public class SimpleStack<T> {
+
+    Object[] storage;
+    int index = -1;
+
+    public SimpleStack() {
+        this(16);
+    }
+
+    public SimpleStack(int length) {
+        this.storage = new Object[length];
+    }
+
+    public boolean isEmpty() {
+        return index == -1;
+    }
+
+    public void push(T value) {
+        if (++index == storage.length) {
+            storage = ArrayUtils.grow(storage, storage.length * 2);
+        }
+        storage[index] = value;
+    }
+
+    public T peek() {
+        return (T) storage[index];
+    }
+
+    public T pop() {
+        T out = (T) storage[index];
+        storage[index] = null;
+        --index;
+        return out;
+    }
+
+    public int size() {
+        return index + 1;
+    }
+}

src/main/java/org/truffleruby/core/array/ArrayNodes.java

@@ -13,7 +13,6 @@
 import static org.truffleruby.core.array.ArrayHelpers.setStoreAndSize;
 import static org.truffleruby.language.dispatch.DispatchNode.PUBLIC;
 
-import java.util.ArrayDeque;
 import java.util.Arrays;
 
 import com.oracle.truffle.api.dsl.CachedLanguage;
@@ -29,6 +28,7 @@
 import org.truffleruby.builtins.Primitive;
 import org.truffleruby.builtins.PrimitiveArrayArgumentsNode;
 import org.truffleruby.builtins.YieldingCoreMethodNode;
+import org.truffleruby.collections.SimpleStack;
 import org.truffleruby.core.CoreLibrary;
 import org.truffleruby.core.Hashing;
 import org.truffleruby.core.array.ArrayBuilderNode.BuilderState;
@@ -2364,7 +2364,6 @@ protected boolean flattenHelper(RubyArray array, RubyArray out, int maxLevels,
 
             boolean modified = false;
             final EconomicSet<RubyArray> visited = EconomicSet.create(Equivalence.IDENTITY);
-            // final HashSet<RubyArray> visited = new HashSet<>();
             class Entry {
                 final RubyArray array;
                 final int index;
@@ -2374,7 +2373,7 @@ class Entry {
                     this.index = index;
                 }
             }
-            final ArrayDeque<Entry> workStack = new ArrayDeque<>();
+            final SimpleStack<Entry> workStack = new SimpleStack<>();
             workStack.push(new Entry(array, 0));
 
             while (!workStack.isEmpty()) {