linera_core/environment/wallet/
memory.rs

1// Copyright (c) Zefchain Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4use futures::{Stream, StreamExt as _};
5use linera_base::identifiers::ChainId;
6use serde::{ser::SerializeMap, Serialize, Serializer};
7
8use super::{Chain, Wallet};
9
10/// A basic implementation of `Wallet` that doesn't persist anything and merely tracks the
11/// chains in memory.
12///
13/// This can be used as-is as an ephemeral wallet for testing or ephemeral clients, or as
14/// a building block for more complex wallets that layer persistence on top of it.
15#[derive(Default, Clone, serde::Deserialize)]
16pub struct Memory(papaya::HashMap<ChainId, Chain>);
17
18/// Custom Serialize implementation that ensures stable ordering by sorting entries by ChainId.
19impl Serialize for Memory {
20    fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
21        let guard = self.0.pin();
22        let mut items: Vec<_> = guard.iter().collect();
23        items.sort_by_key(|(k, _)| *k);
24        let mut map = serializer.serialize_map(Some(items.len()))?;
25        for (k, v) in items {
26            map.serialize_entry(k, v)?;
27        }
28        map.end()
29    }
30}
31
32impl Memory {
33    pub fn get(&self, id: ChainId) -> Option<Chain> {
34        self.0.pin().get(&id).cloned()
35    }
36
37    pub fn insert(&self, id: ChainId, chain: Chain) -> Option<Chain> {
38        self.0.pin().insert(id, chain).cloned()
39    }
40
41    pub fn try_insert(&self, id: ChainId, chain: Chain) -> Option<Chain> {
42        match self.0.pin().try_insert(id, chain) {
43            Ok(_inserted) => None,
44            Err(error) => Some(error.not_inserted),
45        }
46    }
47
48    pub fn remove(&self, id: ChainId) -> Option<Chain> {
49        self.0.pin().remove(&id).cloned()
50    }
51
52    pub fn items(&self) -> Vec<(ChainId, Chain)> {
53        self.0
54            .pin()
55            .iter()
56            .map(|(id, chain)| (*id, chain.clone()))
57            .collect::<Vec<_>>()
58    }
59
60    pub fn chain_ids(&self) -> Vec<ChainId> {
61        self.0.pin().keys().copied().collect::<Vec<_>>()
62    }
63
64    pub fn owned_chain_ids(&self) -> Vec<ChainId> {
65        self.0
66            .pin()
67            .iter()
68            .filter_map(|(id, chain)| chain.owner.as_ref().map(|_| *id))
69            .collect::<Vec<_>>()
70    }
71
72    pub fn mutate<R>(&self, chain_id: ChainId, mutate: impl Fn(&mut Chain) -> R) -> Option<R> {
73        use papaya::Operation::*;
74
75        let mut outcome = None;
76        self.0.pin().compute(chain_id, |chain| {
77            if let Some((_, chain)) = chain {
78                let mut chain = chain.clone();
79                outcome = Some(mutate(&mut chain));
80                Insert(chain)
81            } else {
82                Abort(())
83            }
84        });
85
86        outcome
87    }
88}
89
90impl Extend<(ChainId, Chain)> for Memory {
91    fn extend<It: IntoIterator<Item = (ChainId, Chain)>>(&mut self, chains: It) {
92        let map = self.0.pin();
93        for (id, chain) in chains {
94            map.insert(id, chain);
95        }
96    }
97}
98
99impl Wallet for Memory {
100    type Error = std::convert::Infallible;
101
102    async fn get(&self, id: ChainId) -> Result<Option<Chain>, Self::Error> {
103        Ok(self.get(id))
104    }
105
106    async fn insert(&self, id: ChainId, chain: Chain) -> Result<Option<Chain>, Self::Error> {
107        Ok(self.insert(id, chain))
108    }
109
110    async fn try_insert(&self, id: ChainId, chain: Chain) -> Result<Option<Chain>, Self::Error> {
111        Ok(self.try_insert(id, chain))
112    }
113
114    async fn remove(&self, id: ChainId) -> Result<Option<Chain>, Self::Error> {
115        Ok(self.remove(id))
116    }
117
118    fn items(&self) -> impl Stream<Item = Result<(ChainId, Chain), Self::Error>> {
119        futures::stream::iter(self.items()).map(Ok)
120    }
121
122    async fn modify(
123        &self,
124        id: ChainId,
125        f: impl Fn(&mut Chain) + Send,
126    ) -> Result<Option<()>, Self::Error> {
127        Ok(self.mutate(id, f))
128    }
129}
130
131#[cfg(test)]
132mod tests {
133    use linera_base::{crypto::CryptoHash, data_types::Timestamp};
134
135    use super::*;
136
137    fn make_chain(height: u64) -> Chain {
138        Chain {
139            owner: None,
140            block_hash: None,
141            next_block_height: height.into(),
142            timestamp: Timestamp::from(0),
143            pending_fast_proposal: None,
144            epoch: None,
145        }
146    }
147
148    #[test]
149    fn test_memory_serialization_roundtrip() {
150        let memory = Memory::default();
151
152        // Insert chains in non-sorted order using different hashes
153        let id1 = ChainId(CryptoHash::test_hash("chain1"));
154        let id2 = ChainId(CryptoHash::test_hash("chain2"));
155        let id3 = ChainId(CryptoHash::test_hash("chain3"));
156
157        memory.insert(id2, make_chain(2));
158        memory.insert(id1, make_chain(1));
159        memory.insert(id3, make_chain(3));
160
161        // Serialize to JSON
162        let json = serde_json::to_string_pretty(&memory).unwrap();
163
164        // Deserialize back
165        let restored: Memory = serde_json::from_str(&json).unwrap();
166
167        // Verify data matches
168        assert_eq!(restored.get(id1).unwrap().next_block_height, 1.into());
169        assert_eq!(restored.get(id2).unwrap().next_block_height, 2.into());
170        assert_eq!(restored.get(id3).unwrap().next_block_height, 3.into());
171    }
172
173    #[test]
174    fn test_memory_serialization_is_sorted() {
175        let memory = Memory::default();
176
177        let id1 = ChainId(CryptoHash::test_hash("a"));
178        let id2 = ChainId(CryptoHash::test_hash("b"));
179        let id3 = ChainId(CryptoHash::test_hash("c"));
180
181        // Insert in non-sorted order
182        memory.insert(id3, make_chain(3));
183        memory.insert(id1, make_chain(1));
184        memory.insert(id2, make_chain(2));
185
186        // Serialize and verify output keys are sorted
187        let json = serde_json::to_string(&memory).unwrap();
188        let value: serde_json::Value = serde_json::from_str(&json).unwrap();
189        let keys: Vec<_> = value.as_object().unwrap().keys().collect();
190        let mut sorted_keys = keys.clone();
191        sorted_keys.sort();
192        assert_eq!(keys, sorted_keys);
193    }
194}
linera_core/environment/wallet/memory.rs

linera_core/environment/wallet/
memory.rs
