feat: Type storage placeholders and support for templated maps

objects/src/accounts/component/template/mod.rs

@@ -1,5 +1,5 @@
 use alloc::{
-    collections::BTreeSet,
+    collections::{btree_map::Entry, BTreeMap, BTreeSet},
     string::{String, ToString},
     vec::Vec,
 };
@@ -188,14 +188,14 @@ impl AccountComponentMetadata {
     /// be used for initializing storage slot values, or storage map entries.
     /// For a full example on how a placeholder may be utilized, refer to the docs
     /// for [AccountComponentMetadata].
-    pub fn get_storage_placeholders(&self) -> BTreeSet<StoragePlaceholder> {
-        let mut placeholder_set = BTreeSet::new();
+    pub fn get_storage_placeholders(&self) -> BTreeMap<StoragePlaceholder, PlaceholderType> {
+        let mut placeholder_map = BTreeMap::new();
         for storage_entry in &self.storage {
-            for key in storage_entry.placeholders() {
-                placeholder_set.insert(key.clone());
+            for (placeholder, placeholder_type) in storage_entry.placeholders() {
+                placeholder_map.insert(placeholder.clone(), placeholder_type);
             }
         }
-        placeholder_set
+        placeholder_map
     }
 
     /// Returns the name of the account component.
@@ -254,6 +254,28 @@ impl AccountComponentMetadata {
                 return Err(AccountComponentTemplateError::NonContiguousSlots(slots[0], slots[1]));
             }
         }
+
+        // Check that placeholders do not appear more than once with a different type
+        let mut placeholders = BTreeMap::new();
+        for storage_entry in &self.storage {
+            for (placeholder, placeholder_type) in storage_entry.placeholders() {
+                match placeholders.entry(placeholder.clone()) {
+                    Entry::Occupied(entry) => {
+                        if *entry.get() != placeholder_type {
+                            return Err(
+                                AccountComponentTemplateError::StoragePlaceholderDuplicate(
+                                    placeholder.clone(),
+                                ),
+                            );
+                        }
+                    },
+                    Entry::Vacant(slot) => {
+                        slot.insert(placeholder_type);
+                    },
+                }
+            }
+        }
+
         Ok(())
     }
 }
@@ -290,7 +312,6 @@ impl Deserializable for AccountComponentMetadata {
 
 #[cfg(test)]
 mod tests {
-
     use assembly::Assembler;
     use assert_matches::assert_matches;
     use storage::WordRepresentation;

objects/src/accounts/component/template/storage/mod.rs

@@ -13,7 +13,7 @@ use super::AccountComponentTemplateError;
 use crate::accounts::{StorageMap, StorageSlot};
 
 mod placeholder;
-pub use placeholder::{StoragePlaceholder, StorageValue};
+pub use placeholder::{PlaceholderType, StoragePlaceholder, StorageValue};
 
 mod init_storage_data;
 pub use init_storage_data::InitStorageData;
@@ -54,7 +54,7 @@ pub enum StorageEntry {
         /// The numeric index of this map slot in the component's storage.
         slot: u8,
         /// A list of key-value pairs to initialize in this map slot.
-        map_entries: Vec<MapEntry>,
+        map_entries: MapRepresentation,
     },
 
     /// A multi-slot entry, representing a single logical value across multiple slots.
@@ -91,13 +91,13 @@ impl StorageEntry {
         name: impl Into<String>,
         description: Option<impl Into<String>>,
         slot: u8,
-        map_entries: Vec<MapEntry>,
+        map_representation: MapRepresentation,
     ) -> Self {
         StorageEntry::Map {
             name: name.into(),
             description: description.map(Into::<String>::into),
             slot,
-            map_entries,
+            map_entries: map_representation,
         }
     }
 
@@ -167,24 +167,14 @@ impl StorageEntry {
         }
     }
 
-    /// Returns the map entries for a `Map` variant as a slice.
-    /// Returns an empty slice for non-map variants.
-    pub fn map_entries(&self) -> &[MapEntry] {
-        match self {
-            StorageEntry::Map { map_entries: values, .. } => values.as_slice(),
-            StorageEntry::Value { .. } => &[],
-            StorageEntry::MultiSlot { .. } => &[],
-        }
-    }
-
     /// Returns an iterator over all of the storage entries's placeholder keys.
     // TODO: Should placeholders be typed?
-    pub fn placeholders(&self) -> Box<dyn Iterator<Item = &StoragePlaceholder> + '_> {
+    pub fn placeholders(
+        &self,
+    ) -> Box<dyn Iterator<Item = (&StoragePlaceholder, PlaceholderType)> + '_> {
         match self {
             StorageEntry::Value { value, .. } => value.placeholders(),
-            StorageEntry::Map { map_entries: values, .. } => {
-                Box::new(values.iter().flat_map(|word| word.placeholders()))
-            },
+            StorageEntry::Map { map_entries, .. } => map_entries.placeholders(),
             StorageEntry::MultiSlot { values, .. } => {
                 Box::new(values.iter().flat_map(|word| word.placeholders()))
             },
@@ -209,14 +199,7 @@ impl StorageEntry {
                 Ok(vec![StorageSlot::Value(slot)])
             },
             StorageEntry::Map { map_entries: values, .. } => {
-                let entries = values
-                    .iter()
-                    .map(|map_entry| {
-                        let key = map_entry.key().try_build_word(init_storage_data)?;
-                        let value = map_entry.value().try_build_word(init_storage_data)?;
-                        Ok((key.into(), value))
-                    })
-                    .collect::<Result<Vec<(Digest, Word)>, AccountComponentTemplateError>>()?; // Collect into a Vec and propagate errors
+                let entries = values.try_build_map(init_storage_data)?;
 
                 let storage_map = StorageMap::with_entries(entries);
                 Ok(vec![StorageSlot::Map(storage_map)])
@@ -285,13 +268,13 @@ impl Deserializable for StorageEntry {
             // Map
             1 => {
                 let slot = source.read_u8()?;
-                let values: Vec<MapEntry> = source.read()?;
+                let map_representation: MapRepresentation = source.read()?;
 
                 Ok(StorageEntry::Map {
                     name,
                     description,
                     slot,
-                    map_entries: values,
+                    map_entries: map_representation,
                 })
             },
 
@@ -335,7 +318,7 @@ impl MapEntry {
         &self.value
     }
 
-    pub fn placeholders(&self) -> impl Iterator<Item = &StoragePlaceholder> {
+    pub fn placeholders(&self) -> impl Iterator<Item = (&StoragePlaceholder, PlaceholderType)> {
         self.key.placeholders().chain(self.value.placeholders())
     }
 
@@ -365,7 +348,8 @@ impl Deserializable for MapEntry {
 
 #[cfg(test)]
 mod tests {
-    use std::collections::BTreeSet;
+    use core::panic;
+    use std::{collections::BTreeSet, string::ToString};
 
     use assembly::Assembler;
     use assert_matches::assert_matches;
@@ -380,7 +364,7 @@ mod tests {
         },
         digest,
         testing::account_code::CODE,
-        AccountError,
+        AccountError, MAX_ASSETS_PER_NOTE,
     };
 
     #[test]
@@ -402,7 +386,7 @@ mod tests {
                 name: "map".into(),
                 description: Some("A storage map entry".into()),
                 slot: 1,
-                map_entries: vec![
+                map_entries: MapRepresentation::List(vec![
                     MapEntry {
                         key: WordRepresentation::Template(
                             StoragePlaceholder::new("foo.bar").unwrap(),
@@ -419,7 +403,7 @@ mod tests {
                         key: WordRepresentation::Hexadecimal(digest!("0x3").into()),
                         value: WordRepresentation::Hexadecimal(digest!("0x4").into()),
                     },
-                ],
+                ]),
             },
             StorageEntry::MultiSlot {
                 name: "multi".into(),
@@ -470,7 +454,7 @@ mod tests {
             slot = 0
             values = [
                 { key = "0x1", value = ["{{value.test}}", "0x1", "0x2", "0x3"] },
-                { key = "{{key.test}}", value = "0x3" },
+                { key = "{{map.key.test}}", value = "0x3" },
                 { key = "0x3", value = "0x4" }
             ]
 
@@ -488,12 +472,25 @@ mod tests {
                 "{{word.test}}",
                 ["1", "0", "0", "0"],
             ]
+
+            [[storage]]
+            name = "map-template"
+            description = "a templated map"
+            slot = 4
+            values = "{{map.template}}"
         "#;
 
         let component_metadata = AccountComponentMetadata::from_toml(toml_text).unwrap();
-        let library = Assembler::default().assemble_library([CODE]).unwrap();
+        for (key, placeholder_type) in component_metadata.get_storage_placeholders() {
+            match key.inner() {
+                "map.key.test" | "word.test" => assert_eq!(placeholder_type, PlaceholderType::Word),
+                "value.test" => assert_eq!(placeholder_type, PlaceholderType::Felt),
+                "map.template" => assert_eq!(placeholder_type, PlaceholderType::Map),
+                _ => panic!("all cases are covered"),
+            }
+        }
 
-        assert_eq!(component_metadata.storage_entries().first().unwrap().map_entries().len(), 3);
+        let library = Assembler::default().assemble_library([CODE]).unwrap();
 
         let template = AccountComponentTemplate::new(component_metadata, library);
 
@@ -504,7 +501,7 @@ mod tests {
 
         let storage_placeholders = InitStorageData::new([
             (
-                StoragePlaceholder::new("key.test").unwrap(),
+                StoragePlaceholder::new("map.key.test").unwrap(),
                 StorageValue::Word(Default::default()),
             ),
             (
@@ -515,6 +512,10 @@ mod tests {
                 StoragePlaceholder::new("word.test").unwrap(),
                 StorageValue::Word([Felt::ZERO, Felt::ZERO, Felt::ZERO, Felt::new(128)]),
             ),
+            (
+                StoragePlaceholder::new("map.template").unwrap(),
+                StorageValue::Map(vec![(Digest::default(), Word::default())]),
+            ),
         ]);
 
         let component = AccountComponent::from_template(&template, &storage_placeholders).unwrap();
@@ -542,4 +543,33 @@ mod tests {
             ))
         );
     }
+
+    #[test]
+    pub fn fail_duplicate_key() {
+        let toml_text = r#"
+            name = "Test Component"
+            description = "This is a test component"
+            version = "1.0.1"
+            targets = ["FungibleFaucet"]
+
+            [[storage]]
+            name = "map"
+            description = "A storage map entry"
+            slot = 0
+            values = [
+                { key = "0x1", value = ["{{value.test}}", "0x1", "0x2", "0x3"] },
+            ]
+
+            [[storage]]
+            name = "word"
+            slot = 1
+            value = "{{value.test}}"
+        "#;
+        let component_metadata = AccountComponentMetadata::from_toml(toml_text);
+        let expected_failure_key = StoragePlaceholder::new("value.test").unwrap();
+        assert_matches!(
+            component_metadata,
+            Err(AccountComponentTemplateError::StoragePlaceholderDuplicate(expected_failure_key))
+        );
+    }
 }

objects/src/accounts/component/template/storage/placeholder.rs

@@ -28,6 +28,13 @@ pub struct StoragePlaceholder {
     key: String,
 }
 
+#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
+pub enum PlaceholderType {
+    Felt,
+    Map,
+    Word,
+}
+
 impl StoragePlaceholder {
     /// Creates a new [StoragePlaceholder] from the provided string.
     ///

objects/src/accounts/component/template/storage/toml.rs

@@ -8,7 +8,9 @@ use serde::{de::Error, Deserialize, Deserializer, Serialize, Serializer};
 use vm_core::Felt;
 use vm_processor::Digest;
 
-use super::{FeltRepresentation, StorageEntry, StoragePlaceholder, WordRepresentation};
+use super::{
+    FeltRepresentation, MapRepresentation, StorageEntry, StoragePlaceholder, WordRepresentation,
+};
 use crate::{
     accounts::{component::template::MapEntry, AccountComponentMetadata},
     errors::AccountComponentTemplateError,
@@ -40,6 +42,7 @@ impl AccountComponentMetadata {
         Ok(toml)
     }
 }
+
 // WORD REPRESENTATION SERIALIZATION
 // ================================================================================================
 
@@ -207,41 +210,35 @@ enum StorageValues {
     /// List of individual words (for multi-slot entries).
     Words(Vec<WordRepresentation>),
     /// List of key-value entries (for map storage slots).
-    MapEntries(Vec<MapEntry>),
-    /// A placeholder value, represented as "{{key}}".
-    Template(StoragePlaceholder),
+    MapEntries(MapRepresentation),
 }
 
 impl StorageValues {
     pub fn is_list_of_words(&self) -> bool {
         match self {
             StorageValues::Words(_) => true,
             StorageValues::MapEntries(_) => false,
-            StorageValues::Template(_) => false,
         }
     }
 
     pub fn into_words(self) -> Option<Vec<WordRepresentation>> {
         match self {
             StorageValues::Words(vec) => Some(vec),
             StorageValues::MapEntries(_) => None,
-            StorageValues::Template(_) => None,
         }
     }
 
-    pub fn into_map_entries(self) -> Option<Vec<MapEntry>> {
+    pub fn into_map_entries(self) -> Option<MapRepresentation> {
         match self {
             StorageValues::Words(_) => None,
-            StorageValues::MapEntries(vec) => Some(vec),
-            StorageValues::Template(_) => None,
+            StorageValues::MapEntries(map) => Some(map),
         }
     }
 
     pub fn len(&self) -> Option<usize> {
         match self {
             StorageValues::Words(vec) => Some(vec.len()),
-            StorageValues::MapEntries(vec) => Some(vec.len()),
-            StorageValues::Template(_) => None,
+            StorageValues::MapEntries(map) => map.len(),
         }
     }
 }