1 pub mod fwht;
2 pub mod h264;
3 
4 use log::error;
5 use std::io;
6 
7 /// Trait for classes able to iterate an encoded stream over chunks of decodable units (typically
8 /// frames).
9 pub trait StreamSplitter: Iterator<Item = Vec<u8>> {}
10 
11 /// Splits a stream at each encounter of a given pattern. Useful to extract decodable units (or
12 /// frames from an encoded stream.
13 struct PatternSplitter<S: io::Read> {
14     /// The pattern to split at.
15     pattern: Vec<u8>,
16     stream: io::Bytes<S>,
17 }
18 
19 impl<S: io::Read> PatternSplitter<S> {
20     /// Create a new splitter that will split `stream` at each instance of `pattern`.
21     ///
22     /// `stream` must start with `pattern`, otherwise the input is considered invalid and `None` is
23     /// returned.
new(pattern: impl Into<Vec<u8>>, stream: S) -> Option<Self>24     fn new(pattern: impl Into<Vec<u8>>, stream: S) -> Option<Self> {
25         let mut stream = stream.bytes();
26         let pattern = pattern.into();
27 
28         // The stream must begin by our header, or it is invalid.
29         let stream_start = (0..pattern.len())
30             .map(|_| stream.next().unwrap_or(Ok(0)).unwrap_or(0))
31             .collect::<Vec<_>>();
32 
33         if stream_start == pattern {
34             Some(PatternSplitter { pattern, stream })
35         } else {
36             None
37         }
38     }
39 }
40 
41 impl<S: io::Read> Iterator for PatternSplitter<S> {
42     type Item = Vec<u8>;
43 
44     /// Returns the next frame in the stream, header included.
next(&mut self) -> Option<Self::Item>45     fn next(&mut self) -> Option<Self::Item> {
46         let mut frame_data: Vec<u8> = Vec::with_capacity(0x10000);
47         // Add the header of the frame since we are already past it.
48         frame_data.extend(&self.pattern);
49         // Window used to try and detect the next frame header.
50         let mut header_window: Vec<u8> = Vec::with_capacity(self.pattern.len());
51         // Iterator to the next character expected for a frame header.
52         let mut header_iter = self.pattern.iter().peekable();
53 
54         loop {
55             let b = match self.stream.next() {
56                 Some(Ok(b)) => b,
57                 // End of stream, commit all read data and return.
58                 None => {
59                     frame_data.extend(&header_window);
60                     // If the only data is our pre-filled header, then we haven't
61                     // read anything and thus there is no frame to return.
62                     return if frame_data.len() <= self.pattern.len() {
63                         None
64                     } else {
65                         Some(frame_data)
66                     };
67                 }
68                 Some(Err(e)) => {
69                     error!("Error while reading stream: {}", e);
70                     return None;
71                 }
72             };
73 
74             // Add the read byte to the header candidate buffer.
75             header_window.push(b);
76 
77             // Possibly a header?
78             if Some(&&b) == header_iter.peek() {
79                 header_iter.next();
80                 // Found next frame's header, return read data.
81                 if header_iter.peek().is_none() {
82                     return Some(frame_data);
83                 }
84             } else {
85                 // Not a header, commit header window data to current frame and reset.
86                 frame_data.extend(&header_window);
87                 if header_window.len() > 1 {
88                     header_iter = self.pattern.iter().peekable();
89                 }
90                 header_window.clear();
91             }
92         }
93     }
94 }
95