Support iterable unpacking in SetLiteralNode.

Author: fangerer

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

@@ -293,3 +293,13 @@ def test_set_delete():
     assert s == {'a', 'b', 'c'}
     s.discard('c')
     assert s == {'a', 'b'}
+
+
+def test_literal():
+    d = {"a": 1, "b": 2, "c": 3}
+    e = {"uff": "foo"}
+    assert {*d, *e} == {"a", "b", "c", "uff"}
+
+    d = {}
+    assert {*d} == set()
+

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/nodes/literal/SetLiteralNode.java

@@ -28,35 +28,100 @@
 import com.oracle.graal.python.builtins.objects.PNone;
 import com.oracle.graal.python.builtins.objects.common.HashingStorage;
 import com.oracle.graal.python.builtins.objects.common.HashingStorageNodes;
+import com.oracle.graal.python.builtins.objects.common.SequenceStorageNodes;
 import com.oracle.graal.python.builtins.objects.dict.PDict;
+import com.oracle.graal.python.builtins.objects.set.PSet;
+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.SequenceStorage;
 import com.oracle.truffle.api.CompilerDirectives;
 import com.oracle.truffle.api.frame.VirtualFrame;
 import com.oracle.truffle.api.nodes.ExplodeLoop;
 
 public final class SetLiteralNode extends LiteralNode {
     @Child private PythonObjectFactory factory = PythonObjectFactory.create();
     @Children private final ExpressionNode[] values;
-    @Child HashingStorageNodes.SetItemNode setItemNode;
+    @Child private HashingStorageNodes.SetItemNode setItemNode;
+    @Child private SequenceStorageNodes.LenNode lenNode;
+    @Child private SequenceStorageNodes.GetItemNode getItemNode;
+
+    private final boolean hasStarredExpressions;
 
     public SetLiteralNode(ExpressionNode[] values) {
         this.values = values;
+        for (PNode v : values) {
+            if (v instanceof StarredExpressionNode) {
+                hasStarredExpressions = true;
+                return;
+            }
+        }
+        hasStarredExpressions = false;
     }
 
     @Override
+    public PSet execute(VirtualFrame frame) {
+        if (!hasStarredExpressions) {
+            return directSet(frame);
+        } else {
+            return expandingSet(frame);
+        }
+    }
+
     @ExplodeLoop
-    public Object execute(VirtualFrame frame) {
+    private PSet expandingSet(VirtualFrame frame) {
+        // we will usually have more than 'values.length' elements
         HashingStorage storage = PDict.createNewStorage(true, values.length);
+        for (ExpressionNode n : values) {
+            if (n instanceof StarredExpressionNode) {
+                storage = addAllElement(frame, storage, ((StarredExpressionNode) n).getStorage(frame));
+            } else {
+                Object element = n.execute(frame);
+                storage = ensureSetItemNode().execute(frame, storage, element, PNone.NO_VALUE);
+            }
+        }
+        return factory.createSet(storage);
+    }
 
-        if (setItemNode == null && values.length > 0) {
-            CompilerDirectives.transferToInterpreterAndInvalidate();
-            setItemNode = insert(HashingStorageNodes.SetItemNode.create());
+    private HashingStorage addAllElement(VirtualFrame frame, HashingStorage setStorage, SequenceStorage sequenceStorage) {
+        int n = ensureLenNode().execute(sequenceStorage);
+        for (int i = 0; i < n; i++) {
+            Object element = ensureGetItemNode().execute(sequenceStorage, i);
+            setStorage = ensureSetItemNode().execute(frame, setStorage, element, PNone.NO_VALUE);
         }
+        return setStorage;
+    }
+
+    @ExplodeLoop
+    private PSet directSet(VirtualFrame frame) {
+        HashingStorage storage = PDict.createNewStorage(true, values.length);
         for (ExpressionNode v : this.values) {
-            storage = setItemNode.execute(frame, storage, v.execute(frame), PNone.NO_VALUE);
+            storage = ensureSetItemNode().execute(frame, storage, v.execute(frame), PNone.NO_VALUE);
         }
-
         return factory.createSet(storage);
     }
+
+    private SequenceStorageNodes.LenNode ensureLenNode() {
+        if (lenNode == null) {
+            CompilerDirectives.transferToInterpreterAndInvalidate();
+            lenNode = insert(SequenceStorageNodes.LenNode.create());
+        }
+        return lenNode;
+    }
+
+    private SequenceStorageNodes.GetItemNode ensureGetItemNode() {
+        if (getItemNode == null) {
+            CompilerDirectives.transferToInterpreterAndInvalidate();
+            getItemNode = insert(SequenceStorageNodes.GetItemNode.create());
+        }
+        return getItemNode;
+    }
+
+    private HashingStorageNodes.SetItemNode ensureSetItemNode() {
+        if (setItemNode == null) {
+            CompilerDirectives.transferToInterpreterAndInvalidate();
+            setItemNode = insert(HashingStorageNodes.SetItemNode.create());
+        }
+        return setItemNode;
+    }
 }

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/nodes/literal/TupleLiteralNode.java

@@ -27,6 +27,7 @@
 
 import com.oracle.graal.python.builtins.objects.common.SequenceStorageNodes;
 import com.oracle.graal.python.builtins.objects.common.SequenceStorageNodes.NoGeneralizationNode;
+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;
@@ -68,7 +69,7 @@ public Object execute(VirtualFrame frame) {
     }
 
     @ExplodeLoop
-    private Object expandingTuple(VirtualFrame frame) {
+    private PTuple expandingTuple(VirtualFrame frame) {
         // we will usually have more than 'values.length' elements
         SequenceStorage storage = new ObjectSequenceStorage(values.length);
         for (ExpressionNode n : values) {
@@ -84,7 +85,7 @@ private Object expandingTuple(VirtualFrame frame) {
     }
 
     @ExplodeLoop
-    private Object directTuple(VirtualFrame frame) {
+    private PTuple directTuple(VirtualFrame frame) {
         final Object[] elements = new Object[values.length];
         for (int i = 0; i < values.length; i++) {
             elements[i] = values[i].execute(frame);