Optimize Offset/Pad/Prep: use cached head and slicing, reduce casting (#8808)

This commit is contained in:
Jacob Abrams
2025-12-01 16:00:20 -08:00
committed by GitHub
parent a577050817
commit 597e76a268

View File

@@ -166,7 +166,7 @@ class Builder(object):
def Clear(self): def Clear(self):
## @cond FLATBUFFERS_INTERNAL ## @cond FLATBUFFERS_INTERNAL
self.current_vtable = None self.current_vtable = None
self.head = UOffsetTFlags.py_type(len(self.Bytes)) self.head = len(self.Bytes)
self.minalign = 1 self.minalign = 1
self.objectEnd = None self.objectEnd = None
self.vtables = {} self.vtables = {}
@@ -192,7 +192,7 @@ class Builder(object):
if not self.finished: if not self.finished:
raise BuilderNotFinishedError() raise BuilderNotFinishedError()
return self.Bytes[self.Head() :] return self.Bytes[self.head :]
## @cond FLATBUFFERS_INTERNAL ## @cond FLATBUFFERS_INTERNAL
def StartObject(self, numfields): def StartObject(self, numfields):
@@ -321,7 +321,7 @@ class Builder(object):
self.nested = False self.nested = False
return self.WriteVtable() return self.WriteVtable()
def growByteBuffer(self): def GrowByteBuffer(self):
"""Doubles the size of the byteslice, and copies the old data towards """Doubles the size of the byteslice, and copies the old data towards
the end of the new buffer (since we build the buffer backwards). the end of the new buffer (since we build the buffer backwards).
@@ -352,12 +352,15 @@ class Builder(object):
## @cond FLATBUFFERS_INTERNAL ## @cond FLATBUFFERS_INTERNAL
def Offset(self): def Offset(self):
"""Offset relative to the end of the buffer.""" """Offset relative to the end of the buffer."""
return UOffsetTFlags.py_type(len(self.Bytes) - self.Head()) return len(self.Bytes) - self.head
def Pad(self, n): def Pad(self, n):
"""Pad places zeros at the current offset.""" """Pad places zeros at the current offset."""
for i in range_func(n): if n <= 0:
self.Place(0, N.Uint8Flags) return
new_head = self.head - n
self.Bytes[new_head : self.head] = b"\x00" * n
self.head = new_head
def Prep(self, size, additionalBytes): def Prep(self, size, additionalBytes):
"""Prep prepares to write an element of `size` after `additional_bytes` """Prep prepares to write an element of `size` after `additional_bytes`
@@ -374,15 +377,19 @@ class Builder(object):
# Find the amount of alignment needed such that `size` is properly # Find the amount of alignment needed such that `size` is properly
# aligned after `additionalBytes`: # aligned after `additionalBytes`:
alignSize = (~(len(self.Bytes) - self.Head() + additionalBytes)) + 1 head = self.head
buf_len = len(self.Bytes)
alignSize = (~(buf_len - head + additionalBytes)) + 1
alignSize &= size - 1 alignSize &= size - 1
# Reallocate the buffer if needed: # Reallocate the buffer if needed:
while self.Head() < alignSize + size + additionalBytes: needed = alignSize + size + additionalBytes
oldBufSize = len(self.Bytes) while head < needed:
self.growByteBuffer() oldBufSize = buf_len
updated_head = self.head + len(self.Bytes) - oldBufSize self.GrowByteBuffer()
self.head = UOffsetTFlags.py_type(updated_head) buf_len = len(self.Bytes)
head += buf_len - oldBufSize
self.head = head
self.Pad(alignSize) self.Pad(alignSize)
def PrependSOffsetTRelative(self, off): def PrependSOffsetTRelative(self, off):
@@ -482,16 +489,15 @@ class Builder(object):
else: else:
raise TypeError("non-string passed to CreateString") raise TypeError("non-string passed to CreateString")
self.Prep(N.UOffsetTFlags.bytewidth, (len(x) + 1) * N.Uint8Flags.bytewidth) payload_len = len(x)
self.Prep(N.UOffsetTFlags.bytewidth, (payload_len + 1) * N.Uint8Flags.bytewidth)
self.Place(0, N.Uint8Flags) self.Place(0, N.Uint8Flags)
l = UOffsetTFlags.py_type(len(s)) new_head = self.head - payload_len
## @cond FLATBUFFERS_INTERNAL self.head = new_head
self.head = UOffsetTFlags.py_type(self.Head() - l) self.Bytes[new_head : new_head + payload_len] = x
## @endcond
self.Bytes[self.Head() : self.Head() + l] = x
self.vectorNumElems = len(x) self.vectorNumElems = payload_len
return self.EndVector() return self.EndVector()
def CreateByteVector(self, x): def CreateByteVector(self, x):
@@ -505,15 +511,13 @@ class Builder(object):
if not isinstance(x, compat.binary_types): if not isinstance(x, compat.binary_types):
raise TypeError("non-byte vector passed to CreateByteVector") raise TypeError("non-byte vector passed to CreateByteVector")
self.Prep(N.UOffsetTFlags.bytewidth, len(x) * N.Uint8Flags.bytewidth) data_len = len(x)
self.Prep(N.UOffsetTFlags.bytewidth, data_len * N.Uint8Flags.bytewidth)
new_head = self.head - data_len
self.head = new_head
self.Bytes[new_head : new_head + data_len] = x
l = UOffsetTFlags.py_type(len(x)) self.vectorNumElems = data_len
## @cond FLATBUFFERS_INTERNAL
self.head = UOffsetTFlags.py_type(self.Head() - l)
## @endcond
self.Bytes[self.Head() : self.Head() + l] = x
self.vectorNumElems = len(x)
return self.EndVector() return self.EndVector()
def CreateNumpyVector(self, x): def CreateNumpyVector(self, x):
@@ -540,14 +544,14 @@ class Builder(object):
else: else:
x_lend = x.byteswap(inplace=False) x_lend = x.byteswap(inplace=False)
# Calculate total length
l = UOffsetTFlags.py_type(x_lend.itemsize * x_lend.size)
## @cond FLATBUFFERS_INTERNAL
self.head = UOffsetTFlags.py_type(self.Head() - l)
## @endcond
# tobytes ensures c_contiguous ordering # tobytes ensures c_contiguous ordering
self.Bytes[self.Head() : self.Head() + l] = x_lend.tobytes(order="C") payload = x_lend.tobytes(order="C")
# Calculate total length
payload_len = len(payload)
new_head = self.head - payload_len
self.head = new_head
self.Bytes[new_head : new_head + payload_len] = payload
self.vectorNumElems = x.size self.vectorNumElems = x.size
return self.EndVector() return self.EndVector()
@@ -617,11 +621,11 @@ class Builder(object):
self.PrependUOffsetTRelative(rootTable) self.PrependUOffsetTRelative(rootTable)
if sizePrefix: if sizePrefix:
size = len(self.Bytes) - self.Head() size = len(self.Bytes) - self.head
N.enforce_number(size, N.Int32Flags) N.enforce_number(size, N.Int32Flags)
self.PrependInt32(size) self.PrependInt32(size)
self.finished = True self.finished = True
return self.Head() return self.head
def Finish(self, rootTable, file_identifier=None): def Finish(self, rootTable, file_identifier=None):
"""Finish finalizes a buffer, pointing to the given `rootTable`.""" """Finish finalizes a buffer, pointing to the given `rootTable`."""
@@ -815,8 +819,9 @@ class Builder(object):
""" """
N.enforce_number(x, flags) N.enforce_number(x, flags)
self.head = self.head - flags.bytewidth new_head = self.head - flags.bytewidth
encode.Write(flags.packer_type, self.Bytes, self.Head(), x) self.head = new_head
encode.Write(flags.packer_type, self.Bytes, new_head, x)
def PlaceVOffsetT(self, x): def PlaceVOffsetT(self, x):
"""PlaceVOffsetT prepends a VOffsetT to the Builder, without checking """PlaceVOffsetT prepends a VOffsetT to the Builder, without checking
@@ -824,8 +829,9 @@ class Builder(object):
for space. for space.
""" """
N.enforce_number(x, N.VOffsetTFlags) N.enforce_number(x, N.VOffsetTFlags)
self.head = self.head - N.VOffsetTFlags.bytewidth new_head = self.head - N.VOffsetTFlags.bytewidth
encode.Write(packer.voffset, self.Bytes, self.Head(), x) self.head = new_head
encode.Write(packer.voffset, self.Bytes, new_head, x)
def PlaceSOffsetT(self, x): def PlaceSOffsetT(self, x):
"""PlaceSOffsetT prepends a SOffsetT to the Builder, without checking """PlaceSOffsetT prepends a SOffsetT to the Builder, without checking
@@ -833,8 +839,9 @@ class Builder(object):
for space. for space.
""" """
N.enforce_number(x, N.SOffsetTFlags) N.enforce_number(x, N.SOffsetTFlags)
self.head = self.head - N.SOffsetTFlags.bytewidth new_head = self.head - N.SOffsetTFlags.bytewidth
encode.Write(packer.soffset, self.Bytes, self.Head(), x) self.head = new_head
encode.Write(packer.soffset, self.Bytes, new_head, x)
def PlaceUOffsetT(self, x): def PlaceUOffsetT(self, x):
"""PlaceUOffsetT prepends a UOffsetT to the Builder, without checking """PlaceUOffsetT prepends a UOffsetT to the Builder, without checking
@@ -842,33 +849,10 @@ class Builder(object):
for space. for space.
""" """
N.enforce_number(x, N.UOffsetTFlags) N.enforce_number(x, N.UOffsetTFlags)
self.head = self.head - N.UOffsetTFlags.bytewidth new_head = self.head - N.UOffsetTFlags.bytewidth
encode.Write(packer.uoffset, self.Bytes, self.Head(), x) self.head = new_head
encode.Write(packer.uoffset, self.Bytes, new_head, x)
## @endcond ## @endcond
## @cond FLATBUFFERS_INTERNAL
def vtableEqual(a, objectStart, b):
"""vtableEqual compares an unwritten vtable to a written vtable."""
N.enforce_number(objectStart, N.UOffsetTFlags)
if len(a) * N.VOffsetTFlags.bytewidth != len(b):
return False
for i, elem in enumerate(a):
x = encode.Get(packer.voffset, b, i * N.VOffsetTFlags.bytewidth)
# Skip vtable entries that indicate a default value.
if x == 0 and elem == 0:
pass
else:
y = objectStart - elem
if x != y:
return False
return True
## @endcond
## @} ## @}