expectrl/
needle.rs

1//! A module which contains a [Needle] trait and a list of implementations.
2//! [Needle] is used for seach in a byte stream.
3//!
4//! The list of provided implementations can be found in the documentation.
5
6use crate::error::Error;
7
8/// Needle an interface for search of a match in a buffer.
9pub trait Needle {
10    /// Function returns all matches that were occured.
11    fn check(&self, buf: &[u8], eof: bool) -> Result<Vec<Match>, Error>;
12}
13
14/// Match structure represent a range of bytes where match was found.
15#[derive(Debug, Clone, PartialEq, Eq)]
16pub struct Match {
17    start: usize,
18    end: usize,
19}
20
21impl Match {
22    /// New construct's an intanse of a Match.
23    pub fn new(start: usize, end: usize) -> Self {
24        Self { start, end }
25    }
26
27    /// Start returns a start index of a match.
28    pub fn start(&self) -> usize {
29        self.start
30    }
31
32    /// End returns an end index of a match.
33    pub fn end(&self) -> usize {
34        self.end
35    }
36}
37
38impl From<regex::bytes::Match<'_>> for Match {
39    fn from(m: regex::bytes::Match<'_>) -> Self {
40        Self::new(m.start(), m.end())
41    }
42}
43
44/// Regex tries to look up a match by a regex.
45#[derive(Debug)]
46pub struct Regex<Re: AsRef<str>>(pub Re);
47
48impl<Re: AsRef<str>> Needle for Regex<Re> {
49    fn check(&self, buf: &[u8], _: bool) -> Result<Vec<Match>, Error> {
50        let regex = regex::bytes::Regex::new(self.0.as_ref()).map_err(|_| Error::RegexParsing)?;
51        let matches = regex
52            .captures_iter(buf)
53            .flat_map(|c| c.iter().flatten().map(|m| m.into()).collect::<Vec<Match>>())
54            .collect();
55        Ok(matches)
56    }
57}
58
59/// Eof consider a match when an EOF is reached.
60#[derive(Debug)]
61pub struct Eof;
62
63impl Needle for Eof {
64    fn check(&self, buf: &[u8], eof: bool) -> Result<Vec<Match>, Error> {
65        match eof {
66            true => Ok(vec![Match::new(0, buf.len())]),
67            false => Ok(Vec::new()),
68        }
69    }
70}
71
72/// NBytes matches N bytes from the stream.
73#[derive(Debug)]
74pub struct NBytes(pub usize);
75
76impl NBytes {
77    fn count(&self) -> usize {
78        self.0
79    }
80}
81
82impl Needle for NBytes {
83    fn check(&self, buf: &[u8], _: bool) -> Result<Vec<Match>, Error> {
84        match buf.len() >= self.count() {
85            true => Ok(vec![Match::new(0, self.count())]),
86            false => Ok(Vec::new()),
87        }
88    }
89}
90
91impl Needle for [u8] {
92    fn check(&self, buf: &[u8], _: bool) -> Result<Vec<Match>, Error> {
93        if buf.len() < self.len() {
94            return Ok(Vec::new());
95        }
96
97        for l_bound in 0..buf.len() {
98            let r_bound = l_bound + self.len();
99            if r_bound > buf.len() {
100                return Ok(Vec::new());
101            }
102
103            if self == &buf[l_bound..r_bound] {
104                return Ok(vec![Match::new(l_bound, r_bound)]);
105            }
106        }
107
108        Ok(Vec::new())
109    }
110}
111
112impl Needle for &[u8] {
113    fn check(&self, buf: &[u8], eof: bool) -> Result<Vec<Match>, Error> {
114        (*self).check(buf, eof)
115    }
116}
117
118impl Needle for str {
119    fn check(&self, buf: &[u8], eof: bool) -> Result<Vec<Match>, Error> {
120        self.as_bytes().check(buf, eof)
121    }
122}
123
124impl Needle for &str {
125    fn check(&self, buf: &[u8], eof: bool) -> Result<Vec<Match>, Error> {
126        self.as_bytes().check(buf, eof)
127    }
128}
129
130impl Needle for String {
131    fn check(&self, buf: &[u8], eof: bool) -> Result<Vec<Match>, Error> {
132        self.as_bytes().check(buf, eof)
133    }
134}
135
136impl Needle for u8 {
137    fn check(&self, buf: &[u8], eof: bool) -> Result<Vec<Match>, Error> {
138        ([*self][..]).check(buf, eof)
139    }
140}
141
142impl Needle for char {
143    fn check(&self, buf: &[u8], eof: bool) -> Result<Vec<Match>, Error> {
144        char::to_string(self).check(buf, eof)
145    }
146}
147
148/// Any matches uses all provided lookups and returns a match
149/// from a first successfull match.
150///
151/// It does checks lookups in order they were provided.
152///
153/// # Example
154///
155/// ```no_run,ignore
156/// use expectrl::{spawn, Any};
157///
158/// let mut p = spawn("cat").unwrap();
159/// p.expect(Any(["we", "are", "here"])).unwrap();
160/// ```
161///
162/// To be able to combine different types of lookups you can call [Any::boxed].
163///
164/// ```no_run,ignore
165/// use expectrl::{spawn, Any, NBytes};
166///
167/// let mut p = spawn("cat").unwrap();
168/// p.expect(Any::boxed(vec![Box::new("we"), Box::new(NBytes(3))])).unwrap();
169/// ```
170#[derive(Debug)]
171pub struct Any<I>(pub I);
172
173impl Any<Vec<Box<dyn Needle>>> {
174    /// Boxed expectes a list of [Box]ed lookups.
175    pub fn boxed(v: Vec<Box<dyn Needle>>) -> Self {
176        Self(v)
177    }
178}
179
180impl<T> Needle for Any<&[T]>
181where
182    T: Needle,
183{
184    fn check(&self, buf: &[u8], eof: bool) -> Result<Vec<Match>, Error> {
185        for needle in self.0.iter() {
186            let found = needle.check(buf, eof)?;
187            if !found.is_empty() {
188                return Ok(found);
189            }
190        }
191
192        Ok(Vec::new())
193    }
194}
195
196impl<T> Needle for Any<Vec<T>>
197where
198    T: Needle,
199{
200    fn check(&self, buf: &[u8], eof: bool) -> Result<Vec<Match>, Error> {
201        Any(self.0.as_slice()).check(buf, eof)
202    }
203}
204
205impl<T, const N: usize> Needle for Any<[T; N]>
206where
207    T: Needle,
208{
209    fn check(&self, buf: &[u8], eof: bool) -> Result<Vec<Match>, Error> {
210        Any(&self.0[..]).check(buf, eof)
211    }
212}
213
214impl<T, const N: usize> Needle for Any<&'_ [T; N]>
215where
216    T: Needle,
217{
218    fn check(&self, buf: &[u8], eof: bool) -> Result<Vec<Match>, Error> {
219        Any(&self.0[..]).check(buf, eof)
220    }
221}
222
223impl<T: Needle> Needle for &T {
224    fn check(&self, buf: &[u8], eof: bool) -> Result<Vec<Match>, Error> {
225        T::check(self, buf, eof)
226    }
227}
228
229impl Needle for Box<dyn Needle + '_> {
230    fn check(&self, buf: &[u8], eof: bool) -> Result<Vec<Match>, Error> {
231        self.as_ref().check(buf, eof)
232    }
233}
234
235#[cfg(test)]
236mod tests {
237    use super::*;
238
239    #[test]
240    fn test_regex() {
241        assert_eq!(
242            Regex("[0-9]+").check(b"+012345", false).unwrap(),
243            vec![Match::new(1, 7)]
244        );
245        assert_eq!(
246            Regex(r"\w+").check(b"What's Up Boys", false).unwrap(),
247            vec![
248                Match::new(0, 4),
249                Match::new(5, 6),
250                Match::new(7, 9),
251                Match::new(10, 14)
252            ]
253        );
254        assert_eq!(
255            Regex(r"((?:\w|')+)")
256                .check(b"What's Up Boys", false)
257                .unwrap(),
258            vec![
259                Match::new(0, 6),
260                Match::new(0, 6),
261                Match::new(7, 9),
262                Match::new(7, 9),
263                Match::new(10, 14),
264                Match::new(10, 14)
265            ]
266        );
267        assert_eq!(
268            Regex(r"(\w+)=(\w+)").check(b"asd=123", false).unwrap(),
269            vec![Match::new(0, 7), Match::new(0, 3), Match::new(4, 7)]
270        );
271    }
272
273    #[test]
274    fn test_eof() {
275        assert_eq!(Eof.check(b"qwe", true).unwrap(), vec![Match::new(0, 3)]);
276        assert_eq!(Eof.check(b"qwe", false).unwrap(), vec![]);
277    }
278
279    #[test]
280    fn test_n_bytes() {
281        assert_eq!(
282            NBytes(1).check(b"qwe", false).unwrap(),
283            vec![Match::new(0, 1)]
284        );
285        assert_eq!(
286            NBytes(0).check(b"qwe", false).unwrap(),
287            vec![Match::new(0, 0)]
288        );
289        assert_eq!(NBytes(10).check(b"qwe", false).unwrap(), vec![]);
290    }
291
292    #[test]
293    fn test_str() {
294        assert_eq!(
295            "wer".check(b"qwerty", false).unwrap(),
296            vec![Match::new(1, 4)]
297        );
298        assert_eq!("123".check(b"qwerty", false).unwrap(), vec![]);
299        assert_eq!("".check(b"qwerty", false).unwrap(), vec![Match::new(0, 0)]);
300    }
301
302    #[test]
303    fn test_bytes() {
304        assert_eq!(
305            b"wer".check(b"qwerty", false).unwrap(),
306            vec![Match::new(1, 4)]
307        );
308        assert_eq!(b"123".check(b"qwerty", false).unwrap(), vec![]);
309        assert_eq!(b"".check(b"qwerty", false).unwrap(), vec![Match::new(0, 0)]);
310    }
311
312    #[allow(clippy::needless_borrow)]
313    #[test]
314    #[allow(clippy::needless_borrow)]
315    fn test_bytes_ref() {
316        assert_eq!(
317            (&[b'q', b'w', b'e']).check(b"qwerty", false).unwrap(),
318            vec![Match::new(0, 3)]
319        );
320        assert_eq!(
321            (&[b'1', b'2', b'3']).check(b"qwerty", false).unwrap(),
322            vec![]
323        );
324        assert_eq!(
325            (&[]).check(b"qwerty", false).unwrap(),
326            vec![Match::new(0, 0)]
327        );
328    }
329
330    #[test]
331    fn test_byte() {
332        assert_eq!(
333            (b'3').check(b"1234", false).unwrap(),
334            vec![Match::new(2, 3)]
335        );
336        assert_eq!((b'3').check(b"1234", true).unwrap(), vec![Match::new(2, 3)]);
337    }
338
339    #[test]
340    fn test_char() {
341        for eof in [false, true] {
342            assert_eq!(
343                ('😘').check("😁😄😅😓😠😘😌".as_bytes(), eof).unwrap(),
344                vec![Match::new(20, 24)]
345            );
346        }
347    }
348
349    #[test]
350    fn test_any() {
351        assert_eq!(
352            Any::<Vec<Box<dyn Needle>>>(vec![Box::new("we"), Box::new(NBytes(3))])
353                .check(b"qwerty", false)
354                .unwrap(),
355            vec![Match::new(1, 3)]
356        );
357        assert_eq!(
358            Any::boxed(vec![Box::new("123"), Box::new(NBytes(100))])
359                .check(b"qwerty", false)
360                .unwrap(),
361            vec![],
362        );
363        assert_eq!(
364            Any(["123", "234", "rty"]).check(b"qwerty", false).unwrap(),
365            vec![Match::new(3, 6)]
366        );
367        assert_eq!(
368            Any(&["123", "234", "rty"][..])
369                .check(b"qwerty", false)
370                .unwrap(),
371            vec![Match::new(3, 6)]
372        );
373        assert_eq!(
374            Any(&["123", "234", "rty"]).check(b"qwerty", false).unwrap(),
375            vec![Match::new(3, 6)]
376        );
377    }
378}