From 09767a4a0406e3f3471209613cf40658c7d41ab5 Mon Sep 17 00:00:00 2001
From: Andrei Marinica <andrei.marinica@multiversx.com>
Date: Mon, 18 Nov 2024 13:19:08 +0200
Subject: [PATCH] substitution preprocessor - storage mapper from address fix

---
 .../storage_mapper_address_to_id.scen.json    | 29 +++++++++++++++++--
 .../src/storage_mapper_get_at_address.rs      | 12 ++++++++
 .../src/preprocessing/substitution_list.rs    |  1 +
 3 files changed, 40 insertions(+), 2 deletions(-)

diff --git a/contracts/feature-tests/basic-features/scenarios/storage_mapper_address_to_id.scen.json b/contracts/feature-tests/basic-features/scenarios/storage_mapper_address_to_id.scen.json
index 91493a8d0c..559492d048 100644
--- a/contracts/feature-tests/basic-features/scenarios/storage_mapper_address_to_id.scen.json
+++ b/contracts/feature-tests/basic-features/scenarios/storage_mapper_address_to_id.scen.json
@@ -1,6 +1,5 @@
 {
     "name": "storage",
-    "gasSchedule": "v3",
     "steps": [
         {
             "step": "setState",
@@ -13,7 +12,10 @@
                 "sc:extra-instance": {
                     "nonce": "0",
                     "balance": "0",
-                    "code": "mxsc:../output/basic-features.mxsc.json"
+                    "code": "mxsc:../output/basic-features.mxsc.json",
+                    "storage": {
+                        "str:contract_address": "sc:basic-features"
+                    }
                 },
                 "address:an_account": {
                     "nonce": "0",
@@ -67,6 +69,29 @@
                 "refund": "*"
             }
         },
+        {
+            "step": "scCall",
+            "id": "address_to_id_mapper_get_id_from_address",
+            "tx": {
+                "from": "address:an_account",
+                "to": "sc:extra-instance",
+                "function": "address_to_id_mapper_get_id_from_address",
+                "arguments": [
+                    "address:an_account"
+                ],
+                "gasLimit": "50,000,000",
+                "gasPrice": "0"
+            },
+            "expect": {
+                "out": [
+                    "1"
+                ],
+                "status": "",
+                "logs": "*",
+                "gas": "*",
+                "refund": "*"
+            }
+        },
         {
             "step": "scCall",
             "id": "get_id_non_zero",
diff --git a/contracts/feature-tests/basic-features/src/storage_mapper_get_at_address.rs b/contracts/feature-tests/basic-features/src/storage_mapper_get_at_address.rs
index 324df37021..6a29f41461 100644
--- a/contracts/feature-tests/basic-features/src/storage_mapper_get_at_address.rs
+++ b/contracts/feature-tests/basic-features/src/storage_mapper_get_at_address.rs
@@ -146,4 +146,16 @@ pub trait StorageMapperGetAtAddress {
         self.single_value_from_address_with_keys(address, extra_key)
             .get()
     }
+
+    #[storage_mapper_from_address("address_ids")]
+    fn address_ids_from_address(
+        &self,
+        address: ManagedAddress,
+    ) -> AddressToIdMapper<ManagedAddress>;
+
+    #[view]
+    fn address_to_id_mapper_get_id_from_address(&self, address_arg: ManagedAddress) -> AddressId {
+        let address = self.contract_address().get();
+        self.address_ids_from_address(address).get_id(&address_arg)
+    }
 }
diff --git a/framework/derive/src/preprocessing/substitution_list.rs b/framework/derive/src/preprocessing/substitution_list.rs
index a04da2a4c1..c03bd3b29e 100644
--- a/framework/derive/src/preprocessing/substitution_list.rs
+++ b/framework/derive/src/preprocessing/substitution_list.rs
@@ -90,6 +90,7 @@ fn add_storage_mapper_single_generic_arg(
     substitutions: &mut SubstitutionsMap,
     mapper_name: &proc_macro2::TokenStream,
 ) {
+    add_managed_type_with_generics(substitutions, mapper_name);
     substitutions.add_substitution(
         quote!(#mapper_name<Self::Api>),
         quote!(#mapper_name<Self::Api>),