Add NonZeroType (#1527)

Author: franciszekjob

docs/migration_guide.rst

@@ -1,6 +1,17 @@
 Migration guide
 ===============
 
+******************************
+[unreleased] Migration guide
+******************************
+
+[unreleased] Minor changes
+--------------------------
+
+.. currentmodule:: starknet_py.cairo.data_types
+
+1. Added :class:`NonZeroType` in order to fix parsing ABI which contains Cairo`s `core::zeroable::NonZero <https://github.com/starkware-libs/cairo/blob/a2b9dddeb3212c8d529538454745b27d7a34a6cd/corelib/src/zeroable.cairo#L78>`_
+
 ******************************
 0.24.3 Migration guide
 ******************************

starknet_py/abi/v2/parser_transformer.py

@@ -1,5 +1,5 @@
 from math import log2
-from typing import Any, List, Optional
+from typing import Any, List, Optional, Union
 
 import lark
 from lark import Token, Transformer
@@ -9,6 +9,7 @@
     BoolType,
     CairoType,
     FeltType,
+    NonZeroType,
     OptionType,
     TupleType,
     TypeIdentifier,
@@ -31,6 +32,7 @@
         | type_class_hash
         | type_storage_address
         | type_option
+        | type_non_zero
         | type_array
         | type_span
         | tuple
@@ -49,6 +51,7 @@
     type_option: "core::option::Option::<" (type | type_identifier) ">"
     type_array: "core::array::Array::<" (type | type_identifier) ">"
     type_span: "core::array::Span::<" (type | type_identifier) ">"
+    type_non_zero: "core::zeroable::NonZero::<" (type | type_identifier) ">"
     
     tuple: "(" type? ("," type?)* ")"
     
@@ -185,6 +188,12 @@ def tuple(self, types: List[CairoType]) -> TupleType:
         """
         return TupleType(types)
 
+    def type_non_zero(self, value: List[Union[FeltType, UintType]]) -> NonZeroType:
+        """
+        NonZero contains value which is never zero.
+        """
+        return NonZeroType(value[0])
+
 
 def parse(
     code: str,

starknet_py/cairo/data_types.py

@@ -122,3 +122,12 @@ class EventType(CairoType):
     name: str
     types: OrderedDict[str, CairoType]
     keys: List[str]
+
+
+@dataclass
+class NonZeroType(CairoType):
+    """
+    Type representation of Cairo NonZero.
+    """
+
+    type: CairoType

starknet_py/serialization/data_serializers/non_zero_serializer.py

@@ -0,0 +1,39 @@
+from dataclasses import dataclass
+from typing import Any, Generator
+
+from starknet_py.serialization._context import (
+    Context,
+    DeserializationContext,
+    SerializationContext,
+)
+from starknet_py.serialization.data_serializers.cairo_data_serializer import (
+    CairoDataSerializer,
+)
+
+
+@dataclass
+class NonZeroSerializer(CairoDataSerializer[Any, int]):
+    """
+    Serializer for NonZero type.
+    Can serialize Cairo types which are non-zero.
+    Deserializes data to int.
+    """
+
+    serializer: CairoDataSerializer
+
+    def deserialize_with_context(self, context: DeserializationContext) -> int:
+        deserialized_value = self.serializer.deserialize_with_context(context)
+        self._ensure_valid_nonzero(deserialized_value, context)
+        return deserialized_value
+
+    def serialize_with_context(
+        self,
+        context: SerializationContext,
+        value: Any,
+    ) -> Generator[int, None, None]:
+        self._ensure_valid_nonzero(value, context)
+        return self.serializer.serialize_with_context(context, value)
+
+    @staticmethod
+    def _ensure_valid_nonzero(value: int, context: Context):
+        context.ensure_valid_value(value != 0, "expected value to be non-zero")

starknet_py/serialization/factory.py

@@ -14,6 +14,7 @@
     EventType,
     FeltType,
     NamedTupleType,
+    NonZeroType,
     OptionType,
     StructType,
     TupleType,
@@ -33,6 +34,9 @@
 from starknet_py.serialization.data_serializers.named_tuple_serializer import (
     NamedTupleSerializer,
 )
+from starknet_py.serialization.data_serializers.non_zero_serializer import (
+    NonZeroSerializer,
+)
 from starknet_py.serialization.data_serializers.option_serializer import (
     OptionSerializer,
 )
@@ -119,6 +123,9 @@ def serializer_for_type(cairo_type: CairoType) -> CairoDataSerializer:
     if isinstance(cairo_type, OptionType):
         return OptionSerializer(serializer_for_type(cairo_type.type))
 
+    if isinstance(cairo_type, NonZeroType):
+        return NonZeroSerializer(serializer_for_type(cairo_type.type))
+
     if isinstance(cairo_type, UnitType):
         return UnitSerializer()
 

starknet_py/tests/e2e/mock/contracts_v2/src/abi_types.cairo

@@ -1,5 +1,9 @@
 use core::serde::Serde;
 use starknet::ContractAddress;
+use core::zeroable::NonZero;
+use core::zeroable::IsZeroResult;
+use core::zeroable::NonZeroIntoImpl;
+use core::integer::{u8_try_as_non_zero};
 
 #[derive(Drop, Serde)]
 enum ExampleEnum {
@@ -13,6 +17,8 @@ struct ExampleStruct {
     field_b: felt252,
     field_c: ExampleEnum,
     field_d: (),
+    field_e: NonZero<felt252>,
+    field_f: NonZero<u8>,
 }
 
 #[starknet::interface]
@@ -23,13 +29,26 @@ trait IAbiTest<TContractState> {
     ) -> ExampleStruct;
 }
 
+pub fn felt_to_nonzero(value: felt252) -> NonZero<felt252> {
+    match felt252_is_zero(value) {
+        IsZeroResult::Zero(()) => panic(ArrayTrait::new()),
+        IsZeroResult::NonZero(x) => x,
+    }
+}
+
+pub fn u8_to_nonzero(value: u8) -> NonZero<u8> {
+    match u8_try_as_non_zero(value) {
+        Option::Some(x) => x,
+        Option::None => panic(ArrayTrait::new()),
+    }
+}
 
 #[starknet::contract]
 mod AbiTypes {
     use core::array::ArrayTrait;
     use core::traits::Into;
     use starknet::ContractAddress;
-    use super::{ExampleEnum, ExampleStruct};
+    use super::{ExampleEnum, ExampleStruct, felt_to_nonzero, u8_to_nonzero};
 
     #[storage]
     struct Storage {}
@@ -46,7 +65,12 @@ mod AbiTypes {
             ref self: ContractState, recipient: ContractAddress, amount: u256
         ) -> ExampleStruct {
             ExampleStruct {
-                field_a: 200, field_b: 300, field_c: ExampleEnum::variant_b(400.into()), field_d: ()
+                field_a: 200,
+                field_b: 300,
+                field_c: ExampleEnum::variant_b(400.into()),
+                field_d: (),
+                field_e: felt_to_nonzero(100),
+                field_f: u8_to_nonzero(100),
             }
         }
     }

starknet_py/tests/unit/serialization/data_serializers/non_zero_serializer.py

@@ -0,0 +1,39 @@
+import pytest
+
+from starknet_py.constants import FIELD_PRIME
+from starknet_py.serialization import FeltSerializer
+from starknet_py.serialization.data_serializers.non_zero_serializer import (
+    NonZeroSerializer,
+)
+from starknet_py.serialization.data_serializers.uint_serializer import UintSerializer
+
+
+@pytest.mark.parametrize(
+    "serializer, value, serialized_value",
+    [
+        (NonZeroSerializer(UintSerializer(128)), 123, [123]),
+        (NonZeroSerializer(UintSerializer(256)), 1, [1, 0]),
+        (NonZeroSerializer(FeltSerializer()), 10, [10]),
+        (NonZeroSerializer(FeltSerializer()), FIELD_PRIME - 1, [FIELD_PRIME - 1]),
+    ],
+)
+def test_valid_values(serializer, value, serialized_value):
+    deserialized = serializer.deserialize(serialized_value)
+    assert deserialized == value
+
+    serialized = serializer.serialize(value)
+    assert serialized == serialized_value
+
+
+@pytest.mark.parametrize(
+    "serializer, value, serialized_value",
+    [
+        (NonZeroSerializer(UintSerializer(128)), 0, [0]),
+        (NonZeroSerializer(UintSerializer(256)), 0, [0, 0]),
+        (NonZeroSerializer(FeltSerializer()), 0, [0]),
+    ],
+)
+def test_invalid_values(serializer, value, serialized_value):
+    with pytest.raises(ValueError, match="expected value to be non-zero"):
+        serializer.deserialize(serialized_value)
+        serializer.serialize(value)