verbs_rs/agent/
agent_vec.rs

1//! Vector data structure of simulation agents
2//!
3//! Data structure that stores a vector of agents
4//! of one type. Implements functionality to
5//! iterate over and update agents, and
6//! record and retrieve simulated agent data.
7//!
8
9use crate::agent::traits::{Agent, AgentSet, RecordedAgent, RecordedAgentSet};
10use crate::contract::Transaction;
11use crate::env::{Env, Validator};
12use crate::DB;
13use alloy_primitives::Address;
14use rand::RngCore;
15use std::mem;
16
17/// Implementation of agent set storing agents as a vector
18///
19/// Stores a vector of agents of a single type,
20/// and stores records of their state.
21///
22/// # Examples
23///
24/// ```
25/// use rand::RngCore;
26/// use alloy_primitives::Address;
27/// use verbs_rs::{DB, env::{Env, Validator}};
28/// use verbs_rs::agent::{Agent, RecordedAgent, AgentVec, AgentSet};
29/// use verbs_rs::contract::Transaction;
30///
31/// struct DummyAgent{}
32///
33/// impl Agent for DummyAgent {
34///     fn update<D: DB, V: Validator, R: RngCore>(
35///         &mut self, rng: &mut R, network: &mut Env<D, V>
36///     ) -> Vec<Transaction> {
37///         Vec::default()
38///     }
39///
40///     fn get_address(&self) -> Address {
41///         Address::ZERO
42///     }
43/// }
44///
45/// impl RecordedAgent<bool> for DummyAgent {
46///     fn record<D: DB, V: Validator>(&mut self, _env: &mut Env<D, V>) -> bool {
47///         true
48///     }
49/// }
50///
51/// let agent_vec = AgentVec::<bool, DummyAgent>::from(
52///     vec![DummyAgent{}, DummyAgent{}]
53/// );
54///
55/// let addresses = agent_vec.get_addresses();
56/// ```
57pub struct AgentVec<R, A: Agent + RecordedAgent<R>> {
58    /// Vector of agents of a single type
59    agents: Vec<A>,
60    /// Records of agent states over the course of the simulation
61    records: Vec<Vec<R>>,
62}
63
64impl<R, A: Agent + RecordedAgent<R>> Default for AgentVec<R, A> {
65    fn default() -> Self {
66        AgentVec {
67            agents: Vec::<A>::new(),
68            records: Vec::<Vec<R>>::new(),
69        }
70    }
71}
72
73impl<R, A: Agent + RecordedAgent<R>> AgentVec<R, A> {
74    /// Initialise an empty vector agent-set
75    pub fn new() -> Self {
76        AgentVec {
77            agents: Vec::<A>::new(),
78            records: Vec::<Vec<R>>::new(),
79        }
80    }
81    /// Initialise an agent-vec from an existing vector of agents
82    ///
83    /// # Arguments
84    ///
85    /// * `agents` - Vector af agents of this type
86    ///
87    pub fn from(agents: Vec<A>) -> Self {
88        AgentVec {
89            agents,
90            records: Vec::<Vec<R>>::new(),
91        }
92    }
93    /// Insert an agent into the set.
94    ///
95    /// # Arguments
96    ///
97    /// * `agent` - Agents of this type
98    ///
99    pub fn add_agent(&mut self, agent: A) {
100        self.agents.push(agent);
101    }
102    /// Get the recorded history of agents in this set.
103    pub fn get_records(&self) -> &Vec<Vec<R>> {
104        &self.records
105    }
106}
107
108impl<R, A: Agent + RecordedAgent<R>> RecordedAgentSet<R> for AgentVec<R, A> {
109    /// Take the vector of agent records from the set
110    fn take_records(&mut self) -> Vec<Vec<R>> {
111        mem::take(&mut self.records)
112    }
113}
114
115/// Implementations of agent updates and recording.
116impl<R: 'static, A: Agent + RecordedAgent<R> + 'static> AgentSet for AgentVec<R, A> {
117    /// Call the agents in the set and collect any returned EVM transactions
118    ///
119    /// This is called during the simulation, updating the state of
120    /// the agents, and collecting any submitted transactions into
121    /// a single vector.
122    ///
123    /// # Arguments
124    ///
125    /// * `rng` - Random generator
126    /// * `network` - Protocol deployment(s)
127    ///
128    fn call<D: DB, V: Validator, RG: RngCore>(
129        &mut self,
130        rng: &mut RG,
131        network: &mut Env<D, V>,
132    ) -> Vec<Transaction> {
133        self.agents
134            .iter_mut()
135            .flat_map(|x| x.update(rng, network))
136            .collect()
137    }
138    /// Record the current state of the agents in this set
139    fn record<D: DB, V: Validator>(&mut self, env: &mut Env<D, V>) {
140        let records: Vec<R> = self.agents.iter_mut().map(|x| x.record(env)).collect();
141        self.records.push(records);
142    }
143    /// Get the addresses of the agents in this set.
144    fn get_addresses(&self) -> Vec<Address> {
145        self.agents.iter().map(|x| x.get_address()).collect()
146    }
147}
148
149#[cfg(test)]
150mod tests {
151    use super::*;
152    use crate::LocalDB;
153    use crate::{agent::traits, env::RandomValidator};
154    use alloy_primitives::{Uint, U256};
155    use rand::SeedableRng;
156    use rstest::*;
157
158    struct TestAgent {
159        address: Address,
160        value: u64,
161    }
162
163    impl traits::Agent for TestAgent {
164        fn update<D: DB, V: Validator, RG: RngCore>(
165            &mut self,
166            _rng: &mut RG,
167            _network: &mut crate::env::Env<D, V>,
168        ) -> Vec<crate::contract::Transaction> {
169            self.value += 1;
170            vec![
171                Transaction {
172                    function_selector: [0, 0, 0, 0],
173                    callee: Address::ZERO,
174                    transact_to: Address::ZERO,
175                    args: Vec::default(),
176                    value: U256::ZERO,
177                    checked: false,
178                    gas_priority_fee: None,
179                    nonce: None,
180                },
181                Transaction {
182                    function_selector: [0, 0, 0, 0],
183                    callee: Address::ZERO,
184                    transact_to: Address::ZERO,
185                    args: Vec::default(),
186                    value: U256::ZERO,
187                    checked: false,
188                    gas_priority_fee: None,
189                    nonce: None,
190                },
191            ]
192        }
193
194        fn get_address(&self) -> Address {
195            self.address
196        }
197    }
198
199    impl traits::RecordedAgent<u64> for TestAgent {
200        fn record<D: DB, V: Validator>(&mut self, _env: &mut Env<D, V>) -> u64 {
201            self.value
202        }
203    }
204
205    #[fixture]
206    fn env() -> Env<LocalDB, RandomValidator> {
207        Env::<LocalDB, RandomValidator>::init(U256::ZERO, U256::ZERO, RandomValidator {})
208    }
209
210    #[fixture]
211    fn rng() -> rand_xoshiro::Xoroshiro128StarStar {
212        rand_xoshiro::Xoroshiro128StarStar::seed_from_u64(101)
213    }
214
215    #[rstest]
216    fn test_agent_vec(
217        mut env: Env<LocalDB, RandomValidator>,
218        mut rng: rand_xoshiro::Xoroshiro128StarStar,
219    ) {
220        let a = Address::from(Uint::from(101u128));
221        let b = Address::from(Uint::from(202u128));
222
223        let agents = vec![
224            TestAgent {
225                address: a,
226                value: 0,
227            },
228            TestAgent {
229                address: b,
230                value: 1,
231            },
232        ];
233
234        let mut agent_vec = AgentVec::from(agents);
235
236        assert_eq!(agent_vec.get_addresses(), vec![a, b]);
237
238        agent_vec.record(&mut env);
239        assert_eq!(agent_vec.records.len(), 1);
240
241        let calls = agent_vec.call(&mut rng, &mut env);
242        assert_eq!(calls.len(), 4);
243
244        agent_vec.record(&mut env);
245        assert_eq!(agent_vec.records.len(), 2);
246
247        let records = agent_vec.take_records();
248
249        assert_eq!(records[0], vec![0, 1]);
250        assert_eq!(records[1], vec![1, 2]);
251    }
252}