Add support for segment (de)serialization.

capnp/includes/schema_cpp.pxd

@@ -667,6 +667,8 @@ cdef extern from "capnp/message.h" namespace " ::capnp":
         DynamicStruct_Builder initRootDynamicStruct'initRoot< ::capnp::DynamicStruct>'(StructSchema)
         void setRootDynamicStruct'setRoot< ::capnp::DynamicStruct::Reader>'(DynamicStruct.Reader)
 
+        ConstWordArrayArrayPtr getSegmentsForOutput'getSegmentsForOutput'()
+
         AnyPointer.Builder getRootAnyPointer'getRoot< ::capnp::AnyPointer>'()
 
         DynamicOrphan newOrphan'getOrphanage().newOrphan'(StructSchema)
@@ -691,6 +693,10 @@ cdef extern from "capnp/message.h" namespace " ::capnp":
         MallocMessageBuilder()
         MallocMessageBuilder(int)
 
+    cdef cppclass SegmentArrayMessageReader(MessageReader):
+        SegmentArrayMessageReader(ConstWordArrayArrayPtr array) except +reraise_kj_exception
+        SegmentArrayMessageReader(ConstWordArrayArrayPtr array, ReaderOptions) except +reraise_kj_exception
+
     cdef cppclass FlatMessageBuilder(MessageBuilder):
         FlatMessageBuilder(WordArrayPtr array)
         FlatMessageBuilder(WordArrayPtr array, ReaderOptions)
@@ -714,6 +720,16 @@ cdef extern from "kj/common.h" namespace " ::kj":
         ByteArrayPtr(byte *, size_t size)
         size_t size()
         byte& operator[](size_t index)
+    cdef cppclass ConstWordArrayPtr " ::kj::ArrayPtr< const ::capnp::word>":
+        ConstWordArrayPtr()
+        ConstWordArrayPtr(word *, size_t size)
+        size_t size()
+        const word* begin()
+    cdef cppclass ConstWordArrayArrayPtr " ::kj::ArrayPtr< const ::kj::ArrayPtr< const ::capnp::word>>":
+        ConstWordArrayArrayPtr()
+        ConstWordArrayArrayPtr(ConstWordArrayPtr*, size_t size)
+        size_t size()
+        ConstWordArrayPtr& operator[](size_t index)
 
 cdef extern from "kj/array.h" namespace " ::kj":
     # Cython can't handle Array[word] as a function argument

capnp/lib/capnp.pxd

@@ -53,6 +53,7 @@ cdef class _DynamicStructBuilder:
 
     cdef _check_write(self)
     cpdef to_bytes(_DynamicStructBuilder self) except +reraise_kj_exception
+    cpdef to_segments(_DynamicStructBuilder self) except +reraise_kj_exception
     cpdef _to_bytes_packed_helper(_DynamicStructBuilder self, word_count) except +reraise_kj_exception
     cpdef to_bytes_packed(_DynamicStructBuilder self) except +reraise_kj_exception
 

capnp/lib/capnp.pyx

@@ -1192,6 +1192,12 @@ cdef class _DynamicStructBuilder:
         self._is_written = True
         return ret
 
+    cpdef to_segments(_DynamicStructBuilder self) except +reraise_kj_exception:
+        self._check_write()
+        cdef _MessageBuilder builder = self._parent
+        segments = builder.get_segments_for_output()
+        return segments
+
     cpdef _to_bytes_packed_helper(_DynamicStructBuilder self, word_count) except +reraise_kj_exception:
         cdef _MessageBuilder builder = self._parent
         array = helpers.messageToPackedBytes(deref(builder.thisptr), word_count)
@@ -2989,6 +2995,9 @@ class _StructModule(object):
         else:
             message = _FlatArrayMessageReader(buf, traversal_limit_in_words, nesting_limit)
             return message.get_root(self.schema)
+    def from_segments(self, segments):
+        message = _SegmentArrayMessageReader(segments)
+        return message.get_root(self.schema)
     def from_bytes_packed(self, buf, traversal_limit_in_words = None, nesting_limit = None):
         """Returns a Reader for the packed object in buf.
 
@@ -3285,6 +3294,18 @@ cdef class _MessageBuilder:
             self.thisptr.setRootDynamicStruct((<_DynamicStructReader>value).thisptr)
             return self.get_root(value.schema)
 
+    cpdef get_segments_for_output(self) except +reraise_kj_exception:
+        segments = self.thisptr.getSegmentsForOutput()
+        res = []
+        cdef const char* ptr
+        cdef bytes segment_bytes
+        for i in range(0, segments.size()):
+            segment = segments[i]
+            ptr = <const char *> segment.begin()
+            segment_bytes = ptr[:8*segment.size()]
+            res.append(segment_bytes)
+        return res
+
     cpdef new_orphan(self, schema) except +reraise_kj_exception:
         """A method for instantiating Cap'n Proto orphans
 
@@ -3738,6 +3759,36 @@ cdef class _FlatArrayMessageReader(_MessageReader):
         del self.thisptr
 
 
+@cython.internal
+cdef class _SegmentArrayMessageReader(_MessageReader):
+
+    cdef object _objects_to_pin
+
+    def __init__(self, segments):
+        # takes a Python array of bytes and constructs a ConstWordArrayArrayPtr
+        num_segments = len(segments)
+        cdef char* ptr
+        cdef schema_cpp.ConstWordArrayPtr seg_ptr
+        cdef schema_cpp.ConstWordArrayPtr* seg_ptrs = <schema_cpp.ConstWordArrayPtr*>malloc(num_segments * sizeof(schema_cpp.ConstWordArrayPtr))
+        self._objects_to_pin = []
+        for i in range(0, num_segments):
+            segment = bytes(segments[i])
+            ptr = segment
+            if (<uintptr_t>ptr) % 8 != 0:
+                aligned = _AlignedBuffer(segment)
+                ptr = aligned.buf
+                self._objects_to_pin.append(aligned)
+            else:
+                self._objects_to_pin.append(segment)
+            seg_ptr = schema_cpp.ConstWordArrayPtr(<schema_cpp.word*>ptr, len(segment)//8)
+            seg_ptrs[i] = seg_ptr
+        self.thisptr = new schema_cpp.SegmentArrayMessageReader(
+            schema_cpp.ConstWordArrayArrayPtr(seg_ptrs, num_segments))
+
+    def __dealloc__(self):
+        del self.thisptr
+
+
 @cython.internal
 cdef class _FlatMessageBuilder(_MessageBuilder):
     cdef object _object_to_pin