Skip to content

Commit 72bd97c

Browse files
gui1117gui1117
committed
use u32
1 parent d90e4ee commit 72bd97c

1 file changed

Lines changed: 50 additions & 36 deletions

File tree

‎src/counted_input.rs‎

Lines changed: 50 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@
1616
#[derive(Eq, PartialEq, Clone, Debug)]
1717
pub enum Count {
1818
/// The counter has an exact value.
19-
Exact(u64),
19+
Exact(u32),
2020
/// The counter has reached its maximum countable value.
2121
MaxCountReached,
2222
}
@@ -25,10 +25,10 @@ pub enum Count {
2525
///
2626
/// If inner `Input` fails to read, the counter is not incremented.
2727
///
28-
/// It can count until `u64::MAX - 1` accurately.
28+
/// It can count until `u32::MAX - 1` accurately.
2929
pub struct CountedInput<'a, I: crate::Input> {
3030
input: &'a mut I,
31-
counter: u64,
31+
counter: u32,
3232
}
3333

3434
impl<'a, I: crate::Input> CountedInput<'a, I> {
@@ -38,9 +38,9 @@ impl<'a, I: crate::Input> CountedInput<'a, I> {
3838
}
3939

4040
/// Get the number of bytes successfully read.
41-
/// Count until `u64::MAX - 1` accurately.
41+
/// Count until `u32::MAX - 1` accurately.
4242
pub fn count(&self) -> Count {
43-
if self.counter == u64::MAX {
43+
if self.counter == u32::MAX {
4444
Count::MaxCountReached
4545
} else {
4646
Count::Exact(self.counter)
@@ -57,7 +57,7 @@ impl<I: crate::Input> crate::Input for CountedInput<'_, I> {
5757
self.input.read(into)
5858
.map(|r| {
5959
self.counter = self.counter.saturating_add(
60-
into.len().try_into().unwrap_or(u64::MAX)
60+
into.len().try_into().unwrap_or(u32::MAX)
6161
);
6262
r
6363
})
@@ -122,55 +122,69 @@ mod test {
122122
assert_eq!(counted_input.count(), Count::Exact(5));
123123
}
124124

125-
#[test]
126-
fn test_counted_input_max_count_read_byte() {
127-
let max_exact_count = u64::MAX - 1;
128-
129-
let mut input = &[0u8; 1000][..];
130-
let mut counted_input = CountedInput::new(&mut input);
131-
132-
counted_input.counter = max_exact_count - 2;
133-
134-
assert_eq!(counted_input.count(), Count::Exact(max_exact_count - 2));
135-
136-
counted_input.read_byte().unwrap();
137125

138-
assert_eq!(counted_input.count(), Count::Exact(max_exact_count - 1));
139-
140-
counted_input.read_byte().unwrap();
141-
142-
assert_eq!(counted_input.count(), Count::Exact(max_exact_count));
126+
struct BigInput;
127+
impl Input for BigInput {
128+
fn remaining_len(&mut self) -> Result<Option<usize>, crate::Error> {
129+
Ok(None)
130+
}
143131

144-
counted_input.read_byte().unwrap();
132+
fn read(&mut self, _into: &mut [u8]) -> Result<(), crate::Error> {
133+
Ok(())
134+
}
145135

146-
assert_eq!(counted_input.count(), Count::MaxCountReached);
136+
fn read_byte(&mut self) -> Result<u8, crate::Error> {
137+
Ok(0)
138+
}
147139

148-
counted_input.read_byte().unwrap();
140+
fn ascend_ref(&mut self) {}
149141

150-
assert_eq!(counted_input.count(), Count::MaxCountReached);
142+
fn descend_ref(&mut self) -> Result<(), crate::Error> {
143+
Ok(())
144+
}
151145
}
152146

153147
#[test]
154-
fn test_counted_input_max_count_read() {
155-
let max_exact_count = u64::MAX - 1;
148+
fn test_counted_input_max_count() {
149+
let max = u32::MAX - 1; // 2^32 - 2
156150

157-
let mut input = &[0u8; 1000][..];
151+
let mut input = BigInput;
158152
let mut counted_input = CountedInput::new(&mut input);
159153

160-
counted_input.counter = max_exact_count - 1;
154+
assert_eq!(counted_input.count(), Count::Exact(0));
161155

162-
assert_eq!(counted_input.count(), Count::Exact(max_exact_count - 1));
156+
// Test will read continuously into this page.
157+
let page_size = 1024 * 1024; // 1 MB
158+
let mut page = vec![0u8; page_size];
163159

164-
counted_input.read_byte().unwrap();
160+
// Calculate the number of full pages and the remaining bytes
161+
let total_bytes = max as usize;
162+
let full_pages = total_bytes / page_size;
163+
let remaining_bytes = total_bytes % page_size;
164+
165+
// Read full pages
166+
for _ in 0..full_pages {
167+
counted_input.read(&mut page[..]).unwrap();
168+
}
165169

166-
assert_eq!(counted_input.count(), Count::Exact(max_exact_count));
170+
// Read remaining bytes
171+
if remaining_bytes > 0 {
172+
counted_input.read(&mut page[..remaining_bytes]).unwrap();
173+
}
167174

168-
counted_input.read(&mut [0u8; 2][..]).unwrap();
175+
// Max is reached exactly
176+
assert_eq!(counted_input.count(), Count::Exact(max));
169177

178+
// Perform one additional read
179+
counted_input.read_byte().unwrap();
180+
181+
// Count is more than max.
170182
assert_eq!(counted_input.count(), Count::MaxCountReached);
171183

172-
counted_input.read(&mut [0u8; 2][..]).unwrap();
184+
// Perform one additional read
185+
counted_input.read_byte().unwrap();
173186

187+
// Count is still more than max.
174188
assert_eq!(counted_input.count(), Count::MaxCountReached);
175189
}
176190
}

0 commit comments

Comments
 (0)