1use 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
17pub struct AgentVec<R, A: Agent + RecordedAgent<R>> {
58 agents: Vec<A>,
60 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 pub fn new() -> Self {
76 AgentVec {
77 agents: Vec::<A>::new(),
78 records: Vec::<Vec<R>>::new(),
79 }
80 }
81 pub fn from(agents: Vec<A>) -> Self {
88 AgentVec {
89 agents,
90 records: Vec::<Vec<R>>::new(),
91 }
92 }
93 pub fn add_agent(&mut self, agent: A) {
100 self.agents.push(agent);
101 }
102 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 fn take_records(&mut self) -> Vec<Vec<R>> {
111 mem::take(&mut self.records)
112 }
113}
114
115impl<R: 'static, A: Agent + RecordedAgent<R> + 'static> AgentSet for AgentVec<R, A> {
117 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 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 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}