From fc15e07470bd4748fea9eb4e1cf5f6d356a68fcb Mon Sep 17 00:00:00 2001
From: Jon Michael Aanes <jonjmaa@gmail.com>
Date: Wed, 9 Apr 2025 21:37:26 +0200
Subject: [PATCH] Add support for AvlTreeIds

---
 pbcabi/model.py | 27 +++++++++++++++++++++------
 1 file changed, 21 insertions(+), 6 deletions(-)

diff --git a/pbcabi/model.py b/pbcabi/model.py
index b25c023..204477d 100644
--- a/pbcabi/model.py
+++ b/pbcabi/model.py
@@ -237,6 +237,9 @@ class MapTypeSpec(TypeSpec):
     type_key: TypeSpec
     type_value: TypeSpec
 
+
+
+
 @enforce_typing.enforce_types
 @dataclasses.dataclass(frozen=True, slots=True)
 class AvlTreeTypeSpec(MapTypeSpec):
@@ -252,8 +255,18 @@ class AvlTreeTypeSpec(MapTypeSpec):
         '''
         type_key = TypeSpec.read_from(reader)
         type_value = TypeSpec.read_from(reader)
-        type_spec = AvlTreeTypeSpec(type_key, type_value)
-        return type_spec
+        return AvlTreeTypeSpec(type_key, type_value)
+
+    def read_element_from(self, reader: BinaryReader, type_env,
+                          mode: SerializeMode) -> 'AvlTreeId':
+        avl_tree_id = SIZE_SIMPLE_TYPE_SPEC.read_element_from(reader, type_env, mode)
+        return AvlTreeId(avl_tree_id, self)
+
+@enforce_typing.enforce_types
+@dataclasses.dataclass(frozen=True, slots=True)
+class AvlTreeId:
+    avl_tree_id: int
+    type_spec: AvlTreeTypeSpec
 
 @enforce_typing.enforce_types
 @dataclasses.dataclass(frozen=True, slots=True)
@@ -695,11 +708,13 @@ class ContractAbi:
         return self.named_types_by_id[type_name].read_element_from(
             rpc, self.named_types_by_idx, SerializeMode.RPC)
 
-    def read_state(self, reader: BinaryReader):
+    def read_state(self, reader: BinaryReader | bytes, explicit_type: TypeSpec | None = None):
         if not isinstance(reader, BinaryReader):
             reader = BinaryReader(reader)
 
-        return self.state_type.read_element_from(reader, self.named_types_by_idx,
+        explicit_type = explicit_type or self.state_type
+
+        return explicit_type.read_element_from(reader, self.named_types_by_idx,
                                                  SerializeMode.STATE)
 
     def get_fn_abi_from_shortname(self, shortname: int) -> FnAbi:
@@ -707,7 +722,7 @@ class ContractAbi:
             if derp.shortname == shortname:
                 return derp
 
-    def read_invocation(self, reader: BinaryReader) -> ParsedInvocation:
+    def read_invocation(self, reader: BinaryReader | bytes) -> ParsedInvocation:
         if not isinstance(reader, BinaryReader):
             reader = BinaryReader(reader)
         shortname = reader.readLeb128()
@@ -730,7 +745,7 @@ class FileAbi:
     contract: ContractAbi
 
     @staticmethod
-    def read_from(reader: BinaryReader):
+    def read_from(reader: BinaryReader | bytes):
         '''
         Deserialize 'FileAbi' from 'BinaryReader'.
         '''