1use crate::adapter::StripBytes;
2use crate::IsTerminal;
3use crate::Lockable;
4use crate::RawStream;
5
6/// Only pass printable data to the inner `Write`
7#[derive(Debug)]
8pub struct StripStream<S> {
9 raw: S,
10 state: StripBytes,
11}
12
13impl<S> StripStream<S>
14where
15 S: RawStream,
16{
17 /// Only pass printable data to the inner `Write`
18 #[inline]
19 pub fn new(raw: S) -> Self {
20 Self {
21 raw,
22 state: Default::default(),
23 }
24 }
25
26 /// Get the wrapped [`RawStream`]
27 #[inline]
28 pub fn into_inner(self) -> S {
29 self.raw
30 }
31
32 #[inline]
33 pub fn is_terminal(&self) -> bool {
34 self.raw.is_terminal()
35 }
36}
37
38impl<S> IsTerminal for StripStream<S>
39where
40 S: RawStream,
41{
42 #[inline]
43 fn is_terminal(&self) -> bool {
44 self.is_terminal()
45 }
46}
47
48impl<S> std::io::Write for StripStream<S>
49where
50 S: RawStream,
51{
52 #[inline]
53 fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
54 write(&mut self.raw, &mut self.state, buf)
55 }
56
57 #[inline]
58 fn flush(&mut self) -> std::io::Result<()> {
59 self.raw.flush()
60 }
61
62 // Provide explicit implementations of trait methods
63 // - To reduce bookkeeping
64 // - Avoid acquiring / releasing locks in a loop
65
66 #[inline]
67 fn write_all(&mut self, buf: &[u8]) -> std::io::Result<()> {
68 write_all(&mut self.raw, &mut self.state, buf)
69 }
70
71 // Not bothering with `write_fmt` as it just calls `write_all`
72}
73
74fn write(
75 raw: &mut dyn std::io::Write,
76 state: &mut StripBytes,
77 buf: &[u8],
78) -> std::io::Result<usize> {
79 let initial_state: StripBytes = state.clone();
80
81 for printable: &[u8] in state.strip_next(bytes:buf) {
82 let possible: usize = printable.len();
83 let written: usize = raw.write(buf:printable)?;
84 if possible != written {
85 let divergence: &[u8] = &printable[written..];
86 let offset: usize = offset_to(total:buf, subslice:divergence);
87 let consumed: &[u8] = &buf[offset..];
88 *state = initial_state;
89 state.strip_next(bytes:consumed).last();
90 return Ok(offset);
91 }
92 }
93 Ok(buf.len())
94}
95
96fn write_all(
97 raw: &mut dyn std::io::Write,
98 state: &mut StripBytes,
99 buf: &[u8],
100) -> std::io::Result<()> {
101 for printable: &[u8] in state.strip_next(bytes:buf) {
102 raw.write_all(buf:printable)?;
103 }
104 Ok(())
105}
106
107#[inline]
108fn offset_to(total: &[u8], subslice: &[u8]) -> usize {
109 let total: *const u8 = total.as_ptr();
110 let subslice: *const u8 = subslice.as_ptr();
111
112 debug_assert!(
113 total <= subslice,
114 "`Offset::offset_to` only accepts slices of `self`"
115 );
116 subslice as usize - total as usize
117}
118
119impl<S> Lockable for StripStream<S>
120where
121 S: Lockable,
122{
123 type Locked = StripStream<<S as Lockable>::Locked>;
124
125 #[inline]
126 fn lock(self) -> Self::Locked {
127 Self::Locked {
128 raw: self.raw.lock(),
129 state: self.state,
130 }
131 }
132}
133
134#[cfg(test)]
135mod test {
136 use super::*;
137 use proptest::prelude::*;
138 use std::io::Write as _;
139
140 proptest! {
141 #[test]
142 #[cfg_attr(miri, ignore)] // See https://github.com/AltSysrq/proptest/issues/253
143 fn write_all_no_escapes(s in "\\PC*") {
144 let buffer = crate::Buffer::new();
145 let mut stream = StripStream::new(buffer);
146 stream.write_all(s.as_bytes()).unwrap();
147 let buffer = stream.into_inner();
148 let actual = std::str::from_utf8(buffer.as_ref()).unwrap();
149 assert_eq!(s, actual);
150 }
151
152 #[test]
153 #[cfg_attr(miri, ignore)] // See https://github.com/AltSysrq/proptest/issues/253
154 fn write_byte_no_escapes(s in "\\PC*") {
155 let buffer = crate::Buffer::new();
156 let mut stream = StripStream::new(buffer);
157 for byte in s.as_bytes() {
158 stream.write_all(&[*byte]).unwrap();
159 }
160 let buffer = stream.into_inner();
161 let actual = std::str::from_utf8(buffer.as_ref()).unwrap();
162 assert_eq!(s, actual);
163 }
164
165 #[test]
166 #[cfg_attr(miri, ignore)] // See https://github.com/AltSysrq/proptest/issues/253
167 fn write_all_random(s in any::<Vec<u8>>()) {
168 let buffer = crate::Buffer::new();
169 let mut stream = StripStream::new(buffer);
170 stream.write_all(s.as_slice()).unwrap();
171 let buffer = stream.into_inner();
172 if let Ok(actual) = std::str::from_utf8(buffer.as_ref()) {
173 for char in actual.chars() {
174 assert!(!char.is_ascii() || !char.is_control() || char.is_ascii_whitespace(), "{:?} -> {:?}: {:?}", String::from_utf8_lossy(&s), actual, char);
175 }
176 }
177 }
178
179 #[test]
180 #[cfg_attr(miri, ignore)] // See https://github.com/AltSysrq/proptest/issues/253
181 fn write_byte_random(s in any::<Vec<u8>>()) {
182 let buffer = crate::Buffer::new();
183 let mut stream = StripStream::new(buffer);
184 for byte in s.as_slice() {
185 stream.write_all(&[*byte]).unwrap();
186 }
187 let buffer = stream.into_inner();
188 if let Ok(actual) = std::str::from_utf8(buffer.as_ref()) {
189 for char in actual.chars() {
190 assert!(!char.is_ascii() || !char.is_control() || char.is_ascii_whitespace(), "{:?} -> {:?}: {:?}", String::from_utf8_lossy(&s), actual, char);
191 }
192 }
193 }
194 }
195}
196