Skip to content

Commit

Permalink
add method to encode tuple element
Browse files Browse the repository at this point in the history
  • Loading branch information
barnjamin committed Dec 23, 2022
1 parent 721532f commit dcbcc5a
Showing 1 changed file with 101 additions and 1 deletion.
102 changes: 101 additions & 1 deletion pyteal/ast/abi/tuple.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@
from pyteal.ast.int import Int
from pyteal.ast.bytes import Bytes
from pyteal.ast.unaryexpr import Len
from pyteal.ast.binaryexpr import ExtractUint16
from pyteal.ast.binaryexpr import ExtractUint16, GetBit
from pyteal.ast.naryexpr import Concat
from pyteal.ast.abstractvar import alloc_abstract_var

Expand Down Expand Up @@ -117,6 +117,99 @@ def _encode_tuple(values: Sequence[BaseType]) -> Expr:
return Concat(*toConcat)


def _index_tuple_bytes(
value_types: Sequence[TypeSpec], encoded: Expr, index: int
) -> Expr:
if not (0 <= index < len(value_types)):
raise ValueError("Index outside of range")

offset = 0
ignoreNext = 0
lastBoolStart = 0
lastBoolLength = 0
for i, typeBefore in enumerate(value_types[:index]):
if ignoreNext > 0:
ignoreNext -= 1
continue

if typeBefore == BoolTypeSpec():
lastBoolStart = offset
lastBoolLength = _consecutive_bool_type_spec_num(value_types, i)
offset += _bool_sequence_length(lastBoolLength)
ignoreNext = lastBoolLength - 1
continue

if typeBefore.is_dynamic():
offset += 2
continue

offset += typeBefore.byte_length_static()

valueType = value_types[index]

if type(valueType) is Bool:
if ignoreNext > 0:
# value is in the middle of a bool sequence
bitOffsetInBoolSeq = lastBoolLength - ignoreNext
bitOffsetInEncoded = lastBoolStart * NUM_BITS_IN_BYTE + bitOffsetInBoolSeq
else:
# value is the beginning of a bool sequence (or a single bool)
bitOffsetInEncoded = offset * NUM_BITS_IN_BYTE
return GetBit(encoded, Int(bitOffsetInEncoded))

if valueType.is_dynamic():
hasNextDynamicValue = False
nextDynamicValueOffset = offset + 2
ignoreNext = 0
for i, typeAfter in enumerate(value_types[index + 1 :], start=index + 1):
if ignoreNext > 0:
ignoreNext -= 1
continue

if type(typeAfter) is BoolTypeSpec:
boolLength = _consecutive_bool_type_spec_num(value_types, i)
nextDynamicValueOffset += _bool_sequence_length(boolLength)
ignoreNext = boolLength - 1
continue

if typeAfter.is_dynamic():
hasNextDynamicValue = True
break

nextDynamicValueOffset += typeAfter.byte_length_static()

start_index = ExtractUint16(encoded, Int(offset))
if not hasNextDynamicValue:
# This is the final dynamic value, so decode the substring from start_index to the end of
# encoded
return substring_for_decoding(encoded, start_index=start_index)

# There is a dynamic value after this one, and end_index is where its tail starts, so decode
# the substring from start_index to end_index
end_index = ExtractUint16(encoded, Int(nextDynamicValueOffset))
return substring_for_decoding(
encoded, start_index=start_index, end_index=end_index
)

start_index = Int(offset)
length = Int(valueType.byte_length_static())

if index + 1 == len(value_types):
if offset == 0:
# This is the first and only value in the tuple, so decode all of encoded
return encoded
# This is the last value in the tuple, so decode the substring from start_index to the end of
# encoded
return substring_for_decoding(encoded, start_index=start_index)

if offset == 0:
# This is the first value in the tuple, so decode the substring from 0 with length length
return substring_for_decoding(encoded, length=length)

# This is not the first or last value, so decode the substring from start_index with length length
return substring_for_decoding(encoded, start_index=start_index, length=length)


def _index_tuple(
value_types: Sequence[TypeSpec], encoded: Expr, index: int, output: BaseType
) -> Expr:
Expand Down Expand Up @@ -405,6 +498,13 @@ def store_into(self, output: T) -> Expr:
output,
)

def encode(self) -> Expr:
return _index_tuple_bytes(
self.tuple.type_spec().value_type_specs(),
self.tuple.encode(),
self.index,
)


TupleElement.__module__ = "pyteal.abi"

Expand Down

0 comments on commit dcbcc5a

Please sign in to comment.