[GR-17149] Implemented startswith call with tuple on byte arrays

PullRequest: graalpython/584

Author: timfel

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

@@ -443,6 +443,16 @@ def test_startswith():
     assert not b.startswith(b"hellow")
     assert not b.startswith(b"ha")
 
+    assert b.startswith((b"hellow", b"he"))
+    assert not b.startswith((b"hellow", b"ha"))
+
+    try:
+        assert b.startswith(("hel", "hello"))
+    except TypeError:
+        assert True
+    else:
+        assert False
+
 
 def test_endswith():
     b = b'hello'

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/bytes/ByteArrayBuiltins.java

@@ -67,14 +67,18 @@
 import com.oracle.graal.python.builtins.objects.tuple.PTuple;
 import com.oracle.graal.python.nodes.SpecialMethodNames;
 import com.oracle.graal.python.nodes.call.special.LookupAndCallUnaryNode;
+import com.oracle.graal.python.nodes.control.GetIteratorExpressionNode.GetIteratorNode;
+import com.oracle.graal.python.nodes.control.GetNextNode;
 import com.oracle.graal.python.nodes.expression.BinaryComparisonNode;
 import com.oracle.graal.python.nodes.function.PythonBuiltinBaseNode;
 import com.oracle.graal.python.nodes.function.PythonBuiltinNode;
 import com.oracle.graal.python.nodes.function.builtins.PythonBinaryBuiltinNode;
 import com.oracle.graal.python.nodes.function.builtins.PythonTernaryBuiltinNode;
 import com.oracle.graal.python.nodes.function.builtins.PythonUnaryBuiltinNode;
 import com.oracle.graal.python.nodes.object.GetLazyClassNode;
+import com.oracle.graal.python.nodes.object.IsBuiltinClassProfile;
 import com.oracle.graal.python.nodes.truffle.PythonArithmeticTypes;
+import com.oracle.graal.python.runtime.exception.PException;
 import com.oracle.graal.python.runtime.sequence.PSequence;
 import com.oracle.graal.python.runtime.sequence.storage.ByteSequenceStorage;
 import com.oracle.graal.python.runtime.sequence.storage.IntSequenceStorage;
@@ -477,6 +481,31 @@ public Object iter(PByteArray byteArray) {
     abstract static class StartsWithNode extends PythonBuiltinNode {
         @Child private SequenceStorageNodes.LenNode lenNode;
 
+        @Specialization
+        boolean startswith(VirtualFrame frame, PByteArray self, PTuple prefixes, @SuppressWarnings("unused") PNone start, @SuppressWarnings("unused") PNone end,
+                        @Cached GetIteratorNode getIteratorNode,
+                        @Cached IsBuiltinClassProfile errorProfile,
+                        @Cached GetNextNode getNextNode,
+                        @Cached BytesNodes.FindNode findNode) {
+            Object iterator = getIteratorNode.executeWith(frame, prefixes);
+            while (true) {
+                try {
+                    Object arrayObj = getNextNode.execute(frame, iterator);
+                    if (arrayObj instanceof PIBytesLike) {
+                        PIBytesLike array = (PIBytesLike) arrayObj;
+                        if (startswith(self, array, start, end, findNode)) {
+                            return true;
+                        }
+                    } else {
+                        throw raise(PythonBuiltinClassType.TypeError, "a bytes-like object is required, not '%p'", arrayObj);
+                    }
+                } catch (PException e) {
+                    e.expectStopIteration(errorProfile);
+                    return false;
+                }
+            }
+        }
+
         @Specialization
         boolean startswith(PByteArray self, PIBytesLike prefix, @SuppressWarnings("unused") PNone start, @SuppressWarnings("unused") PNone end,
                         @Cached("create()") BytesNodes.FindNode findNode) {

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/bytes/BytesNodes.java

@@ -224,9 +224,14 @@ int find(PIBytesLike primary, PIBytesLike sub, Object starting, Object ending) {
             // TODO implement a more efficient algorithm
             outer: for (int i = start; i < end; i++) {
                 for (int j = 0; j < len2; j++) {
+                    if (i + j >= end) {
+                        continue outer;
+                    }
+
                     int hb = getGetLeftItemNode().executeInt(haystack, i + j);
                     int nb = getGetRightItemNode().executeInt(needle, j);
-                    if (nb != hb || i + j >= end) {
+
+                    if (nb != hb) {
                         continue outer;
                     }
                 }