[pypy-commit] pypy faster-rstruct: make sure that struct.unpack_from uses the fast path as well

[antocuni]

Author: Antonio Cuni <anto.cuni at gmail.com>
Branch: faster-rstruct
Changeset: r80808:f2275be59406
Date: 2015-11-20 17:37 +0100
http://bitbucket.org/pypy/pypy/changeset/f2275be59406/

Log:	make sure that struct.unpack_from uses the fast path as well

diff --git a/pypy/module/struct/formatiterator.py b/pypy/module/struct/formatiterator.py
--- a/pypy/module/struct/formatiterator.py
+++ b/pypy/module/struct/formatiterator.py
@@ -154,9 +154,8 @@
         return self.pos
 
     def get_buffer_as_string_maybe(self):
-        # XXX: if self.buf is something different that StringBuffer, this has
-        # the effect to copy the whole string at each unpack!
-        return self.buf.as_str()
+        string, pos = self.buf.as_str_and_offset_maybe()
+        return string, pos+self.pos
 
     def skip(self, size):
         self.read(size) # XXX, could avoid taking the slice
diff --git a/pypy/module/struct/test/test_struct.py b/pypy/module/struct/test/test_struct.py
--- a/pypy/module/struct/test/test_struct.py
+++ b/pypy/module/struct/test/test_struct.py
@@ -462,3 +462,29 @@
         assert self.struct.unpack_from("ii", b, 2) == (17, 42)
         b[:sz] = self.struct.pack("ii", 18, 43)
         assert self.struct.unpack_from("ii", b) == (18, 43)
+
+
+class AppTestFastPath(object):
+    spaceconfig = dict(usemodules=['struct', '__pypy__'])
+
+    def setup_class(cls):
+        from rpython.rlib.rstruct import standardfmttable
+        standardfmttable.ALLOW_SLOWPATH = False
+        #
+        cls.w_struct = cls.space.appexec([], """():
+            import struct
+            return struct
+        """)
+        cls.w_bytebuffer = cls.space.appexec([], """():
+            import __pypy__
+            return __pypy__.bytebuffer
+        """)
+
+    def teardown_class(cls):
+        from rpython.rlib.rstruct import standardfmttable
+        standardfmttable.ALLOW_SLOWPATH = True
+
+    def test_unpack_from(self):
+        buf = self.struct.pack("iii", 0, 42, 43)
+        offset = self.struct.calcsize("i")
+        assert self.struct.unpack_from("ii", buf, offset) == (42, 43)
diff --git a/rpython/rlib/buffer.py b/rpython/rlib/buffer.py
--- a/rpython/rlib/buffer.py
+++ b/rpython/rlib/buffer.py
@@ -22,6 +22,14 @@
         # May be overridden.
         return self.getslice(0, self.getlength(), 1, self.getlength())
 
+    def as_str_and_offset_maybe(self):
+        """
+        If the buffer is backed by a string, return a pair (string, offset), where
+        offset is the offset inside the string where the buffer start.
+        Else, return (None, 0).
+        """
+        return None, 0
+
     def getitem(self, index):
         "Returns the index'th character in the buffer."
         raise NotImplementedError   # Must be overriden.  No bounds checks.
@@ -66,6 +74,9 @@
     def as_str(self):
         return self.value
 
+    def as_str_and_offset_maybe(self):
+        return self.value, 0
+
     def getitem(self, index):
         return self.value[index]
 
@@ -99,6 +110,12 @@
         else:
             return 0
 
+    def as_str_and_offset_maybe(self):
+        string, offset = self.buffer.as_str_and_offset_maybe()
+        if string is not None:
+            return string, offset+self.offset
+        return None, 0
+
     def getitem(self, index):
         return self.buffer.getitem(self.offset + index)
 
diff --git a/rpython/rlib/rstruct/standardfmttable.py b/rpython/rlib/rstruct/standardfmttable.py
--- a/rpython/rlib/rstruct/standardfmttable.py
+++ b/rpython/rlib/rstruct/standardfmttable.py
@@ -130,7 +130,8 @@
 
 # ____________________________________________________________
 
-ALLOW_FASTPATH = True # set to False by TestNoFastPath
+USE_FASTPATH = True    # set to False by some tests
+ALLOW_SLOWPATH = True  # set to False by some tests
 
 class CannotUnpack(Exception):
     pass
@@ -139,7 +140,7 @@
 def unpack_fastpath(TYPE, fmtiter):
     size = rffi.sizeof(TYPE)
     strbuf, pos = fmtiter.get_buffer_as_string_maybe()
-    if pos % size != 0 or strbuf is None or not ALLOW_FASTPATH:
+    if strbuf is None or pos % size != 0 or not USE_FASTPATH:
         raise CannotUnpack
     fmtiter.skip(size)
     return str_storage_getitem(TYPE, strbuf, pos)
@@ -226,6 +227,9 @@
         if unpack_int_fastpath_maybe(fmtiter):
             return
         # slow path
+        if not ALLOW_SLOWPATH:
+            # we enter here only on some tests
+            raise ValueError("fastpath not taken :(")
         intvalue = inttype(0)
         s = fmtiter.read(size)
         idx = 0
diff --git a/rpython/rlib/rstruct/test/test_runpack.py b/rpython/rlib/rstruct/test/test_runpack.py
--- a/rpython/rlib/rstruct/test/test_runpack.py
+++ b/rpython/rlib/rstruct/test/test_runpack.py
@@ -94,7 +94,7 @@
 class TestNoFastPath(TestRStruct):
 
     def setup_method(self, meth):
-        standardfmttable.ALLOW_FASTPATH = False
+        standardfmttable.USE_FASTPATH = False
 
     def teardown_method(self, meth):
-        standardfmttable.ALLOW_FASTPATH = True
+        standardfmttable.USE_FASTPATH = True
diff --git a/rpython/rlib/test/test_buffer.py b/rpython/rlib/test/test_buffer.py
--- a/rpython/rlib/test/test_buffer.py
+++ b/rpython/rlib/test/test_buffer.py
@@ -32,3 +32,16 @@
     a = RPythonAnnotator()
     s = a.build_types(func, [int])
     assert s == SomeInteger(nonneg=True)
+
+
+def test_as_str_and_offset_maybe():
+    buf = StringBuffer('hello world')
+    assert buf.as_str_and_offset_maybe() == ('hello world', 0)
+    #
+    sbuf = SubBuffer(buf, 6, 5)
+    assert sbuf.getslice(0, 5, 1, 5) == 'world'
+    assert sbuf.as_str_and_offset_maybe() == ('hello world', 6)
+    #
+    ssbuf = SubBuffer(sbuf, 3, 2)
+    assert ssbuf.getslice(0, 2, 1, 2) == 'ld'
+    assert ssbuf.as_str_and_offset_maybe() == ('hello world', 9)