Support iterable unpacking in ListLiteralNode.

fangerer · fangerer · commit 2206833358d1 · 2019-10-10T21:03:22.000+02:00
diff --git a/graalpython/com.oracle.graal.python.test/src/tests/test_list.py b/graalpython/com.oracle.graal.python.test/src/tests/test_list.py
@@ -54,6 +54,18 @@ def test_basic(self):
         l[0] = "hello"
         self.assertEqual(l, ["hello", 1, 2, 3, 4])
 
+    def test_literal(self):
+        self.assertEqual([1,2,3], [*[1,2,3]])
+        self.assertEqual([1,2,3], [*(1,2,3)])
+        self.assertEqual([1,2,3,4,5,6], [*(1,2,3), *(4,5,6)])
+
+        # this will certainly create a list where the capacity of the storage is not exhausted, i.e., 'length < cap',
+        # and so the storage contains null elements
+        l = []
+        for c in "abcdefghijk":
+            l.append(c)
+        self.assertEqual(['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k'], [*l])
+
     def test_truth(self):
         super().test_truth()
         self.assertTrue(not [])
diff --git a/graalpython/com.oracle.graal.python.test/src/tests/test_tuple.py b/graalpython/com.oracle.graal.python.test/src/tests/test_tuple.py
@@ -22,6 +22,9 @@ def test_constructors(self):
         self.assertEqual(tuple(''), ())
         self.assertEqual(tuple('spam'), ('s', 'p', 'a', 'm'))
 
+    def test_literal(self):
+        self.assertEqual((1,2,3), (*[1,2,3],))
+
     def test_truth(self):
         super().test_truth()
         self.assertTrue(not ())
diff --git a/graalpython/com.oracle.graal.python/src/com/oracle/graal/python/nodes/expression/CastToListExpressionNode.java b/graalpython/com.oracle.graal.python/src/com/oracle/graal/python/nodes/expression/CastToListExpressionNode.java
@@ -43,6 +43,7 @@
 import static com.oracle.graal.python.runtime.exception.PythonErrorType.TypeError;
 
 import com.oracle.graal.python.PythonLanguage;
+import com.oracle.graal.python.builtins.objects.common.SequenceNodes.GetSequenceStorageNode;
 import com.oracle.graal.python.builtins.objects.common.SequenceStorageNodes;
 import com.oracle.graal.python.builtins.objects.list.PList;
 import com.oracle.graal.python.builtins.objects.object.PythonObject;
@@ -78,15 +79,15 @@
 
 public abstract class CastToListExpressionNode extends UnaryOpNode {
 
-    public Object[] getArray(VirtualFrame frame) {
+    @Child private GetSequenceStorageNode getSequenceStorageNode;
+
+    public final SequenceStorage getStorage(VirtualFrame frame) {
         Object result = execute(frame);
-        if (result instanceof PTuple) {
-            return ((PTuple) result).getSequenceStorage().getInternalArray();
-        } else if (result instanceof PList) {
-            return ((PList) result).getSequenceStorage().getInternalArray();
-        } else {
-            throw new RuntimeException("Got an unexpected result when casting to a list");
+        if (getSequenceStorageNode == null) {
+            CompilerDirectives.transferToInterpreterAndInvalidate();
+            getSequenceStorageNode = insert(GetSequenceStorageNode.create());
         }
+        return getSequenceStorageNode.execute(result);
     }
 
     @Specialization
diff --git a/graalpython/com.oracle.graal.python/src/com/oracle/graal/python/nodes/literal/ListLiteralNode.java b/graalpython/com.oracle.graal.python/src/com/oracle/graal/python/nodes/literal/ListLiteralNode.java
@@ -27,8 +27,11 @@
 
 import java.lang.reflect.Array;
 
+import com.oracle.graal.python.builtins.objects.common.SequenceStorageNodes;
+import com.oracle.graal.python.builtins.objects.common.SequenceStorageNodes.ListGeneralizationNode;
 import com.oracle.graal.python.builtins.objects.list.PList;
 import com.oracle.graal.python.builtins.objects.tuple.PTuple;
+import com.oracle.graal.python.nodes.PNode;
 import com.oracle.graal.python.nodes.expression.ExpressionNode;
 import com.oracle.graal.python.runtime.object.PythonObjectFactory;
 import com.oracle.graal.python.runtime.sequence.storage.DoubleSequenceStorage;
@@ -40,6 +43,7 @@
 import com.oracle.graal.python.runtime.sequence.storage.SequenceStorage.ListStorageType;
 import com.oracle.graal.python.runtime.sequence.storage.SequenceStorageFactory;
 import com.oracle.graal.python.runtime.sequence.storage.TupleSequenceStorage;
+import com.oracle.truffle.api.CompilerDirectives;
 import com.oracle.truffle.api.CompilerDirectives.CompilationFinal;
 import com.oracle.truffle.api.frame.VirtualFrame;
 import com.oracle.truffle.api.nodes.ExplodeLoop;
@@ -48,20 +52,55 @@
 public final class ListLiteralNode extends LiteralNode {
     @Child private PythonObjectFactory factory = PythonObjectFactory.create();
     @Children protected final ExpressionNode[] values;
+    @Child private SequenceStorageNodes.ConcatNode concatStoragesNode;
+    @Child private SequenceStorageNodes.AppendNode appendNode;
+
+    private final boolean hasStarredExpressions;
 
     @CompilationFinal private ListStorageType type = ListStorageType.Uninitialized;
 
     public ListLiteralNode(ExpressionNode[] values) {
         this.values = values;
+        for (PNode v : values) {
+            if (v instanceof StarredExpressionNode) {
+                hasStarredExpressions = true;
+                return;
+            }
+        }
+        hasStarredExpressions = false;
     }
 
     public ExpressionNode[] getValues() {
         return values;
     }
 
+    @ExplodeLoop
+    private PList expandingList(VirtualFrame frame) {
+        // we will usually have more than 'values.length' elements
+        SequenceStorage storage = new ObjectSequenceStorage(values.length);
+        for (ExpressionNode n : values) {
+            if (n instanceof StarredExpressionNode) {
+                SequenceStorage addElements = ((StarredExpressionNode) n).getStorage(frame);
+                storage = ensureConcatStoragesNode().execute(storage, addElements);
+            } else {
+                Object element = n.execute(frame);
+                storage = ensureAppendNode().execute(storage, element, ListGeneralizationNode.SUPPLIER);
+            }
+        }
+        return factory.createList(storage);
+    }
+
     @Override
+    public PList execute(VirtualFrame frame) {
+        if (!hasStarredExpressions) {
+            return directList(frame);
+        } else {
+            return expandingList(frame);
+        }
+    }
+
     @ExplodeLoop
-    public Object execute(VirtualFrame frame) {
+    private PList directList(VirtualFrame frame) {
         SequenceStorage storage;
         if (type == ListStorageType.Uninitialized) {
             try {
@@ -169,6 +208,22 @@ private SequenceStorage genericFallback(VirtualFrame frame, Object array, int co
         return new ObjectSequenceStorage(elements);
     }
 
+    private SequenceStorageNodes.ConcatNode ensureConcatStoragesNode() {
+        if (concatStoragesNode == null) {
+            CompilerDirectives.transferToInterpreterAndInvalidate();
+            concatStoragesNode = insert(SequenceStorageNodes.ConcatNode.create(ListGeneralizationNode::create));
+        }
+        return concatStoragesNode;
+    }
+
+    private SequenceStorageNodes.AppendNode ensureAppendNode() {
+        if (appendNode == null) {
+            CompilerDirectives.transferToInterpreterAndInvalidate();
+            appendNode = insert(SequenceStorageNodes.AppendNode.create());
+        }
+        return appendNode;
+    }
+
     public static ListLiteralNode create(ExpressionNode[] values) {
         return new ListLiteralNode(values);
     }
diff --git a/graalpython/com.oracle.graal.python/src/com/oracle/graal/python/nodes/literal/StarredExpressionNode.java b/graalpython/com.oracle.graal.python/src/com/oracle/graal/python/nodes/literal/StarredExpressionNode.java
@@ -43,6 +43,7 @@
 import com.oracle.graal.python.nodes.expression.CastToListExpressionNode;
 import com.oracle.graal.python.nodes.expression.CastToListExpressionNodeGen;
 import com.oracle.graal.python.nodes.expression.ExpressionNode;
+import com.oracle.graal.python.runtime.sequence.storage.SequenceStorage;
 import com.oracle.truffle.api.frame.VirtualFrame;
 
 public final class StarredExpressionNode extends LiteralNode {
@@ -60,8 +61,8 @@ public ExpressionNode getValue() {
         return childNode.getOperand();
     }
 
-    public Object[] getArray(VirtualFrame frame) {
-        return childNode.getArray(frame);
+    public SequenceStorage getStorage(VirtualFrame frame) {
+        return childNode.getStorage(frame);
     }
 
     @Override
diff --git a/graalpython/com.oracle.graal.python/src/com/oracle/graal/python/nodes/literal/TupleLiteralNode.java b/graalpython/com.oracle.graal.python/src/com/oracle/graal/python/nodes/literal/TupleLiteralNode.java
@@ -25,21 +25,23 @@
  */
 package com.oracle.graal.python.nodes.literal;
 
-import java.util.ArrayList;
-import java.util.Arrays;
-import java.util.List;
-
+import com.oracle.graal.python.builtins.objects.common.SequenceStorageNodes;
+import com.oracle.graal.python.builtins.objects.common.SequenceStorageNodes.NoGeneralizationNode;
 import com.oracle.graal.python.nodes.PNode;
 import com.oracle.graal.python.nodes.expression.ExpressionNode;
 import com.oracle.graal.python.runtime.object.PythonObjectFactory;
-import com.oracle.truffle.api.CompilerDirectives.TruffleBoundary;
+import com.oracle.graal.python.runtime.sequence.storage.ObjectSequenceStorage;
+import com.oracle.graal.python.runtime.sequence.storage.SequenceStorage;
+import com.oracle.truffle.api.CompilerDirectives;
 import com.oracle.truffle.api.frame.VirtualFrame;
 import com.oracle.truffle.api.nodes.ExplodeLoop;
 
 public final class TupleLiteralNode extends LiteralNode {
     @Child private PythonObjectFactory factory = PythonObjectFactory.create();
     @Children private final ExpressionNode[] values;
-    protected final boolean hasStarredExpressions;
+    @Child private SequenceStorageNodes.ConcatNode concatStoragesNode;
+    @Child private SequenceStorageNodes.AppendNode appendNode;
+    private final boolean hasStarredExpressions;
 
     public ExpressionNode[] getValues() {
         return values;
@@ -67,42 +69,18 @@ public Object execute(VirtualFrame frame) {
 
     @ExplodeLoop
     private Object expandingTuple(VirtualFrame frame) {
-        List<Object> elements = makeList();
+        // we will usually have more than 'values.length' elements
+        SequenceStorage storage = new ObjectSequenceStorage(values.length);
         for (ExpressionNode n : values) {
             if (n instanceof StarredExpressionNode) {
-                Object[] array = ((StarredExpressionNode) n).getArray(frame);
-                addAllElements(elements, array);
+                SequenceStorage addElements = ((StarredExpressionNode) n).getStorage(frame);
+                storage = ensureConcatStoragesNode().execute(storage, addElements);
             } else {
                 Object element = n.execute(frame);
-                addElement(elements, element);
+                storage = ensureAppendNode().execute(storage, element, NoGeneralizationNode.DEFAULT);
             }
         }
-        return factory.createTuple(listToArray(elements));
-    }
-
-    @TruffleBoundary
-    private static Object[] listToArray(List<Object> elements) {
-        return elements.toArray();
-    }
-
-    @TruffleBoundary
-    private static void addElement(List<Object> elements, Object element) {
-        elements.add(element);
-    }
-
-    @TruffleBoundary
-    private static void addAllElements(List<Object> elements, Object[] array) {
-        elements.addAll(asList(array));
-    }
-
-    @TruffleBoundary
-    private static List<Object> asList(Object[] results) {
-        return Arrays.asList(results);
-    }
-
-    @TruffleBoundary
-    private ArrayList<Object> makeList() {
-        return new ArrayList<>(values.length);
+        return factory.createTuple(storage);
     }
 
     @ExplodeLoop
@@ -113,4 +91,20 @@ private Object directTuple(VirtualFrame frame) {
         }
         return factory.createTuple(elements);
     }
+
+    private SequenceStorageNodes.ConcatNode ensureConcatStoragesNode() {
+        if (concatStoragesNode == null) {
+            CompilerDirectives.transferToInterpreterAndInvalidate();
+            concatStoragesNode = insert(SequenceStorageNodes.ConcatNode.create());
+        }
+        return concatStoragesNode;
+    }
+
+    private SequenceStorageNodes.AppendNode ensureAppendNode() {
+        if (appendNode == null) {
+            CompilerDirectives.transferToInterpreterAndInvalidate();
+            appendNode = insert(SequenceStorageNodes.AppendNode.create());
+        }
+        return appendNode;
+    }
 }
