shadow_rs/utility/
give.rs1use std::io::Write;
2
3pub struct Give<T> {
11 inner: T,
12 limit: u64,
13}
14
15impl<T> Give<T> {
16 pub fn new(writer: T, limit: u64) -> Self {
17 Self {
18 inner: writer,
19 limit,
20 }
21 }
22
23 pub fn limit(&self) -> u64 {
26 self.limit
27 }
28
29 pub fn set_limit(&mut self, limit: u64) {
33 self.limit = limit;
34 }
35
36 pub fn into_inner(self) -> T {
38 self.inner
39 }
40
41 pub fn get_ref(&self) -> &T {
43 &self.inner
44 }
45
46 pub fn get_mut(&mut self) -> &mut T {
48 &mut self.inner
49 }
50}
51
52impl<T: Write> Write for Give<T> {
53 fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
54 if self.limit == 0 {
56 return Ok(0);
57 }
58
59 let max = std::cmp::min(buf.len() as u64, self.limit) as usize;
60 let n = self.inner.write(&buf[..max])?;
61 self.limit -= n as u64;
62 Ok(n)
63 }
64
65 fn flush(&mut self) -> std::io::Result<()> {
66 self.inner.flush()
67 }
68}
69
70#[cfg(test)]
71mod tests {
72 use super::*;
73
74 #[test]
75 fn test_give_write_1() {
76 let mut buf = vec![];
77 let mut give = Give::new(&mut buf, 5);
78 assert_eq!(give.write(&[0u8; 20]).unwrap(), 5);
79 assert_eq!(give.write(&[0u8; 20]).unwrap(), 0);
80 assert_eq!(buf.len(), 5);
81 }
82
83 #[test]
84 fn test_give_write_2() {
85 let mut buf = vec![];
86 let mut give = Give::new(&mut buf, 5);
87 assert_eq!(give.write(&[0u8; 3]).unwrap(), 3);
88 assert_eq!(give.write(&[0u8; 3]).unwrap(), 2);
89 assert_eq!(give.write(&[0u8; 3]).unwrap(), 0);
90 assert_eq!(buf.len(), 5);
91 }
92
93 #[test]
94 fn test_give_write_all_1() {
95 let mut buf = vec![];
96 let mut give = Give::new(&mut buf, 5);
97 assert!(give.write_all(&[0u8; 3]).is_ok());
98 assert!(give.write_all(&[0u8; 2]).is_ok());
99 assert!(give.write_all(&[0u8; 1]).is_err());
100 assert_eq!(buf.len(), 5);
101 }
102
103 #[test]
104 fn test_give_write_all_2() {
105 let mut buf = vec![];
106 let mut give = Give::new(&mut buf, 5);
107 assert!(give.write_all(&[0u8; 7]).is_err());
108 assert!(give.write_all(&[0u8; 1]).is_err());
109 assert!(give.write_all(&[0u8; 7]).is_err());
110 assert!(give.write_all(&[0u8; 1]).is_err());
111 assert!(give.write_all(&[0u8; 3]).is_err());
112 assert_eq!(buf.len(), 5);
113 }
114}