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    /// Returns the state of the chain with the given ID, if it is tracked.
34    pub fn get(&self, id: ChainId) -> Option<Chain> {
35        self.0.pin().get(&id).cloned()
36    }
37
38    /// Inserts or replaces the state of the given chain, returning the previous state if any.
39    pub fn insert(&self, id: ChainId, chain: Chain) -> Option<Chain> {
40        self.0.pin().insert(id, chain).cloned()
41    }
42
43    /// Inserts the given chain only if it is not already tracked, returning the existing state otherwise.
44    pub fn try_insert(&self, id: ChainId, chain: Chain) -> Option<Chain> {
45        match self.0.pin().try_insert(id, chain) {
46            Ok(_inserted) => None,
47            Err(error) => Some(error.not_inserted),
48        }
49    }
50
51    /// Removes the chain with the given ID, returning its previous state if any.
52    pub fn remove(&self, id: ChainId) -> Option<Chain> {
53        self.0.pin().remove(&id).cloned()
54    }
55
56    /// Returns all tracked chains and their states.
57    pub fn items(&self) -> Vec<(ChainId, Chain)> {
58        self.0
59            .pin()
60            .iter()
61            .map(|(id, chain)| (*id, chain.clone()))
62            .collect::<Vec<_>>()
63    }
64
65    /// Returns the IDs of all tracked chains.
66    pub fn chain_ids(&self) -> Vec<ChainId> {
67        self.0.pin().keys().copied().collect::<Vec<_>>()
68    }
69
70    /// Returns the IDs of the tracked chains that have an owner configured.
71    pub fn owned_chain_ids(&self) -> Vec<ChainId> {
72        self.0
73            .pin()
74            .iter()
75            .filter_map(|(id, chain)| chain.owner.as_ref().map(|_| *id))
76            .collect::<Vec<_>>()
77    }
78
79    /// Applies a closure to the given chain if it is tracked, returning the closure's result.
80    pub fn mutate<R>(&self, chain_id: ChainId, mutate: impl Fn(&mut Chain) -> R) -> Option<R> {
81        use papaya::Operation::*;
82
83        let mut outcome = None;
84        self.0.pin().compute(chain_id, |chain| {
85            if let Some((_, chain)) = chain {
86                let mut chain = chain.clone();
87                outcome = Some(mutate(&mut chain));
88                Insert(chain)
89            } else {
90                Abort(())
91            }
92        });
93
94        outcome
95    }
96}
97
98impl Extend<(ChainId, Chain)> for Memory {
99    fn extend<It: IntoIterator<Item = (ChainId, Chain)>>(&mut self, chains: It) {
100        let map = self.0.pin();
101        for (id, chain) in chains {
102            map.insert(id, chain);
103        }
104    }
105}
106
107impl Wallet for Memory {
108    type Error = std::convert::Infallible;
109
110    async fn get(&self, id: ChainId) -> Result<Option<Chain>, Self::Error> {
111        Ok(self.get(id))
112    }
113
114    async fn insert(&self, id: ChainId, chain: Chain) -> Result<Option<Chain>, Self::Error> {
115        Ok(self.insert(id, chain))
116    }
117
118    async fn try_insert(&self, id: ChainId, chain: Chain) -> Result<Option<Chain>, Self::Error> {
119        Ok(self.try_insert(id, chain))
120    }
121
122    async fn remove(&self, id: ChainId) -> Result<Option<Chain>, Self::Error> {
123        Ok(self.remove(id))
124    }
125
126    fn items(&self) -> impl Stream<Item = Result<(ChainId, Chain), Self::Error>> {
127        futures::stream::iter(self.items()).map(Ok)
128    }
129
130    async fn modify(
131        &self,
132        id: ChainId,
133        f: impl Fn(&mut Chain) + Send,
134    ) -> Result<Option<()>, Self::Error> {
135        Ok(self.mutate(id, f))
136    }
137}
138
139#[cfg(test)]
140mod tests {
141    use linera_base::{crypto::CryptoHash, data_types::Timestamp};
142
143    use super::*;
144
145    fn make_chain(height: u64) -> Chain {
146        Chain {
147            owner: None,
148            block_hash: None,
149            next_block_height: height.into(),
150            timestamp: Timestamp::from(0),
151            pending_fast_proposal: None,
152            epoch: None,
153        }
154    }
155
156    #[test]
157    fn test_memory_serialization_roundtrip() {
158        let memory = Memory::default();
159
160        // Insert chains in non-sorted order using different hashes
161        let id1 = ChainId(CryptoHash::test_hash("chain1"));
162        let id2 = ChainId(CryptoHash::test_hash("chain2"));
163        let id3 = ChainId(CryptoHash::test_hash("chain3"));
164
165        memory.insert(id2, make_chain(2));
166        memory.insert(id1, make_chain(1));
167        memory.insert(id3, make_chain(3));
168
169        // Serialize to JSON
170        let json = serde_json::to_string_pretty(&memory).unwrap();
171
172        // Deserialize back
173        let restored: Memory = serde_json::from_str(&json).unwrap();
174
175        // Verify data matches
176        assert_eq!(restored.get(id1).unwrap().next_block_height, 1.into());
177        assert_eq!(restored.get(id2).unwrap().next_block_height, 2.into());
178        assert_eq!(restored.get(id3).unwrap().next_block_height, 3.into());
179    }
180
181    #[test]
182    fn test_memory_serialization_is_sorted() {
183        let memory = Memory::default();
184
185        let id1 = ChainId(CryptoHash::test_hash("a"));
186        let id2 = ChainId(CryptoHash::test_hash("b"));
187        let id3 = ChainId(CryptoHash::test_hash("c"));
188
189        // Insert in non-sorted order
190        memory.insert(id3, make_chain(3));
191        memory.insert(id1, make_chain(1));
192        memory.insert(id2, make_chain(2));
193
194        // Serialize and verify output keys are sorted
195        let json = serde_json::to_string(&memory).unwrap();
196        let value: serde_json::Value = serde_json::from_str(&json).unwrap();
197        let keys: Vec<_> = value.as_object().unwrap().keys().collect();
198        let mut sorted_keys = keys.clone();
199        sorted_keys.sort();
200        assert_eq!(keys, sorted_keys);
201    }
202}
linera_core/environment/wallet/memory.rs

linera_core/environment/wallet/
memory.rs
