[pypy-commit] pypy faster-rstruct-2: add MutableStringBuffer, which will be used by struct.pack to incrementally build the desired string

[antocuni]

Author: Antonio Cuni <anto.cuni at gmail.com>
Branch: faster-rstruct-2
Changeset: r91218:ac79316e004f
Date: 2017-05-10 01:28 +0200
http://bitbucket.org/pypy/pypy/changeset/ac79316e004f/

Log:	add MutableStringBuffer, which will be used by struct.pack to
	incrementally build the desired string

diff --git a/rpython/rlib/mutbuffer.py b/rpython/rlib/mutbuffer.py
new file mode 100644
--- /dev/null
+++ b/rpython/rlib/mutbuffer.py
@@ -0,0 +1,40 @@
+from rpython.rtyper.lltypesystem import lltype, llmemory
+from rpython.rtyper.lltypesystem.lloperation import llop
+from rpython.rtyper.lltypesystem.rstr import STR
+from rpython.rtyper.annlowlevel import llstr, hlstr
+from rpython.rlib.buffer import Buffer
+
+class MutableStringBuffer(Buffer):
+    """
+    A writeable buffer to incrementally fill a string of a fixed size.
+
+    You can fill the string by calling setitem, setslice and typed_write, and
+    get the result by calling finish().
+
+    After you call finish(), you can no longer modify the buffer. There is no
+    check, you will probably get a segfault after translation.
+
+    You can call finish() only once.
+    """
+    _attrs_ = ['readonly', 'll_val']
+    _immutable_ = True
+
+    def __init__(self, size):
+        self.readonly = False
+        # rstr.mallocstr does not pass zero=True, so we call lltype.malloc
+        # directly
+        self.ll_val = lltype.malloc(STR, size, zero=True)
+
+    def finish(self):
+        if not self.ll_val:
+            raise ValueError("Cannot call finish() twice")
+        result = hlstr(self.ll_val)
+        self.ll_val = lltype.nullptr(STR)
+        self.readonly = True
+        return result
+
+    def as_str(self):
+        raise ValueError('as_str() is not supported. Use finish() instead')
+
+    def setitem(self, index, char):
+        self.ll_val.chars[index] = char
diff --git a/rpython/rlib/rstruct/formatiterator.py b/rpython/rlib/rstruct/formatiterator.py
--- a/rpython/rlib/rstruct/formatiterator.py
+++ b/rpython/rlib/rstruct/formatiterator.py
@@ -116,6 +116,7 @@
 def table2desclist(table):
     items = table.items()
     items.sort()
+    import pdb;pdb.set_trace()
     lst = [FmtDesc(key, attrs) for key, attrs in items]
     return unrolling_iterable(lst)
 
diff --git a/rpython/rlib/test/test_mutbuffer.py b/rpython/rlib/test/test_mutbuffer.py
new file mode 100644
--- /dev/null
+++ b/rpython/rlib/test/test_mutbuffer.py
@@ -0,0 +1,24 @@
+import pytest
+from rpython.rlib.mutbuffer import MutableStringBuffer
+
+class TestMutableStringBuffer(object):
+
+    def test_finish(self):
+        buf = MutableStringBuffer(4)
+        pytest.raises(ValueError, "buf.as_str()")
+        s = buf.finish()
+        assert s == '\x00' * 4
+        pytest.raises(ValueError, "buf.finish()")
+
+    def test_setitem(self):
+        buf = MutableStringBuffer(4)
+        buf.setitem(0, 'A')
+        buf.setitem(1, 'B')
+        buf.setitem(2, 'C')
+        buf.setitem(3, 'D')
+        assert buf.finish() == 'ABCD'
+
+    def test_setslice(self):
+        buf = MutableStringBuffer(6)
+        buf.setslice(2, 'ABCD')
+        assert buf.finish() == '\x00\x00ABCD'