[pypy-commit] pypy null_byte_after_str: rffi.scoped_nonmovingbuffer_final_null

[arigo]

Author: Armin Rigo <arigo at tunes.org>
Branch: null_byte_after_str
Changeset: r85915:4496557d38d9
Date: 2016-07-29 19:25 +0200
http://bitbucket.org/pypy/pypy/changeset/4496557d38d9/

Log:	rffi.scoped_nonmovingbuffer_final_null

diff --git a/rpython/rlib/rgc.py b/rpython/rlib/rgc.py
--- a/rpython/rlib/rgc.py
+++ b/rpython/rlib/rgc.py
@@ -1279,7 +1279,7 @@
     """
     from rpython.rtyper.lltypesystem import rffi
     PSTR = lltype.typeOf(s)
-    _check_final_null_char(PSTR)
+    assert has_final_null_char(PSTR) == 1
     n = llmemory.offsetof(PSTR.TO, 'chars')
     n += llmemory.itemoffsetof(PSTR.TO.chars, 0)
     n = llmemory.raw_malloc_usage(n)
@@ -1289,5 +1289,5 @@
     ptr[n] = '\x00'
 
 @specialize.memo()
-def _check_final_null_char(PSTR):
-    assert PSTR.TO.chars._hints.get('extra_item_after_alloc', 0) == 1
+def has_final_null_char(PSTR):
+    return PSTR.TO.chars._hints.get('extra_item_after_alloc', 0)
diff --git a/rpython/rtyper/lltypesystem/rffi.py b/rpython/rtyper/lltypesystem/rffi.py
--- a/rpython/rtyper/lltypesystem/rffi.py
+++ b/rpython/rtyper/lltypesystem/rffi.py
@@ -819,6 +819,11 @@
         First bool returned indicates if 'data' was pinned. Second bool returned
         indicates if we did a raw alloc because pinning failed. Both bools
         should never be true at the same time.
+
+        For strings (not unicodes), the len()th character of the resulting
+        raw buffer is available, but not initialized.  Use
+        get_nonmovingbuffer_final_null() instead of get_nonmovingbuffer()
+        to get a regular null-terminated "char *".
         """
 
         lldata = llstrtype(data)
@@ -829,6 +834,7 @@
             if rgc.pin(data):
                 pinned = True
             else:
+                count += has_final_null_char(TYPEP)
                 buf = lltype.malloc(TYPEP.TO, count, flavor='raw')
                 copy_string_to_raw(lldata, buf, 0, count)
                 return buf, pinned, True
@@ -846,6 +852,14 @@
     get_nonmovingbuffer._always_inline_ = 'try' # get rid of the returned tuple
     get_nonmovingbuffer._annenforceargs_ = [strtype]
 
+    @jit.dont_look_inside
+    def get_nonmovingbuffer_final_null(data):
+        buf, is_pinned, is_raw = get_nonmovingbuffer(data)
+        buf[len(data)] = lastchar
+        return buf, is_pinned, is_raw
+    get_nonmovingbuffer_final_null._always_inline_ = 'try'
+    get_nonmovingbuffer_final_null._annenforceargs_ = [strtype]
+
     # (str, char*, bool, bool) -> None
     # Can't inline this because of the raw address manipulation.
     @jit.dont_look_inside
@@ -947,18 +961,19 @@
 
     return (str2charp, free_charp, charp2str,
             get_nonmovingbuffer, free_nonmovingbuffer,
+            get_nonmovingbuffer_final_null,
             alloc_buffer, str_from_buffer, keep_buffer_alive_until_here,
             charp2strn, charpsize2str, str2chararray, str2rawmem,
             )
 
 (str2charp, free_charp, charp2str,
- get_nonmovingbuffer, free_nonmovingbuffer,
+ get_nonmovingbuffer, free_nonmovingbuffer, get_nonmovingbuffer_final_null,
  alloc_buffer, str_from_buffer, keep_buffer_alive_until_here,
  charp2strn, charpsize2str, str2chararray, str2rawmem,
  ) = make_string_mappings(str)
 
 (unicode2wcharp, free_wcharp, wcharp2unicode,
- get_nonmoving_unicodebuffer, free_nonmoving_unicodebuffer,
+ get_nonmoving_unicodebuffer, free_nonmoving_unicodebuffer, __not_usable,
  alloc_unicodebuffer, unicode_from_buffer, keep_unicodebuffer_alive_until_here,
  wcharp2unicoden, wcharpsize2unicode, unicode2wchararray, unicode2rawmem,
  ) = make_string_mappings(unicode)
@@ -1202,6 +1217,19 @@
     __enter__._always_inline_ = 'try'
     __exit__._always_inline_ = 'try'
 
+class scoped_nonmovingbuffer_final_null:
+    def __init__(self, data):
+        self.data = data
+    def __enter__(self):
+        self.buf, self.pinned, self.is_raw = (
+            get_nonmovingbuffer_final_null(self.data))
+        return self.buf
+    def __exit__(self, *args):
+        free_nonmovingbuffer(self.data, self.buf, self.pinned, self.is_raw)
+    __init__._always_inline_ = 'try'
+    __enter__._always_inline_ = 'try'
+    __exit__._always_inline_ = 'try'
+
 class scoped_nonmoving_unicodebuffer:
     def __init__(self, data):
         self.data = data
diff --git a/rpython/rtyper/lltypesystem/test/test_rffi.py b/rpython/rtyper/lltypesystem/test/test_rffi.py
--- a/rpython/rtyper/lltypesystem/test/test_rffi.py
+++ b/rpython/rtyper/lltypesystem/test/test_rffi.py
@@ -835,3 +835,11 @@
     if hasattr(rffi, '__INT128_T'):
         value = 0xAAAABBBBCCCCDDDD
         assert cast(rffi.__INT128_T, r_uint64(value)) == value
+
+def test_scoped_nonmovingbuffer_final_null():
+    s = 'bar'
+    with scoped_nonmovingbuffer_final_null(s) as buf:
+        assert buf[0] == 'b'
+        assert buf[1] == 'a'
+        assert buf[2] == 'r'
+        assert buf[3] == '\x00'