diff --git a/src/puya/awst_build/eb/_utils.py b/src/puya/awst_build/eb/_utils.py index 019a43c849..add10bd439 100644 --- a/src/puya/awst_build/eb/_utils.py +++ b/src/puya/awst_build/eb/_utils.py @@ -5,9 +5,18 @@ import structlog from puya.awst import wtypes -from puya.awst.nodes import BoolConstant, Expression, IntrinsicCall, Literal +from puya.awst.nodes import ( + BigUIntBinaryOperator, + BoolConstant, + Expression, + IntrinsicCall, + Literal, + UInt64BinaryOperator, +) +from puya.awst_build.eb.base import BuilderBinaryOp from puya.awst_build.eb.var_factory import var_expression from puya.awst_build.utils import expect_operand_wtype +from puya.errors import CodeError if TYPE_CHECKING: from puya.awst_build.eb.base import ExpressionBuilder @@ -37,3 +46,41 @@ def uint64_to_biguint( stack_args=[arg], ) return itob_call + + +def translate_uint64_math_operator( + operator: BuilderBinaryOp, loc: SourceLocation +) -> UInt64BinaryOperator: + if operator is BuilderBinaryOp.div: + logger.error( + ( + "To maintain semantic compatibility with Python, " + "only the truncating division operator (//) is supported " + ), + location=loc, + ) + # continue traversing code to generate any further errors + operator = BuilderBinaryOp.floor_div + try: + return UInt64BinaryOperator(operator.value) + except ValueError as ex: + raise CodeError(f"Unsupported UInt64 math operator {operator.value}", loc) from ex + + +def translate_biguint_math_operator( + operator: BuilderBinaryOp, loc: SourceLocation +) -> BigUIntBinaryOperator: + if operator is BuilderBinaryOp.div: + logger.error( + ( + "To maintain semantic compatibility with Python, " + "only the truncating division operator (//) is supported " + ), + location=loc, + ) + # continue traversing code to generate any further errors + operator = BuilderBinaryOp.floor_div + try: + return BigUIntBinaryOperator(operator.value) + except ValueError as ex: + raise CodeError(f"Unsupported BigUInt math operator {operator.value}", loc) from ex diff --git a/src/puya/awst_build/eb/arc4/numeric.py b/src/puya/awst_build/eb/arc4/numeric.py index 6fb4bd4eca..954aa398cf 100644 --- a/src/puya/awst_build/eb/arc4/numeric.py +++ b/src/puya/awst_build/eb/arc4/numeric.py @@ -7,7 +7,9 @@ from puya.awst import wtypes from puya.awst.nodes import ( + ARC4Decode, ARC4Encode, + BigUIntBinaryOperation, DecimalConstant, Expression, IntegerConstant, @@ -15,15 +17,21 @@ NumericComparison, NumericComparisonExpression, ReinterpretCast, + Statement, + UInt64BinaryOperation, +) +from puya.awst_build.eb._utils import ( + translate_biguint_math_operator, + translate_uint64_math_operator, + uint64_to_biguint, ) -from puya.awst_build.eb._utils import uint64_to_biguint from puya.awst_build.eb.arc4.base import ( ARC4ClassExpressionBuilder, ARC4EncodedExpressionBuilder, arc4_bool_bytes, get_integer_literal_value, ) -from puya.awst_build.eb.base import BuilderComparisonOp, ExpressionBuilder +from puya.awst_build.eb.base import BuilderBinaryOp, BuilderComparisonOp, ExpressionBuilder from puya.awst_build.eb.var_factory import var_expression from puya.awst_build.utils import convert_literal_to_expr from puya.errors import CodeError, InternalError, TodoError @@ -185,6 +193,14 @@ def bool_eval(self, location: SourceLocation, *, negate: bool = False) -> Expres negate=negate, ) + def unary_plus(self, location: SourceLocation) -> ExpressionBuilder: + # unary + is allowed, but for the current types it has no real impact + # so just expand the existing expression to include the unary operator + raise TodoError(location) + + def bitwise_invert(self, location: SourceLocation) -> ExpressionBuilder: + raise TodoError(location) + def compare( self, other: ExpressionBuilder | Literal, op: BuilderComparisonOp, location: SourceLocation ) -> ExpressionBuilder: @@ -216,6 +232,96 @@ def compare( ) return var_expression(cmp_expr) + def binary_op( + self, + other: ExpressionBuilder | Literal, + op: BuilderBinaryOp, + location: SourceLocation, + *, + reverse: bool, + ) -> ExpressionBuilder: + other_expr = convert_literal_to_expr(other, self.wtype) + if self.wtype.n <= 64: + result_expr = self._uint64_binary_op(other_expr, op, location, reverse=reverse) + else: + result_expr = self._biguint_binary_op(other_expr, op, location, reverse=reverse) + encoded_result = ARC4Encode(value=result_expr, source_location=location, wtype=self.wtype) + return var_expression(encoded_result) + + def _uint64_binary_op( + self, other: Expression, op: BuilderBinaryOp, location: SourceLocation, *, reverse: bool + ) -> Expression: + if other.wtype == self.wtype: + other = ARC4Decode( + value=other, + wtype=wtypes.uint64_wtype, + source_location=location, + ) + elif isinstance(other.wtype, wtypes.ARC4UIntN): + raise TodoError(location, "TODO: support mixed size operators with arc4 numerics") + elif other.wtype == wtypes.uint64_wtype: + pass + elif other.wtype == wtypes.bool_wtype: + raise TodoError(location, "TODO: support upcast from bool to arc4.UIntN") + else: + return NotImplemented + lhs: Expression = ARC4Decode( + value=self.expr, + wtype=wtypes.uint64_wtype, + source_location=self.source_location, + ) + rhs = other + if reverse: + (lhs, rhs) = (rhs, lhs) + uint64_op = translate_uint64_math_operator(op, location) + bin_op_expr = UInt64BinaryOperation( + source_location=location, left=lhs, op=uint64_op, right=rhs + ) + return bin_op_expr + + def _biguint_binary_op( + self, other: Expression, op: BuilderBinaryOp, location: SourceLocation, *, reverse: bool + ) -> Expression: + if other.wtype == self.wtype: + other = ReinterpretCast( + expr=other, + wtype=wtypes.biguint_wtype, + source_location=other.source_location, + ) + elif isinstance(other.wtype, wtypes.ARC4UIntN): + raise TodoError(location, "TODO: support mixed size operators with arc4 numerics") + elif other.wtype == wtypes.uint64_wtype: + other = uint64_to_biguint(other, location) + elif other.wtype == wtypes.biguint_wtype: + pass + elif other.wtype == wtypes.bool_wtype: + raise TodoError(location, "TODO: support upcast from bool to arc4.UIntN") + else: + return NotImplemented + lhs: Expression = ReinterpretCast( + expr=self.expr, + wtype=wtypes.biguint_wtype, + source_location=self.source_location, + ) + rhs = other + if reverse: + (lhs, rhs) = (rhs, lhs) + biguint_op = translate_biguint_math_operator(op, location) + bin_op_expr = BigUIntBinaryOperation( + source_location=location, left=lhs, op=biguint_op, right=rhs + ) + return bin_op_expr + + def augmented_assignment( + self, op: BuilderBinaryOp, rhs: ExpressionBuilder | Literal, location: SourceLocation + ) -> Statement: + raise TodoError(location) + # rhs_expr = convert_literal_to_expr(rhs, self.wtype) + # if self.wtype.n <= 64: + # return self._uint64_augmented_assignment(rhs_expr, op, location) + # else: + # return self._biguint_augmented_assignment(rhs_expr, op, location) + class UFixedNxMExpressionBuilder(ARC4EncodedExpressionBuilder): def __init__(self, expr: Expression): diff --git a/src/puya/awst_build/eb/biguint.py b/src/puya/awst_build/eb/biguint.py index 7d90db0829..75a76bbc00 100644 --- a/src/puya/awst_build/eb/biguint.py +++ b/src/puya/awst_build/eb/biguint.py @@ -10,7 +10,6 @@ from puya.awst.nodes import ( BigUIntAugmentedAssignment, BigUIntBinaryOperation, - BigUIntBinaryOperator, BigUIntConstant, Literal, NumericComparison, @@ -18,7 +17,7 @@ ReinterpretCast, Statement, ) -from puya.awst_build.eb._utils import uint64_to_biguint +from puya.awst_build.eb._utils import translate_biguint_math_operator, uint64_to_biguint from puya.awst_build.eb.base import ( BuilderBinaryOp, BuilderComparisonOp, @@ -124,7 +123,7 @@ def binary_op( if other_expr.wtype == self.wtype: pass elif other_expr.wtype == wtypes.uint64_wtype: - other_expr = uint64_to_biguint(other, location) + other_expr = uint64_to_biguint(other_expr, location) elif other_expr.wtype == wtypes.bool_wtype: raise TodoError(location, "TODO: support upcast from bool to biguint") else: @@ -133,7 +132,7 @@ def binary_op( rhs = other_expr if reverse: (lhs, rhs) = (rhs, lhs) - biguint_op = _translate_biguint_math_operator(op, location) + biguint_op = translate_biguint_math_operator(op, location) bin_op_expr = BigUIntBinaryOperation( source_location=location, left=lhs, op=biguint_op, right=rhs ) @@ -154,29 +153,10 @@ def augmented_assignment( f"Invalid operand type {value.wtype} for {op.value}= with {self.wtype}", location ) target = self.lvalue() - biguint_op = _translate_biguint_math_operator(op, location) + biguint_op = translate_biguint_math_operator(op, location) return BigUIntAugmentedAssignment( source_location=location, target=target, value=value, op=biguint_op, ) - - -def _translate_biguint_math_operator( - operator: BuilderBinaryOp, loc: SourceLocation -) -> BigUIntBinaryOperator: - if operator is BuilderBinaryOp.div: - logger.error( - ( - "To maintain semantic compatibility with Python, " - "only the truncating division operator (//) is supported " - ), - location=loc, - ) - # continue traversing code to generate any further errors - operator = BuilderBinaryOp.floor_div - try: - return BigUIntBinaryOperator(operator.value) - except ValueError as ex: - raise CodeError(f"Unsupported BigUInt math operator {operator.value}", loc) from ex diff --git a/src/puya/awst_build/eb/uint64.py b/src/puya/awst_build/eb/uint64.py index 1d65b977b6..2016b05081 100644 --- a/src/puya/awst_build/eb/uint64.py +++ b/src/puya/awst_build/eb/uint64.py @@ -17,11 +17,11 @@ Statement, UInt64AugmentedAssignment, UInt64BinaryOperation, - UInt64BinaryOperator, UInt64Constant, UInt64UnaryOperation, UInt64UnaryOperator, ) +from puya.awst_build.eb._utils import translate_uint64_math_operator from puya.awst_build.eb.base import ( BuilderBinaryOp, BuilderComparisonOp, @@ -130,7 +130,7 @@ def binary_op( rhs = other_expr if reverse: (lhs, rhs) = (rhs, lhs) - uint64_op = _translate_uint64_math_operator(op, location) + uint64_op = translate_uint64_math_operator(op, location) bin_op_expr = UInt64BinaryOperation( source_location=location, left=lhs, op=uint64_op, right=rhs ) @@ -149,26 +149,7 @@ def augmented_assignment( f"Invalid operand type {value.wtype} for {op.value}= with {self.wtype}", location ) target = self.lvalue() - uint64_op = _translate_uint64_math_operator(op, location) + uint64_op = translate_uint64_math_operator(op, location) return UInt64AugmentedAssignment( source_location=location, target=target, value=value, op=uint64_op ) - - -def _translate_uint64_math_operator( - operator: BuilderBinaryOp, loc: SourceLocation -) -> UInt64BinaryOperator: - if operator is BuilderBinaryOp.div: - logger.error( - ( - "To maintain semantic compatibility with Python, " - "only the truncating division operator (//) is supported " - ), - location=loc, - ) - # continue traversing code to generate any further errors - operator = BuilderBinaryOp.floor_div - try: - return UInt64BinaryOperator(operator.value) - except ValueError as ex: - raise CodeError(f"Unsupported UInt64 math operator {operator.value}", loc) from ex diff --git a/src/puyapy-stubs/arc4.pyi b/src/puyapy-stubs/arc4.pyi index e06a81245d..575fe358f6 100644 --- a/src/puyapy-stubs/arc4.pyi +++ b/src/puyapy-stubs/arc4.pyi @@ -82,6 +82,10 @@ class UIntN(typing.Generic[_TBitSize], _ABIEncoded[puyapy.UInt64]): Max Size: 64 bits""" + # TODO: match support + # __match_value__: int + # __match_args__ = ("__match_value__",) + def __init__(self, value: int | puyapy.UInt64 | puyapy.BigUInt) -> None: ... # ~~~ https://docs.python.org/3/reference/datamodel.html#basic-customization ~~~ @@ -123,8 +127,79 @@ class UIntN(typing.Generic[_TBitSize], _ABIEncoded[puyapy.UInt64]): | puyapy.BigUInt | int, ) -> bool: ... + # truthiness def __bool__(self) -> bool: """Returns `True` if not equal to zero""" + # ~~~ https://docs.python.org/3/reference/datamodel.html#emulating-numeric-types ~~~ + # + + # TODO: these ops should support operating on different sizes and return the larger size, + # or whatever makes sense. Impossible to do purely through type annotations in a + # single class though + def __add__(self, other: typing.Self | puyapy.UInt64 | int) -> typing.Self: + """This will error on overflow""" + def __radd__(self, other: typing.Self | puyapy.UInt64 | int) -> typing.Self: + """This will error on overflow""" + def __iadd__(self, other: typing.Self | puyapy.UInt64 | int) -> typing.Self: + """This will error on overflow""" + # - + def __sub__(self, other: typing.Self | puyapy.UInt64 | int) -> typing.Self: + """This will error on underflow""" + def __rsub__(self, other: typing.Self | puyapy.UInt64 | int) -> typing.Self: + """This will error on underflow""" + def __isub__(self, other: typing.Self | puyapy.UInt64 | int) -> typing.Self: + """This will error on underflow""" + # * + def __mul__(self, other: typing.Self | puyapy.UInt64 | int) -> typing.Self: + """This will error on overflow""" + def __rmul__(self, other: typing.Self | puyapy.UInt64 | int) -> typing.Self: + """This will error on overflow""" + def __imul__(self, other: typing.Self | puyapy.UInt64 | int) -> typing.Self: + """This will error on overflow""" + # // + def __floordiv__(self, other: typing.Self | puyapy.UInt64 | int) -> typing.Self: + """This will error on divide by zero""" + def __rfloordiv__(self, other: typing.Self | puyapy.UInt64 | int) -> typing.Self: + """This will error on divide by zero""" + def __ifloordiv__(self, other: typing.Self | puyapy.UInt64 | int) -> typing.Self: + """This will error on divide by zero""" + # % + def __mod__(self, other: typing.Self | puyapy.UInt64 | int) -> typing.Self: + """This will error on mod by zero""" + def __rmod__(self, other: typing.Self | puyapy.UInt64 | int) -> typing.Self: + """This will error on mod by zero""" + def __imod__(self, other: typing.Self | puyapy.UInt64 | int) -> typing.Self: + """This will error on mod by zero""" + # TODO: __divmod__? only supported as single op via divmodw though 🤔 + # **, pow + def __pow__(self, power: typing.Self | puyapy.UInt64 | int) -> typing.Self: + """This will error on overflow""" + def __rpow__(self, power: typing.Self | puyapy.UInt64 | int) -> typing.Self: + """This will error on overflow""" + def __ipow__(self, power: typing.Self | puyapy.UInt64 | int) -> typing.Self: + """This will error on overflow""" + # TODO: the below may operator differently to just a .decode() -> op -> .encode() + # # << + # def __lshift__(self, other: typing.Self | puyapy.UInt64 | int) -> typing.Self: ... + # def __rlshift__(self, other: typing.Self | puyapy.UInt64 | int) -> typing.Self: ... + # def __ilshift__(self, other: typing.Self | puyapy.UInt64 | int) -> typing.Self: ... + # # >> + # def __rshift__(self, other: typing.Self | puyapy.UInt64 | int) -> typing.Self: ... + # def __rrshift__(self, other: typing.Self | puyapy.UInt64 | int) -> typing.Self: ... + # def __irshift__(self, other: typing.Self | puyapy.UInt64 | int) -> typing.Self: ... + # # & + # def __and__(self, other: typing.Self | puyapy.UInt64 | int) -> typing.Self: ... + # def __rand__(self, other: typing.Self | puyapy.UInt64 | int) -> typing.Self: ... + # def __iand__(self, other: typing.Self | puyapy.UInt64 | int) -> typing.Self: ... + # # ^ + # def __xor__(self, other: typing.Self | puyapy.UInt64 | int) -> typing.Self: ... + # def __rxor__(self, other: typing.Self | puyapy.UInt64 | int) -> typing.Self: ... + # def __ixor__(self, other: typing.Self | puyapy.UInt64 | int) -> typing.Self: ... + # # | + # def __or__(self, other: typing.Self | puyapy.UInt64 | int) -> typing.Self: ... + # def __ror__(self, other: typing.Self | puyapy.UInt64 | int) -> typing.Self: ... + # def __ior__(self, other: typing.Self | puyapy.UInt64 | int) -> typing.Self: ... + # # ~ + # def __invert__(self) -> typing.Self: ... class BigUIntN(typing.Generic[_TBitSize], _ABIEncoded[puyapy.BigUInt]): """An ARC4 UInt consisting of the number of bits specified. @@ -172,6 +247,7 @@ class BigUIntN(typing.Generic[_TBitSize], _ABIEncoded[puyapy.BigUInt]): | puyapy.BigUInt | int, ) -> bool: ... + # truthiness def __bool__(self) -> bool: """Returns `True` if not equal to zero"""