Skip to content

Commit 08289db

Browse files
fangererfangerer
committed
Implement _deque_iterator.__reduce__
1 parent 4b1d0ec commit 08289db

File tree

2 files changed

+76
-16
lines changed

2 files changed

+76
-16
lines changed

‎graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/modules/CollectionsModuleBuiltins.java

Lines changed: 57 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -46,15 +46,21 @@
4646
import com.oracle.graal.python.builtins.CoreFunctions;
4747
import com.oracle.graal.python.builtins.PythonBuiltinClassType;
4848
import com.oracle.graal.python.builtins.PythonBuiltins;
49+
import com.oracle.graal.python.builtins.objects.PNone;
50+
import com.oracle.graal.python.builtins.objects.deque.DequeIterBuiltins.DequeIterNextNode;
4951
import com.oracle.graal.python.builtins.objects.deque.PDeque;
5052
import com.oracle.graal.python.builtins.objects.deque.PDequeIter;
5153
import com.oracle.graal.python.builtins.objects.function.PKeyword;
54+
import com.oracle.graal.python.lib.PyNumberIndexNode;
5255
import com.oracle.graal.python.nodes.BuiltinNames;
5356
import com.oracle.graal.python.nodes.ErrorMessages;
5457
import com.oracle.graal.python.nodes.function.PythonBuiltinBaseNode;
55-
import com.oracle.graal.python.nodes.function.builtins.PythonBinaryBuiltinNode;
58+
import com.oracle.graal.python.nodes.function.builtins.PythonTernaryBuiltinNode;
5659
import com.oracle.graal.python.nodes.function.builtins.PythonVarargsBuiltinNode;
60+
import com.oracle.graal.python.nodes.util.CastToJavaIntExactNode;
5761
import com.oracle.truffle.api.CompilerDirectives;
62+
import com.oracle.truffle.api.dsl.Cached;
63+
import com.oracle.truffle.api.dsl.Cached.Shared;
5864
import com.oracle.truffle.api.dsl.GenerateNodeFactory;
5965
import com.oracle.truffle.api.dsl.NodeFactory;
6066
import com.oracle.truffle.api.dsl.Specialization;
@@ -89,38 +95,74 @@ PDeque doGeneric(Object cls, Object[] args, PKeyword[] kwargs) {
8995
}
9096

9197
// _collections._deque_iterator
92-
@Builtin(name = BuiltinNames.DEQUE_ITER, minNumOfPositionalArgs = 2, constructsClass = PythonBuiltinClassType.PDequeIter)
98+
@Builtin(name = BuiltinNames.DEQUE_ITER, constructsClass = PythonBuiltinClassType.PDequeIter, //
99+
minNumOfPositionalArgs = 2, parameterNames = {"$self", "iterable", "index"})
93100
@GenerateNodeFactory
94-
abstract static class DequeIterNode extends PythonBinaryBuiltinNode {
101+
abstract static class DequeIterNode extends PythonTernaryBuiltinNode {
95102

96-
@Specialization
97-
PDequeIter doDeque(@SuppressWarnings("unused") Object cls, PDeque deque) {
103+
@Specialization(guards = "isNoValue(index)")
104+
PDequeIter doDeque(@SuppressWarnings("unused") Object cls, PDeque deque, @SuppressWarnings("unused") PNone index) {
98105
return factory().createDequeIter(deque);
99106
}
100107

101-
@Specialization(replaces = "doDeque")
102-
PDequeIter doGeneric(@SuppressWarnings("unused") Object cls, Object deque) {
108+
@Specialization
109+
PDequeIter doDequeInt(@SuppressWarnings("unused") Object cls, PDeque deque, int index,
110+
@Shared("getNextNode") @Cached DequeIterNextNode getNextNode) {
111+
PDequeIter dequeIter = factory().createDequeIter(deque);
112+
for (int i = 0; i < index; i++) {
113+
getNextNode.execute(dequeIter);
114+
}
115+
return dequeIter;
116+
}
117+
118+
@Specialization(replaces = {"doDeque", "doDequeInt"})
119+
PDequeIter doGeneric(VirtualFrame frame, @SuppressWarnings("unused") Object cls, Object deque, Object indexObj,
120+
@Cached PyNumberIndexNode toIndexNode,
121+
@Cached CastToJavaIntExactNode castToJavaIntExactNode,
122+
@Shared("getNextNode") @Cached DequeIterNextNode getNextNode) {
103123
if (deque instanceof PDeque) {
104-
return factory().createDequeIter((PDeque) deque);
124+
if (indexObj != PNone.NO_VALUE) {
125+
int index = castToJavaIntExactNode.execute(toIndexNode.execute(frame, indexObj));
126+
return doDequeInt(cls, (PDeque) deque, index, getNextNode);
127+
}
128+
return doDeque(cls, (PDeque) deque, PNone.NO_VALUE);
105129
}
106130
throw raise(PythonBuiltinClassType.TypeError, ErrorMessages.EXPECTED_OBJ_TYPE_S_GOT_P, BuiltinNames.DEQUE, deque);
107131
}
108132
}
109133

110134
// _collections._deque_reverse_iterator
111-
@Builtin(name = BuiltinNames.DEQUE_REV_ITER, minNumOfPositionalArgs = 2, constructsClass = PythonBuiltinClassType.PDequeRevIter)
135+
@Builtin(name = BuiltinNames.DEQUE_REV_ITER, constructsClass = PythonBuiltinClassType.PDequeRevIter, //
136+
minNumOfPositionalArgs = 2, parameterNames = {"$self", "iterable", "index"})
112137
@GenerateNodeFactory
113-
abstract static class DequeRevIterNode extends PythonBinaryBuiltinNode {
138+
abstract static class DequeRevIterNode extends PythonTernaryBuiltinNode {
114139

115-
@Specialization
116-
PDequeIter doDeque(@SuppressWarnings("unused") Object cls, PDeque deque) {
140+
@Specialization(guards = "isNoValue(index)")
141+
PDequeIter doDeque(@SuppressWarnings("unused") Object cls, PDeque deque, @SuppressWarnings("unused") PNone index) {
117142
return factory().createDequeRevIter(deque);
118143
}
119144

120-
@Specialization(replaces = "doDeque")
121-
PDequeIter doGeneric(@SuppressWarnings("unused") Object cls, Object deque) {
145+
@Specialization
146+
PDequeIter doDequeInt(@SuppressWarnings("unused") Object cls, PDeque deque, int index,
147+
@Shared("getNextNode") @Cached DequeIterNextNode getNextNode) {
148+
PDequeIter dequeIter = factory().createDequeRevIter(deque);
149+
for (int i = 0; i < index; i++) {
150+
getNextNode.execute(dequeIter);
151+
}
152+
return dequeIter;
153+
}
154+
155+
@Specialization(replaces = {"doDeque", "doDequeInt"})
156+
PDequeIter doGeneric(VirtualFrame frame, @SuppressWarnings("unused") Object cls, Object deque, Object indexObj,
157+
@Cached PyNumberIndexNode toIndexNode,
158+
@Cached CastToJavaIntExactNode castToJavaIntExactNode,
159+
@Shared("getNextNode") @Cached DequeIterNextNode getNextNode) {
122160
if (deque instanceof PDeque) {
123-
return factory().createDequeRevIter((PDeque) deque);
161+
if (indexObj != PNone.NO_VALUE) {
162+
int index = castToJavaIntExactNode.execute(toIndexNode.execute(frame, indexObj));
163+
return doDequeInt(cls, (PDeque) deque, index, getNextNode);
164+
}
165+
return doDeque(cls, (PDeque) deque, PNone.NO_VALUE);
124166
}
125167
throw raise(PythonBuiltinClassType.TypeError, ErrorMessages.EXPECTED_OBJ_TYPE_S_GOT_P, BuiltinNames.DEQUE, deque);
126168
}

‎graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/deque/DequeIterBuiltins.java

Lines changed: 19 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -51,13 +51,16 @@
5151
import com.oracle.graal.python.builtins.CoreFunctions;
5252
import com.oracle.graal.python.builtins.PythonBuiltinClassType;
5353
import com.oracle.graal.python.builtins.PythonBuiltins;
54+
import com.oracle.graal.python.builtins.objects.tuple.PTuple;
5455
import com.oracle.graal.python.nodes.ErrorMessages;
5556
import com.oracle.graal.python.nodes.PRaiseNode;
5657
import com.oracle.graal.python.nodes.SpecialMethodNames;
5758
import com.oracle.graal.python.nodes.function.PythonBuiltinBaseNode;
5859
import com.oracle.graal.python.nodes.function.builtins.PythonUnaryBuiltinNode;
60+
import com.oracle.graal.python.nodes.object.GetClassNode;
5961
import com.oracle.truffle.api.CompilerDirectives;
6062
import com.oracle.truffle.api.CompilerDirectives.TruffleBoundary;
63+
import com.oracle.truffle.api.dsl.Cached;
6164
import com.oracle.truffle.api.dsl.GenerateNodeFactory;
6265
import com.oracle.truffle.api.dsl.NodeFactory;
6366
import com.oracle.truffle.api.dsl.Specialization;
@@ -85,6 +88,8 @@ static PDequeIter doGeneric(PDequeIter self) {
8588
@Builtin(name = SpecialMethodNames.__NEXT__, minNumOfPositionalArgs = 1)
8689
@GenerateNodeFactory
8790
public abstract static class DequeIterNextNode extends PythonUnaryBuiltinNode {
91+
92+
public abstract Object execute(PDequeIter self);
8893

8994
@Specialization
9095
@TruffleBoundary
@@ -107,7 +112,7 @@ Object doGeneric(PDequeIter self) {
107112
}
108113
}
109114

110-
// deque.__length_hint__()
115+
// _deque_iterator.__length_hint__()
111116
@Builtin(name = SpecialMethodNames.__LENGTH_HINT__, minNumOfPositionalArgs = 1)
112117
@GenerateNodeFactory
113118
public abstract static class DequeIterLengthHintNode extends PythonUnaryBuiltinNode {
@@ -117,4 +122,17 @@ static int doGeneric(PDequeIter self) {
117122
return self.lengthHint();
118123
}
119124
}
125+
126+
// _deque_iterator.__reduce__()
127+
@Builtin(name = SpecialMethodNames.__REDUCE__, minNumOfPositionalArgs = 1)
128+
@GenerateNodeFactory
129+
public abstract static class DequeIterReduceNode extends PythonUnaryBuiltinNode {
130+
131+
@Specialization
132+
PTuple doGeneric(PDequeIter self,
133+
@Cached GetClassNode getClassNode) {
134+
Object clazz = getClassNode.execute(self);
135+
return factory().createTuple(new Object[]{clazz, factory().createTuple(new Object[]{self.deque, self.deque.getSize() - self.lengthHint()})});
136+
}
137+
}
120138
}

0 commit comments

Comments
 (0)