From 2e10080ec5673ec6c378dc25648e4921cf35b357 Mon Sep 17 00:00:00 2001
From: saeed-zil <saeed@zilliqa.com>
Date: Thu, 21 Nov 2024 19:13:58 +0330
Subject: [PATCH] Fix parsing complex address types in transitions.

---
 src/simplified_representation/emitter.rs | 54 ++++++++++++++----------
 tests/contracts/ByStr.scilla             |  6 ++-
 tests/full_contract_tests.rs             | 14 +++++-
 3 files changed, 50 insertions(+), 24 deletions(-)

diff --git a/src/simplified_representation/emitter.rs b/src/simplified_representation/emitter.rs
index f74f59f..bec032f 100644
--- a/src/simplified_representation/emitter.rs
+++ b/src/simplified_representation/emitter.rs
@@ -283,38 +283,48 @@ impl AstConverting for SrEmitter {
     }
     fn emit_address_type_field(
         &mut self,
-        _mode: TreeTraversalMode,
+        mode: TreeTraversalMode,
         node: &NodeAddressTypeField,
     ) -> Result<TraversalResult, String> {
-        if let NodeVariableIdentifier::VariableName(n) = &node.identifier.node {
-            node.type_name.visit(self)?;
-            let typename = self.pop_type_definition()?;
-            let s = StackObject::VariableDeclaration(Field::new(&n.node, typename.into()));
-            self.stack.push(s);
+        match mode {
+            TreeTraversalMode::Enter => {
+                if let NodeVariableIdentifier::VariableName(n) = &node.identifier.node {
+                    node.type_name.visit(self)?;
+                    let typename = self.pop_type_definition()?;
+                    let s = StackObject::VariableDeclaration(Field::new(&n.node, typename.into()));
+                    self.stack.push(s);
+                }
+            }
+            TreeTraversalMode::Exit => (),
         }
         Ok(TraversalResult::SkipChildren)
     }
     fn emit_address_type(
         &mut self,
-        _mode: TreeTraversalMode,
+        mode: TreeTraversalMode,
         node: &NodeAddressType,
     ) -> Result<TraversalResult, String> {
-        node.identifier.visit(self)?;
-        let identifier = self.pop_ir_identifier()?;
-        self.stack
-            .push(StackObject::TypeDefinition(identifier.into()));
-        let mut main_type = self.pop_type_definition()?;
-        let mut fields = vec![];
-        for field in &node.address_fields {
-            field.visit(self)?;
-            let field = self.pop_variable_declaration()?;
-            fields.push(field);
+        match mode {
+            TreeTraversalMode::Enter => {
+                node.identifier.visit(self)?;
+                let identifier = self.pop_ir_identifier()?;
+                self.stack
+                    .push(StackObject::TypeDefinition(identifier.into()));
+                let mut main_type = self.pop_type_definition()?;
+                let mut fields = vec![];
+                for field in &node.address_fields {
+                    field.visit(self)?;
+                    let field = self.pop_variable_declaration()?;
+                    fields.push(field);
+                }
+                main_type.address_type = Some(AddressType {
+                    type_name: node.type_name.node.clone(),
+                    fields: FieldList(fields),
+                });
+                self.stack.push(StackObject::TypeDefinition(main_type));
+            }
+            TreeTraversalMode::Exit => (),
         }
-        main_type.address_type = Some(AddressType {
-            type_name: node.type_name.node.clone(),
-            fields: FieldList(fields),
-        });
-        self.stack.push(StackObject::TypeDefinition(main_type));
         Ok(TraversalResult::SkipChildren)
     }
 
diff --git a/tests/contracts/ByStr.scilla b/tests/contracts/ByStr.scilla
index 31090e3..c07c2e2 100644
--- a/tests/contracts/ByStr.scilla
+++ b/tests/contracts/ByStr.scilla
@@ -12,4 +12,8 @@ contract AllByStrVariants
         field balances : Map ByStr20 Uint128,
         field total_supply : Uint128
         end
-  )
\ No newline at end of file
+  )
+
+transition ArbitrageFromXCAD
+  (token : ByStr20 with contract field balances : Map ByStr20 Uint128 end)
+end
diff --git a/tests/full_contract_tests.rs b/tests/full_contract_tests.rs
index 350294b..3ed4511 100644
--- a/tests/full_contract_tests.rs
+++ b/tests/full_contract_tests.rs
@@ -67,7 +67,19 @@ fn test_bystr_contract_parse() {
                 )
             ]),
             fields: FieldList::default(),
-            transitions: TransitionList::default(),
+            transitions: TransitionList(vec![Transition::new(
+                "ArbitrageFromXCAD",
+                FieldList(vec![Field::new(
+                    "token",
+                    Type::ByStr20With {
+                        type_name: "contract".to_string(),
+                        fields: FieldList(vec![Field::new(
+                            "balances",
+                            Type::Map(Box::new(Type::ByStr20), Box::new(Type::Uint128))
+                        ),],),
+                    },
+                )]),
+            )])
         }
     );
 }