[pypy-commit] pypy faster-rstruct: introduce rawstorage.str_storage_getitem{, _unaligned}, which works as rawstorage.raw_storage_* but operates on RPython strings instead of raw buffers: since strings are immutable, we provide only the functions for reading out of them, not writing

[antocuni]

Author: Antonio Cuni <anto.cuni at gmail.com>
Branch: faster-rstruct
Changeset: r80668:84c49eed1754
Date: 2015-05-15 22:38 +0200
http://bitbucket.org/pypy/pypy/changeset/84c49eed1754/

Log:	introduce rawstorage.str_storage_getitem{,_unaligned}, which works
	as rawstorage.raw_storage_* but operates on RPython strings instead
	of raw buffers: since strings are immutable, we provide only the
	functions for reading out of them, not writing

diff --git a/rpython/rlib/rawstorage.py b/rpython/rlib/rawstorage.py
--- a/rpython/rlib/rawstorage.py
+++ b/rpython/rlib/rawstorage.py
@@ -1,10 +1,13 @@
-from rpython.rlib.objectmodel import we_are_translated
+from rpython.rlib.objectmodel import we_are_translated, keepalive_until_here
 from rpython.rtyper.extregistry import ExtRegistryEntry
 from rpython.rtyper.lltypesystem import lltype, rffi, llmemory
 from rpython.annotator import model as annmodel
 from rpython.rtyper.llannotation import lltype_to_annotation
+from rpython.rlib import rgc
 from rpython.rlib.rgc import lltype_is_gc
 from rpython.rlib.objectmodel import specialize
+from rpython.rtyper.lltypesystem.rstr import STR, _get_raw_str_buf
+from rpython.rtyper.annlowlevel import llstr
 
 RAW_STORAGE = rffi.CCHARP.TO
 RAW_STORAGE_PTR = rffi.CCHARP
@@ -40,6 +43,29 @@
 def free_raw_storage(storage, track_allocation=True):
     lltype.free(storage, flavor='raw', track_allocation=track_allocation)
 
+
+ at rgc.no_collect
+def str_storage_getitem(TP, s, index):
+    lls = llstr(s)
+    # from here, no GC operations can happen
+    buf = _get_raw_str_buf(STR, lls, 0)
+    storage = rffi.cast(RAW_STORAGE_PTR, buf)
+    res = raw_storage_getitem(TP, storage, index)
+    # end of "no GC" section
+    keepalive_until_here(lls)
+    return res
+
+ at rgc.no_collect
+def str_storage_getitem_unaligned(TP, s, index):
+    lls = llstr(s)
+    # from here, no GC operations can happen
+    buf = _get_raw_str_buf(STR, lls, 0)
+    storage = rffi.cast(RAW_STORAGE_PTR, buf)
+    res = raw_storage_getitem_unaligned(TP, storage, index)
+    # end of "no GC" section
+    keepalive_until_here(lls)
+    return res
+
 # ____________________________________________________________
 #
 # Support for possibly-unaligned accesses
diff --git a/rpython/rlib/test/test_rawstorage.py b/rpython/rlib/test/test_rawstorage.py
--- a/rpython/rlib/test/test_rawstorage.py
+++ b/rpython/rlib/test/test_rawstorage.py
@@ -4,7 +4,8 @@
 from rpython.rlib import rawstorage
 from rpython.rlib.rawstorage import alloc_raw_storage, free_raw_storage,\
      raw_storage_setitem, raw_storage_getitem, AlignmentError,\
-     raw_storage_setitem_unaligned, raw_storage_getitem_unaligned
+     raw_storage_setitem_unaligned, raw_storage_getitem_unaligned,\
+     str_storage_getitem, str_storage_getitem_unaligned
 from rpython.rtyper.test.tool import BaseRtypingTest
 from rpython.translator.c.test.test_genc import compile
 
@@ -32,6 +33,26 @@
     assert res == 3.14
     free_raw_storage(r)
 
+def test_untranslated_str_storage():
+    import struct
+    buf = struct.pack('@lld', 42, 43, 123.0)
+    size = struct.calcsize('@l')
+    res = str_storage_getitem(lltype.Signed, buf, 0)
+    assert res == 42
+    res = str_storage_getitem(lltype.Signed, buf, size)
+    assert res == 43
+    res = str_storage_getitem(lltype.Float, buf, size*2)
+    assert res == 123.0
+
+def test_untranslated_str_storage_unaligned(monkeypatch):
+    import struct
+    monkeypatch.setattr(rawstorage, 'misaligned_is_fine', False)
+    buf = 'foo' + struct.pack('@ll', 42, 43)
+    size = struct.calcsize('@l')
+    res = str_storage_getitem_unaligned(lltype.Signed, buf, 3)
+    assert res == 42
+    res = str_storage_getitem_unaligned(lltype.Signed, buf, size+3)
+    assert res == 43
 
 class TestRawStorage(BaseRtypingTest):
 
@@ -46,6 +67,21 @@
         x = self.interpret(f, [1<<30])
         assert x == 1 << 30
 
+    def test_str_storage_int(self):
+        import struct
+        buf = struct.pack('@ll', 42, 43)
+        size = struct.calcsize('@l')
+
+        def f(i):
+            res = str_storage_getitem(lltype.Signed, buf, i)
+            return res
+
+        x = self.interpret(f, [0])
+        assert x == 42
+        x = self.interpret(f, [8])
+        assert x == 43
+
+
     def test_storage_float_unaligned(self, monkeypatch):
         def f(v):
             r = alloc_raw_storage(24)
diff --git a/rpython/rtyper/lltypesystem/rstr.py b/rpython/rtyper/lltypesystem/rstr.py
--- a/rpython/rtyper/lltypesystem/rstr.py
+++ b/rpython/rtyper/lltypesystem/rstr.py
@@ -60,6 +60,13 @@
     @signature(types.any(), types.any(), types.int(), returns=types.any())
     @specialize.arg(0)
     def _get_raw_buf(TP, src, ofs):
+        """
+        WARNING: dragons ahead.
+        Return the address of the internal char* buffer of the low level
+        string. The return value is valid as long as no GC operation occur, so
+        you must ensure that it will be used inside a "GC safe" section, for
+        example by marking your function with @rgc.no_collect
+        """
         assert typeOf(src).TO == TP
         assert ofs >= 0
         return llmemory.cast_ptr_to_adr(src) + _str_ofs(TP, ofs)
@@ -129,11 +136,17 @@
     copy_raw_to_string = func_with_new_name(copy_raw_to_string,
                                               'copy_raw_to_%s' % name)
 
-    return copy_string_to_raw, copy_raw_to_string, copy_string_contents
+    return _get_raw_buf, copy_string_to_raw, copy_raw_to_string, copy_string_contents
 
-copy_string_to_raw, copy_raw_to_string, copy_string_contents = _new_copy_contents_fun(STR, STR, Char, 'string')
-copy_unicode_to_raw, copy_raw_to_unicode, copy_unicode_contents = _new_copy_contents_fun(UNICODE, UNICODE,
-                                                                    UniChar, 'unicode')
+(_get_raw_str_buf,
+ copy_string_to_raw,
+ copy_raw_to_string,
+ copy_string_contents) = _new_copy_contents_fun(STR, STR, Char, 'string')
+
+(_get_raw_unicode_buf,
+ copy_unicode_to_raw,
+ copy_raw_to_unicode,
+ copy_unicode_contents) = _new_copy_contents_fun(UNICODE, UNICODE, UniChar, 'unicode')
 
 CONST_STR_CACHE = WeakValueDictionary()
 CONST_UNICODE_CACHE = WeakValueDictionary()