[GR-12855] Implement or for frozenset.

PullRequest: graalpython/317

Author: ppisl

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

@@ -80,6 +80,12 @@ def test_set_or_union():
     assert or_result == {1, 2, 3}
     assert union_result == {1, 2, 3}
 
+    assert frozenset((1,2)) | {1:2}.items() == {1, 2, (1, 2)}
+    assert frozenset((1,2)) | {1:2}.keys() == {1, 2}
+
+def test_set_and():
+    assert frozenset((1,2)) & {1:2}.items() == set()
+    assert frozenset((1,2)) & {1:2}.keys() == {1}
 
 def test_set_union():
     assert {1, 2, 3}.union({1: 'a', 2: 'b', 4: 'd'}) == {1, 2, 3, 4}

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/set/FrozenSetBuiltins.java

@@ -26,6 +26,7 @@
 package com.oracle.graal.python.builtins.objects.set;
 
 import static com.oracle.graal.python.nodes.SpecialMethodNames.__AND__;
+import static com.oracle.graal.python.nodes.SpecialMethodNames.__OR__;
 import static com.oracle.graal.python.nodes.SpecialMethodNames.__CONTAINS__;
 import static com.oracle.graal.python.nodes.SpecialMethodNames.__EQ__;
 import static com.oracle.graal.python.nodes.SpecialMethodNames.__GE__;
@@ -51,6 +52,7 @@
 import com.oracle.graal.python.builtins.objects.common.HashingStorageNodes;
 import com.oracle.graal.python.builtins.objects.common.HashingStorageNodes.PythonEquivalence;
 import com.oracle.graal.python.builtins.objects.common.PHashingCollection;
+import com.oracle.graal.python.builtins.objects.dict.PDictView;
 import com.oracle.graal.python.builtins.objects.set.FrozenSetBuiltinsFactory.BinaryUnionNodeGen;
 import com.oracle.graal.python.builtins.objects.tuple.PTuple;
 import com.oracle.graal.python.nodes.PNodeWithContext;
@@ -73,6 +75,7 @@
 import com.oracle.truffle.api.dsl.GenerateNodeFactory;
 import com.oracle.truffle.api.dsl.NodeFactory;
 import com.oracle.truffle.api.dsl.Specialization;
+import com.oracle.truffle.api.nodes.Node;
 import com.oracle.truffle.api.profiles.ConditionProfile;
 import com.oracle.truffle.api.profiles.ValueProfile;
 
@@ -206,9 +209,101 @@ PBaseSet doPBaseSet(PFrozenSet left, PBaseSet right) {
             return factory().createFrozenSet(intersectedStorage);
         }
 
+        @Specialization
+        PBaseSet doPBaseSet(PSet left, PDictView right,
+                        @Cached("create()") SetNodes.ConstructSetNode constructSetNode) {
+            PSet rightSet = constructSetNode.executeWith(right);
+            HashingStorage intersectedStorage = getIntersectNode().execute(left.getDictStorage(), rightSet.getDictStorage());
+            return factory().createSet(intersectedStorage);
+        }
+
+        @Specialization
+        PBaseSet doPBaseSet(PFrozenSet left, PDictView right,
+                        @Cached("create()") SetNodes.ConstructSetNode constructSetNode) {
+            PSet rightSet = constructSetNode.executeWith(right);
+            HashingStorage intersectedStorage = getIntersectNode().execute(left.getDictStorage(), rightSet.getDictStorage());
+            return factory().createSet(intersectedStorage);
+        }
+
         @Fallback
         Object doAnd(Object self, Object other) {
-            throw raise(PythonErrorType.TypeError, "unsupported operand type(s) for &=: '%p' and '%p'", self, other);
+            throw raise(PythonErrorType.TypeError, "unsupported operand type(s) for &: '%p' and '%p'", self, other);
+        }
+    }
+
+    @Builtin(name = __OR__, fixedNumOfPositionalArgs = 2)
+    @GenerateNodeFactory
+    abstract static class OrNode extends PythonBinaryBuiltinNode {
+        @Node.Child private HashingStorageNodes.UnionNode unionNode;
+        @Node.Child private HashingStorageNodes.SetItemNode setItemNode;
+
+        private HashingStorageNodes.SetItemNode getSetItemNode() {
+            if (setItemNode == null) {
+                CompilerDirectives.transferToInterpreterAndInvalidate();
+                setItemNode = insert(HashingStorageNodes.SetItemNode.create());
+            }
+            return setItemNode;
+        }
+
+        @TruffleBoundary
+        private HashingStorage getStringAsHashingStorage(String str) {
+            HashingStorage storage = EconomicMapStorage.create(str.length(), true);
+            for (int i = 0; i < str.length(); i++) {
+                String key = String.valueOf(str.charAt(i));
+                getSetItemNode().execute(storage, key, PNone.NO_VALUE);
+            }
+            return storage;
+        }
+
+        @Specialization
+        PBaseSet doPBaseSet(PSet left, String right) {
+            return factory().createSet(getUnionNode().execute(left.getDictStorage(), getStringAsHashingStorage(right)));
+        }
+
+        @Specialization
+        PBaseSet doPBaseSet(PFrozenSet left, String right) {
+            return factory().createFrozenSet(getUnionNode().execute(left.getDictStorage(), getStringAsHashingStorage(right)));
+        }
+
+        private HashingStorageNodes.UnionNode getUnionNode() {
+            if (unionNode == null) {
+                CompilerDirectives.transferToInterpreterAndInvalidate();
+                unionNode = insert(HashingStorageNodes.UnionNode.create());
+            }
+            return unionNode;
+        }
+
+        @Specialization
+        PBaseSet doPBaseSet(PSet left, PBaseSet right) {
+            HashingStorage intersectedStorage = getUnionNode().execute(left.getDictStorage(), right.getDictStorage());
+            return factory().createSet(intersectedStorage);
+        }
+
+        @Specialization
+        PBaseSet doPBaseSet(PFrozenSet left, PBaseSet right) {
+            HashingStorage intersectedStorage = getUnionNode().execute(left.getDictStorage(), right.getDictStorage());
+            return factory().createFrozenSet(intersectedStorage);
+        }
+
+        @Specialization
+        PBaseSet doPBaseSet(PSet left, PDictView right,
+                        @Cached("create()") SetNodes.ConstructSetNode constructSetNode) {
+            PSet rightSet = constructSetNode.executeWith(right);
+            HashingStorage intersectedStorage = getUnionNode().execute(left.getDictStorage(), rightSet.getDictStorage());
+            return factory().createSet(intersectedStorage);
+        }
+
+        @Specialization
+        PBaseSet doPBaseSet(PFrozenSet left, PDictView right,
+                        @Cached("create()") SetNodes.ConstructSetNode constructSetNode) {
+            PSet rightSet = constructSetNode.executeWith(right);
+            HashingStorage intersectedStorage = getUnionNode().execute(left.getDictStorage(), rightSet.getDictStorage());
+            return factory().createSet(intersectedStorage);
+        }
+
+        @Fallback
+        Object doOr(Object self, Object other) {
+            throw raise(PythonErrorType.TypeError, "unsupported operand type(s) for |: '%p' and '%p'", self, other);
         }
     }