1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155

use std::{collections::HashMap, default::Default, fmt::Debug, io, path::Path};

use fuel_tx::{Bytes32, StorageSlot};
use fuels_core::types::errors::{error, Result};

/// Configuration for contract storage
#[derive(Debug, Clone)]
pub struct StorageConfiguration {
    autoload_storage: bool,
    slot_overrides: StorageSlots,
}

impl Default for StorageConfiguration {
    fn default() -> Self {
        Self {
            autoload_storage: true,
            slot_overrides: Default::default(),
        }
    }
}

impl StorageConfiguration {
    pub fn new(autoload_enabled: bool, slots: impl IntoIterator<Item = StorageSlot>) -> Self {
        let config = Self {
            autoload_storage: autoload_enabled,
            slot_overrides: Default::default(),
        };

        config.add_slot_overrides(slots)
    }

    /// If enabled will try to automatically discover and load the storage configuration from the
    /// storage config json file.
    pub fn with_autoload(mut self, enabled: bool) -> Self {
        self.autoload_storage = enabled;
        self
    }

    pub fn autoload_enabled(&self) -> bool {
        self.autoload_storage
    }

    /// Slots added via [`add_slot_overrides`] will override any
    /// existing slots with matching keys.
    pub fn add_slot_overrides(
        mut self,
        storage_slots: impl IntoIterator<Item = StorageSlot>,
    ) -> Self {
        self.slot_overrides.add_overrides(storage_slots);
        self
    }

    /// Slots added via [`add_slot_overrides_from_file`] will override any
    /// existing slots with matching keys.
    ///
    /// `path` - path to a JSON file containing the storage slots.
    pub fn add_slot_overrides_from_file(mut self, path: impl AsRef<Path>) -> Result<Self> {
        let slots = StorageSlots::load_from_file(path.as_ref())?;
        self.slot_overrides.add_overrides(slots.into_iter());
        Ok(self)
    }

    pub fn into_slots(self) -> impl Iterator<Item = StorageSlot> {
        self.slot_overrides.into_iter()
    }
}

#[derive(Debug, Clone, Default)]
pub(crate) struct StorageSlots {
    storage_slots: HashMap<Bytes32, StorageSlot>,
}

impl StorageSlots {
    fn from(storage_slots: impl IntoIterator<Item = StorageSlot>) -> Self {
        let pairs = storage_slots.into_iter().map(|slot| (*slot.key(), slot));
        Self {
            storage_slots: pairs.collect(),
        }
    }

    pub(crate) fn add_overrides(
        &mut self,
        storage_slots: impl IntoIterator<Item = StorageSlot>,
    ) -> &mut Self {
        let pairs = storage_slots.into_iter().map(|slot| (*slot.key(), slot));
        self.storage_slots.extend(pairs);
        self
    }

    pub(crate) fn load_from_file(storage_path: impl AsRef<Path>) -> Result<Self> {
        let storage_path = storage_path.as_ref();
        validate_path_and_extension(storage_path, "json")?;

        let storage_json_string = std::fs::read_to_string(storage_path).map_err(|e| {
            io::Error::new(
                e.kind(),
                format!("failed to read storage slots from: {storage_path:?}: {e}"),
            )
        })?;

        let decoded_slots = serde_json::from_str::<Vec<StorageSlot>>(&storage_json_string)?;

        Ok(StorageSlots::from(decoded_slots))
    }

    pub(crate) fn into_iter(self) -> impl Iterator<Item = StorageSlot> {
        self.storage_slots.into_values()
    }
}

pub(crate) fn validate_path_and_extension(file_path: &Path, extension: &str) -> Result<()> {
    if !file_path.exists() {
        return Err(error!(IO, "file {file_path:?} does not exist"));
    }

    let path_extension = file_path
        .extension()
        .ok_or_else(|| error!(Other, "could not extract extension from: {file_path:?}"))?;

    if extension != path_extension {
        return Err(error!(
            Other,
            "expected {file_path:?} to have '.{extension}' extension"
        ));
    }

    Ok(())
}

#[cfg(test)]
mod tests {
    use std::collections::HashSet;

    use super::*;

    #[test]
    fn merging_overrides_storage_slots() {
        // given
        let make_slot = |id, value| StorageSlot::new([id; 32].into(), [value; 32].into());

        let slots = (1..3).map(|id| make_slot(id, 100));
        let original_config = StorageConfiguration::new(false, slots);

        let overlapping_slots = (2..4).map(|id| make_slot(id, 200));

        // when
        let original_config = original_config.add_slot_overrides(overlapping_slots);

        // then
        assert_eq!(
            HashSet::from_iter(original_config.slot_overrides.into_iter()),
            HashSet::from([make_slot(1, 100), make_slot(2, 200), make_slot(3, 200)])
        );
    }
}