1use alloy_primitives::{Address, B256};
4use nectar_primitives::SwarmAddress;
5
6use crate::{StampError, StampIndex, calculate_bucket};
7
8pub type BatchId = B256;
10
11#[derive(Debug, Clone, PartialEq, Eq)]
13#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
14pub struct BatchParams {
15 pub owner: Address,
17 pub depth: u8,
19 pub bucket_depth: u8,
21 pub immutable: bool,
27 pub amount: u128,
29}
30
31impl BatchParams {
32 pub const fn new(owner: Address, depth: u8, bucket_depth: u8, amount: u128) -> Self {
34 Self {
35 owner,
36 depth,
37 bucket_depth,
38 immutable: false,
39 amount,
40 }
41 }
42
43 #[must_use]
45 pub const fn immutable(mut self, immutable: bool) -> Self {
46 self.immutable = immutable;
47 self
48 }
49}
50
51#[derive(Debug, Clone, PartialEq, Eq)]
57#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
58pub struct Batch {
59 id: BatchId,
61 value: u128,
63 start: u64,
65 owner: Address,
67 depth: u8,
69 bucket_depth: u8,
71 immutable: bool,
77}
78
79impl Batch {
80 #[inline]
82 pub const fn new(
83 id: BatchId,
84 value: u128,
85 start: u64,
86 owner: Address,
87 depth: u8,
88 bucket_depth: u8,
89 immutable: bool,
90 ) -> Self {
91 Self {
92 id,
93 value,
94 start,
95 owner,
96 depth,
97 bucket_depth,
98 immutable,
99 }
100 }
101
102 #[inline]
104 pub const fn id(&self) -> BatchId {
105 self.id
106 }
107
108 #[inline]
110 pub const fn value(&self) -> u128 {
111 self.value
112 }
113
114 #[inline]
116 pub const fn start(&self) -> u64 {
117 self.start
118 }
119
120 #[inline]
122 pub const fn owner(&self) -> Address {
123 self.owner
124 }
125
126 #[inline]
130 pub const fn depth(&self) -> u8 {
131 self.depth
132 }
133
134 #[inline]
138 pub const fn bucket_depth(&self) -> u8 {
139 self.bucket_depth
140 }
141
142 #[inline]
148 pub const fn immutable(&self) -> bool {
149 self.immutable
150 }
151
152 #[inline]
156 pub const fn bucket_upper_bound(&self) -> u32 {
157 1u32 << (self.depth - self.bucket_depth)
158 }
159
160 #[inline]
164 pub const fn bucket_count(&self) -> u32 {
165 1u32 << self.bucket_depth
166 }
167
168 #[inline]
170 pub const fn set_value(&mut self, value: u128) {
171 self.value = value;
172 }
173
174 #[inline]
176 pub const fn set_depth(&mut self, depth: u8) {
177 self.depth = depth;
178 }
179
180 #[inline]
182 pub const fn is_expired(&self, total_amount: u128) -> bool {
183 self.value <= total_amount
184 }
185
186 #[inline]
188 pub const fn is_usable(&self, current_block: u64, threshold: u64) -> bool {
189 current_block >= self.start.saturating_add(threshold)
190 }
191
192 pub const fn validate_index(&self, index: &StampIndex) -> Result<(), StampError> {
206 if index.bucket() >= self.bucket_count() {
208 return Err(StampError::InvalidIndex);
209 }
210
211 if index.index() >= self.bucket_upper_bound() {
213 return Err(StampError::InvalidIndex);
214 }
215
216 Ok(())
217 }
218
219 #[inline]
224 pub fn bucket_for_address(&self, address: &SwarmAddress) -> u32 {
225 calculate_bucket(address, self.bucket_depth)
226 }
227
228 pub fn validate_bucket(
234 &self,
235 index: &StampIndex,
236 address: &SwarmAddress,
237 ) -> Result<(), StampError> {
238 let expected_bucket = self.bucket_for_address(address);
239 if index.bucket() != expected_bucket {
240 return Err(StampError::BucketMismatch);
241 }
242 Ok(())
243 }
244}
245
246#[cfg(feature = "arbitrary")]
249impl<'a> arbitrary::Arbitrary<'a> for BatchParams {
250 fn arbitrary(u: &mut arbitrary::Unstructured<'a>) -> arbitrary::Result<Self> {
251 let depth: u8 = u.int_in_range(1..=32)?;
253 let bucket_depth: u8 = u.int_in_range(1..=depth)?;
254
255 Ok(Self {
256 owner: Address::arbitrary(u)?,
257 depth,
258 bucket_depth,
259 immutable: u.arbitrary()?,
260 amount: u.arbitrary()?,
261 })
262 }
263}
264
265#[cfg(feature = "arbitrary")]
266impl<'a> arbitrary::Arbitrary<'a> for Batch {
267 fn arbitrary(u: &mut arbitrary::Unstructured<'a>) -> arbitrary::Result<Self> {
268 let depth: u8 = u.int_in_range(1..=32)?;
270 let bucket_depth: u8 = u.int_in_range(1..=depth)?;
271
272 Ok(Self::new(
273 B256::arbitrary(u)?,
274 u.arbitrary()?,
275 u.arbitrary()?,
276 Address::arbitrary(u)?,
277 depth,
278 bucket_depth,
279 u.arbitrary()?,
280 ))
281 }
282}
283
284#[cfg(test)]
285mod tests {
286 use super::*;
287
288 #[test]
289 fn test_batch_creation() {
290 let id = B256::ZERO;
291 let batch = Batch::new(id, 1000, 100, Address::ZERO, 18, 16, false);
292
293 assert_eq!(batch.id(), id);
294 assert_eq!(batch.value(), 1000);
295 assert_eq!(batch.start(), 100);
296 assert_eq!(batch.owner(), Address::ZERO);
297 assert_eq!(batch.depth(), 18);
298 assert_eq!(batch.bucket_depth(), 16);
299 assert!(!batch.immutable());
300 }
301
302 #[test]
303 fn test_bucket_calculations() {
304 let batch = Batch::new(B256::ZERO, 0, 0, Address::ZERO, 18, 16, false);
305
306 assert_eq!(batch.bucket_upper_bound(), 4);
308 assert_eq!(batch.bucket_count(), 65536);
310 }
311
312 #[test]
313 fn test_batch_expiry() {
314 let batch = Batch::new(B256::ZERO, 1000, 0, Address::ZERO, 18, 16, false);
315
316 assert!(!batch.is_expired(999));
317 assert!(batch.is_expired(1000));
318 assert!(batch.is_expired(1001));
319 }
320
321 #[test]
322 fn test_batch_usability() {
323 let batch = Batch::new(B256::ZERO, 1000, 100, Address::ZERO, 18, 16, false);
324
325 assert!(!batch.is_usable(100, 10)); assert!(!batch.is_usable(109, 10)); assert!(batch.is_usable(110, 10)); assert!(batch.is_usable(111, 10)); }
330
331 #[test]
332 fn test_batch_params_builder() {
333 let params = BatchParams::new(Address::ZERO, 20, 16, 1000).immutable(true);
334
335 assert_eq!(params.owner, Address::ZERO);
336 assert_eq!(params.depth, 20);
337 assert_eq!(params.bucket_depth, 16);
338 assert_eq!(params.amount, 1000);
339 assert!(params.immutable);
340 }
341}