fuels_programs/contract/
storage.rs

1use std::{
2    collections::HashMap,
3    default::Default,
4    fmt::Debug,
5    io,
6    path::{Path, PathBuf},
7};
8
9use fuel_tx::{Bytes32, StorageSlot};
10use fuels_core::types::errors::{error, Result};
11
12/// Configuration for contract storage
13#[derive(Debug, Clone)]
14pub struct StorageConfiguration {
15    autoload_storage: bool,
16    slot_overrides: StorageSlots,
17}
18
19impl Default for StorageConfiguration {
20    fn default() -> Self {
21        Self {
22            autoload_storage: true,
23            slot_overrides: Default::default(),
24        }
25    }
26}
27
28impl StorageConfiguration {
29    pub fn new(autoload_enabled: bool, slots: impl IntoIterator<Item = StorageSlot>) -> Self {
30        let config = Self {
31            autoload_storage: autoload_enabled,
32            slot_overrides: Default::default(),
33        };
34
35        config.add_slot_overrides(slots)
36    }
37
38    /// If enabled will try to automatically discover and load the storage configuration from the
39    /// storage config json file.
40    pub fn with_autoload(mut self, enabled: bool) -> Self {
41        self.autoload_storage = enabled;
42        self
43    }
44
45    pub fn autoload_enabled(&self) -> bool {
46        self.autoload_storage
47    }
48
49    /// Slots added via [`add_slot_overrides`] will override any
50    /// existing slots with matching keys.
51    pub fn add_slot_overrides(
52        mut self,
53        storage_slots: impl IntoIterator<Item = StorageSlot>,
54    ) -> Self {
55        self.slot_overrides.add_overrides(storage_slots);
56        self
57    }
58
59    /// Slots added via [`add_slot_overrides_from_file`] will override any
60    /// existing slots with matching keys.
61    ///
62    /// `path` - path to a JSON file containing the storage slots.
63    pub fn add_slot_overrides_from_file(mut self, path: impl AsRef<Path>) -> Result<Self> {
64        let slots = StorageSlots::load_from_file(path.as_ref())?;
65        self.slot_overrides.add_overrides(slots.into_iter());
66        Ok(self)
67    }
68
69    pub fn into_slots(self) -> impl Iterator<Item = StorageSlot> {
70        self.slot_overrides.into_iter()
71    }
72}
73
74#[derive(Debug, Clone, Default)]
75pub(crate) struct StorageSlots {
76    storage_slots: HashMap<Bytes32, StorageSlot>,
77}
78
79impl StorageSlots {
80    fn from(storage_slots: impl IntoIterator<Item = StorageSlot>) -> Self {
81        let pairs = storage_slots.into_iter().map(|slot| (*slot.key(), slot));
82        Self {
83            storage_slots: pairs.collect(),
84        }
85    }
86
87    pub(crate) fn add_overrides(
88        &mut self,
89        storage_slots: impl IntoIterator<Item = StorageSlot>,
90    ) -> &mut Self {
91        let pairs = storage_slots.into_iter().map(|slot| (*slot.key(), slot));
92        self.storage_slots.extend(pairs);
93        self
94    }
95
96    pub(crate) fn load_from_file(storage_path: impl AsRef<Path>) -> Result<Self> {
97        let storage_path = storage_path.as_ref();
98        validate_path_and_extension(storage_path, "json")?;
99
100        let storage_json_string = std::fs::read_to_string(storage_path).map_err(|e| {
101            io::Error::new(
102                e.kind(),
103                format!("failed to read storage slots from: {storage_path:?}: {e}"),
104            )
105        })?;
106
107        let decoded_slots = serde_json::from_str::<Vec<StorageSlot>>(&storage_json_string)?;
108
109        Ok(StorageSlots::from(decoded_slots))
110    }
111
112    pub(crate) fn into_iter(self) -> impl Iterator<Item = StorageSlot> {
113        self.storage_slots.into_values()
114    }
115}
116
117pub(crate) fn determine_storage_slots(
118    storage_config: StorageConfiguration,
119    binary_filepath: &Path,
120) -> Result<Vec<StorageSlot>> {
121    let autoload_enabled = storage_config.autoload_enabled();
122    let user_overrides = storage_config.into_slots().collect::<Vec<_>>();
123    let slots = if autoload_enabled {
124        let mut slots = autoload_storage_slots(binary_filepath)?;
125        slots.add_overrides(user_overrides);
126        slots.into_iter().collect()
127    } else {
128        user_overrides
129    };
130
131    Ok(slots)
132}
133
134pub(crate) fn autoload_storage_slots(contract_binary: &Path) -> Result<StorageSlots> {
135    let storage_file = expected_storage_slots_filepath(contract_binary)
136        .ok_or_else(|| error!(Other, "could not determine storage slots file"))?;
137
138    StorageSlots::load_from_file(&storage_file)
139                .map_err(|_| error!(Other, "could not autoload storage slots from file: {storage_file:?}. \
140                                    Either provide the file or disable autoloading in `StorageConfiguration`"))
141}
142
143pub(crate) fn expected_storage_slots_filepath(contract_binary: &Path) -> Option<PathBuf> {
144    let dir = contract_binary.parent()?;
145
146    let binary_filename = contract_binary.file_stem()?.to_str()?;
147
148    Some(dir.join(format!("{binary_filename}-storage_slots.json")))
149}
150pub(crate) fn validate_path_and_extension(file_path: &Path, extension: &str) -> Result<()> {
151    if !file_path.exists() {
152        return Err(error!(IO, "file {file_path:?} does not exist"));
153    }
154
155    let path_extension = file_path
156        .extension()
157        .ok_or_else(|| error!(Other, "could not extract extension from: {file_path:?}"))?;
158
159    if extension != path_extension {
160        return Err(error!(
161            Other,
162            "expected {file_path:?} to have '.{extension}' extension"
163        ));
164    }
165
166    Ok(())
167}
168
169#[cfg(test)]
170mod tests {
171    use std::collections::HashSet;
172
173    use super::*;
174
175    #[test]
176    fn merging_overrides_storage_slots() {
177        // given
178        let make_slot = |id, value| StorageSlot::new([id; 32].into(), [value; 32].into());
179
180        let slots = (1..3).map(|id| make_slot(id, 100));
181        let original_config = StorageConfiguration::new(false, slots);
182
183        let overlapping_slots = (2..4).map(|id| make_slot(id, 200));
184
185        // when
186        let original_config = original_config.add_slot_overrides(overlapping_slots);
187
188        // then
189        assert_eq!(
190            HashSet::from_iter(original_config.slot_overrides.into_iter()),
191            HashSet::from([make_slot(1, 100), make_slot(2, 200), make_slot(3, 200)])
192        );
193    }
194}