pycapnp supports lazy zero segment allocation strategy

capnp/includes/schema_cpp.pxd

@@ -2,7 +2,9 @@
 # distutils: language = c++
 
 from libc.stdint cimport *
+from libcpp.set cimport set
 from capnp.helpers.non_circular cimport c_reraise_kj_exception as reraise_kj_exception
+from capnp.includes.capnp_cpp cimport TypeWhich
 
 from capnp.includes.types cimport *
 
@@ -649,6 +651,13 @@ cdef extern from "capnp/message.h" namespace " ::capnp":
         uint64_t traversalLimitInWords
         uint nestingLimit
 
+    cdef cppclass LazyZeroSegmentAlloc "BuilderOptions::LazyZeroSegmentAlloc":
+        bint enableLazyZero
+        set[TypeWhich] skipLazyZeroTypes
+
+    cdef cppclass BuilderOptions:
+        LazyZeroSegmentAlloc lazyZeroSegmentAlloc
+
     cdef cppclass MessageBuilder nogil:
         CodeGeneratorRequest.Builder getRootCodeGeneratorRequest'getRoot< ::capnp::schema::CodeGeneratorRequest>'()
         CodeGeneratorRequest.Builder initRootCodeGeneratorRequest'initRoot< ::capnp::schema::CodeGeneratorRequest>'()
@@ -683,6 +692,9 @@ cdef extern from "capnp/message.h" namespace " ::capnp":
 
         DynamicOrphan newOrphan'getOrphanage().newOrphan'(StructSchema)
 
+        void setOptions'setOptions'(BuilderOptions options)
+        BuilderOptions getOptions'getOptions'() const
+
     cdef cppclass MessageReader nogil:
         CodeGeneratorRequest.Reader getRootCodeGeneratorRequest'getRoot< ::capnp::schema::CodeGeneratorRequest>'()
         InterfaceNode.Reader getRootInterfaceNode'getRoot< ::capnp::schema::InterfaceNode>'()

capnp/lib/capnp.pxd

@@ -151,6 +151,15 @@ cdef class _MessageBuilder:
     cpdef set_root(self, value)
     cpdef get_segments_for_output(self)
     cpdef new_orphan(self, schema)
+    cpdef get_options(self)
+    cpdef set_options(self, BuilderOptions py_opts)
+
+cdef class LazyZeroSegmentAlloc:
+    cdef public bint enableLazyZero
+    cdef public object skipLazyZeroTypes
+
+cdef class BuilderOptions:
+    cdef public LazyZeroSegmentAlloc lazyZeroSegmentAlloc
 
 cdef to_python_reader(C_DynamicValue.Reader self, object parent)
 cdef to_python_builder(C_DynamicValue.Builder self, object parent)

capnp/lib/capnp.pyx

@@ -3196,12 +3196,14 @@ class _StructABCMeta(type):
         return isinstance(obj, cls.__base__) and obj.schema == cls._schema
 
 
-cdef _new_message(self, kwargs, num_first_segment_words, allocate_seg_callable):
+cdef _new_message(self, kwargs, num_first_segment_words, allocate_seg_callable, builder_options):
     cdef _MessageBuilder builder
     if allocate_seg_callable is None:
         builder = _MallocMessageBuilder(num_first_segment_words)
+        builder.set_options(builder_options)
     else:
         builder = _PyCustomMessageBuilder(allocate_seg_callable, num_first_segment_words)
+        builder.set_options(builder_options)
     msg = builder.init_root(self.schema)
     if kwargs is not None:
         msg.from_dict(kwargs)
@@ -3442,7 +3444,7 @@ class _StructModule(object):
     def __call__(self, num_first_segment_words=None, **kwargs):
         return self.new_message(num_first_segment_words=num_first_segment_words, **kwargs)
 
-    def new_message(self, num_first_segment_words=None, allocate_seg_callable=None, **kwargs):
+    def new_message(self, num_first_segment_words=None, allocate_seg_callable=None, builder_options=BuilderOptions(), **kwargs):
         """Returns a newly allocated builder message.
 
         :type num_first_segment_words: int
@@ -3453,6 +3455,9 @@ class _StructModule(object):
         words to allocate (as an `int`) and returns a `bytearray`. This is used to customize the memory
         allocation strategy.
 
+        :type builder_options: BuilderOptions
+        :param builder_options: This is for configuring message builder options, such as lazy zero segment alloc.
+
         :type kwargs: dict
         :param kwargs: A list of fields and their values to initialize in the struct.
 
@@ -3461,7 +3466,7 @@ class _StructModule(object):
 
         :rtype: :class:`_DynamicStructBuilder`
         """
-        return _new_message(self, kwargs, num_first_segment_words, allocate_seg_callable)
+        return _new_message(self, kwargs, num_first_segment_words, allocate_seg_callable, builder_options)
 
 
 class _InterfaceModule(object):
@@ -3795,6 +3800,51 @@ cdef class _MessageBuilder:
         ptr = s._thisptr()
         return _DynamicOrphan()._init(self.thisptr.newOrphan(ptr), self)
 
+    cpdef set_options(self, BuilderOptions py_opts):
+        if py_opts is None:
+            raise ValueError("options must be an BuilderOptions instance, got None")
+
+        cdef schema_cpp.BuilderOptions opts
+        opts.lazyZeroSegmentAlloc.enableLazyZero = <bint> py_opts.lazyZeroSegmentAlloc.enableLazyZero
+
+        if py_opts.lazyZeroSegmentAlloc is not None:
+            py_set = py_opts.lazyZeroSegmentAlloc.skipLazyZeroTypes
+            if py_set is not None:
+                try:
+                    if types.Data in py_set:
+                        opts.lazyZeroSegmentAlloc.skipLazyZeroTypes.insert(capnp.TypeWhichDATA)
+                except TypeError:
+                    raise TypeError("skipLazyZeroTypes must be an iterable (e.g. set)")
+
+            self.thisptr.setOptions(opts)
+
+
+    cpdef get_options(self):
+        cdef schema_cpp.BuilderOptions opts = self.thisptr.getOptions()
+
+        cdef BuilderOptions py_opts = BuilderOptions()
+        py_opts.lazyZeroSegmentAlloc = LazyZeroSegmentAlloc()
+        py_opts.lazyZeroSegmentAlloc.enableLazyZero = opts.lazyZeroSegmentAlloc.enableLazyZero
+        py_opts.lazyZeroSegmentAlloc.skipLazyZeroTypes = set()
+
+        if opts.lazyZeroSegmentAlloc.skipLazyZeroTypes.count(capnp.TypeWhichDATA):
+            py_opts.lazyZeroSegmentAlloc.skipLazyZeroTypes.add(types.Data)
+
+        return py_opts
+
+cdef class LazyZeroSegmentAlloc:
+    def __init__(self, enableLazyZero=False, skipLazyZeroTypes=None):
+        self.enableLazyZero = enableLazyZero
+        self.skipLazyZeroTypes = skipLazyZeroTypes
+        if skipLazyZeroTypes == None:
+            self.skipLazyZeroTypes = set()
+
+cdef class BuilderOptions:
+    def __init__(self, lazyZeroSegmentAlloc=None):
+        self.lazyZeroSegmentAlloc = lazyZeroSegmentAlloc
+        if lazyZeroSegmentAlloc == None:
+            self.lazyZeroSegmentAlloc = LazyZeroSegmentAlloc()
+
 
 cdef class _MallocMessageBuilder(_MessageBuilder):
     """The main class for building Cap'n Proto messages

examples/lazy_zero_allocation.py

@@ -0,0 +1,41 @@
+import capnp
+import random
+
+import addressbook_capnp
+
+
+WORD_SIZE = 8
+
+
+class Allocator:
+    def __init__(self):
+        return
+
+    def __call__(self, minimum_size: int) -> bytearray:
+        # dirty the memory on purpose, to simulate reusing shared memory
+        return bytearray(random.getrandbits(8) for _ in range(minimum_size * WORD_SIZE))
+
+
+lazy_zero = capnp.LazyZeroSegmentAlloc(enableLazyZero=True, skipLazyZeroTypes={capnp.types.Data})
+builder_options=capnp.BuilderOptions(lazyZeroSegmentAlloc=lazy_zero)
+
+# message creation method 1
+person = addressbook_capnp.Person.new_message(allocate_seg_callable=Allocator(), builder_options=builder_options)
+print(person.name) # guaranteed empty string
+person.name = "test name 1"
+print(person.name)
+person.init("extraData", 100) # random dirty bytes
+print(bytes(person.extraData))
+print()
+
+# message creation method 2
+builder = capnp._PyCustomMessageBuilder(allocate_seg_callable=Allocator())
+builder.set_options(builder_options)
+person = builder.init_root(addressbook_capnp.Person)
+print(person.name) # guaranteed empty string
+person.name = "test name 2"
+print(person.name)
+person.init("extraData", 100) # random dirty bytes
+print(bytes(person.extraData))
+builder.get_options()
+print(len(builder.get_options().lazyZeroSegmentAlloc.skipLazyZeroTypes)) # return a set of length 1