Skip to content

Commit e5fa2f0

Browse files
committed
intrinsified itertools.zip_logest()
1 parent 769a513 commit e5fa2f0

File tree

7 files changed

+342
-63
lines changed

7 files changed

+342
-63
lines changed

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/Python3Core.java

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -212,6 +212,7 @@
212212
import com.oracle.graal.python.builtins.objects.itertools.TakewhileBuiltins;
213213
import com.oracle.graal.python.builtins.objects.itertools.TeeBuiltins;
214214
import com.oracle.graal.python.builtins.objects.itertools.TeeDataObjectBuiltins;
215+
import com.oracle.graal.python.builtins.objects.itertools.ZipLongestBuiltins;
215216
import com.oracle.graal.python.builtins.objects.keywrapper.KeyWrapperBuiltins;
216217
import com.oracle.graal.python.builtins.objects.list.ListBuiltins;
217218
import com.oracle.graal.python.builtins.objects.map.MapBuiltins;
@@ -529,6 +530,7 @@ private static PythonBuiltins[] initializeBuiltins(boolean nativeAccessAllowed)
529530
new TakewhileBuiltins(),
530531
new TeeBuiltins(),
531532
new TeeDataObjectBuiltins(),
533+
new ZipLongestBuiltins(),
532534

533535
// zlib
534536
new ZLibModuleBuiltins(),

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/PythonBuiltinClassType.java

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -227,6 +227,7 @@ public enum PythonBuiltinClassType implements TruffleObject {
227227
PIslice("islice", "itertools"),
228228
PStarmap("starmap", "itertools"),
229229
PTakewhile("takewhile", "itertools"),
230+
PZipLongest("zip_longest", "itertools"),
230231

231232
// json
232233
JSONScanner("Scanner", "_json", Flags.PUBLIC_BASE_WODICT),

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/modules/ItertoolsModuleBuiltins.java

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,7 @@
6060
import com.oracle.graal.python.builtins.objects.itertools.PStarmap;
6161
import com.oracle.graal.python.builtins.objects.itertools.PTakewhile;
6262
import com.oracle.graal.python.builtins.objects.itertools.PTeeDataObject;
63+
import com.oracle.graal.python.builtins.objects.itertools.PZipLongest;
6364
import com.oracle.graal.python.builtins.objects.type.TypeNodes.IsTypeNode;
6465
import com.oracle.graal.python.lib.PyCallableCheckNode;
6566
import com.oracle.graal.python.lib.PyNumberAsSizeNode;
@@ -993,4 +994,48 @@ protected Object construct(Object cls, Object iterable, Object[] args,
993994
}
994995
}
995996

997+
@Builtin(name = "zip_longest", minNumOfPositionalArgs = 1, takesVarArgs = true, constructsClass = PythonBuiltinClassType.PZipLongest, keywordOnlyNames = {
998+
"fillvalue"}, doc = "zip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> zip_longest object\n\n" +
999+
"Return a zip_longest object whose .next() method returns a tuple where\n" +
1000+
"the i-th element comes from the i-th iterable argument. The .next()\n" +
1001+
"method continues until the longest iterable in the argument sequence\n" +
1002+
"is exhausted and then it raises StopIteration. When the shorter iterables\n" +
1003+
"are exhausted, the fillvalue is substituted in their place. The fillvalue\n" +
1004+
"defaults to None or can be specified by a keyword argument.")
1005+
@GenerateNodeFactory
1006+
public abstract static class ZipLongestNode extends PythonBuiltinNode {
1007+
@Specialization(guards = "isTypeNode.execute(cls)")
1008+
Object constructNoFillValue(VirtualFrame frame, Object cls, Object[] args, PNone fillValue,
1009+
@Cached PyObjectGetIter getIterNode,
1010+
@Cached LoopConditionProfile loopProfile,
1011+
@SuppressWarnings("unused") @Cached IsTypeNode isTypeNode) {
1012+
return construct(frame, cls, args, null, getIterNode, loopProfile, isTypeNode);
1013+
}
1014+
1015+
@Specialization(guards = {"isTypeNode.execute(cls)", "!isNoValue(fillValue)"})
1016+
Object construct(VirtualFrame frame, Object cls, Object[] args, Object fillValue,
1017+
@Cached PyObjectGetIter getIterNode,
1018+
@Cached LoopConditionProfile loopProfile,
1019+
@SuppressWarnings("unused") @Cached IsTypeNode isTypeNode) {
1020+
PZipLongest self = factory().createZipLongest(cls);
1021+
self.setFillValue(fillValue);
1022+
self.setNumActive(args.length);
1023+
1024+
Object[] itTuple = new Object[args.length];
1025+
loopProfile.profileCounted(itTuple.length);
1026+
for (int i = 0; loopProfile.inject(i < itTuple.length); i++) {
1027+
itTuple[i] = getIterNode.execute(frame, args[i]);
1028+
}
1029+
self.setItTuple(itTuple);
1030+
return self;
1031+
}
1032+
1033+
@Specialization(guards = "!isTypeNode.execute(cls)")
1034+
@SuppressWarnings("unused")
1035+
protected Object construct(Object cls, Object iterables, Object repeat,
1036+
@Cached IsTypeNode isTypeNode) {
1037+
throw raise(TypeError, ErrorMessages.IS_NOT_TYPE_OBJ, "'cls'", cls);
1038+
}
1039+
}
1040+
9961041
}
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,80 @@
1+
/*
2+
* Copyright (c) 2021, Oracle and/or its affiliates. All rights reserved.
3+
* DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
4+
*
5+
* The Universal Permissive License (UPL), Version 1.0
6+
*
7+
* Subject to the condition set forth below, permission is hereby granted to any
8+
* person obtaining a copy of this software, associated documentation and/or
9+
* data (collectively the "Software"), free of charge and under any and all
10+
* copyright rights in the Software, and any and all patent rights owned or
11+
* freely licensable by each licensor hereunder covering either (i) the
12+
* unmodified Software as contributed to or provided by such licensor, or (ii)
13+
* the Larger Works (as defined below), to deal in both
14+
*
15+
* (a) the Software, and
16+
*
17+
* (b) any piece of software and/or hardware listed in the lrgrwrks.txt file if
18+
* one is included with the Software each a "Larger Work" to which the Software
19+
* is contributed by such licensors),
20+
*
21+
* without restriction, including without limitation the rights to copy, create
22+
* derivative works of, display, perform, and distribute the Software and make,
23+
* use, sell, offer for sale, import, export, have made, and have sold the
24+
* Software and the Larger Work(s), and to sublicense the foregoing rights on
25+
* either these or other terms.
26+
*
27+
* This license is subject to the following condition:
28+
*
29+
* The above copyright notice and either this complete permission notice or at a
30+
* minimum a reference to the UPL must be included in all copies or substantial
31+
* portions of the Software.
32+
*
33+
* THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
34+
* IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
35+
* FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
36+
* AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
37+
* LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
38+
* OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
39+
* SOFTWARE.
40+
*/
41+
package com.oracle.graal.python.builtins.objects.itertools;
42+
43+
import com.oracle.graal.python.builtins.objects.object.PythonBuiltinObject;
44+
import com.oracle.truffle.api.object.Shape;
45+
46+
public final class PZipLongest extends PythonBuiltinObject {
47+
48+
private Object fillValue;
49+
private int numActive;
50+
private Object[] itTuple;
51+
52+
public PZipLongest(Object cls, Shape instanceShape) {
53+
super(cls, instanceShape);
54+
}
55+
56+
public Object getFillValue() {
57+
return fillValue;
58+
}
59+
60+
public void setFillValue(Object fillValue) {
61+
this.fillValue = fillValue;
62+
}
63+
64+
public int getNumActive() {
65+
return numActive;
66+
}
67+
68+
public void setNumActive(int numActive) {
69+
this.numActive = numActive;
70+
}
71+
72+
public Object[] getItTuple() {
73+
return itTuple;
74+
}
75+
76+
public void setItTuple(Object[] itTuple) {
77+
this.itTuple = itTuple;
78+
}
79+
80+
}
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,209 @@
1+
/*
2+
* Copyright (c) 2021, Oracle and/or its affiliates. All rights reserved.
3+
* DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
4+
*
5+
* The Universal Permissive License (UPL), Version 1.0
6+
*
7+
* Subject to the condition set forth below, permission is hereby granted to any
8+
* person obtaining a copy of this software, associated documentation and/or
9+
* data (collectively the "Software"), free of charge and under any and all
10+
* copyright rights in the Software, and any and all patent rights owned or
11+
* freely licensable by each licensor hereunder covering either (i) the
12+
* unmodified Software as contributed to or provided by such licensor, or (ii)
13+
* the Larger Works (as defined below), to deal in both
14+
*
15+
* (a) the Software, and
16+
*
17+
* (b) any piece of software and/or hardware listed in the lrgrwrks.txt file if
18+
* one is included with the Software each a "Larger Work" to which the Software
19+
* is contributed by such licensors),
20+
*
21+
* without restriction, including without limitation the rights to copy, create
22+
* derivative works of, display, perform, and distribute the Software and make,
23+
* use, sell, offer for sale, import, export, have made, and have sold the
24+
* Software and the Larger Work(s), and to sublicense the foregoing rights on
25+
* either these or other terms.
26+
*
27+
* This license is subject to the following condition:
28+
*
29+
* The above copyright notice and either this complete permission notice or at a
30+
* minimum a reference to the UPL must be included in all copies or substantial
31+
* portions of the Software.
32+
*
33+
* THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
34+
* IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
35+
* FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
36+
* AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
37+
* LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
38+
* OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
39+
* SOFTWARE.
40+
*/
41+
package com.oracle.graal.python.builtins.objects.itertools;
42+
43+
import static com.oracle.graal.python.builtins.PythonBuiltinClassType.StopIteration;
44+
import static com.oracle.graal.python.nodes.SpecialMethodNames.__ITER__;
45+
import static com.oracle.graal.python.nodes.SpecialMethodNames.__NEXT__;
46+
import static com.oracle.graal.python.nodes.SpecialMethodNames.__REDUCE__;
47+
import static com.oracle.graal.python.nodes.SpecialMethodNames.__SETSTATE__;
48+
49+
import com.oracle.graal.python.builtins.Builtin;
50+
import java.util.List;
51+
52+
import com.oracle.graal.python.builtins.CoreFunctions;
53+
import com.oracle.graal.python.builtins.PythonBuiltinClassType;
54+
import com.oracle.graal.python.builtins.PythonBuiltins;
55+
import com.oracle.graal.python.builtins.modules.BuiltinFunctions;
56+
import com.oracle.graal.python.builtins.objects.PNone;
57+
import com.oracle.graal.python.builtins.objects.tuple.PTuple;
58+
import com.oracle.graal.python.nodes.function.PythonBuiltinBaseNode;
59+
import com.oracle.graal.python.nodes.function.builtins.PythonBinaryBuiltinNode;
60+
import com.oracle.graal.python.nodes.function.builtins.PythonUnaryBuiltinNode;
61+
import com.oracle.graal.python.nodes.object.GetClassNode;
62+
import com.oracle.graal.python.nodes.object.IsBuiltinClassProfile;
63+
import com.oracle.graal.python.runtime.exception.PException;
64+
import com.oracle.truffle.api.dsl.Cached;
65+
import com.oracle.truffle.api.dsl.GenerateNodeFactory;
66+
import com.oracle.truffle.api.dsl.NodeFactory;
67+
import com.oracle.truffle.api.dsl.Specialization;
68+
import com.oracle.truffle.api.frame.VirtualFrame;
69+
import com.oracle.truffle.api.profiles.ConditionProfile;
70+
import com.oracle.truffle.api.profiles.LoopConditionProfile;
71+
72+
@CoreFunctions(extendClasses = {PythonBuiltinClassType.PZipLongest})
73+
public final class ZipLongestBuiltins extends PythonBuiltins {
74+
75+
@Override
76+
protected List<? extends NodeFactory<? extends PythonBuiltinBaseNode>> getNodeFactories() {
77+
return ZipLongestBuiltinsFactory.getFactories();
78+
}
79+
80+
@Builtin(name = __ITER__, minNumOfPositionalArgs = 1)
81+
@GenerateNodeFactory
82+
public abstract static class IterNode extends PythonUnaryBuiltinNode {
83+
@Specialization
84+
static Object iter(PZipLongest self) {
85+
return self;
86+
}
87+
}
88+
89+
@Builtin(name = __NEXT__, minNumOfPositionalArgs = 1)
90+
@GenerateNodeFactory
91+
public abstract static class NextNode extends PythonUnaryBuiltinNode {
92+
@SuppressWarnings("unused")
93+
@Specialization(guards = "zeroSize(self)")
94+
Object nextNoFillValue(VirtualFrame frame, PZipLongest self) {
95+
throw raise(StopIteration);
96+
}
97+
98+
@Specialization(guards = {"!zeroSize(self)", "isNullFillValue(self)"})
99+
Object nextNoFillValue(VirtualFrame frame, PZipLongest self,
100+
@Cached BuiltinFunctions.NextNode nextNode,
101+
@Cached IsBuiltinClassProfile isStopIterationProfile,
102+
@Cached ConditionProfile noItProfile,
103+
@Cached ConditionProfile noActiveProfile,
104+
@Cached LoopConditionProfile loopProfile) {
105+
return next(frame, self, PNone.NONE, nextNode, isStopIterationProfile, loopProfile, noItProfile, noActiveProfile);
106+
}
107+
108+
@Specialization(guards = {"!zeroSize(self)", "!isNullFillValue(self)"})
109+
Object next(VirtualFrame frame, PZipLongest self,
110+
@Cached BuiltinFunctions.NextNode nextNode,
111+
@Cached IsBuiltinClassProfile isStopIterationProfile,
112+
@Cached ConditionProfile noItProfile,
113+
@Cached ConditionProfile noActiveProfile,
114+
@Cached LoopConditionProfile loopProfile) {
115+
return next(frame, self, self.getFillValue(), nextNode, isStopIterationProfile, loopProfile, noItProfile, noActiveProfile);
116+
}
117+
118+
private Object next(VirtualFrame frame, PZipLongest self, Object fillValue, BuiltinFunctions.NextNode nextNode, IsBuiltinClassProfile isStopIterationProfile, LoopConditionProfile loopProfile,
119+
ConditionProfile noItProfile, ConditionProfile noActiveProfile) {
120+
Object[] result = new Object[self.getItTuple().length];
121+
loopProfile.profileCounted(result.length);
122+
for (int i = 0; loopProfile.inject(i < result.length); i++) {
123+
Object it = self.getItTuple()[i];
124+
Object item;
125+
if (noItProfile.profile(it == PNone.NONE)) {
126+
item = fillValue;
127+
} else {
128+
try {
129+
item = nextNode.execute(frame, it, PNone.NO_VALUE);
130+
} catch (PException e) {
131+
if (isStopIterationProfile.profileException(e, StopIteration)) {
132+
self.setNumActive(self.getNumActive() - 1);
133+
if (noActiveProfile.profile(self.getNumActive() == 0)) {
134+
throw raise(StopIteration);
135+
} else {
136+
item = fillValue;
137+
self.getItTuple()[i] = PNone.NONE;
138+
}
139+
} else {
140+
self.setNumActive(0);
141+
throw e;
142+
}
143+
}
144+
}
145+
result[i] = item;
146+
}
147+
return factory().createTuple(result);
148+
}
149+
150+
protected boolean isNullFillValue(PZipLongest self) {
151+
return self.getFillValue() == null;
152+
}
153+
154+
protected boolean zeroSize(PZipLongest self) {
155+
return self.getItTuple().length == 0 || self.getNumActive() == 0;
156+
}
157+
}
158+
159+
@Builtin(name = __REDUCE__, minNumOfPositionalArgs = 1)
160+
@GenerateNodeFactory
161+
public abstract static class ReduceNode extends PythonUnaryBuiltinNode {
162+
@Specialization(guards = "isNullFillValue(self)")
163+
Object reduceNoFillValue(PZipLongest self,
164+
@Cached GetClassNode getClass,
165+
@Cached LoopConditionProfile loopProfile,
166+
@Cached ConditionProfile noItProfile) {
167+
return reduce(getClass, self, PNone.NONE, loopProfile, noItProfile);
168+
}
169+
170+
@Specialization(guards = "!isNullFillValue(self)")
171+
Object reducePos(PZipLongest self,
172+
@Cached GetClassNode getClass,
173+
@Cached LoopConditionProfile loopProfile,
174+
@Cached ConditionProfile noItProfile) {
175+
return reduce(getClass, self, self.getFillValue(), loopProfile, noItProfile);
176+
}
177+
178+
private Object reduce(GetClassNode getClass, PZipLongest self, Object fillValue, LoopConditionProfile loopProfile, ConditionProfile noItProfile) {
179+
Object type = getClass.execute(self);
180+
Object[] its = new Object[self.getItTuple().length];
181+
loopProfile.profileCounted(its.length);
182+
for (int i = 0; loopProfile.profile(i < its.length); i++) {
183+
Object it = self.getItTuple()[i];
184+
if (noItProfile.profile(it == PNone.NONE)) {
185+
its[i] = factory().createEmptyTuple();
186+
} else {
187+
its[i] = it;
188+
}
189+
}
190+
PTuple tuple = factory().createTuple(its);
191+
return factory().createTuple(new Object[]{type, tuple, fillValue});
192+
}
193+
194+
protected boolean isNullFillValue(PZipLongest self) {
195+
return self.getFillValue() == null;
196+
}
197+
}
198+
199+
@Builtin(name = __SETSTATE__, minNumOfPositionalArgs = 2)
200+
@GenerateNodeFactory
201+
public abstract static class SetStateNode extends PythonBinaryBuiltinNode {
202+
@Specialization
203+
Object setState(PZipLongest self, Object state) {
204+
self.setFillValue(state);
205+
return PNone.NONE;
206+
}
207+
}
208+
209+
}

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/runtime/object/PythonObjectFactory.java

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -130,6 +130,7 @@
130130
import com.oracle.graal.python.builtins.objects.itertools.PTakewhile;
131131
import com.oracle.graal.python.builtins.objects.itertools.PTee;
132132
import com.oracle.graal.python.builtins.objects.itertools.PTeeDataObject;
133+
import com.oracle.graal.python.builtins.objects.itertools.PZipLongest;
133134
import com.oracle.graal.python.builtins.objects.keywrapper.PKeyWrapper;
134135
import com.oracle.graal.python.builtins.objects.list.PList;
135136
import com.oracle.graal.python.builtins.objects.map.PMap;
@@ -1248,6 +1249,10 @@ public final PTeeDataObject createTeeDataObject(Object it) {
12481249
return trace(new PTeeDataObject(it, PythonBuiltinClassType.PTeeDataObject, PythonBuiltinClassType.PTeeDataObject.getInstanceShape(getLanguage())));
12491250
}
12501251

1252+
public final PZipLongest createZipLongest(Object cls) {
1253+
return trace(new PZipLongest(cls, getShape(cls)));
1254+
}
1255+
12511256
public final PTextIO createTextIO(Object clazz) {
12521257
return trace(new PTextIO(clazz, getShape(clazz)));
12531258
}

0 commit comments

Comments
 (0)