Optimize Offset/Pad/Prep: use cached head and slicing, reduce casting

Author: satur9nine

python/flatbuffers/builder.py

@@ -166,7 +166,7 @@ def __init__(self, initialSize=1024):
   def Clear(self):
     ## @cond FLATBUFFERS_INTERNAL
     self.current_vtable = None
-    self.head = UOffsetTFlags.py_type(len(self.Bytes))
+    self.head = len(self.Bytes)
     self.minalign = 1
     self.objectEnd = None
     self.vtables = {}
@@ -192,7 +192,7 @@ def Output(self):
     if not self.finished:
       raise BuilderNotFinishedError()
 
-    return self.Bytes[self.Head() :]
+    return self.Bytes[self.head :]
 
   ## @cond FLATBUFFERS_INTERNAL
   def StartObject(self, numfields):
@@ -321,7 +321,7 @@ def EndObject(self):
     self.nested = False
     return self.WriteVtable()
 
-  def growByteBuffer(self):
+  def GrowByteBuffer(self):
     """Doubles the size of the byteslice, and copies the old data towards
 
     the end of the new buffer (since we build the buffer backwards).
@@ -352,12 +352,15 @@ def Head(self):
   ## @cond FLATBUFFERS_INTERNAL
   def Offset(self):
     """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):
     """Pad places zeros at the current offset."""
-    for i in range_func(n):
-      self.Place(0, N.Uint8Flags)
+    if n <= 0:
+      return
+    new_head = self.head - n
+    self.Bytes[new_head : self.head] = b"\x00" * n
+    self.head = new_head
 
   def Prep(self, size, additionalBytes):
     """Prep prepares to write an element of `size` after `additional_bytes`
@@ -374,15 +377,19 @@ def Prep(self, size, additionalBytes):
 
     # Find the amount of alignment needed such that `size` is properly
     # 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
 
     # Reallocate the buffer if needed:
-    while self.Head() < alignSize + size + additionalBytes:
-      oldBufSize = len(self.Bytes)
-      self.growByteBuffer()
-      updated_head = self.head + len(self.Bytes) - oldBufSize
-      self.head = UOffsetTFlags.py_type(updated_head)
+    needed = alignSize + size + additionalBytes
+    while head < needed:
+      oldBufSize = buf_len
+      self.GrowByteBuffer()
+      buf_len = len(self.Bytes)
+      head += buf_len - oldBufSize
+    self.head = head
     self.Pad(alignSize)
 
   def PrependSOffsetTRelative(self, off):
@@ -482,16 +489,15 @@ def CreateString(self, s, encoding="utf-8", errors="strict"):
     else:
       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)
 
-    l = UOffsetTFlags.py_type(len(s))
-    ## @cond FLATBUFFERS_INTERNAL
-    self.head = UOffsetTFlags.py_type(self.Head() - l)
-    ## @endcond
-    self.Bytes[self.Head() : self.Head() + l] = x
+    new_head = self.head - payload_len
+    self.head = new_head
+    self.Bytes[new_head : new_head + payload_len] = x
 
-    self.vectorNumElems = len(x)
+    self.vectorNumElems = payload_len
     return self.EndVector()
 
   def CreateByteVector(self, x):
@@ -505,15 +511,13 @@ def CreateByteVector(self, x):
     if not isinstance(x, compat.binary_types):
       raise TypeError("non-byte vector passed to CreateByteVector")
 
-    self.Prep(N.UOffsetTFlags.bytewidth, len(x) * N.Uint8Flags.bytewidth)
-
-    l = UOffsetTFlags.py_type(len(x))
-    ## @cond FLATBUFFERS_INTERNAL
-    self.head = UOffsetTFlags.py_type(self.Head() - l)
-    ## @endcond
-    self.Bytes[self.Head() : self.Head() + l] = x
+    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
 
-    self.vectorNumElems = len(x)
+    self.vectorNumElems = data_len
     return self.EndVector()
 
   def CreateNumpyVector(self, x):
@@ -540,14 +544,14 @@ def CreateNumpyVector(self, x):
     else:
       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
-    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
     return self.EndVector()
@@ -617,11 +621,11 @@ def __Finish(self, rootTable, sizePrefix, file_identifier=None):
 
     self.PrependUOffsetTRelative(rootTable)
     if sizePrefix:
-      size = len(self.Bytes) - self.Head()
+      size = len(self.Bytes) - self.head
       N.enforce_number(size, N.Int32Flags)
       self.PrependInt32(size)
     self.finished = True
-    return self.Head()
+    return self.head
 
   def Finish(self, rootTable, file_identifier=None):
     """Finish finalizes a buffer, pointing to the given `rootTable`."""
@@ -815,35 +819,39 @@ def Place(self, x, flags):
     """
 
     N.enforce_number(x, flags)
-    self.head = self.head - flags.bytewidth
-    encode.Write(flags.packer_type, self.Bytes, self.Head(), x)
+    new_head = self.head - flags.bytewidth
+    self.head = new_head
+    encode.Write(flags.packer_type, self.Bytes, new_head, x)
 
   def PlaceVOffsetT(self, x):
     """PlaceVOffsetT prepends a VOffsetT to the Builder, without checking
 
     for space.
     """
     N.enforce_number(x, N.VOffsetTFlags)
-    self.head = self.head - N.VOffsetTFlags.bytewidth
-    encode.Write(packer.voffset, self.Bytes, self.Head(), x)
+    new_head = self.head - N.VOffsetTFlags.bytewidth
+    self.head = new_head
+    encode.Write(packer.voffset, self.Bytes, new_head, x)
 
   def PlaceSOffsetT(self, x):
     """PlaceSOffsetT prepends a SOffsetT to the Builder, without checking
 
     for space.
     """
     N.enforce_number(x, N.SOffsetTFlags)
-    self.head = self.head - N.SOffsetTFlags.bytewidth
-    encode.Write(packer.soffset, self.Bytes, self.Head(), x)
+    new_head = self.head - N.SOffsetTFlags.bytewidth
+    self.head = new_head
+    encode.Write(packer.soffset, self.Bytes, new_head, x)
 
   def PlaceUOffsetT(self, x):
     """PlaceUOffsetT prepends a UOffsetT to the Builder, without checking
 
     for space.
     """
     N.enforce_number(x, N.UOffsetTFlags)
-    self.head = self.head - N.UOffsetTFlags.bytewidth
-    encode.Write(packer.uoffset, self.Bytes, self.Head(), x)
+    new_head = self.head - N.UOffsetTFlags.bytewidth
+    self.head = new_head
+    encode.Write(packer.uoffset, self.Bytes, new_head, x)
 
   ## @endcond