1 // Copyright (c) 2023 Huawei Device Co., Ltd.
2 // Licensed under the Apache License, Version 2.0 (the "License");
3 // you may not use this file except in compliance with the License.
4 // You may obtain a copy of the License at
5 //
6 //     http://www.apache.org/licenses/LICENSE-2.0
7 //
8 // Unless required by applicable law or agreed to in writing, software
9 // distributed under the License is distributed on an "AS IS" BASIS,
10 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11 // See the License for the specific language governing permissions and
12 // limitations under the License.
13 
14 use super::{BytesReader, Cacheable, Position};
15 use std::io::{Error, ErrorKind, Read, Result};
16 
17 /// Reader for reading I\O. This reader implements `BytesReader` trait and
18 /// `Cacheable` trait.
19 ///
20 /// # Examples
21 /// ```not run
22 /// use std::fs::File;
23 /// use ylong_bytes_reader::{IoReader, BytesReader};
24 ///
25 /// let file = File::open("./test.txt").unwrap();
26 /// let mut io_reader = IoReader::new(file);
27 /// let char = io_reader.next();
28 /// let char = io_reader.peek();
29 /// ```
30 pub(crate) struct IoReader<R: Read> {
31     io: R,
32     buf: Vec<u8>, // Buffer for storing read bytes.
33     cur: usize,   // The position of the cursor in the current buf.
34     idx: usize,   // A counter of all bytes that have been read.
35     pos: Position,
36     cache: Option<Cache>,
37 }
38 
39 // A simple cache implementation for `IoReader`.
40 struct Cache {
41     cache: Vec<u8>,
42     pre: usize, // Last cached location.
43 }
44 
45 impl Cache {
46     /// Create a new `Cache`.
new() -> Self47     fn new() -> Self {
48         Self {
49             cache: Vec::new(),
50             pre: 0,
51         }
52     }
53 }
54 
55 impl<R: Read> IoReader<R> {
56     /// Create a new `IoReader` from the given I\O.
new(io: R) -> Self57     pub(crate) fn new(io: R) -> Self {
58         Self {
59             io,
60             buf: Vec::with_capacity(1024), // Default size is 1024.
61             cur: 0,
62             idx: 0,
63             pos: Position::new(1, 1),
64             cache: None,
65         }
66     }
67 
68     // Try to read some bytes from io to fill buf.
read_bytes(&mut self) -> Result<bool>69     fn read_bytes(&mut self) -> Result<bool> {
70         unsafe {
71             self.buf.set_len(1024);
72         }
73         loop {
74             return match self.io.read(self.buf.as_mut_slice()) {
75                 Ok(0) => unsafe {
76                     self.buf.set_len(0);
77                     Ok(false)
78                 },
79                 Ok(n) => unsafe {
80                     self.buf.set_len(n);
81                     Ok(true)
82                 },
83                 Err(ref e) if e.kind() == ErrorKind::WouldBlock => continue,
84                 Err(e) => Err(e),
85             };
86         }
87     }
88 
89     // If there is not enough bytes in buf, try to read some bytes from io and
90     // reset some parameters inside.
load(&mut self) -> Result<bool>91     fn load(&mut self) -> Result<bool> {
92         if let Some(ref mut cacher) = self.cache {
93             cacher.cache.extend_from_slice(&self.buf[cacher.pre..]);
94             cacher.pre = 0;
95         }
96         let result = self.read_bytes();
97         if let Ok(true) = result {
98             self.cur = 0;
99         }
100         result
101     }
102 
103     // Every time a user calls a cache-related interface, the cache content
104     // needs to be updated in time.
update_cache(&mut self)105     fn update_cache(&mut self) {
106         if let Some(ref mut cacher) = self.cache {
107             if self.cur > cacher.pre {
108                 cacher
109                     .cache
110                     .extend_from_slice(&self.buf[cacher.pre..self.cur]);
111             }
112             cacher.pre = self.cur;
113         }
114     }
115 }
116 
117 impl<R: Read> BytesReader for IoReader<R> {
118     type Error = Error;
119 
next(&mut self) -> Result<Option<u8>>120     fn next(&mut self) -> Result<Option<u8>> {
121         if self.cur == self.buf.len() {
122             match self.load() {
123                 Ok(true) => {}
124                 Ok(false) => return Ok(None),
125                 Err(e) => return Err(e),
126             }
127         }
128 
129         let ch = self.buf[self.cur];
130         self.cur += 1;
131         self.idx += 1;
132 
133         if ch == b'\n' {
134             self.pos.line += 1;
135             self.pos.column = 1;
136         } else {
137             self.pos.column += 1;
138         }
139 
140         Ok(Some(ch))
141     }
142 
peek(&mut self) -> Result<Option<u8>>143     fn peek(&mut self) -> Result<Option<u8>> {
144         if self.cur == self.buf.len() {
145             match self.load() {
146                 Ok(true) => {}
147                 Ok(false) => return Ok(None),
148                 Err(e) => return Err(e),
149             }
150         }
151 
152         Ok(Some(self.buf[self.cur]))
153     }
154 
discard(&mut self)155     fn discard(&mut self) {
156         if self.cur == self.buf.len() {
157             match self.load() {
158                 Ok(true) => {}
159                 Ok(false) => return,
160                 Err(_) => return,
161             }
162         }
163 
164         let ch = self.buf[self.cur];
165         self.cur += 1;
166         self.idx += 1;
167 
168         if ch == b'\n' {
169             self.pos.line += 1;
170             self.pos.column = 1;
171         } else {
172             self.pos.column += 1;
173         }
174     }
175 
176     #[inline]
index(&self) -> usize177     fn index(&self) -> usize {
178         self.idx
179     }
180 
181     #[inline]
position(&self) -> Position182     fn position(&self) -> Position {
183         self.pos.clone()
184     }
185 }
186 
187 impl<R: Read> Cacheable for IoReader<R> {
start_caching(&mut self)188     fn start_caching(&mut self) {
189         if let Some(ref mut cacher) = self.cache {
190             cacher.cache.clear();
191             cacher.pre = self.cur;
192         } else {
193             let mut cache = Cache::new();
194             cache.pre = self.cur;
195             self.cache = Some(cache);
196         }
197     }
198 
cached_len(&mut self) -> Option<usize>199     fn cached_len(&mut self) -> Option<usize> {
200         self.update_cache();
201         self.cache.as_ref().map(|c| c.cache.len())
202     }
203 
cached_slice(&mut self) -> Option<&[u8]>204     fn cached_slice(&mut self) -> Option<&[u8]> {
205         self.update_cache();
206         self.cache.as_ref().map(|c| c.cache.as_slice())
207     }
208 
cached_data(&mut self) -> Option<Vec<u8>>209     fn cached_data(&mut self) -> Option<Vec<u8>> {
210         self.update_cache();
211         self.cache.as_ref().map(|c| c.cache.clone())
212     }
213 
end_caching(&mut self)214     fn end_caching(&mut self) {
215         self.cache = None;
216     }
217 
take_cached_data(&mut self) -> Option<Vec<u8>>218     fn take_cached_data(&mut self) -> Option<Vec<u8>> {
219         self.update_cache();
220         self.cache.take().map(|c| c.cache)
221     }
222 }
223 
224 #[cfg(test)]
225 mod ut_io_reader {
226     use super::{BytesReader, Cacheable, IoReader};
227     use std::cmp;
228     use std::io::{ErrorKind, Read};
229 
230     struct TestIo {
231         vec: Vec<u8>,
232         idx: usize,
233     }
234 
235     impl TestIo {
new(vec: Vec<u8>) -> Self236         fn new(vec: Vec<u8>) -> Self {
237             Self { vec, idx: 0 }
238         }
239     }
240 
241     impl Read for TestIo {
read(&mut self, buf: &mut [u8]) -> std::io::Result<usize>242         fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
243             if self.idx == self.vec.len() {
244                 return Ok(0);
245             }
246             let last = cmp::min(self.idx + buf.len(), self.vec.len());
247             let len = last - self.idx;
248             buf[..len].copy_from_slice(&self.vec[self.idx..last]);
249             self.idx = last;
250             Ok(len)
251         }
252     }
253 
254     struct TestWouldBlockIo {
255         cnt: usize,
256     }
257 
258     impl TestWouldBlockIo {
new() -> Self259         fn new() -> Self {
260             Self { cnt: 0 }
261         }
262 
is_finished(&self) -> bool263         fn is_finished(&self) -> bool {
264             self.cnt == 10
265         }
266     }
267 
268     impl Read for TestWouldBlockIo {
read(&mut self, _buf: &mut [u8]) -> std::io::Result<usize>269         fn read(&mut self, _buf: &mut [u8]) -> std::io::Result<usize> {
270             if self.cnt < 10 {
271                 self.cnt += 1;
272                 return Err(ErrorKind::WouldBlock.into());
273             }
274             Ok(0)
275         }
276     }
277 
278     struct TestErrIo;
279 
280     impl TestErrIo {
new() -> Self281         fn new() -> Self {
282             Self
283         }
284     }
285 
286     impl Read for TestErrIo {
read(&mut self, _buf: &mut [u8]) -> std::io::Result<usize>287         fn read(&mut self, _buf: &mut [u8]) -> std::io::Result<usize> {
288             Err(ErrorKind::AddrInUse.into())
289         }
290     }
291 
292     /// UT test case for `IoReader::new`.
293     ///
294     /// # Title
295     /// ut_io_reader_new
296     ///
297     /// # Brief
298     /// 1. Call `IoReader::new`.
299     /// 2. Check that parts of the return value are default values.
300     #[test]
ut_io_reader_new()301     fn ut_io_reader_new() {
302         let io = TestIo::new(Vec::new());
303         let io_reader = IoReader::new(io);
304 
305         assert_eq!(io_reader.cur, 0);
306         assert_eq!(io_reader.idx, 0);
307         assert_eq!(io_reader.pos.line, 1);
308         assert_eq!(io_reader.pos.column, 1);
309         assert_eq!(io_reader.buf.capacity(), 1024);
310         assert!(io_reader.buf.is_empty());
311         assert!(io_reader.cache.is_none());
312     }
313 
314     /// UT test case for `IoReader::next`.
315     ///
316     /// # Title
317     /// ut_test_case_io_reader_next
318     ///
319     /// # Brief
320     /// 1. Create a `IoReader`.
321     /// 2. Call `IoReader::next`.
322     /// 3. Check the return value against the following conditions:
323     ///     - If the end is not read, it returns `Ok(Some(..))`, and the index
324     ///     is moved backward; if the end is read, it returns `Ok(None)`, and
325     ///     the index is not moved.
326     #[test]
ut_io_reader_next()327     fn ut_io_reader_next() {
328         // Use TestIo.
329         let io = TestIo::new(vec![1u8; 1025]);
330         let mut io_reader = IoReader::new(io);
331         assert_eq!(io_reader.next().unwrap(), Some(1));
332         for _ in 0..1023 {
333             let _ = io_reader.next().unwrap();
334         }
335         assert_eq!(io_reader.next().unwrap(), Some(1));
336         assert_eq!(io_reader.next().unwrap(), None);
337 
338         // Use TestWouldBlockIo.
339         let io = TestWouldBlockIo::new();
340         let mut io_reader = IoReader::new(io);
341         assert_eq!(io_reader.next().unwrap(), None);
342         assert!(io_reader.io.is_finished());
343 
344         // Use TestErrIo
345         let io = TestErrIo::new();
346         let mut io_reader = IoReader::new(io);
347         assert!(io_reader.next().is_err());
348     }
349 
350     /// UT test case for `IoReader::peek`.
351     ///
352     /// # Title
353     /// ut_io_reader_peek
354     ///
355     /// # Brief
356     /// 1. Create a `IoReader`.
357     /// 2. Call `IoReader::peek`.
358     /// 3. Check the return value against the following conditions:
359     ///     - If the end is not read, it returns `Ok(Some(..))`; if the end is
360     ///     read, it returns `Ok(None)`.
361     #[test]
ut_io_reader_peek()362     fn ut_io_reader_peek() {
363         // Use TestIo.
364         let io = TestIo::new(vec![1u8; 1]);
365         let mut io_reader = IoReader::new(io);
366         assert_eq!(io_reader.peek().unwrap(), Some(1));
367         assert_eq!(io_reader.peek().unwrap(), Some(1));
368         assert_eq!(io_reader.next().unwrap(), Some(1));
369         assert_eq!(io_reader.peek().unwrap(), None);
370 
371         // Use TestWouldBlockIo.
372         let io = TestWouldBlockIo::new();
373         let mut io_reader = IoReader::new(io);
374         assert_eq!(io_reader.peek().unwrap(), None);
375         assert!(io_reader.io.is_finished());
376 
377         // Use TestErrorIo.
378         let io = TestErrIo::new();
379         let mut io_reader = IoReader::new(io);
380         assert!(io_reader.peek().is_err());
381     }
382 
383     /// UT test case for `IoReader::discard`.
384     ///
385     /// # Title
386     /// ut_io_reader_discard
387     ///
388     /// # Brief
389     /// 1. Create a `IoReader`.
390     /// 2. Call `IoReader::discard`.
391     /// 3. Check `index` against the following conditions:
392     ///     - If the end is not read, the index is moved backward; if the end is
393     ///     read, the index is not moved.
394     #[test]
ut_io_reader_discard()395     fn ut_io_reader_discard() {
396         let io = TestIo::new(vec![1u8; 1]);
397         let mut io_reader = IoReader::new(io);
398         assert_eq!(io_reader.index(), 0);
399         io_reader.discard();
400         assert_eq!(io_reader.index(), 1);
401         io_reader.discard();
402         assert_eq!(io_reader.index(), 1);
403     }
404 
405     /// UT test case for `IoReader::index`.
406     ///
407     /// # Title
408     /// ut_io_reader_index
409     ///
410     /// # Brief
411     /// 1. Create a `IoReader`.
412     /// 2. Call `IoReader::index`.
413     /// 3. Check if the `index` is correct.
414     #[test]
ut_io_reader_index()415     fn ut_io_reader_index() {
416         let io = TestIo::new(vec![1u8; 1]);
417         let io_reader = IoReader::new(io);
418         assert_eq!(io_reader.index(), 0);
419     }
420 
421     /// UT test case for `IoReader::position`.
422     ///
423     /// # Title
424     /// ut_io_reader_position
425     ///
426     /// # Brief
427     /// 1. Create a `IoReader`.
428     /// 2. Call `IoReader::position`.
429     /// 3. Check the return value against the following conditions:
430     ///     - If `'\n'` is read, the line number will increase and the column
431     ///     number will return to 1; if other characters are read, the line
432     ///     number will remain unchanged and the column number will increase.
433     #[test]
ut_io_reader_position()434     fn ut_io_reader_position() {
435         let io = TestIo::new(vec![1u8, b'\n', 2, b'\n', 3]);
436         let mut io_reader = IoReader::new(io);
437         let position = io_reader.position();
438         assert_eq!(position.line(), 1);
439         assert_eq!(position.column(), 1);
440         assert_eq!(io_reader.next().unwrap(), Some(1));
441 
442         // Use `next()`.
443         let position = io_reader.position();
444         assert_eq!(position.line(), 1);
445         assert_eq!(position.column(), 2);
446         assert_eq!(io_reader.next().unwrap(), Some(b'\n'));
447 
448         let position = io_reader.position();
449         assert_eq!(position.line(), 2);
450         assert_eq!(position.column(), 1);
451         assert_eq!(io_reader.next().unwrap(), Some(2));
452 
453         // Use `peek()` and `discard()`.
454         let position = io_reader.position();
455         assert_eq!(position.line(), 2);
456         assert_eq!(position.column(), 2);
457         assert_eq!(io_reader.peek().unwrap(), Some(b'\n'));
458         io_reader.discard();
459 
460         let position = io_reader.position();
461         assert_eq!(position.line(), 3);
462         assert_eq!(position.column(), 1);
463         assert_eq!(io_reader.peek().unwrap(), Some(3));
464         io_reader.discard();
465 
466         let position = io_reader.position();
467         assert_eq!(position.line(), 3);
468         assert_eq!(position.column(), 2);
469         assert_eq!(io_reader.peek().unwrap(), None);
470     }
471 
472     /// UT test case for `IoReader::start_caching`.
473     ///
474     /// # Title
475     /// ut_io_reader_start_caching
476     ///
477     /// # Brief
478     /// 1. Create a `IoReader`.
479     /// 2. Call `IoReader::start_caching`.
480     /// 3. Check if `cache` is correct.
481     #[test]
ut_io_reader_start_caching()482     fn ut_io_reader_start_caching() {
483         let io = TestIo::new(vec![1]);
484         let mut io_reader = IoReader::new(io);
485         assert!(io_reader.cache.is_none());
486         io_reader.start_caching();
487         assert!(io_reader.cache.is_some());
488         assert_eq!(io_reader.cached_len(), Some(0));
489 
490         assert_eq!(io_reader.next().unwrap(), Some(1));
491         assert_eq!(io_reader.cached_len(), Some(1));
492         io_reader.start_caching();
493         assert_eq!(io_reader.cached_len(), Some(0));
494     }
495 
496     /// UT test case for `IoReader::cached_len`.
497     ///
498     /// # Title
499     /// ut_io_reader_cached_len
500     ///
501     /// # Brief
502     /// 1. Create a `IoReader`.
503     /// 2. Call `IoReader::cached_len`.
504     /// 3. Check the return value against the following conditions:
505     ///     - Returns `None` if caching is not enabled, otherwise returns
506     ///     `Some(..)`.
507     #[test]
ut_io_reader_cached_len()508     fn ut_io_reader_cached_len() {
509         let io = TestIo::new(Vec::new());
510         let mut io_reader = IoReader::new(io);
511         assert_eq!(io_reader.cached_len(), None);
512         io_reader.start_caching();
513         assert_eq!(io_reader.cached_len(), Some(0));
514     }
515 
516     /// UT test case for `IoReader::cached_slice`.
517     ///
518     /// # Title
519     /// ut_io_reader_cached_slice
520     ///
521     /// # Brief
522     /// 1. Create a `IoReader`.
523     /// 2. Call `IoReader::cached_slice`.
524     /// 3. Check the return value against the following conditions:
525     ///     - Returns `None` if caching is not enabled, otherwise returns
526     ///     `Some(..)`.
527     #[test]
ut_io_reader_cached_slice()528     fn ut_io_reader_cached_slice() {
529         let io = TestIo::new(Vec::new());
530         let mut io_reader = IoReader::new(io);
531         assert_eq!(io_reader.cached_slice(), None);
532         io_reader.start_caching();
533         assert_eq!(io_reader.cached_slice(), Some([].as_slice()));
534 
535         // Test 1025 bytes.
536         let mut input = vec![0; 1024];
537         input.push(1);
538         let io = TestIo::new(input);
539         let mut io_reader = IoReader::new(io);
540         for _ in 0..1023 {
541             let _ = io_reader.next();
542         }
543         io_reader.start_caching();
544         assert_eq!(io_reader.next().unwrap(), Some(0));
545         assert_eq!(io_reader.next().unwrap(), Some(1));
546         assert_eq!(io_reader.cached_len(), Some(2));
547     }
548 
549     /// UT test case for `IoReader::cached_data`.
550     ///
551     /// # Title
552     /// ut_io_reader_cached_slice
553     ///
554     /// # Brief
555     /// 1. Create a `IoReader`.
556     /// 2. Call `IoReader::cached_data`.
557     /// 3. Check the return value against the following conditions:
558     ///     - Returns `None` if caching is not enabled, otherwise returns
559     ///     `Some(..)`.
560     #[test]
ut_io_reader_cached_data()561     fn ut_io_reader_cached_data() {
562         let io = TestIo::new(Vec::new());
563         let mut io_reader = IoReader::new(io);
564         assert_eq!(io_reader.cached_data(), None);
565         io_reader.start_caching();
566         assert_eq!(io_reader.cached_data(), Some(Vec::new()));
567     }
568 
569     /// UT test case for `IoReader::end_caching`.
570     ///
571     /// # Title
572     /// ut_io_reader_end_caching
573     ///
574     /// # Brief
575     /// 1. Create a `IoReader`.
576     /// 2. Call `IoReader::end_caching`.
577     /// 3. Check if `cache` is correct.
578     #[test]
ut_io_reader_end_caching()579     fn ut_io_reader_end_caching() {
580         let io = TestIo::new(Vec::new());
581         let mut io_reader = IoReader::new(io);
582         io_reader.start_caching();
583         assert!(io_reader.cache.is_some());
584         io_reader.end_caching();
585         assert!(io_reader.cache.is_none());
586     }
587 
588     /// UT test case for `IoReader::take_cached_data`.
589     ///
590     /// # Title
591     /// ut_io_reader_take_cached_data
592     ///
593     /// # Brief
594     /// 1. Create a `IoReader`.
595     /// 2. Call `IoReader::take_cached_data`.
596     /// 3. Check if the return value is correct.
597     #[test]
ut_io_reader_take_cached_data()598     fn ut_io_reader_take_cached_data() {
599         let io = TestIo::new(Vec::new());
600         let mut io_reader = IoReader::new(io);
601         io_reader.start_caching();
602         assert!(io_reader.cache.is_some());
603         assert_eq!(io_reader.take_cached_data(), Some(Vec::new()));
604         assert!(io_reader.cache.is_none());
605     }
606 }
607