diff --git a/pyteal/ast/abi/tuple.py b/pyteal/ast/abi/tuple.py index 0b1ed0076..e01d1b036 100644 --- a/pyteal/ast/abi/tuple.py +++ b/pyteal/ast/abi/tuple.py @@ -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 @@ -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: @@ -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"