# Lint as: python3
# Copyright 2020 Google Inc. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Implementation of FlexBuffers binary format.

For more info check https://google.github.io/flatbuffers/flexbuffers.html and
corresponding C++ implementation at
https://github.com/google/flatbuffers/blob/master/include/flatbuffers/flexbuffers.h
"""

# pylint: disable=invalid-name
# TODO(dkovalev): Add type hints everywhere, so tools like pytypes could work.

import array
import contextlib
import enum
import struct

__all__ = ('Type', 'Builder', 'GetRoot', 'Dumps', 'Loads')


class BitWidth(enum.IntEnum):
  """Supported bit widths of value types.

  These are used in the lower 2 bits of a type field to determine the size of
  the elements (and or size field) of the item pointed to (e.g. vector).
  """
  W8 = 0  # 2^0 = 1 byte
  W16 = 1  # 2^1 = 2 bytes
  W32 = 2  # 2^2 = 4 bytes
  W64 = 3  # 2^3 = 8 bytes

  @staticmethod
  def U(value):
    """Returns the minimum `BitWidth` to encode unsigned integer value."""
    assert value >= 0

    if value < (1 << 8):
      return BitWidth.W8
    elif value < (1 << 16):
      return BitWidth.W16
    elif value < (1 << 32):
      return BitWidth.W32
    elif value < (1 << 64):
      return BitWidth.W64
    else:
      raise ValueError('value is too big to encode: %s' % value)

  @staticmethod
  def I(value):
    """Returns the minimum `BitWidth` to encode signed integer value."""
    # -2^(n-1) <=     value < 2^(n-1)
    # -2^n     <= 2 * value < 2^n
    # 2 * value < 2^n, when value >= 0 or 2 * (-value) <= 2^n, when value < 0
    # 2 * value < 2^n, when value >= 0 or 2 * (-value) - 1 < 2^n, when value < 0
    #
    # if value >= 0:
    #   return BitWidth.U(2 * value)
    # else:
    #   return BitWidth.U(2 * (-value) - 1)  # ~x = -x - 1
    value *= 2
    return BitWidth.U(value if value >= 0 else ~value)

  @staticmethod
  def F(value):
    """Returns the `BitWidth` to encode floating point value."""
    if struct.unpack('f', struct.pack('f', value))[0] == value:
      return BitWidth.W32
    return BitWidth.W64

  @staticmethod
  def B(byte_width):
    return {
        1: BitWidth.W8,
        2: BitWidth.W16,
        4: BitWidth.W32,
        8: BitWidth.W64
    }[byte_width]


I = {1: 'b', 2: 'h', 4: 'i', 8: 'q'}  # Integer formats
U = {1: 'B', 2: 'H', 4: 'I', 8: 'Q'}  # Unsigned integer formats
F = {4: 'f', 8: 'd'}  # Floating point formats


def _Unpack(fmt, buf):
  return struct.unpack(fmt[len(buf)], buf)[0]


def _UnpackVector(fmt, buf, length):
  byte_width = len(buf) // length
  return struct.unpack('%d%s' % (length, fmt[byte_width]), buf)


def _Pack(fmt, value, byte_width):
  return struct.pack(fmt[byte_width], value)


def _PackVector(fmt, values, byte_width):
  return struct.pack('%d%s' % (len(values), fmt[byte_width]), *values)


def _Mutate(fmt, buf, value, byte_width, value_bit_width):
  if (1 << value_bit_width) <= byte_width:
    buf[:byte_width] = _Pack(fmt, value, byte_width)
    return True
  return False


# Computes how many bytes you'd have to pad to be able to write an
# "scalar_size" scalar if the buffer had grown to "buf_size",
# "scalar_size" is a power of two.
def _PaddingBytes(buf_size, scalar_size):
  # ((buf_size + (scalar_size - 1)) // scalar_size) * scalar_size - buf_size
  return -buf_size & (scalar_size - 1)


def _ShiftSlice(s, offset, length):
  start = offset + (0 if s.start is None else s.start)
  stop = offset + (length if s.stop is None else s.stop)
  return slice(start, stop, s.step)


# https://en.cppreference.com/w/cpp/algorithm/lower_bound
def _LowerBound(values, value, pred):
  """Implementation of C++ std::lower_bound() algorithm."""
  first, last = 0, len(values)
  count = last - first
  while count > 0:
    i = first
    step = count // 2
    i += step
    if pred(values[i], value):
      i += 1
      first = i
      count -= step + 1
    else:
      count = step
  return first


# https://en.cppreference.com/w/cpp/algorithm/binary_search
def _BinarySearch(values, value, pred=lambda x, y: x < y):
  """Implementation of C++ std::binary_search() algorithm."""
  index = _LowerBound(values, value, pred)
  if index != len(values) and not pred(value, values[index]):
    return index
  return -1


class Type(enum.IntEnum):
  """Supported types of encoded data.

  These are used as the upper 6 bits of a type field to indicate the actual
  type.
  """
  NULL = 0
  INT = 1
  UINT = 2
  FLOAT = 3
  # Types above stored inline, types below store an offset.
  KEY = 4
  STRING = 5
  INDIRECT_INT = 6
  INDIRECT_UINT = 7
  INDIRECT_FLOAT = 8
  MAP = 9
  VECTOR = 10  # Untyped.

  VECTOR_INT = 11  # Typed any size (stores no type table).
  VECTOR_UINT = 12
  VECTOR_FLOAT = 13
  VECTOR_KEY = 14
  # DEPRECATED, use VECTOR or VECTOR_KEY instead.
  # Read test.cpp/FlexBuffersDeprecatedTest() for details on why.
  VECTOR_STRING_DEPRECATED = 15

  VECTOR_INT2 = 16  # Typed tuple (no type table, no size field).
  VECTOR_UINT2 = 17
  VECTOR_FLOAT2 = 18
  VECTOR_INT3 = 19  # Typed triple (no type table, no size field).
  VECTOR_UINT3 = 20
  VECTOR_FLOAT3 = 21
  VECTOR_INT4 = 22  # Typed quad (no type table, no size field).
  VECTOR_UINT4 = 23
  VECTOR_FLOAT4 = 24

  BLOB = 25
  BOOL = 26
  VECTOR_BOOL = 36  # To do the same type of conversion of type to vector type

  @staticmethod
  def Pack(type_, bit_width):
    return (int(type_) << 2) | bit_width

  @staticmethod
  def Unpack(packed_type):
    return 1 << (packed_type & 0b11), Type(packed_type >> 2)

  @staticmethod
  def IsInline(type_):
    return type_ <= Type.FLOAT or type_ == Type.BOOL

  @staticmethod
  def IsTypedVector(type_):
    return Type.VECTOR_INT <= type_ <= Type.VECTOR_STRING_DEPRECATED or \
           type_ == Type.VECTOR_BOOL

  @staticmethod
  def IsTypedVectorElementType(type_):
    return Type.INT <= type_ <= Type.STRING or type_ == Type.BOOL

  @staticmethod
  def ToTypedVectorElementType(type_):
    if not Type.IsTypedVector(type_):
      raise ValueError('must be typed vector type')

    return Type(type_ - Type.VECTOR_INT + Type.INT)

  @staticmethod
  def IsFixedTypedVector(type_):
    return Type.VECTOR_INT2 <= type_ <= Type.VECTOR_FLOAT4

  @staticmethod
  def IsFixedTypedVectorElementType(type_):
    return Type.INT <= type_ <= Type.FLOAT

  @staticmethod
  def ToFixedTypedVectorElementType(type_):
    if not Type.IsFixedTypedVector(type_):
      raise ValueError('must be fixed typed vector type')

    # 3 types each, starting from length 2.
    fixed_type = type_ - Type.VECTOR_INT2
    return Type(fixed_type % 3 + Type.INT), fixed_type // 3 + 2

  @staticmethod
  def ToTypedVector(element_type, fixed_len=0):
    """Converts element type to corresponding vector type.

    Args:
      element_type: vector element type
      fixed_len: number of elements: 0 for typed vector; 2, 3, or 4 for fixed
        typed vector.

    Returns:
      Typed vector type or fixed typed vector type.
    """
    if fixed_len == 0:
      if not Type.IsTypedVectorElementType(element_type):
        raise ValueError('must be typed vector element type')
    else:
      if not Type.IsFixedTypedVectorElementType(element_type):
        raise ValueError('must be fixed typed vector element type')

    offset = element_type - Type.INT
    if fixed_len == 0:
      return Type(offset + Type.VECTOR_INT)  # TypedVector
    elif fixed_len == 2:
      return Type(offset + Type.VECTOR_INT2)  # FixedTypedVector
    elif fixed_len == 3:
      return Type(offset + Type.VECTOR_INT3)  # FixedTypedVector
    elif fixed_len == 4:
      return Type(offset + Type.VECTOR_INT4)  # FixedTypedVector
    else:
      raise ValueError('unsupported fixed_len: %s' % fixed_len)


class Buf:
  """Class to access underlying buffer object starting from the given offset."""

  def __init__(self, buf, offset):
    self._buf = buf
    self._offset = offset if offset >= 0 else len(buf) + offset
    self._length = len(buf) - self._offset

  def __getitem__(self, key):
    if isinstance(key, slice):
      return self._buf[_ShiftSlice(key, self._offset, self._length)]
    elif isinstance(key, int):
      return self._buf[self._offset + key]
    else:
      raise TypeError('invalid key type')

  def __setitem__(self, key, value):
    if isinstance(key, slice):
      self._buf[_ShiftSlice(key, self._offset, self._length)] = value
    elif isinstance(key, int):
      self._buf[self._offset + key] = key
    else:
      raise TypeError('invalid key type')

  def __repr__(self):
    return 'buf[%d:]' % self._offset

  def Find(self, sub):
    """Returns the lowest index where the sub subsequence is found."""
    return self._buf[self._offset:].find(sub)

  def Slice(self, offset):
    """Returns new `Buf` which starts from the given offset."""
    return Buf(self._buf, self._offset + offset)

  def Indirect(self, offset, byte_width):
    """Return new `Buf` based on the encoded offset (indirect encoding)."""
    return self.Slice(offset - _Unpack(U, self[offset:offset + byte_width]))


class Object:
  """Base class for all non-trivial data accessors."""
  __slots__ = '_buf', '_byte_width'

  def __init__(self, buf, byte_width):
    self._buf = buf
    self._byte_width = byte_width

  @property
  def ByteWidth(self):
    return self._byte_width


class Sized(Object):
  """Base class for all data accessors which need to read encoded size."""
  __slots__ = '_size',

  def __init__(self, buf, byte_width, size=0):
    super().__init__(buf, byte_width)
    if size == 0:
      self._size = _Unpack(U, self.SizeBytes)
    else:
      self._size = size

  @property
  def SizeBytes(self):
    return self._buf[-self._byte_width:0]

  def __len__(self):
    return self._size


class Blob(Sized):
  """Data accessor for the encoded blob bytes."""
  __slots__ = ()

  @property
  def Bytes(self):
    return self._buf[0:len(self)]

  def __repr__(self):
    return 'Blob(%s, size=%d)' % (self._buf, len(self))


class String(Sized):
  """Data accessor for the encoded string bytes."""
  __slots__ = ()

  @property
  def Bytes(self):
    return self._buf[0:len(self)]

  def Mutate(self, value):
    """Mutates underlying string bytes in place.

    Args:
      value: New string to replace the existing one. New string must have less
        or equal UTF-8-encoded bytes than the existing one to successfully
        mutate underlying byte buffer.

    Returns:
      Whether the value was mutated or not.
    """
    encoded = value.encode('utf-8')
    n = len(encoded)
    if n <= len(self):
      self._buf[-self._byte_width:0] = _Pack(U, n, self._byte_width)
      self._buf[0:n] = encoded
      self._buf[n:len(self)] = bytearray(len(self) - n)
      return True
    return False

  def __str__(self):
    return self.Bytes.decode('utf-8')

  def __repr__(self):
    return 'String(%s, size=%d)' % (self._buf, len(self))


class Key(Object):
  """Data accessor for the encoded key bytes."""
  __slots__ = ()

  def __init__(self, buf, byte_width):
    assert byte_width == 1
    super().__init__(buf, byte_width)

  @property
  def Bytes(self):
    return self._buf[0:len(self)]

  def __len__(self):
    return self._buf.Find(0)

  def __str__(self):
    return self.Bytes.decode('ascii')

  def __repr__(self):
    return 'Key(%s, size=%d)' % (self._buf, len(self))


class Vector(Sized):
  """Data accessor for the encoded vector bytes."""
  __slots__ = ()

  def __getitem__(self, index):
    if index < 0 or index >= len(self):
      raise IndexError('vector index %s is out of [0, %d) range' % \
          (index, len(self)))

    packed_type = self._buf[len(self) * self._byte_width + index]
    buf = self._buf.Slice(index * self._byte_width)
    return Ref.PackedType(buf, self._byte_width, packed_type)

  @property
  def Value(self):
    """Returns the underlying encoded data as a list object."""
    return [e.Value for e in self]

  def __repr__(self):
    return 'Vector(%s, byte_width=%d, size=%d)' % \
        (self._buf, self._byte_width, self._size)


class TypedVector(Sized):
  """Data accessor for the encoded typed vector or fixed typed vector bytes."""
  __slots__ = '_element_type', '_size'

  def __init__(self, buf, byte_width, element_type, size=0):
    super().__init__(buf, byte_width, size)

    if element_type == Type.STRING:
      # These can't be accessed as strings, since we don't know the bit-width
      # of the size field, see the declaration of
      # FBT_VECTOR_STRING_DEPRECATED above for details.
      # We change the type here to be keys, which are a subtype of strings,
      # and will ignore the size field. This will truncate strings with
      # embedded nulls.
      element_type = Type.KEY

    self._element_type = element_type

  @property
  def Bytes(self):
    return self._buf[:self._byte_width * len(self)]

  @property
  def ElementType(self):
    return self._element_type

  def __getitem__(self, index):
    if index < 0 or index >= len(self):
      raise IndexError('vector index %s is out of [0, %d) range' % \
          (index, len(self)))

    buf = self._buf.Slice(index * self._byte_width)
    return Ref(buf, self._byte_width, 1, self._element_type)

  @property
  def Value(self):
    """Returns underlying data as list object."""
    if not self:
      return []

    if self._element_type is Type.BOOL:
      return [bool(e) for e in _UnpackVector(U, self.Bytes, len(self))]
    elif self._element_type is Type.INT:
      return list(_UnpackVector(I, self.Bytes, len(self)))
    elif self._element_type is Type.UINT:
      return list(_UnpackVector(U, self.Bytes, len(self)))
    elif self._element_type is Type.FLOAT:
      return list(_UnpackVector(F, self.Bytes, len(self)))
    elif self._element_type is Type.KEY:
      return [e.AsKey for e in self]
    elif self._element_type is Type.STRING:
      return [e.AsString for e in self]
    else:
      raise TypeError('unsupported element_type: %s' % self._element_type)

  def __repr__(self):
    return 'TypedVector(%s, byte_width=%d, element_type=%s, size=%d)' % \
        (self._buf, self._byte_width, self._element_type, self._size)


class Map(Vector):
  """Data accessor for the encoded map bytes."""

  @staticmethod
  def CompareKeys(a, b):
    if isinstance(a, Ref):
      a = a.AsKeyBytes
    if isinstance(b, Ref):
      b = b.AsKeyBytes
    return a < b

  def __getitem__(self, key):
    if isinstance(key, int):
      return super().__getitem__(key)

    index = _BinarySearch(self.Keys, key.encode('ascii'), self.CompareKeys)
    if index != -1:
      return super().__getitem__(index)

    raise KeyError(key)

  @property
  def Keys(self):
    byte_width = _Unpack(U, self._buf[-2 * self._byte_width:-self._byte_width])
    buf = self._buf.Indirect(-3 * self._byte_width, self._byte_width)
    return TypedVector(buf, byte_width, Type.KEY)

  @property
  def Values(self):
    return Vector(self._buf, self._byte_width)

  @property
  def Value(self):
    return {k.Value: v.Value for k, v in zip(self.Keys, self.Values)}

  def __repr__(self):
    return 'Map(%s, size=%d)' % (self._buf, len(self))


class Ref:
  """Data accessor for the encoded data bytes."""
  __slots__ = '_buf', '_parent_width', '_byte_width', '_type'

  @staticmethod
  def PackedType(buf, parent_width, packed_type):
    byte_width, type_ = Type.Unpack(packed_type)
    return Ref(buf, parent_width, byte_width, type_)

  def __init__(self, buf, parent_width, byte_width, type_):
    self._buf = buf
    self._parent_width = parent_width
    self._byte_width = byte_width
    self._type = type_

  def __repr__(self):
    return 'Ref(%s, parent_width=%d, byte_width=%d, type_=%s)' % \
            (self._buf, self._parent_width, self._byte_width, self._type)

  @property
  def _Bytes(self):
    return self._buf[:self._parent_width]

  def _ConvertError(self, target_type):
    raise TypeError('cannot convert %s to %s' % (self._type, target_type))

  def _Indirect(self):
    return self._buf.Indirect(0, self._parent_width)

  @property
  def IsNull(self):
    return self._type is Type.NULL

  @property
  def IsBool(self):
    return self._type is Type.BOOL

  @property
  def AsBool(self):
    if self._type is Type.BOOL:
      return bool(_Unpack(U, self._Bytes))
    else:
      return self.AsInt != 0

  def MutateBool(self, value):
    """Mutates underlying boolean value bytes in place.

    Args:
      value: New boolean value.

    Returns:
      Whether the value was mutated or not.
    """
    return self.IsBool and \
           _Mutate(U, self._buf, value, self._parent_width, BitWidth.W8)

  @property
  def IsNumeric(self):
    return self.IsInt or self.IsFloat

  @property
  def IsInt(self):
    return self._type in (Type.INT, Type.INDIRECT_INT, Type.UINT,
                          Type.INDIRECT_UINT)

  @property
  def AsInt(self):
    """Returns current reference as integer value."""
    if self.IsNull:
      return 0
    elif self.IsBool:
      return int(self.AsBool)
    elif self._type is Type.INT:
      return _Unpack(I, self._Bytes)
    elif self._type is Type.INDIRECT_INT:
      return _Unpack(I, self._Indirect()[:self._byte_width])
    if self._type is Type.UINT:
      return _Unpack(U, self._Bytes)
    elif self._type is Type.INDIRECT_UINT:
      return _Unpack(U, self._Indirect()[:self._byte_width])
    elif self.IsString:
      return len(self.AsString)
    elif self.IsKey:
      return len(self.AsKey)
    elif self.IsBlob:
      return len(self.AsBlob)
    elif self.IsVector:
      return len(self.AsVector)
    elif self.IsTypedVector:
      return len(self.AsTypedVector)
    elif self.IsFixedTypedVector:
      return len(self.AsFixedTypedVector)
    else:
      raise self._ConvertError(Type.INT)

  def MutateInt(self, value):
    """Mutates underlying integer value bytes in place.

    Args:
      value: New integer value. It must fit to the byte size of the existing
        encoded value.

    Returns:
      Whether the value was mutated or not.
    """
    if self._type is Type.INT:
      return _Mutate(I, self._buf, value, self._parent_width, BitWidth.I(value))
    elif self._type is Type.INDIRECT_INT:
      return _Mutate(I, self._Indirect(), value, self._byte_width,
                     BitWidth.I(value))
    elif self._type is Type.UINT:
      return _Mutate(U, self._buf, value, self._parent_width, BitWidth.U(value))
    elif self._type is Type.INDIRECT_UINT:
      return _Mutate(U, self._Indirect(), value, self._byte_width,
                     BitWidth.U(value))
    else:
      return False

  @property
  def IsFloat(self):
    return self._type in (Type.FLOAT, Type.INDIRECT_FLOAT)

  @property
  def AsFloat(self):
    """Returns current reference as floating point value."""
    if self.IsNull:
      return 0.0
    elif self.IsBool:
      return float(self.AsBool)
    elif self.IsInt:
      return float(self.AsInt)
    elif self._type is Type.FLOAT:
      return _Unpack(F, self._Bytes)
    elif self._type is Type.INDIRECT_FLOAT:
      return _Unpack(F, self._Indirect()[:self._byte_width])
    elif self.IsString:
      return float(self.AsString)
    elif self.IsVector:
      return float(len(self.AsVector))
    elif self.IsTypedVector():
      return float(len(self.AsTypedVector))
    elif self.IsFixedTypedVector():
      return float(len(self.FixedTypedVector))
    else:
      raise self._ConvertError(Type.FLOAT)

  def MutateFloat(self, value):
    """Mutates underlying floating point value bytes in place.

    Args:
      value: New float value. It must fit to the byte size of the existing
        encoded value.

    Returns:
      Whether the value was mutated or not.
    """
    if self._type is Type.FLOAT:
      return _Mutate(F, self._buf, value, self._parent_width,
                     BitWidth.B(self._parent_width))
    elif self._type is Type.INDIRECT_FLOAT:
      return _Mutate(F, self._Indirect(), value, self._byte_width,
                     BitWidth.B(self._byte_width))
    else:
      return False

  @property
  def IsKey(self):
    return self._type is Type.KEY

  @property
  def AsKeyBytes(self):
    if self.IsKey:
      return Key(self._Indirect(), self._byte_width).Bytes
    else:
      raise self._ConvertError(Type.KEY)

  @property
  def AsKey(self):
    if self.IsKey:
      return str(Key(self._Indirect(), self._byte_width))
    else:
      raise self._ConvertError(Type.KEY)

  @property
  def IsString(self):
    return self._type is Type.STRING

  @property
  def AsString(self):
    if self.IsString:
      return str(String(self._Indirect(), self._byte_width))
    elif self.IsKey:
      return self.AsKey
    else:
      raise self._ConvertError(Type.STRING)

  def MutateString(self, value):
    return String(self._Indirect(), self._byte_width).Mutate(value)

  @property
  def IsBlob(self):
    return self._type is Type.BLOB

  @property
  def AsBlob(self):
    if self.IsBlob:
      return Blob(self._Indirect(), self._byte_width).Bytes
    else:
      raise self._ConvertError(Type.BLOB)

  @property
  def IsAnyVector(self):
    return self.IsVector or self.IsTypedVector or self.IsFixedTypedVector()

  @property
  def IsVector(self):
    return self._type in (Type.VECTOR, Type.MAP)

  @property
  def AsVector(self):
    if self.IsVector:
      return Vector(self._Indirect(), self._byte_width)
    else:
      raise self._ConvertError(Type.VECTOR)

  @property
  def IsTypedVector(self):
    return Type.IsTypedVector(self._type)

  @property
  def AsTypedVector(self):
    if self.IsTypedVector:
      return TypedVector(self._Indirect(), self._byte_width,
                         Type.ToTypedVectorElementType(self._type))
    else:
      raise self._ConvertError('TYPED_VECTOR')

  @property
  def IsFixedTypedVector(self):
    return Type.IsFixedTypedVector(self._type)

  @property
  def AsFixedTypedVector(self):
    if self.IsFixedTypedVector:
      element_type, size = Type.ToFixedTypedVectorElementType(self._type)
      return TypedVector(self._Indirect(), self._byte_width, element_type, size)
    else:
      raise self._ConvertError('FIXED_TYPED_VECTOR')

  @property
  def IsMap(self):
    return self._type is Type.MAP

  @property
  def AsMap(self):
    if self.IsMap:
      return Map(self._Indirect(), self._byte_width)
    else:
      raise self._ConvertError(Type.MAP)

  @property
  def Value(self):
    """Converts current reference to value of corresponding type.

    This is equivalent to calling `AsInt` for integer values, `AsFloat` for
    floating point values, etc.

    Returns:
      Value of corresponding type.
    """
    if self.IsNull:
      return None
    elif self.IsBool:
      return self.AsBool
    elif self.IsInt:
      return self.AsInt
    elif self.IsFloat:
      return self.AsFloat
    elif self.IsString:
      return self.AsString
    elif self.IsKey:
      return self.AsKey
    elif self.IsBlob:
      return self.AsBlob
    elif self.IsMap:
      return self.AsMap.Value
    elif self.IsVector:
      return self.AsVector.Value
    elif self.IsTypedVector:
      return self.AsTypedVector.Value
    elif self.IsFixedTypedVector:
      return self.AsFixedTypedVector.Value
    else:
      raise TypeError('cannot convert %r to value' % self)


def _IsIterable(obj):
  try:
    iter(obj)
    return True
  except TypeError:
    return False


class Value:
  """Class to represent given value during the encoding process."""

  @staticmethod
  def Null():
    return Value(0, Type.NULL, BitWidth.W8)

  @staticmethod
  def Bool(value):
    return Value(value, Type.BOOL, BitWidth.W8)

  @staticmethod
  def Int(value, bit_width):
    return Value(value, Type.INT, bit_width)

  @staticmethod
  def UInt(value, bit_width):
    return Value(value, Type.UINT, bit_width)

  @staticmethod
  def Float(value, bit_width):
    return Value(value, Type.FLOAT, bit_width)

  @staticmethod
  def Key(offset):
    return Value(offset, Type.KEY, BitWidth.W8)

  def __init__(self, value, type_, min_bit_width):
    self._value = value
    self._type = type_

    # For scalars: of itself, for vector: of its elements, for string: length.
    self._min_bit_width = min_bit_width

  @property
  def Value(self):
    return self._value

  @property
  def Type(self):
    return self._type

  @property
  def MinBitWidth(self):
    return self._min_bit_width

  def StoredPackedType(self, parent_bit_width=BitWidth.W8):
    return Type.Pack(self._type, self.StoredWidth(parent_bit_width))

  # We have an absolute offset, but want to store a relative offset
  # elem_index elements beyond the current buffer end. Since whether
  # the relative offset fits in a certain byte_width depends on
  # the size of the elements before it (and their alignment), we have
  # to test for each size in turn.
  def ElemWidth(self, buf_size, elem_index=0):
    if Type.IsInline(self._type):
      return self._min_bit_width
    for byte_width in 1, 2, 4, 8:
      offset_loc = buf_size + _PaddingBytes(buf_size, byte_width) + \
                   elem_index * byte_width
      bit_width = BitWidth.U(offset_loc - self._value)
      if byte_width == (1 << bit_width):
        return bit_width
    raise ValueError('relative offset is too big')

  def StoredWidth(self, parent_bit_width=BitWidth.W8):
    if Type.IsInline(self._type):
      return max(self._min_bit_width, parent_bit_width)
    return self._min_bit_width

  def __repr__(self):
    return 'Value(%s, %s, %s)' % (self._value, self._type, self._min_bit_width)

  def __str__(self):
    return str(self._value)


def InMap(func):
  def wrapper(self, *args, **kwargs):
    if isinstance(args[0], str):
      self.Key(args[0])
      func(self, *args[1:], **kwargs)
    else:
      func(self, *args, **kwargs)
  return wrapper


def InMapForString(func):
  def wrapper(self, *args):
    if len(args) == 1:
      func(self, args[0])
    elif len(args) == 2:
      self.Key(args[0])
      func(self, args[1])
    else:
      raise ValueError('invalid number of arguments')
  return wrapper


class Pool:
  """Collection of (data, offset) pairs sorted by data for quick access."""

  def __init__(self):
    self._pool = []  # sorted list of (data, offset) tuples

  def FindOrInsert(self, data, offset):
    do = data, offset
    index = _BinarySearch(self._pool, do, lambda a, b: a[0] < b[0])
    if index != -1:
      _, offset = self._pool[index]
      return offset
    self._pool.insert(index, do)
    return None

  def Clear(self):
    self._pool = []

  @property
  def Elements(self):
    return [data for data, _ in self._pool]


class Builder:
  """Helper class to encode structural data into flexbuffers format."""

  def __init__(self,
               share_strings=False,
               share_keys=True,
               force_min_bit_width=BitWidth.W8):
    self._share_strings = share_strings
    self._share_keys = share_keys
    self._force_min_bit_width = force_min_bit_width

    self._string_pool = Pool()
    self._key_pool = Pool()

    self._finished = False
    self._buf = bytearray()
    self._stack = []

  def __len__(self):
    return len(self._buf)

  @property
  def StringPool(self):
    return self._string_pool

  @property
  def KeyPool(self):
    return self._key_pool

  def Clear(self):
    self._string_pool.Clear()
    self._key_pool.Clear()
    self._finished = False
    self._buf = bytearray()
    self._stack = []

  def Finish(self):
    """Finishes encoding process and returns underlying buffer."""
    if self._finished:
      raise RuntimeError('builder has been already finished')

    # If you hit this exception, you likely have objects that were never
    # included in a parent. You need to have exactly one root to finish a
    # buffer. Check your Start/End calls are matched, and all objects are inside
    # some other object.
    if len(self._stack) != 1:
      raise RuntimeError('internal stack size must be one')

    value = self._stack[0]
    byte_width = self._Align(value.ElemWidth(len(self._buf)))
    self._WriteAny(value, byte_width=byte_width)  # Root value
    self._Write(U, value.StoredPackedType(), byte_width=1)  # Root type
    self._Write(U, byte_width, byte_width=1)  # Root size

    self.finished = True
    return self._buf

  def _ReadKey(self, offset):
    key = self._buf[offset:]
    return key[:key.find(0)]

  def _Align(self, alignment):
    byte_width = 1 << alignment
    self._buf.extend(b'\x00' * _PaddingBytes(len(self._buf), byte_width))
    return byte_width

  def _Write(self, fmt, value, byte_width):
    self._buf.extend(_Pack(fmt, value, byte_width))

  def _WriteVector(self, fmt, values, byte_width):
    self._buf.extend(_PackVector(fmt, values, byte_width))

  def _WriteOffset(self, offset, byte_width):
    relative_offset = len(self._buf) - offset
    assert byte_width == 8 or relative_offset < (1 << (8 * byte_width))
    self._Write(U, relative_offset, byte_width)

  def _WriteAny(self, value, byte_width):
    fmt = {
        Type.NULL: U, Type.BOOL: U, Type.INT: I, Type.UINT: U, Type.FLOAT: F
    }.get(value.Type)
    if fmt:
      self._Write(fmt, value.Value, byte_width)
    else:
      self._WriteOffset(value.Value, byte_width)

  def _WriteBlob(self, data, append_zero, type_):
    bit_width = BitWidth.U(len(data))
    byte_width = self._Align(bit_width)
    self._Write(U, len(data), byte_width)
    loc = len(self._buf)
    self._buf.extend(data)
    if append_zero:
      self._buf.append(0)
    self._stack.append(Value(loc, type_, bit_width))
    return loc

  def _WriteScalarVector(self, element_type, byte_width, elements, fixed):
    """Writes scalar vector elements to the underlying buffer."""
    bit_width = BitWidth.B(byte_width)
    # If you get this exception, you're trying to write a vector with a size
    # field that is bigger than the scalars you're trying to write (e.g. a
    # byte vector > 255 elements). For such types, write a "blob" instead.
    if BitWidth.U(len(elements)) > bit_width:
      raise ValueError('too many elements for the given byte_width')

    self._Align(bit_width)
    if not fixed:
      self._Write(U, len(elements), byte_width)

    loc = len(self._buf)

    fmt = {Type.INT: I, Type.UINT: U, Type.FLOAT: F}.get(element_type)
    if not fmt:
      raise TypeError('unsupported element_type')
    self._WriteVector(fmt, elements, byte_width)

    type_ = Type.ToTypedVector(element_type, len(elements) if fixed else 0)
    self._stack.append(Value(loc, type_, bit_width))
    return loc

  def _CreateVector(self, elements, typed, fixed, keys=None):
    """Writes vector elements to the underlying buffer."""
    length = len(elements)

    if fixed and not typed:
      raise ValueError('fixed vector must be typed')

    # Figure out smallest bit width we can store this vector with.
    bit_width = max(self._force_min_bit_width, BitWidth.U(length))
    prefix_elems = 1  # Vector size
    if keys:
      bit_width = max(bit_width, keys.ElemWidth(len(self._buf)))
      prefix_elems += 2  # Offset to the keys vector and its byte width.

    vector_type = Type.KEY
    # Check bit widths and types for all elements.
    for i, e in enumerate(elements):
      bit_width = max(bit_width, e.ElemWidth(len(self._buf), prefix_elems + i))

      if typed:
        if i == 0:
          vector_type = e.Type
        else:
          if vector_type != e.Type:
            raise RuntimeError('typed vector elements must be of the same type')

    if fixed and not Type.IsFixedTypedVectorElementType(vector_type):
      raise RuntimeError('must be fixed typed vector element type')

    byte_width = self._Align(bit_width)
    # Write vector. First the keys width/offset if available, and size.
    if keys:
      self._WriteOffset(keys.Value, byte_width)
      self._Write(U, 1 << keys.MinBitWidth, byte_width)

    if not fixed:
      self._Write(U, length, byte_width)

    # Then the actual data.
    loc = len(self._buf)
    for e in elements:
      self._WriteAny(e, byte_width)

    # Then the types.
    if not typed:
      for e in elements:
        self._buf.append(e.StoredPackedType(bit_width))

    if keys:
      type_ = Type.MAP
    else:
      if typed:
        type_ = Type.ToTypedVector(vector_type, length if fixed else 0)
      else:
        type_ = Type.VECTOR

    return Value(loc, type_, bit_width)

  def _PushIndirect(self, value, type_, bit_width):
    byte_width = self._Align(bit_width)
    loc = len(self._buf)
    fmt = {
        Type.INDIRECT_INT: I,
        Type.INDIRECT_UINT: U,
        Type.INDIRECT_FLOAT: F
    }[type_]
    self._Write(fmt, value, byte_width)
    self._stack.append(Value(loc, type_, bit_width))

  @InMapForString
  def String(self, value):
    """Encodes string value."""
    reset_to = len(self._buf)
    encoded = value.encode('utf-8')
    loc = self._WriteBlob(encoded, append_zero=True, type_=Type.STRING)
    if self._share_strings:
      prev_loc = self._string_pool.FindOrInsert(encoded, loc)
      if prev_loc is not None:
        del self._buf[reset_to:]
        self._stack[-1]._value = loc = prev_loc  # pylint: disable=protected-access

    return loc

  @InMap
  def Blob(self, value):
    """Encodes binary blob value.

    Args:
      value: A byte/bytearray value to encode

    Returns:
      Offset of the encoded value in underlying the byte buffer.
    """
    return self._WriteBlob(value, append_zero=False, type_=Type.BLOB)

  def Key(self, value):
    """Encodes key value.

    Args:
      value: A byte/bytearray/str value to encode. Byte object must not contain
        zero bytes. String object must be convertible to ASCII.

    Returns:
      Offset of the encoded value in the underlying byte buffer.
    """
    if isinstance(value, (bytes, bytearray)):
      encoded = value
    else:
      encoded = value.encode('ascii')

    if 0 in encoded:
      raise ValueError('key contains zero byte')

    loc = len(self._buf)
    self._buf.extend(encoded)
    self._buf.append(0)
    if self._share_keys:
      prev_loc = self._key_pool.FindOrInsert(encoded, loc)
      if prev_loc is not None:
        del self._buf[loc:]
        loc = prev_loc

    self._stack.append(Value.Key(loc))
    return loc

  def Null(self, key=None):
    """Encodes None value."""
    if key:
      self.Key(key)
    self._stack.append(Value.Null())

  @InMap
  def Bool(self, value):
    """Encodes boolean value.

    Args:
      value: A boolean value.
    """
    self._stack.append(Value.Bool(value))

  @InMap
  def Int(self, value, byte_width=0):
    """Encodes signed integer value.

    Args:
      value: A signed integer value.
      byte_width: Number of bytes to use: 1, 2, 4, or 8.
    """
    bit_width = BitWidth.I(value) if byte_width == 0 else BitWidth.B(byte_width)
    self._stack.append(Value.Int(value, bit_width))

  @InMap
  def IndirectInt(self, value, byte_width=0):
    """Encodes signed integer value indirectly.

    Args:
      value: A signed integer value.
      byte_width: Number of bytes to use: 1, 2, 4, or 8.
    """
    bit_width = BitWidth.I(value) if byte_width == 0 else BitWidth.B(byte_width)
    self._PushIndirect(value, Type.INDIRECT_INT, bit_width)

  @InMap
  def UInt(self, value, byte_width=0):
    """Encodes unsigned integer value.

    Args:
      value: An unsigned integer value.
      byte_width: Number of bytes to use: 1, 2, 4, or 8.
    """
    bit_width = BitWidth.U(value) if byte_width == 0 else BitWidth.B(byte_width)
    self._stack.append(Value.UInt(value, bit_width))

  @InMap
  def IndirectUInt(self, value, byte_width=0):
    """Encodes unsigned integer value indirectly.

    Args:
      value: An unsigned integer value.
      byte_width: Number of bytes to use: 1, 2, 4, or 8.
    """
    bit_width = BitWidth.U(value) if byte_width == 0 else BitWidth.B(byte_width)
    self._PushIndirect(value, Type.INDIRECT_UINT, bit_width)

  @InMap
  def Float(self, value, byte_width=0):
    """Encodes floating point value.

    Args:
      value: A floating point value.
      byte_width: Number of bytes to use: 4 or 8.
    """
    bit_width = BitWidth.F(value) if byte_width == 0 else BitWidth.B(byte_width)
    self._stack.append(Value.Float(value, bit_width))

  @InMap
  def IndirectFloat(self, value, byte_width=0):
    """Encodes floating point value indirectly.

    Args:
      value: A floating point value.
      byte_width: Number of bytes to use: 4 or 8.
    """
    bit_width = BitWidth.F(value) if byte_width == 0 else BitWidth.B(byte_width)
    self._PushIndirect(value, Type.INDIRECT_FLOAT, bit_width)

  def _StartVector(self):
    """Starts vector construction."""
    return len(self._stack)

  def _EndVector(self, start, typed, fixed):
    """Finishes vector construction by encodung its elements."""
    vec = self._CreateVector(self._stack[start:], typed, fixed)
    del self._stack[start:]
    self._stack.append(vec)
    return vec.Value

  @contextlib.contextmanager
  def Vector(self, key=None):
    if key:
      self.Key(key)

    try:
      start = self._StartVector()
      yield self
    finally:
      self._EndVector(start, typed=False, fixed=False)

  @InMap
  def VectorFromElements(self, elements):
    """Encodes sequence of any elements as a vector.

    Args:
      elements: sequence of elements, they may have different types.
    """
    with self.Vector():
      for e in elements:
        self.Add(e)

  @contextlib.contextmanager
  def TypedVector(self, key=None):
    if key:
      self.Key(key)

    try:
      start = self._StartVector()
      yield self
    finally:
      self._EndVector(start, typed=True, fixed=False)

  @InMap
  def TypedVectorFromElements(self, elements, element_type=None):
    """Encodes sequence of elements of the same type as typed vector.

    Args:
      elements: Sequence of elements, they must be of the same type.
      element_type: Suggested element type. Setting it to None means determining
        correct value automatically based on the given elements.
    """
    if isinstance(elements, array.array):
      if elements.typecode == 'f':
        self._WriteScalarVector(Type.FLOAT, 4, elements, fixed=False)
      elif elements.typecode == 'd':
        self._WriteScalarVector(Type.FLOAT, 8, elements, fixed=False)
      elif elements.typecode in ('b', 'h', 'i', 'l', 'q'):
        self._WriteScalarVector(
            Type.INT, elements.itemsize, elements, fixed=False)
      elif elements.typecode in ('B', 'H', 'I', 'L', 'Q'):
        self._WriteScalarVector(
            Type.UINT, elements.itemsize, elements, fixed=False)
      else:
        raise ValueError('unsupported array typecode: %s' % elements.typecode)
    else:
      add = self.Add if element_type is None else self.Adder(element_type)
      with self.TypedVector():
        for e in elements:
          add(e)

  @InMap
  def FixedTypedVectorFromElements(self,
                                   elements,
                                   element_type=None,
                                   byte_width=0):
    """Encodes sequence of elements of the same type as fixed typed vector.

    Args:
      elements: Sequence of elements, they must be of the same type. Allowed
        types are `Type.INT`, `Type.UINT`, `Type.FLOAT`. Allowed number of
        elements are 2, 3, or 4.
      element_type: Suggested element type. Setting it to None means determining
        correct value automatically based on the given elements.
      byte_width: Number of bytes to use per element. For `Type.INT` and
        `Type.UINT`: 1, 2, 4, or 8. For `Type.FLOAT`: 4 or 8. Setting it to 0
        means determining correct value automatically based on the given
        elements.
    """
    if not 2 <= len(elements) <= 4:
      raise ValueError('only 2, 3, or 4 elements are supported')

    types = {type(e) for e in elements}
    if len(types) != 1:
      raise TypeError('all elements must be of the same type')

    type_, = types

    if element_type is None:
      element_type = {int: Type.INT, float: Type.FLOAT}.get(type_)
      if not element_type:
        raise TypeError('unsupported element_type: %s' % type_)

    if byte_width == 0:
      width = {
          Type.UINT: BitWidth.U,
          Type.INT: BitWidth.I,
          Type.FLOAT: BitWidth.F
      }[element_type]
      byte_width = 1 << max(width(e) for e in elements)

    self._WriteScalarVector(element_type, byte_width, elements, fixed=True)

  def _StartMap(self):
    """Starts map construction."""
    return len(self._stack)

  def _EndMap(self, start):
    """Finishes map construction by encodung its elements."""
    # Interleaved keys and values on the stack.
    stack = self._stack[start:]

    if len(stack) % 2 != 0:
      raise RuntimeError('must be even number of keys and values')

    for key in stack[::2]:
      if key.Type is not Type.KEY:
        raise RuntimeError('all map keys must be of %s type' % Type.KEY)

    pairs = zip(stack[::2], stack[1::2])  # [(key, value), ...]
    pairs = sorted(pairs, key=lambda pair: self._ReadKey(pair[0].Value))

    del self._stack[start:]
    for pair in pairs:
      self._stack.extend(pair)

    keys = self._CreateVector(self._stack[start::2], typed=True, fixed=False)
    values = self._CreateVector(
        self._stack[start + 1::2], typed=False, fixed=False, keys=keys)

    del self._stack[start:]
    self._stack.append(values)
    return values.Value

  @contextlib.contextmanager
  def Map(self, key=None):
    if key:
      self.Key(key)

    try:
      start = self._StartMap()
      yield self
    finally:
      self._EndMap(start)

  def MapFromElements(self, elements):
    start = self._StartMap()
    for k, v in elements.items():
      self.Key(k)
      self.Add(v)
    self._EndMap(start)

  def Adder(self, type_):
    return {
        Type.BOOL: self.Bool,
        Type.INT: self.Int,
        Type.INDIRECT_INT: self.IndirectInt,
        Type.UINT: self.UInt,
        Type.INDIRECT_UINT: self.IndirectUInt,
        Type.FLOAT: self.Float,
        Type.INDIRECT_FLOAT: self.IndirectFloat,
        Type.KEY: self.Key,
        Type.BLOB: self.Blob,
        Type.STRING: self.String,
    }[type_]

  @InMapForString
  def Add(self, value):
    """Encodes value of any supported type."""
    if value is None:
      self.Null()
    elif isinstance(value, bool):
      self.Bool(value)
    elif isinstance(value, int):
      self.Int(value)
    elif isinstance(value, float):
      self.Float(value)
    elif isinstance(value, str):
      self.String(value)
    elif isinstance(value, (bytes, bytearray)):
      self.Blob(value)
    elif isinstance(value, dict):
      with self.Map():
        for k, v in value.items():
          self.Key(k)
          self.Add(v)
    elif isinstance(value, array.array):
      self.TypedVectorFromElements(value)
    elif _IsIterable(value):
      self.VectorFromElements(value)
    else:
      raise TypeError('unsupported python type: %s' % type(value))

  @property
  def LastValue(self):
    return self._stack[-1]

  @InMap
  def ReuseValue(self, value):
    self._stack.append(value)


def GetRoot(buf):
  """Returns root `Ref` object for the given buffer."""
  if len(buf) < 3:
    raise ValueError('buffer is too small')
  byte_width = buf[-1]
  return Ref.PackedType(
      Buf(buf, -(2 + byte_width)), byte_width, packed_type=buf[-2])


def Dumps(obj):
  """Returns bytearray with the encoded python object."""
  fbb = Builder()
  fbb.Add(obj)
  return fbb.Finish()


def Loads(buf):
  """Returns python object decoded from the buffer."""
  return GetRoot(buf).Value

# Lint as: python3
# Copyright 2020 Google Inc. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Implementation of FlexBuffers binary format.

For more info check https://google.github.io/flatbuffers/flexbuffers.html and
corresponding C++ implementation at
https://github.com/google/flatbuffers/blob/master/include/flatbuffers/flexbuffers.h
"""

# pylint: disable=invalid-name
# TODO(dkovalev): Add type hints everywhere, so tools like pytypes could work.

import array
import contextlib
import enum
import struct

__all__ = ('Type', 'Builder', 'GetRoot', 'Dumps', 'Loads')


class BitWidth(enum.IntEnum):
  """Supported bit widths of value types.

  These are used in the lower 2 bits of a type field to determine the size of
  the elements (and or size field) of the item pointed to (e.g. vector).
  """
  W8 = 0  # 2^0 = 1 byte
  W16 = 1  # 2^1 = 2 bytes
  W32 = 2  # 2^2 = 4 bytes
  W64 = 3  # 2^3 = 8 bytes

  @staticmethod
  def U(value):
    """Returns the minimum `BitWidth` to encode unsigned integer value."""
    assert value >= 0

    if value < (1 << 8):
      return BitWidth.W8
    elif value < (1 << 16):
      return BitWidth.W16
    elif value < (1 << 32):
      return BitWidth.W32
    elif value < (1 << 64):
      return BitWidth.W64
    else:
      raise ValueError('value is too big to encode: %s' % value)

  @staticmethod
  def I(value):
    """Returns the minimum `BitWidth` to encode signed integer value."""
    # -2^(n-1) <=     value < 2^(n-1)
    # -2^n     <= 2 * value < 2^n
    # 2 * value < 2^n, when value >= 0 or 2 * (-value) <= 2^n, when value < 0
    # 2 * value < 2^n, when value >= 0 or 2 * (-value) - 1 < 2^n, when value < 0
    #
    # if value >= 0:
    #   return BitWidth.U(2 * value)
    # else:
    #   return BitWidth.U(2 * (-value) - 1)  # ~x = -x - 1
    value *= 2
    return BitWidth.U(value if value >= 0 else ~value)

  @staticmethod
  def F(value):
    """Returns the `BitWidth` to encode floating point value."""
    if struct.unpack('f', struct.pack('f', value))[0] == value:
      return BitWidth.W32
    return BitWidth.W64

  @staticmethod
  def B(byte_width):
    return {
        1: BitWidth.W8,
        2: BitWidth.W16,
        4: BitWidth.W32,
        8: BitWidth.W64
    }[byte_width]


I = {1: 'b', 2: 'h', 4: 'i', 8: 'q'}  # Integer formats
U = {1: 'B', 2: 'H', 4: 'I', 8: 'Q'}  # Unsigned integer formats
F = {4: 'f', 8: 'd'}  # Floating point formats


def _Unpack(fmt, buf):
  return struct.unpack(fmt[len(buf)], buf)[0]


def _UnpackVector(fmt, buf, length):
  byte_width = len(buf) // length
  return struct.unpack('%d%s' % (length, fmt[byte_width]), buf)


def _Pack(fmt, value, byte_width):
  return struct.pack(fmt[byte_width], value)


def _PackVector(fmt, values, byte_width):
  return struct.pack('%d%s' % (len(values), fmt[byte_width]), *values)


def _Mutate(fmt, buf, value, byte_width, value_bit_width):
  if (1 << value_bit_width) <= byte_width:
    buf[:byte_width] = _Pack(fmt, value, byte_width)
    return True
  return False


# Computes how many bytes you'd have to pad to be able to write an
# "scalar_size" scalar if the buffer had grown to "buf_size",
# "scalar_size" is a power of two.
def _PaddingBytes(buf_size, scalar_size):
  # ((buf_size + (scalar_size - 1)) // scalar_size) * scalar_size - buf_size
  return -buf_size & (scalar_size - 1)


def _ShiftSlice(s, offset, length):
  start = offset + (0 if s.start is None else s.start)
  stop = offset + (length if s.stop is None else s.stop)
  return slice(start, stop, s.step)


# https://en.cppreference.com/w/cpp/algorithm/lower_bound
def _LowerBound(values, value, pred):
  """Implementation of C++ std::lower_bound() algorithm."""
  first, last = 0, len(values)
  count = last - first
  while count > 0:
    i = first
    step = count // 2
    i += step
    if pred(values[i], value):
      i += 1
      first = i
      count -= step + 1
    else:
      count = step
  return first


# https://en.cppreference.com/w/cpp/algorithm/binary_search
def _BinarySearch(values, value, pred=lambda x, y: x < y):
  """Implementation of C++ std::binary_search() algorithm."""
  index = _LowerBound(values, value, pred)
  if index != len(values) and not pred(value, values[index]):
    return index
  return -1


class Type(enum.IntEnum):
  """Supported types of encoded data.

  These are used as the upper 6 bits of a type field to indicate the actual
  type.
  """
  NULL = 0
  INT = 1
  UINT = 2
  FLOAT = 3
  # Types above stored inline, types below store an offset.
  KEY = 4
  STRING = 5
  INDIRECT_INT = 6
  INDIRECT_UINT = 7
  INDIRECT_FLOAT = 8
  MAP = 9
  VECTOR = 10  # Untyped.

  VECTOR_INT = 11  # Typed any size (stores no type table).
  VECTOR_UINT = 12
  VECTOR_FLOAT = 13
  VECTOR_KEY = 14
  # DEPRECATED, use VECTOR or VECTOR_KEY instead.
  # Read test.cpp/FlexBuffersDeprecatedTest() for details on why.
  VECTOR_STRING_DEPRECATED = 15

  VECTOR_INT2 = 16  # Typed tuple (no type table, no size field).
  VECTOR_UINT2 = 17
  VECTOR_FLOAT2 = 18
  VECTOR_INT3 = 19  # Typed triple (no type table, no size field).
  VECTOR_UINT3 = 20
  VECTOR_FLOAT3 = 21
  VECTOR_INT4 = 22  # Typed quad (no type table, no size field).
  VECTOR_UINT4 = 23
  VECTOR_FLOAT4 = 24

  BLOB = 25
  BOOL = 26
  VECTOR_BOOL = 36  # To do the same type of conversion of type to vector type

  @staticmethod
  def Pack(type_, bit_width):
    return (int(type_) << 2) | bit_width

  @staticmethod
  def Unpack(packed_type):
    return 1 << (packed_type & 0b11), Type(packed_type >> 2)

  @staticmethod
  def IsInline(type_):
    return type_ <= Type.FLOAT or type_ == Type.BOOL

  @staticmethod
  def IsTypedVector(type_):
    return Type.VECTOR_INT <= type_ <= Type.VECTOR_STRING_DEPRECATED or \
           type_ == Type.VECTOR_BOOL

  @staticmethod
  def IsTypedVectorElementType(type_):
    return Type.INT <= type_ <= Type.STRING or type_ == Type.BOOL

  @staticmethod
  def ToTypedVectorElementType(type_):
    if not Type.IsTypedVector(type_):
      raise ValueError('must be typed vector type')

    return Type(type_ - Type.VECTOR_INT + Type.INT)

  @staticmethod
  def IsFixedTypedVector(type_):
    return Type.VECTOR_INT2 <= type_ <= Type.VECTOR_FLOAT4

  @staticmethod
  def IsFixedTypedVectorElementType(type_):
    return Type.INT <= type_ <= Type.FLOAT

  @staticmethod
  def ToFixedTypedVectorElementType(type_):
    if not Type.IsFixedTypedVector(type_):
      raise ValueError('must be fixed typed vector type')

    # 3 types each, starting from length 2.
    fixed_type = type_ - Type.VECTOR_INT2
    return Type(fixed_type % 3 + Type.INT), fixed_type // 3 + 2

  @staticmethod
  def ToTypedVector(element_type, fixed_len=0):
    """Converts element type to corresponding vector type.

    Args:
      element_type: vector element type
      fixed_len: number of elements: 0 for typed vector; 2, 3, or 4 for fixed
        typed vector.

    Returns:
      Typed vector type or fixed typed vector type.
    """
    if fixed_len == 0:
      if not Type.IsTypedVectorElementType(element_type):
        raise ValueError('must be typed vector element type')
    else:
      if not Type.IsFixedTypedVectorElementType(element_type):
        raise ValueError('must be fixed typed vector element type')

    offset = element_type - Type.INT
    if fixed_len == 0:
      return Type(offset + Type.VECTOR_INT)  # TypedVector
    elif fixed_len == 2:
      return Type(offset + Type.VECTOR_INT2)  # FixedTypedVector
    elif fixed_len == 3:
      return Type(offset + Type.VECTOR_INT3)  # FixedTypedVector
    elif fixed_len == 4:
      return Type(offset + Type.VECTOR_INT4)  # FixedTypedVector
    else:
      raise ValueError('unsupported fixed_len: %s' % fixed_len)


class Buf:
  """Class to access underlying buffer object starting from the given offset."""

  def __init__(self, buf, offset):
    self._buf = buf
    self._offset = offset if offset >= 0 else len(buf) + offset
    self._length = len(buf) - self._offset

  def __getitem__(self, key):
    if isinstance(key, slice):
      return self._buf[_ShiftSlice(key, self._offset, self._length)]
    elif isinstance(key, int):
      return self._buf[self._offset + key]
    else:
      raise TypeError('invalid key type')

  def __setitem__(self, key, value):
    if isinstance(key, slice):
      self._buf[_ShiftSlice(key, self._offset, self._length)] = value
    elif isinstance(key, int):
      self._buf[self._offset + key] = key
    else:
      raise TypeError('invalid key type')

  def __repr__(self):
    return 'buf[%d:]' % self._offset

  def Find(self, sub):
    """Returns the lowest index where the sub subsequence is found."""
    return self._buf[self._offset:].find(sub)

  def Slice(self, offset):
    """Returns new `Buf` which starts from the given offset."""
    return Buf(self._buf, self._offset + offset)

  def Indirect(self, offset, byte_width):
    """Return new `Buf` based on the encoded offset (indirect encoding)."""
    return self.Slice(offset - _Unpack(U, self[offset:offset + byte_width]))


class Object:
  """Base class for all non-trivial data accessors."""
  __slots__ = '_buf', '_byte_width'

  def __init__(self, buf, byte_width):
    self._buf = buf
    self._byte_width = byte_width

  @property
  def ByteWidth(self):
    return self._byte_width


class Sized(Object):
  """Base class for all data accessors which need to read encoded size."""
  __slots__ = '_size',

  def __init__(self, buf, byte_width, size=0):
    super().__init__(buf, byte_width)
    if size == 0:
      self._size = _Unpack(U, self.SizeBytes)
    else:
      self._size = size

  @property
  def SizeBytes(self):
    return self._buf[-self._byte_width:0]

  def __len__(self):
    return self._size


class Blob(Sized):
  """Data accessor for the encoded blob bytes."""
  __slots__ = ()

  @property
  def Bytes(self):
    return self._buf[0:len(self)]

  def __repr__(self):
    return 'Blob(%s, size=%d)' % (self._buf, len(self))


class String(Sized):
  """Data accessor for the encoded string bytes."""
  __slots__ = ()

  @property
  def Bytes(self):
    return self._buf[0:len(self)]

  def Mutate(self, value):
    """Mutates underlying string bytes in place.

    Args:
      value: New string to replace the existing one. New string must have less
        or equal UTF-8-encoded bytes than the existing one to successfully
        mutate underlying byte buffer.

    Returns:
      Whether the value was mutated or not.
    """
    encoded = value.encode('utf-8')
    n = len(encoded)
    if n <= len(self):
      self._buf[-self._byte_width:0] = _Pack(U, n, self._byte_width)
      self._buf[0:n] = encoded
      self._buf[n:len(self)] = bytearray(len(self) - n)
      return True
    return False

  def __str__(self):
    return self.Bytes.decode('utf-8')

  def __repr__(self):
    return 'String(%s, size=%d)' % (self._buf, len(self))


class Key(Object):
  """Data accessor for the encoded key bytes."""
  __slots__ = ()

  def __init__(self, buf, byte_width):
    assert byte_width == 1
    super().__init__(buf, byte_width)

  @property
  def Bytes(self):
    return self._buf[0:len(self)]

  def __len__(self):
    return self._buf.Find(0)

  def __str__(self):
    return self.Bytes.decode('ascii')

  def __repr__(self):
    return 'Key(%s, size=%d)' % (self._buf, len(self))


class Vector(Sized):
  """Data accessor for the encoded vector bytes."""
  __slots__ = ()

  def __getitem__(self, index):
    if index < 0 or index >= len(self):
      raise IndexError('vector index %s is out of [0, %d) range' % \
          (index, len(self)))

    packed_type = self._buf[len(self) * self._byte_width + index]
    buf = self._buf.Slice(index * self._byte_width)
    return Ref.PackedType(buf, self._byte_width, packed_type)

  @property
  def Value(self):
    """Returns the underlying encoded data as a list object."""
    return [e.Value for e in self]

  def __repr__(self):
    return 'Vector(%s, byte_width=%d, size=%d)' % \
        (self._buf, self._byte_width, self._size)


class TypedVector(Sized):
  """Data accessor for the encoded typed vector or fixed typed vector bytes."""
  __slots__ = '_element_type', '_size'

  def __init__(self, buf, byte_width, element_type, size=0):
    super().__init__(buf, byte_width, size)

    if element_type == Type.STRING:
      # These can't be accessed as strings, since we don't know the bit-width
      # of the size field, see the declaration of
      # FBT_VECTOR_STRING_DEPRECATED above for details.
      # We change the type here to be keys, which are a subtype of strings,
      # and will ignore the size field. This will truncate strings with
      # embedded nulls.
      element_type = Type.KEY

    self._element_type = element_type

  @property
  def Bytes(self):
    return self._buf[:self._byte_width * len(self)]

  @property
  def ElementType(self):
    return self._element_type

  def __getitem__(self, index):
    if index < 0 or index >= len(self):
      raise IndexError('vector index %s is out of [0, %d) range' % \
          (index, len(self)))

    buf = self._buf.Slice(index * self._byte_width)
    return Ref(buf, self._byte_width, 1, self._element_type)

  @property
  def Value(self):
    """Returns underlying data as list object."""
    if not self:
      return []

    if self._element_type is Type.BOOL:
      return [bool(e) for e in _UnpackVector(U, self.Bytes, len(self))]
    elif self._element_type is Type.INT:
      return list(_UnpackVector(I, self.Bytes, len(self)))
    elif self._element_type is Type.UINT:
      return list(_UnpackVector(U, self.Bytes, len(self)))
    elif self._element_type is Type.FLOAT:
      return list(_UnpackVector(F, self.Bytes, len(self)))
    elif self._element_type is Type.KEY:
      return [e.AsKey for e in self]
    elif self._element_type is Type.STRING:
      return [e.AsString for e in self]
    else:
      raise TypeError('unsupported element_type: %s' % self._element_type)

  def __repr__(self):
    return 'TypedVector(%s, byte_width=%d, element_type=%s, size=%d)' % \
        (self._buf, self._byte_width, self._element_type, self._size)


class Map(Vector):
  """Data accessor for the encoded map bytes."""

  @staticmethod
  def CompareKeys(a, b):
    if isinstance(a, Ref):
      a = a.AsKeyBytes
    if isinstance(b, Ref):
      b = b.AsKeyBytes
    return a < b

  def __getitem__(self, key):
    if isinstance(key, int):
      return super().__getitem__(key)

    index = _BinarySearch(self.Keys, key.encode('ascii'), self.CompareKeys)
    if index != -1:
      return super().__getitem__(index)

    raise KeyError(key)

  @property
  def Keys(self):
    byte_width = _Unpack(U, self._buf[-2 * self._byte_width:-self._byte_width])
    buf = self._buf.Indirect(-3 * self._byte_width, self._byte_width)
    return TypedVector(buf, byte_width, Type.KEY)

  @property
  def Values(self):
    return Vector(self._buf, self._byte_width)

  @property
  def Value(self):
    return {k.Value: v.Value for k, v in zip(self.Keys, self.Values)}

  def __repr__(self):
    return 'Map(%s, size=%d)' % (self._buf, len(self))


class Ref:
  """Data accessor for the encoded data bytes."""
  __slots__ = '_buf', '_parent_width', '_byte_width', '_type'

  @staticmethod
  def PackedType(buf, parent_width, packed_type):
    byte_width, type_ = Type.Unpack(packed_type)
    return Ref(buf, parent_width, byte_width, type_)

  def __init__(self, buf, parent_width, byte_width, type_):
    self._buf = buf
    self._parent_width = parent_width
    self._byte_width = byte_width
    self._type = type_

  def __repr__(self):
    return 'Ref(%s, parent_width=%d, byte_width=%d, type_=%s)' % \
            (self._buf, self._parent_width, self._byte_width, self._type)

  @property
  def _Bytes(self):
    return self._buf[:self._parent_width]

  def _ConvertError(self, target_type):
    raise TypeError('cannot convert %s to %s' % (self._type, target_type))

  def _Indirect(self):
    return self._buf.Indirect(0, self._parent_width)

  @property
  def IsNull(self):
    return self._type is Type.NULL

  @property
  def IsBool(self):
    return self._type is Type.BOOL

  @property
  def AsBool(self):
    if self._type is Type.BOOL:
      return bool(_Unpack(U, self._Bytes))
    else:
      return self.AsInt != 0

  def MutateBool(self, value):
    """Mutates underlying boolean value bytes in place.

    Args:
      value: New boolean value.

    Returns:
      Whether the value was mutated or not.
    """
    return self.IsBool and \
           _Mutate(U, self._buf, value, self._parent_width, BitWidth.W8)

  @property
  def IsNumeric(self):
    return self.IsInt or self.IsFloat

  @property
  def IsInt(self):
    return self._type in (Type.INT, Type.INDIRECT_INT, Type.UINT,
                          Type.INDIRECT_UINT)

  @property
  def AsInt(self):
    """Returns current reference as integer value."""
    if self.IsNull:
      return 0
    elif self.IsBool:
      return int(self.AsBool)
    elif self._type is Type.INT:
      return _Unpack(I, self._Bytes)
    elif self._type is Type.INDIRECT_INT:
      return _Unpack(I, self._Indirect()[:self._byte_width])
    if self._type is Type.UINT:
      return _Unpack(U, self._Bytes)
    elif self._type is Type.INDIRECT_UINT:
      return _Unpack(U, self._Indirect()[:self._byte_width])
    elif self.IsString:
      return len(self.AsString)
    elif self.IsKey:
      return len(self.AsKey)
    elif self.IsBlob:
      return len(self.AsBlob)
    elif self.IsVector:
      return len(self.AsVector)
    elif self.IsTypedVector:
      return len(self.AsTypedVector)
    elif self.IsFixedTypedVector:
      return len(self.AsFixedTypedVector)
    else:
      raise self._ConvertError(Type.INT)

  def MutateInt(self, value):
    """Mutates underlying integer value bytes in place.

    Args:
      value: New integer value. It must fit to the byte size of the existing
        encoded value.

    Returns:
      Whether the value was mutated or not.
    """
    if self._type is Type.INT:
      return _Mutate(I, self._buf, value, self._parent_width, BitWidth.I(value))
    elif self._type is Type.INDIRECT_INT:
      return _Mutate(I, self._Indirect(), value, self._byte_width,
                     BitWidth.I(value))
    elif self._type is Type.UINT:
      return _Mutate(U, self._buf, value, self._parent_width, BitWidth.U(value))
    elif self._type is Type.INDIRECT_UINT:
      return _Mutate(U, self._Indirect(), value, self._byte_width,
                     BitWidth.U(value))
    else:
      return False

  @property
  def IsFloat(self):
    return self._type in (Type.FLOAT, Type.INDIRECT_FLOAT)

  @property
  def AsFloat(self):
    """Returns current reference as floating point value."""
    if self.IsNull:
      return 0.0
    elif self.IsBool:
      return float(self.AsBool)
    elif self.IsInt:
      return float(self.AsInt)
    elif self._type is Type.FLOAT:
      return _Unpack(F, self._Bytes)
    elif self._type is Type.INDIRECT_FLOAT:
      return _Unpack(F, self._Indirect()[:self._byte_width])
    elif self.IsString:
      return float(self.AsString)
    elif self.IsVector:
      return float(len(self.AsVector))
    elif self.IsTypedVector():
      return float(len(self.AsTypedVector))
    elif self.IsFixedTypedVector():
      return float(len(self.FixedTypedVector))
    else:
      raise self._ConvertError(Type.FLOAT)

  def MutateFloat(self, value):
    """Mutates underlying floating point value bytes in place.

    Args:
      value: New float value. It must fit to the byte size of the existing
        encoded value.

    Returns:
      Whether the value was mutated or not.
    """
    if self._type is Type.FLOAT:
      return _Mutate(F, self._buf, value, self._parent_width,
                     BitWidth.B(self._parent_width))
    elif self._type is Type.INDIRECT_FLOAT:
      return _Mutate(F, self._Indirect(), value, self._byte_width,
                     BitWidth.B(self._byte_width))
    else:
      return False

  @property
  def IsKey(self):
    return self._type is Type.KEY

  @property
  def AsKeyBytes(self):
    if self.IsKey:
      return Key(self._Indirect(), self._byte_width).Bytes
    else:
      raise self._ConvertError(Type.KEY)

  @property
  def AsKey(self):
    if self.IsKey:
      return str(Key(self._Indirect(), self._byte_width))
    else:
      raise self._ConvertError(Type.KEY)

  @property
  def IsString(self):
    return self._type is Type.STRING

  @property
  def AsString(self):
    if self.IsString:
      return str(String(self._Indirect(), self._byte_width))
    elif self.IsKey:
      return self.AsKey
    else:
      raise self._ConvertError(Type.STRING)

  def MutateString(self, value):
    return String(self._Indirect(), self._byte_width).Mutate(value)

  @property
  def IsBlob(self):
    return self._type is Type.BLOB

  @property
  def AsBlob(self):
    if self.IsBlob:
      return Blob(self._Indirect(), self._byte_width).Bytes
    else:
      raise self._ConvertError(Type.BLOB)

  @property
  def IsAnyVector(self):
    return self.IsVector or self.IsTypedVector or self.IsFixedTypedVector()

  @property
  def IsVector(self):
    return self._type in (Type.VECTOR, Type.MAP)

  @property
  def AsVector(self):
    if self.IsVector:
      return Vector(self._Indirect(), self._byte_width)
    else:
      raise self._ConvertError(Type.VECTOR)

  @property
  def IsTypedVector(self):
    return Type.IsTypedVector(self._type)

  @property
  def AsTypedVector(self):
    if self.IsTypedVector:
      return TypedVector(self._Indirect(), self._byte_width,
                         Type.ToTypedVectorElementType(self._type))
    else:
      raise self._ConvertError('TYPED_VECTOR')

  @property
  def IsFixedTypedVector(self):
    return Type.IsFixedTypedVector(self._type)

  @property
  def AsFixedTypedVector(self):
    if self.IsFixedTypedVector:
      element_type, size = Type.ToFixedTypedVectorElementType(self._type)
      return TypedVector(self._Indirect(), self._byte_width, element_type, size)
    else:
      raise self._ConvertError('FIXED_TYPED_VECTOR')

  @property
  def IsMap(self):
    return self._type is Type.MAP

  @property
  def AsMap(self):
    if self.IsMap:
      return Map(self._Indirect(), self._byte_width)
    else:
      raise self._ConvertError(Type.MAP)

  @property
  def Value(self):
    """Converts current reference to value of corresponding type.

    This is equivalent to calling `AsInt` for integer values, `AsFloat` for
    floating point values, etc.

    Returns:
      Value of corresponding type.
    """
    if self.IsNull:
      return None
    elif self.IsBool:
      return self.AsBool
    elif self.IsInt:
      return self.AsInt
    elif self.IsFloat:
      return self.AsFloat
    elif self.IsString:
      return self.AsString
    elif self.IsKey:
      return self.AsKey
    elif self.IsBlob:
      return self.AsBlob
    elif self.IsMap:
      return self.AsMap.Value
    elif self.IsVector:
      return self.AsVector.Value
    elif self.IsTypedVector:
      return self.AsTypedVector.Value
    elif self.IsFixedTypedVector:
      return self.AsFixedTypedVector.Value
    else:
      raise TypeError('cannot convert %r to value' % self)


def _IsIterable(obj):
  try:
    iter(obj)
    return True
  except TypeError:
    return False


class Value:
  """Class to represent given value during the encoding process."""

  @staticmethod
  def Null():
    return Value(0, Type.NULL, BitWidth.W8)

  @staticmethod
  def Bool(value):
    return Value(value, Type.BOOL, BitWidth.W8)

  @staticmethod
  def Int(value, bit_width):
    return Value(value, Type.INT, bit_width)

  @staticmethod
  def UInt(value, bit_width):
    return Value(value, Type.UINT, bit_width)

  @staticmethod
  def Float(value, bit_width):
    return Value(value, Type.FLOAT, bit_width)

  @staticmethod
  def Key(offset):
    return Value(offset, Type.KEY, BitWidth.W8)

  def __init__(self, value, type_, min_bit_width):
    self._value = value
    self._type = type_

    # For scalars: of itself, for vector: of its elements, for string: length.
    self._min_bit_width = min_bit_width

  @property
  def Value(self):
    return self._value

  @property
  def Type(self):
    return self._type

  @property
  def MinBitWidth(self):
    return self._min_bit_width

  def StoredPackedType(self, parent_bit_width=BitWidth.W8):
    return Type.Pack(self._type, self.StoredWidth(parent_bit_width))

  # We have an absolute offset, but want to store a relative offset
  # elem_index elements beyond the current buffer end. Since whether
  # the relative offset fits in a certain byte_width depends on
  # the size of the elements before it (and their alignment), we have
  # to test for each size in turn.
  def ElemWidth(self, buf_size, elem_index=0):
    if Type.IsInline(self._type):
      return self._min_bit_width
    for byte_width in 1, 2, 4, 8:
      offset_loc = buf_size + _PaddingBytes(buf_size, byte_width) + \
                   elem_index * byte_width
      bit_width = BitWidth.U(offset_loc - self._value)
      if byte_width == (1 << bit_width):
        return bit_width
    raise ValueError('relative offset is too big')

  def StoredWidth(self, parent_bit_width=BitWidth.W8):
    if Type.IsInline(self._type):
      return max(self._min_bit_width, parent_bit_width)
    return self._min_bit_width

  def __repr__(self):
    return 'Value(%s, %s, %s)' % (self._value, self._type, self._min_bit_width)

  def __str__(self):
    return str(self._value)


def InMap(func):
  def wrapper(self, *args, **kwargs):
    if isinstance(args[0], str):
      self.Key(args[0])
      func(self, *args[1:], **kwargs)
    else:
      func(self, *args, **kwargs)
  return wrapper


def InMapForString(func):
  def wrapper(self, *args):
    if len(args) == 1:
      func(self, args[0])
    elif len(args) == 2:
      self.Key(args[0])
      func(self, args[1])
    else:
      raise ValueError('invalid number of arguments')
  return wrapper


class Pool:
  """Collection of (data, offset) pairs sorted by data for quick access."""

  def __init__(self):
    self._pool = []  # sorted list of (data, offset) tuples

  def FindOrInsert(self, data, offset):
    do = data, offset
    index = _BinarySearch(self._pool, do, lambda a, b: a[0] < b[0])
    if index != -1:
      _, offset = self._pool[index]
      return offset
    self._pool.insert(index, do)
    return None

  def Clear(self):
    self._pool = []

  @property
  def Elements(self):
    return [data for data, _ in self._pool]


class Builder:
  """Helper class to encode structural data into flexbuffers format."""

  def __init__(self,
               share_strings=False,
               share_keys=True,
               force_min_bit_width=BitWidth.W8):
    self._share_strings = share_strings
    self._share_keys = share_keys
    self._force_min_bit_width = force_min_bit_width

    self._string_pool = Pool()
    self._key_pool = Pool()

    self._finished = False
    self._buf = bytearray()
    self._stack = []

  def __len__(self):
    return len(self._buf)

  @property
  def StringPool(self):
    return self._string_pool

  @property
  def KeyPool(self):
    return self._key_pool

  def Clear(self):
    self._string_pool.Clear()
    self._key_pool.Clear()
    self._finished = False
    self._buf = bytearray()
    self._stack = []

  def Finish(self):
    """Finishes encoding process and returns underlying buffer."""
    if self._finished:
      raise RuntimeError('builder has been already finished')

    # If you hit this exception, you likely have objects that were never
    # included in a parent. You need to have exactly one root to finish a
    # buffer. Check your Start/End calls are matched, and all objects are inside
    # some other object.
    if len(self._stack) != 1:
      raise RuntimeError('internal stack size must be one')

    value = self._stack[0]
    byte_width = self._Align(value.ElemWidth(len(self._buf)))
    self._WriteAny(value, byte_width=byte_width)  # Root value
    self._Write(U, value.StoredPackedType(), byte_width=1)  # Root type
    self._Write(U, byte_width, byte_width=1)  # Root size

    self.finished = True
    return self._buf

  def _ReadKey(self, offset):
    key = self._buf[offset:]
    return key[:key.find(0)]

  def _Align(self, alignment):
    byte_width = 1 << alignment
    self._buf.extend(b'\x00' * _PaddingBytes(len(self._buf), byte_width))
    return byte_width

  def _Write(self, fmt, value, byte_width):
    self._buf.extend(_Pack(fmt, value, byte_width))

  def _WriteVector(self, fmt, values, byte_width):
    self._buf.extend(_PackVector(fmt, values, byte_width))

  def _WriteOffset(self, offset, byte_width):
    relative_offset = len(self._buf) - offset
    assert byte_width == 8 or relative_offset < (1 << (8 * byte_width))
    self._Write(U, relative_offset, byte_width)

  def _WriteAny(self, value, byte_width):
    fmt = {
        Type.NULL: U, Type.BOOL: U, Type.INT: I, Type.UINT: U, Type.FLOAT: F
    }.get(value.Type)
    if fmt:
      self._Write(fmt, value.Value, byte_width)
    else:
      self._WriteOffset(value.Value, byte_width)

  def _WriteBlob(self, data, append_zero, type_):
    bit_width = BitWidth.U(len(data))
    byte_width = self._Align(bit_width)
    self._Write(U, len(data), byte_width)
    loc = len(self._buf)
    self._buf.extend(data)
    if append_zero:
      self._buf.append(0)
    self._stack.append(Value(loc, type_, bit_width))
    return loc

  def _WriteScalarVector(self, element_type, byte_width, elements, fixed):
    """Writes scalar vector elements to the underlying buffer."""
    bit_width = BitWidth.B(byte_width)
    # If you get this exception, you're trying to write a vector with a size
    # field that is bigger than the scalars you're trying to write (e.g. a
    # byte vector > 255 elements). For such types, write a "blob" instead.
    if BitWidth.U(len(elements)) > bit_width:
      raise ValueError('too many elements for the given byte_width')

    self._Align(bit_width)
    if not fixed:
      self._Write(U, len(elements), byte_width)

    loc = len(self._buf)

    fmt = {Type.INT: I, Type.UINT: U, Type.FLOAT: F}.get(element_type)
    if not fmt:
      raise TypeError('unsupported element_type')
    self._WriteVector(fmt, elements, byte_width)

    type_ = Type.ToTypedVector(element_type, len(elements) if fixed else 0)
    self._stack.append(Value(loc, type_, bit_width))
    return loc

  def _CreateVector(self, elements, typed, fixed, keys=None):
    """Writes vector elements to the underlying buffer."""
    length = len(elements)

    if fixed and not typed:
      raise ValueError('fixed vector must be typed')

    # Figure out smallest bit width we can store this vector with.
    bit_width = max(self._force_min_bit_width, BitWidth.U(length))
    prefix_elems = 1  # Vector size
    if keys:
      bit_width = max(bit_width, keys.ElemWidth(len(self._buf)))
      prefix_elems += 2  # Offset to the keys vector and its byte width.

    vector_type = Type.KEY
    # Check bit widths and types for all elements.
    for i, e in enumerate(elements):
      bit_width = max(bit_width, e.ElemWidth(len(self._buf), prefix_elems + i))

      if typed:
        if i == 0:
          vector_type = e.Type
        else:
          if vector_type != e.Type:
            raise RuntimeError('typed vector elements must be of the same type')

    if fixed and not Type.IsFixedTypedVectorElementType(vector_type):
      raise RuntimeError('must be fixed typed vector element type')

    byte_width = self._Align(bit_width)
    # Write vector. First the keys width/offset if available, and size.
    if keys:
      self._WriteOffset(keys.Value, byte_width)
      self._Write(U, 1 << keys.MinBitWidth, byte_width)

    if not fixed:
      self._Write(U, length, byte_width)

    # Then the actual data.
    loc = len(self._buf)
    for e in elements:
      self._WriteAny(e, byte_width)

    # Then the types.
    if not typed:
      for e in elements:
        self._buf.append(e.StoredPackedType(bit_width))

    if keys:
      type_ = Type.MAP
    else:
      if typed:
        type_ = Type.ToTypedVector(vector_type, length if fixed else 0)
      else:
        type_ = Type.VECTOR

    return Value(loc, type_, bit_width)

  def _PushIndirect(self, value, type_, bit_width):
    byte_width = self._Align(bit_width)
    loc = len(self._buf)
    fmt = {
        Type.INDIRECT_INT: I,
        Type.INDIRECT_UINT: U,
        Type.INDIRECT_FLOAT: F
    }[type_]
    self._Write(fmt, value, byte_width)
    self._stack.append(Value(loc, type_, bit_width))

  @InMapForString
  def String(self, value):
    """Encodes string value."""
    reset_to = len(self._buf)
    encoded = value.encode('utf-8')
    loc = self._WriteBlob(encoded, append_zero=True, type_=Type.STRING)
    if self._share_strings:
      prev_loc = self._string_pool.FindOrInsert(encoded, loc)
      if prev_loc is not None:
        del self._buf[reset_to:]
        self._stack[-1]._value = loc = prev_loc  # pylint: disable=protected-access

    return loc

  @InMap
  def Blob(self, value):
    """Encodes binary blob value.

    Args:
      value: A byte/bytearray value to encode

    Returns:
      Offset of the encoded value in underlying the byte buffer.
    """
    return self._WriteBlob(value, append_zero=False, type_=Type.BLOB)

  def Key(self, value):
    """Encodes key value.

    Args:
      value: A byte/bytearray/str value to encode. Byte object must not contain
        zero bytes. String object must be convertible to ASCII.

    Returns:
      Offset of the encoded value in the underlying byte buffer.
    """
    if isinstance(value, (bytes, bytearray)):
      encoded = value
    else:
      encoded = value.encode('ascii')

    if 0 in encoded:
      raise ValueError('key contains zero byte')

    loc = len(self._buf)
    self._buf.extend(encoded)
    self._buf.append(0)
    if self._share_keys:
      prev_loc = self._key_pool.FindOrInsert(encoded, loc)
      if prev_loc is not None:
        del self._buf[loc:]
        loc = prev_loc

    self._stack.append(Value.Key(loc))
    return loc

  def Null(self, key=None):
    """Encodes None value."""
    if key:
      self.Key(key)
    self._stack.append(Value.Null())

  @InMap
  def Bool(self, value):
    """Encodes boolean value.

    Args:
      value: A boolean value.
    """
    self._stack.append(Value.Bool(value))

  @InMap
  def Int(self, value, byte_width=0):
    """Encodes signed integer value.

    Args:
      value: A signed integer value.
      byte_width: Number of bytes to use: 1, 2, 4, or 8.
    """
    bit_width = BitWidth.I(value) if byte_width == 0 else BitWidth.B(byte_width)
    self._stack.append(Value.Int(value, bit_width))

  @InMap
  def IndirectInt(self, value, byte_width=0):
    """Encodes signed integer value indirectly.

    Args:
      value: A signed integer value.
      byte_width: Number of bytes to use: 1, 2, 4, or 8.
    """
    bit_width = BitWidth.I(value) if byte_width == 0 else BitWidth.B(byte_width)
    self._PushIndirect(value, Type.INDIRECT_INT, bit_width)

  @InMap
  def UInt(self, value, byte_width=0):
    """Encodes unsigned integer value.

    Args:
      value: An unsigned integer value.
      byte_width: Number of bytes to use: 1, 2, 4, or 8.
    """
    bit_width = BitWidth.U(value) if byte_width == 0 else BitWidth.B(byte_width)
    self._stack.append(Value.UInt(value, bit_width))

  @InMap
  def IndirectUInt(self, value, byte_width=0):
    """Encodes unsigned integer value indirectly.

    Args:
      value: An unsigned integer value.
      byte_width: Number of bytes to use: 1, 2, 4, or 8.
    """
    bit_width = BitWidth.U(value) if byte_width == 0 else BitWidth.B(byte_width)
    self._PushIndirect(value, Type.INDIRECT_UINT, bit_width)

  @InMap
  def Float(self, value, byte_width=0):
    """Encodes floating point value.

    Args:
      value: A floating point value.
      byte_width: Number of bytes to use: 4 or 8.
    """
    bit_width = BitWidth.F(value) if byte_width == 0 else BitWidth.B(byte_width)
    self._stack.append(Value.Float(value, bit_width))

  @InMap
  def IndirectFloat(self, value, byte_width=0):
    """Encodes floating point value indirectly.

    Args:
      value: A floating point value.
      byte_width: Number of bytes to use: 4 or 8.
    """
    bit_width = BitWidth.F(value) if byte_width == 0 else BitWidth.B(byte_width)
    self._PushIndirect(value, Type.INDIRECT_FLOAT, bit_width)

  def _StartVector(self):
    """Starts vector construction."""
    return len(self._stack)

  def _EndVector(self, start, typed, fixed):
    """Finishes vector construction by encodung its elements."""
    vec = self._CreateVector(self._stack[start:], typed, fixed)
    del self._stack[start:]
    self._stack.append(vec)
    return vec.Value

  @contextlib.contextmanager
  def Vector(self, key=None):
    if key:
      self.Key(key)

    try:
      start = self._StartVector()
      yield self
    finally:
      self._EndVector(start, typed=False, fixed=False)

  @InMap
  def VectorFromElements(self, elements):
    """Encodes sequence of any elements as a vector.

    Args:
      elements: sequence of elements, they may have different types.
    """
    with self.Vector():
      for e in elements:
        self.Add(e)

  @contextlib.contextmanager
  def TypedVector(self, key=None):
    if key:
      self.Key(key)

    try:
      start = self._StartVector()
      yield self
    finally:
      self._EndVector(start, typed=True, fixed=False)

  @InMap
  def TypedVectorFromElements(self, elements, element_type=None):
    """Encodes sequence of elements of the same type as typed vector.

    Args:
      elements: Sequence of elements, they must be of the same type.
      element_type: Suggested element type. Setting it to None means determining
        correct value automatically based on the given elements.
    """
    if isinstance(elements, array.array):
      if elements.typecode == 'f':
        self._WriteScalarVector(Type.FLOAT, 4, elements, fixed=False)
      elif elements.typecode == 'd':
        self._WriteScalarVector(Type.FLOAT, 8, elements, fixed=False)
      elif elements.typecode in ('b', 'h', 'i', 'l', 'q'):
        self._WriteScalarVector(
            Type.INT, elements.itemsize, elements, fixed=False)
      elif elements.typecode in ('B', 'H', 'I', 'L', 'Q'):
        self._WriteScalarVector(
            Type.UINT, elements.itemsize, elements, fixed=False)
      else:
        raise ValueError('unsupported array typecode: %s' % elements.typecode)
    else:
      add = self.Add if element_type is None else self.Adder(element_type)
      with self.TypedVector():
        for e in elements:
          add(e)

  @InMap
  def FixedTypedVectorFromElements(self,
                                   elements,
                                   element_type=None,
                                   byte_width=0):
    """Encodes sequence of elements of the same type as fixed typed vector.

    Args:
      elements: Sequence of elements, they must be of the same type. Allowed
        types are `Type.INT`, `Type.UINT`, `Type.FLOAT`. Allowed number of
        elements are 2, 3, or 4.
      element_type: Suggested element type. Setting it to None means determining
        correct value automatically based on the given elements.
      byte_width: Number of bytes to use per element. For `Type.INT` and
        `Type.UINT`: 1, 2, 4, or 8. For `Type.FLOAT`: 4 or 8. Setting it to 0
        means determining correct value automatically based on the given
        elements.
    """
    if not 2 <= len(elements) <= 4:
      raise ValueError('only 2, 3, or 4 elements are supported')

    types = {type(e) for e in elements}
    if len(types) != 1:
      raise TypeError('all elements must be of the same type')

    type_, = types

    if element_type is None:
      element_type = {int: Type.INT, float: Type.FLOAT}.get(type_)
      if not element_type:
        raise TypeError('unsupported element_type: %s' % type_)

    if byte_width == 0:
      width = {
          Type.UINT: BitWidth.U,
          Type.INT: BitWidth.I,
          Type.FLOAT: BitWidth.F
      }[element_type]
      byte_width = 1 << max(width(e) for e in elements)

    self._WriteScalarVector(element_type, byte_width, elements, fixed=True)

  def _StartMap(self):
    """Starts map construction."""
    return len(self._stack)

  def _EndMap(self, start):
    """Finishes map construction by encodung its elements."""
    # Interleaved keys and values on the stack.
    stack = self._stack[start:]

    if len(stack) % 2 != 0:
      raise RuntimeError('must be even number of keys and values')

    for key in stack[::2]:
      if key.Type is not Type.KEY:
        raise RuntimeError('all map keys must be of %s type' % Type.KEY)

    pairs = zip(stack[::2], stack[1::2])  # [(key, value), ...]
    pairs = sorted(pairs, key=lambda pair: self._ReadKey(pair[0].Value))

    del self._stack[start:]
    for pair in pairs:
      self._stack.extend(pair)

    keys = self._CreateVector(self._stack[start::2], typed=True, fixed=False)
    values = self._CreateVector(
        self._stack[start + 1::2], typed=False, fixed=False, keys=keys)

    del self._stack[start:]
    self._stack.append(values)
    return values.Value

  @contextlib.contextmanager
  def Map(self, key=None):
    if key:
      self.Key(key)

    try:
      start = self._StartMap()
      yield self
    finally:
      self._EndMap(start)

  def MapFromElements(self, elements):
    start = self._StartMap()
    for k, v in elements.items():
      self.Key(k)
      self.Add(v)
    self._EndMap(start)

  def Adder(self, type_):
    return {
        Type.BOOL: self.Bool,
        Type.INT: self.Int,
        Type.INDIRECT_INT: self.IndirectInt,
        Type.UINT: self.UInt,
        Type.INDIRECT_UINT: self.IndirectUInt,
        Type.FLOAT: self.Float,
        Type.INDIRECT_FLOAT: self.IndirectFloat,
        Type.KEY: self.Key,
        Type.BLOB: self.Blob,
        Type.STRING: self.String,
    }[type_]

  @InMapForString
  def Add(self, value):
    """Encodes value of any supported type."""
    if value is None:
      self.Null()
    elif isinstance(value, bool):
      self.Bool(value)
    elif isinstance(value, int):
      self.Int(value)
    elif isinstance(value, float):
      self.Float(value)
    elif isinstance(value, str):
      self.String(value)
    elif isinstance(value, (bytes, bytearray)):
      self.Blob(value)
    elif isinstance(value, dict):
      with self.Map():
        for k, v in value.items():
          self.Key(k)
          self.Add(v)
    elif isinstance(value, array.array):
      self.TypedVectorFromElements(value)
    elif _IsIterable(value):
      self.VectorFromElements(value)
    else:
      raise TypeError('unsupported python type: %s' % type(value))

  @property
  def LastValue(self):
    return self._stack[-1]

  @InMap
  def ReuseValue(self, value):
    self._stack.append(value)


def GetRoot(buf):
  """Returns root `Ref` object for the given buffer."""
  if len(buf) < 3:
    raise ValueError('buffer is too small')
  byte_width = buf[-1]
  return Ref.PackedType(
      Buf(buf, -(2 + byte_width)), byte_width, packed_type=buf[-2])


def Dumps(obj):
  """Returns bytearray with the encoded python object."""
  fbb = Builder()
  fbb.Add(obj)
  return fbb.Finish()


def Loads(buf):
  """Returns python object decoded from the buffer."""
  return GetRoot(buf).Value

# Lint as: python3 # Copyright 2020 Google Inc. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # #     http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. Implementation of FlexBuffers binary format.  For more info check https://google.github.io/flatbuffers/flexbuffers.html and corresponding C++ implementation at https://github.com/google/flatbuffers/blob/master/include/flatbuffers/flexbuffers.h # pylint: disable=invalid-name # TODO(dkovalev): Add type hints everywhere, so tools like pytypes could work. Supported bit widths of value types.    These are used in the lower 2 bits of a type field to determine the size of   the elements (and or size field) of the item pointed to (e.g. vector). # 2^0 = 1 byte # 2^1 = 2 bytes # 2^2 = 4 bytes # 2^3 = 8 bytes Returns the minimum `BitWidth` to encode unsigned integer value. Returns the minimum `BitWidth` to encode signed integer value. # -2^(n-1) <=     value < 2^(n-1) # -2^n     <= 2 * value < 2^n # 2 * value < 2^n, when value >= 0 or 2 * (-value) <= 2^n, when value < 0 # 2 * value < 2^n, when value >= 0 or 2 * (-value) - 1 < 2^n, when value < 0 # # if value >= 0: #   return BitWidth.U(2 * value) # else: #   return BitWidth.U(2 * (-value) - 1)  # ~x = -x - 1 Returns the `BitWidth` to encode floating point value. # Integer formats # Unsigned integer formats # Floating point formats # Computes how many bytes you'd have to pad to be able to write an # "scalar_size" scalar if the buffer had grown to "buf_size", # "scalar_size" is a power of two. # ((buf_size + (scalar_size - 1)) // scalar_size) * scalar_size - buf_size # https://en.cppreference.com/w/cpp/algorithm/lower_bound Implementation of C++ std::lower_bound() algorithm. # https://en.cppreference.com/w/cpp/algorithm/binary_search Implementation of C++ std::binary_search() algorithm. Supported types of encoded data.    These are used as the upper 6 bits of a type field to indicate the actual   type. # Types above stored inline, types below store an offset. # Untyped. # Typed any size (stores no type table). # DEPRECATED, use VECTOR or VECTOR_KEY instead. # Read test.cpp/FlexBuffersDeprecatedTest() for details on why. # Typed tuple (no type table, no size field). # Typed triple (no type table, no size field). # Typed quad (no type table, no size field). # To do the same type of conversion of type to vector type # 3 types each, starting from length 2. Converts element type to corresponding vector type.      Args:       element_type: vector element type       fixed_len: number of elements: 0 for typed vector; 2, 3, or 4 for fixed         typed vector.      Returns:       Typed vector type or fixed typed vector type. # TypedVector # FixedTypedVector # FixedTypedVector # FixedTypedVector Class to access underlying buffer object starting from the given offset. Returns the lowest index where the sub subsequence is found. Returns new `Buf` which starts from the given offset. Return new `Buf` based on the encoded offset (indirect encoding). Base class for all non-trivial data accessors. Base class for all data accessors which need to read encoded size. Data accessor for the encoded blob bytes. Data accessor for the encoded string bytes. Mutates underlying string bytes in place.      Args:       value: New string to replace the existing one. New string must have less         or equal UTF-8-encoded bytes than the existing one to successfully         mutate underlying byte buffer.      Returns:       Whether the value was mutated or not. Data accessor for the encoded key bytes. Data accessor for the encoded vector bytes. Returns the underlying encoded data as a list object. Data accessor for the encoded typed vector or fixed typed vector bytes. # These can't be accessed as strings, since we don't know the bit-width # of the size field, see the declaration of # FBT_VECTOR_STRING_DEPRECATED above for details. # We change the type here to be keys, which are a subtype of strings, # and will ignore the size field. This will truncate strings with # embedded nulls. Returns underlying data as list object. Data accessor for the encoded map bytes. Data accessor for the encoded data bytes. Mutates underlying boolean value bytes in place.      Args:       value: New boolean value.      Returns:       Whether the value was mutated or not. Returns current reference as integer value. Mutates underlying integer value bytes in place.      Args:       value: New integer value. It must fit to the byte size of the existing         encoded value.      Returns:       Whether the value was mutated or not. Returns current reference as floating point value. Mutates underlying floating point value bytes in place.      Args:       value: New float value. It must fit to the byte size of the existing         encoded value.      Returns:       Whether the value was mutated or not. Converts current reference to value of corresponding type.      This is equivalent to calling `AsInt` for integer values, `AsFloat` for     floating point values, etc.      Returns:       Value of corresponding type. Class to represent given value during the encoding process. # For scalars: of itself, for vector: of its elements, for string: length. # We have an absolute offset, but want to store a relative offset # elem_index elements beyond the current buffer end. Since whether # the relative offset fits in a certain byte_width depends on # the size of the elements before it (and their alignment), we have # to test for each size in turn. Collection of (data, offset) pairs sorted by data for quick access. # sorted list of (data, offset) tuples Helper class to encode structural data into flexbuffers format. Finishes encoding process and returns underlying buffer. # If you hit this exception, you likely have objects that were never # included in a parent. You need to have exactly one root to finish a # buffer. Check your Start/End calls are matched, and all objects are inside # some other object. # Root value # Root type # Root size Writes scalar vector elements to the underlying buffer. # If you get this exception, you're trying to write a vector with a size # field that is bigger than the scalars you're trying to write (e.g. a # byte vector > 255 elements). For such types, write a "blob" instead. Writes vector elements to the underlying buffer. # Figure out smallest bit width we can store this vector with. # Vector size # Offset to the keys vector and its byte width. # Check bit widths and types for all elements. # Write vector. First the keys width/offset if available, and size. # Then the actual data. # Then the types. Encodes string value. # pylint: disable=protected-access Encodes binary blob value.      Args:       value: A byte/bytearray value to encode      Returns:       Offset of the encoded value in underlying the byte buffer. Encodes key value.      Args:       value: A byte/bytearray/str value to encode. Byte object must not contain         zero bytes. String object must be convertible to ASCII.      Returns:       Offset of the encoded value in the underlying byte buffer. Encodes None value. Encodes boolean value.      Args:       value: A boolean value. Encodes signed integer value.      Args:       value: A signed integer value.       byte_width: Number of bytes to use: 1, 2, 4, or 8. Encodes signed integer value indirectly.      Args:       value: A signed integer value.       byte_width: Number of bytes to use: 1, 2, 4, or 8. Encodes unsigned integer value.      Args:       value: An unsigned integer value.       byte_width: Number of bytes to use: 1, 2, 4, or 8. Encodes unsigned integer value indirectly.      Args:       value: An unsigned integer value.       byte_width: Number of bytes to use: 1, 2, 4, or 8. Encodes floating point value.      Args:       value: A floating point value.       byte_width: Number of bytes to use: 4 or 8. Encodes floating point value indirectly.      Args:       value: A floating point value.       byte_width: Number of bytes to use: 4 or 8. Starts vector construction. Finishes vector construction by encodung its elements. Encodes sequence of any elements as a vector.      Args:       elements: sequence of elements, they may have different types. Encodes sequence of elements of the same type as typed vector.      Args:       elements: Sequence of elements, they must be of the same type.       element_type: Suggested element type. Setting it to None means determining         correct value automatically based on the given elements. Encodes sequence of elements of the same type as fixed typed vector.      Args:       elements: Sequence of elements, they must be of the same type. Allowed         types are `Type.INT`, `Type.UINT`, `Type.FLOAT`. Allowed number of         elements are 2, 3, or 4.       element_type: Suggested element type. Setting it to None means determining         correct value automatically based on the given elements.       byte_width: Number of bytes to use per element. For `Type.INT` and         `Type.UINT`: 1, 2, 4, or 8. For `Type.FLOAT`: 4 or 8. Setting it to 0         means determining correct value automatically based on the given         elements. Starts map construction. Finishes map construction by encodung its elements. # Interleaved keys and values on the stack. # [(key, value), ...] Encodes value of any supported type. Returns root `Ref` object for the given buffer. Returns bytearray with the encoded python object. Returns python object decoded from the buffer.