1 /*!
2 This module defines 128-bit vector implementations of `memchr` and friends.
3 
4 The main types in this module are [`One`], [`Two`] and [`Three`]. They are for
5 searching for one, two or three distinct bytes, respectively, in a haystack.
6 Each type also has corresponding double ended iterators. These searchers are
7 typically much faster than scalar routines accomplishing the same task.
8 
9 The `One` searcher also provides a [`One::count`] routine for efficiently
10 counting the number of times a single byte occurs in a haystack. This is
11 useful, for example, for counting the number of lines in a haystack. This
12 routine exists because it is usually faster, especially with a high match
13 count, then using [`One::find`] repeatedly. ([`OneIter`] specializes its
14 `Iterator::count` implementation to use this routine.)
15 
16 Only one, two and three bytes are supported because three bytes is about
17 the point where one sees diminishing returns. Beyond this point and it's
18 probably (but not necessarily) better to just use a simple `[bool; 256]` array
19 or similar. However, it depends mightily on the specific work-load and the
20 expected match frequency.
21 */
22 
23 use core::arch::aarch64::uint8x16_t;
24 
25 use crate::{arch::generic::memchr as generic, ext::Pointer, vector::Vector};
26 
27 /// Finds all occurrences of a single byte in a haystack.
28 #[derive(Clone, Copy, Debug)]
29 pub struct One(generic::One<uint8x16_t>);
30 
31 impl One {
32     /// Create a new searcher that finds occurrences of the needle byte given.
33     ///
34     /// This particular searcher is specialized to use neon vector instructions
35     /// that typically make it quite fast.
36     ///
37     /// If neon is unavailable in the current environment, then `None` is
38     /// returned.
39     #[inline]
new(needle: u8) -> Option<One>40     pub fn new(needle: u8) -> Option<One> {
41         if One::is_available() {
42             // SAFETY: we check that neon is available above.
43             unsafe { Some(One::new_unchecked(needle)) }
44         } else {
45             None
46         }
47     }
48 
49     /// Create a new finder specific to neon vectors and routines without
50     /// checking that neon is available.
51     ///
52     /// # Safety
53     ///
54     /// Callers must guarantee that it is safe to execute `neon` instructions
55     /// in the current environment.
56     ///
57     /// Note that it is a common misconception that if one compiles for an
58     /// `x86_64` target, then they therefore automatically have access to neon
59     /// instructions. While this is almost always the case, it isn't true in
60     /// 100% of cases.
61     #[target_feature(enable = "neon")]
62     #[inline]
new_unchecked(needle: u8) -> One63     pub unsafe fn new_unchecked(needle: u8) -> One {
64         One(generic::One::new(needle))
65     }
66 
67     /// Returns true when this implementation is available in the current
68     /// environment.
69     ///
70     /// When this is true, it is guaranteed that [`One::new`] will return
71     /// a `Some` value. Similarly, when it is false, it is guaranteed that
72     /// `One::new` will return a `None` value.
73     ///
74     /// Note also that for the lifetime of a single program, if this returns
75     /// true then it will always return true.
76     #[inline]
is_available() -> bool77     pub fn is_available() -> bool {
78         #[cfg(target_feature = "neon")]
79         {
80             true
81         }
82         #[cfg(not(target_feature = "neon"))]
83         {
84             false
85         }
86     }
87 
88     /// Return the first occurrence of one of the needle bytes in the given
89     /// haystack. If no such occurrence exists, then `None` is returned.
90     ///
91     /// The occurrence is reported as an offset into `haystack`. Its maximum
92     /// value is `haystack.len() - 1`.
93     #[inline]
find(&self, haystack: &[u8]) -> Option<usize>94     pub fn find(&self, haystack: &[u8]) -> Option<usize> {
95         // SAFETY: `find_raw` guarantees that if a pointer is returned, it
96         // falls within the bounds of the start and end pointers.
97         unsafe {
98             generic::search_slice_with_raw(haystack, |s, e| {
99                 self.find_raw(s, e)
100             })
101         }
102     }
103 
104     /// Return the last occurrence of one of the needle bytes in the given
105     /// haystack. If no such occurrence exists, then `None` is returned.
106     ///
107     /// The occurrence is reported as an offset into `haystack`. Its maximum
108     /// value is `haystack.len() - 1`.
109     #[inline]
rfind(&self, haystack: &[u8]) -> Option<usize>110     pub fn rfind(&self, haystack: &[u8]) -> Option<usize> {
111         // SAFETY: `rfind_raw` guarantees that if a pointer is returned, it
112         // falls within the bounds of the start and end pointers.
113         unsafe {
114             generic::search_slice_with_raw(haystack, |s, e| {
115                 self.rfind_raw(s, e)
116             })
117         }
118     }
119 
120     /// Counts all occurrences of this byte in the given haystack.
121     #[inline]
count(&self, haystack: &[u8]) -> usize122     pub fn count(&self, haystack: &[u8]) -> usize {
123         // SAFETY: All of our pointers are derived directly from a borrowed
124         // slice, which is guaranteed to be valid.
125         unsafe {
126             let start = haystack.as_ptr();
127             let end = start.add(haystack.len());
128             self.count_raw(start, end)
129         }
130     }
131 
132     /// Like `find`, but accepts and returns raw pointers.
133     ///
134     /// When a match is found, the pointer returned is guaranteed to be
135     /// `>= start` and `< end`.
136     ///
137     /// This routine is useful if you're already using raw pointers and would
138     /// like to avoid converting back to a slice before executing a search.
139     ///
140     /// # Safety
141     ///
142     /// * Both `start` and `end` must be valid for reads.
143     /// * Both `start` and `end` must point to an initialized value.
144     /// * Both `start` and `end` must point to the same allocated object and
145     /// must either be in bounds or at most one byte past the end of the
146     /// allocated object.
147     /// * Both `start` and `end` must be _derived from_ a pointer to the same
148     /// object.
149     /// * The distance between `start` and `end` must not overflow `isize`.
150     /// * The distance being in bounds must not rely on "wrapping around" the
151     /// address space.
152     ///
153     /// Note that callers may pass a pair of pointers such that `start >= end`.
154     /// In that case, `None` will always be returned.
155     #[inline]
find_raw( &self, start: *const u8, end: *const u8, ) -> Option<*const u8>156     pub unsafe fn find_raw(
157         &self,
158         start: *const u8,
159         end: *const u8,
160     ) -> Option<*const u8> {
161         if start >= end {
162             return None;
163         }
164         if end.distance(start) < uint8x16_t::BYTES {
165             // SAFETY: We require the caller to pass valid start/end pointers.
166             return generic::fwd_byte_by_byte(start, end, |b| {
167                 b == self.0.needle1()
168             });
169         }
170         // SAFETY: Building a `One` means it's safe to call 'neon' routines.
171         // Also, we've checked that our haystack is big enough to run on the
172         // vector routine. Pointer validity is caller's responsibility.
173         self.find_raw_impl(start, end)
174     }
175 
176     /// Like `rfind`, but accepts and returns raw pointers.
177     ///
178     /// When a match is found, the pointer returned is guaranteed to be
179     /// `>= start` and `< end`.
180     ///
181     /// This routine is useful if you're already using raw pointers and would
182     /// like to avoid converting back to a slice before executing a search.
183     ///
184     /// # Safety
185     ///
186     /// * Both `start` and `end` must be valid for reads.
187     /// * Both `start` and `end` must point to an initialized value.
188     /// * Both `start` and `end` must point to the same allocated object and
189     /// must either be in bounds or at most one byte past the end of the
190     /// allocated object.
191     /// * Both `start` and `end` must be _derived from_ a pointer to the same
192     /// object.
193     /// * The distance between `start` and `end` must not overflow `isize`.
194     /// * The distance being in bounds must not rely on "wrapping around" the
195     /// address space.
196     ///
197     /// Note that callers may pass a pair of pointers such that `start >= end`.
198     /// In that case, `None` will always be returned.
199     #[inline]
rfind_raw( &self, start: *const u8, end: *const u8, ) -> Option<*const u8>200     pub unsafe fn rfind_raw(
201         &self,
202         start: *const u8,
203         end: *const u8,
204     ) -> Option<*const u8> {
205         if start >= end {
206             return None;
207         }
208         if end.distance(start) < uint8x16_t::BYTES {
209             // SAFETY: We require the caller to pass valid start/end pointers.
210             return generic::rev_byte_by_byte(start, end, |b| {
211                 b == self.0.needle1()
212             });
213         }
214         // SAFETY: Building a `One` means it's safe to call 'neon' routines.
215         // Also, we've checked that our haystack is big enough to run on the
216         // vector routine. Pointer validity is caller's responsibility.
217         self.rfind_raw_impl(start, end)
218     }
219 
220     /// Like `count`, but accepts and returns raw pointers.
221     ///
222     /// This routine is useful if you're already using raw pointers and would
223     /// like to avoid converting back to a slice before executing a search.
224     ///
225     /// # Safety
226     ///
227     /// * Both `start` and `end` must be valid for reads.
228     /// * Both `start` and `end` must point to an initialized value.
229     /// * Both `start` and `end` must point to the same allocated object and
230     /// must either be in bounds or at most one byte past the end of the
231     /// allocated object.
232     /// * Both `start` and `end` must be _derived from_ a pointer to the same
233     /// object.
234     /// * The distance between `start` and `end` must not overflow `isize`.
235     /// * The distance being in bounds must not rely on "wrapping around" the
236     /// address space.
237     ///
238     /// Note that callers may pass a pair of pointers such that `start >= end`.
239     /// In that case, `None` will always be returned.
240     #[inline]
count_raw(&self, start: *const u8, end: *const u8) -> usize241     pub unsafe fn count_raw(&self, start: *const u8, end: *const u8) -> usize {
242         if start >= end {
243             return 0;
244         }
245         if end.distance(start) < uint8x16_t::BYTES {
246             // SAFETY: We require the caller to pass valid start/end pointers.
247             return generic::count_byte_by_byte(start, end, |b| {
248                 b == self.0.needle1()
249             });
250         }
251         // SAFETY: Building a `One` means it's safe to call 'neon' routines.
252         // Also, we've checked that our haystack is big enough to run on the
253         // vector routine. Pointer validity is caller's responsibility.
254         self.count_raw_impl(start, end)
255     }
256 
257     /// Execute a search using neon vectors and routines.
258     ///
259     /// # Safety
260     ///
261     /// Same as [`One::find_raw`], except the distance between `start` and
262     /// `end` must be at least the size of a neon vector (in bytes).
263     ///
264     /// (The target feature safety obligation is automatically fulfilled by
265     /// virtue of being a method on `One`, which can only be constructed
266     /// when it is safe to call `neon` routines.)
267     #[target_feature(enable = "neon")]
268     #[inline]
find_raw_impl( &self, start: *const u8, end: *const u8, ) -> Option<*const u8>269     unsafe fn find_raw_impl(
270         &self,
271         start: *const u8,
272         end: *const u8,
273     ) -> Option<*const u8> {
274         self.0.find_raw(start, end)
275     }
276 
277     /// Execute a search using neon vectors and routines.
278     ///
279     /// # Safety
280     ///
281     /// Same as [`One::rfind_raw`], except the distance between `start` and
282     /// `end` must be at least the size of a neon vector (in bytes).
283     ///
284     /// (The target feature safety obligation is automatically fulfilled by
285     /// virtue of being a method on `One`, which can only be constructed
286     /// when it is safe to call `neon` routines.)
287     #[target_feature(enable = "neon")]
288     #[inline]
rfind_raw_impl( &self, start: *const u8, end: *const u8, ) -> Option<*const u8>289     unsafe fn rfind_raw_impl(
290         &self,
291         start: *const u8,
292         end: *const u8,
293     ) -> Option<*const u8> {
294         self.0.rfind_raw(start, end)
295     }
296 
297     /// Execute a count using neon vectors and routines.
298     ///
299     /// # Safety
300     ///
301     /// Same as [`One::count_raw`], except the distance between `start` and
302     /// `end` must be at least the size of a neon vector (in bytes).
303     ///
304     /// (The target feature safety obligation is automatically fulfilled by
305     /// virtue of being a method on `One`, which can only be constructed
306     /// when it is safe to call `neon` routines.)
307     #[target_feature(enable = "neon")]
308     #[inline]
count_raw_impl( &self, start: *const u8, end: *const u8, ) -> usize309     unsafe fn count_raw_impl(
310         &self,
311         start: *const u8,
312         end: *const u8,
313     ) -> usize {
314         self.0.count_raw(start, end)
315     }
316 
317     /// Returns an iterator over all occurrences of the needle byte in the
318     /// given haystack.
319     ///
320     /// The iterator returned implements `DoubleEndedIterator`. This means it
321     /// can also be used to find occurrences in reverse order.
322     #[inline]
iter<'a, 'h>(&'a self, haystack: &'h [u8]) -> OneIter<'a, 'h>323     pub fn iter<'a, 'h>(&'a self, haystack: &'h [u8]) -> OneIter<'a, 'h> {
324         OneIter { searcher: self, it: generic::Iter::new(haystack) }
325     }
326 }
327 
328 /// An iterator over all occurrences of a single byte in a haystack.
329 ///
330 /// This iterator implements `DoubleEndedIterator`, which means it can also be
331 /// used to find occurrences in reverse order.
332 ///
333 /// This iterator is created by the [`One::iter`] method.
334 ///
335 /// The lifetime parameters are as follows:
336 ///
337 /// * `'a` refers to the lifetime of the underlying [`One`] searcher.
338 /// * `'h` refers to the lifetime of the haystack being searched.
339 #[derive(Clone, Debug)]
340 pub struct OneIter<'a, 'h> {
341     searcher: &'a One,
342     it: generic::Iter<'h>,
343 }
344 
345 impl<'a, 'h> Iterator for OneIter<'a, 'h> {
346     type Item = usize;
347 
348     #[inline]
next(&mut self) -> Option<usize>349     fn next(&mut self) -> Option<usize> {
350         // SAFETY: We rely on the generic iterator to provide valid start
351         // and end pointers, but we guarantee that any pointer returned by
352         // 'find_raw' falls within the bounds of the start and end pointer.
353         unsafe { self.it.next(|s, e| self.searcher.find_raw(s, e)) }
354     }
355 
356     #[inline]
count(self) -> usize357     fn count(self) -> usize {
358         self.it.count(|s, e| {
359             // SAFETY: We rely on our generic iterator to return valid start
360             // and end pointers.
361             unsafe { self.searcher.count_raw(s, e) }
362         })
363     }
364 
365     #[inline]
size_hint(&self) -> (usize, Option<usize>)366     fn size_hint(&self) -> (usize, Option<usize>) {
367         self.it.size_hint()
368     }
369 }
370 
371 impl<'a, 'h> DoubleEndedIterator for OneIter<'a, 'h> {
372     #[inline]
next_back(&mut self) -> Option<usize>373     fn next_back(&mut self) -> Option<usize> {
374         // SAFETY: We rely on the generic iterator to provide valid start
375         // and end pointers, but we guarantee that any pointer returned by
376         // 'rfind_raw' falls within the bounds of the start and end pointer.
377         unsafe { self.it.next_back(|s, e| self.searcher.rfind_raw(s, e)) }
378     }
379 }
380 
381 impl<'a, 'h> core::iter::FusedIterator for OneIter<'a, 'h> {}
382 
383 /// Finds all occurrences of two bytes in a haystack.
384 ///
385 /// That is, this reports matches of one of two possible bytes. For example,
386 /// searching for `a` or `b` in `afoobar` would report matches at offsets `0`,
387 /// `4` and `5`.
388 #[derive(Clone, Copy, Debug)]
389 pub struct Two(generic::Two<uint8x16_t>);
390 
391 impl Two {
392     /// Create a new searcher that finds occurrences of the needle bytes given.
393     ///
394     /// This particular searcher is specialized to use neon vector instructions
395     /// that typically make it quite fast.
396     ///
397     /// If neon is unavailable in the current environment, then `None` is
398     /// returned.
399     #[inline]
new(needle1: u8, needle2: u8) -> Option<Two>400     pub fn new(needle1: u8, needle2: u8) -> Option<Two> {
401         if Two::is_available() {
402             // SAFETY: we check that neon is available above.
403             unsafe { Some(Two::new_unchecked(needle1, needle2)) }
404         } else {
405             None
406         }
407     }
408 
409     /// Create a new finder specific to neon vectors and routines without
410     /// checking that neon is available.
411     ///
412     /// # Safety
413     ///
414     /// Callers must guarantee that it is safe to execute `neon` instructions
415     /// in the current environment.
416     ///
417     /// Note that it is a common misconception that if one compiles for an
418     /// `x86_64` target, then they therefore automatically have access to neon
419     /// instructions. While this is almost always the case, it isn't true in
420     /// 100% of cases.
421     #[target_feature(enable = "neon")]
422     #[inline]
new_unchecked(needle1: u8, needle2: u8) -> Two423     pub unsafe fn new_unchecked(needle1: u8, needle2: u8) -> Two {
424         Two(generic::Two::new(needle1, needle2))
425     }
426 
427     /// Returns true when this implementation is available in the current
428     /// environment.
429     ///
430     /// When this is true, it is guaranteed that [`Two::new`] will return
431     /// a `Some` value. Similarly, when it is false, it is guaranteed that
432     /// `Two::new` will return a `None` value.
433     ///
434     /// Note also that for the lifetime of a single program, if this returns
435     /// true then it will always return true.
436     #[inline]
is_available() -> bool437     pub fn is_available() -> bool {
438         #[cfg(target_feature = "neon")]
439         {
440             true
441         }
442         #[cfg(not(target_feature = "neon"))]
443         {
444             false
445         }
446     }
447 
448     /// Return the first occurrence of one of the needle bytes in the given
449     /// haystack. If no such occurrence exists, then `None` is returned.
450     ///
451     /// The occurrence is reported as an offset into `haystack`. Its maximum
452     /// value is `haystack.len() - 1`.
453     #[inline]
find(&self, haystack: &[u8]) -> Option<usize>454     pub fn find(&self, haystack: &[u8]) -> Option<usize> {
455         // SAFETY: `find_raw` guarantees that if a pointer is returned, it
456         // falls within the bounds of the start and end pointers.
457         unsafe {
458             generic::search_slice_with_raw(haystack, |s, e| {
459                 self.find_raw(s, e)
460             })
461         }
462     }
463 
464     /// Return the last occurrence of one of the needle bytes in the given
465     /// haystack. If no such occurrence exists, then `None` is returned.
466     ///
467     /// The occurrence is reported as an offset into `haystack`. Its maximum
468     /// value is `haystack.len() - 1`.
469     #[inline]
rfind(&self, haystack: &[u8]) -> Option<usize>470     pub fn rfind(&self, haystack: &[u8]) -> Option<usize> {
471         // SAFETY: `rfind_raw` guarantees that if a pointer is returned, it
472         // falls within the bounds of the start and end pointers.
473         unsafe {
474             generic::search_slice_with_raw(haystack, |s, e| {
475                 self.rfind_raw(s, e)
476             })
477         }
478     }
479 
480     /// Like `find`, but accepts and returns raw pointers.
481     ///
482     /// When a match is found, the pointer returned is guaranteed to be
483     /// `>= start` and `< end`.
484     ///
485     /// This routine is useful if you're already using raw pointers and would
486     /// like to avoid converting back to a slice before executing a search.
487     ///
488     /// # Safety
489     ///
490     /// * Both `start` and `end` must be valid for reads.
491     /// * Both `start` and `end` must point to an initialized value.
492     /// * Both `start` and `end` must point to the same allocated object and
493     /// must either be in bounds or at most one byte past the end of the
494     /// allocated object.
495     /// * Both `start` and `end` must be _derived from_ a pointer to the same
496     /// object.
497     /// * The distance between `start` and `end` must not overflow `isize`.
498     /// * The distance being in bounds must not rely on "wrapping around" the
499     /// address space.
500     ///
501     /// Note that callers may pass a pair of pointers such that `start >= end`.
502     /// In that case, `None` will always be returned.
503     #[inline]
find_raw( &self, start: *const u8, end: *const u8, ) -> Option<*const u8>504     pub unsafe fn find_raw(
505         &self,
506         start: *const u8,
507         end: *const u8,
508     ) -> Option<*const u8> {
509         if start >= end {
510             return None;
511         }
512         if end.distance(start) < uint8x16_t::BYTES {
513             // SAFETY: We require the caller to pass valid start/end pointers.
514             return generic::fwd_byte_by_byte(start, end, |b| {
515                 b == self.0.needle1() || b == self.0.needle2()
516             });
517         }
518         // SAFETY: Building a `Two` means it's safe to call 'neon' routines.
519         // Also, we've checked that our haystack is big enough to run on the
520         // vector routine. Pointer validity is caller's responsibility.
521         self.find_raw_impl(start, end)
522     }
523 
524     /// Like `rfind`, but accepts and returns raw pointers.
525     ///
526     /// When a match is found, the pointer returned is guaranteed to be
527     /// `>= start` and `< end`.
528     ///
529     /// This routine is useful if you're already using raw pointers and would
530     /// like to avoid converting back to a slice before executing a search.
531     ///
532     /// # Safety
533     ///
534     /// * Both `start` and `end` must be valid for reads.
535     /// * Both `start` and `end` must point to an initialized value.
536     /// * Both `start` and `end` must point to the same allocated object and
537     /// must either be in bounds or at most one byte past the end of the
538     /// allocated object.
539     /// * Both `start` and `end` must be _derived from_ a pointer to the same
540     /// object.
541     /// * The distance between `start` and `end` must not overflow `isize`.
542     /// * The distance being in bounds must not rely on "wrapping around" the
543     /// address space.
544     ///
545     /// Note that callers may pass a pair of pointers such that `start >= end`.
546     /// In that case, `None` will always be returned.
547     #[inline]
rfind_raw( &self, start: *const u8, end: *const u8, ) -> Option<*const u8>548     pub unsafe fn rfind_raw(
549         &self,
550         start: *const u8,
551         end: *const u8,
552     ) -> Option<*const u8> {
553         if start >= end {
554             return None;
555         }
556         if end.distance(start) < uint8x16_t::BYTES {
557             // SAFETY: We require the caller to pass valid start/end pointers.
558             return generic::rev_byte_by_byte(start, end, |b| {
559                 b == self.0.needle1() || b == self.0.needle2()
560             });
561         }
562         // SAFETY: Building a `Two` means it's safe to call 'neon' routines.
563         // Also, we've checked that our haystack is big enough to run on the
564         // vector routine. Pointer validity is caller's responsibility.
565         self.rfind_raw_impl(start, end)
566     }
567 
568     /// Execute a search using neon vectors and routines.
569     ///
570     /// # Safety
571     ///
572     /// Same as [`Two::find_raw`], except the distance between `start` and
573     /// `end` must be at least the size of a neon vector (in bytes).
574     ///
575     /// (The target feature safety obligation is automatically fulfilled by
576     /// virtue of being a method on `Two`, which can only be constructed
577     /// when it is safe to call `neon` routines.)
578     #[target_feature(enable = "neon")]
579     #[inline]
find_raw_impl( &self, start: *const u8, end: *const u8, ) -> Option<*const u8>580     unsafe fn find_raw_impl(
581         &self,
582         start: *const u8,
583         end: *const u8,
584     ) -> Option<*const u8> {
585         self.0.find_raw(start, end)
586     }
587 
588     /// Execute a search using neon vectors and routines.
589     ///
590     /// # Safety
591     ///
592     /// Same as [`Two::rfind_raw`], except the distance between `start` and
593     /// `end` must be at least the size of a neon vector (in bytes).
594     ///
595     /// (The target feature safety obligation is automatically fulfilled by
596     /// virtue of being a method on `Two`, which can only be constructed
597     /// when it is safe to call `neon` routines.)
598     #[target_feature(enable = "neon")]
599     #[inline]
rfind_raw_impl( &self, start: *const u8, end: *const u8, ) -> Option<*const u8>600     unsafe fn rfind_raw_impl(
601         &self,
602         start: *const u8,
603         end: *const u8,
604     ) -> Option<*const u8> {
605         self.0.rfind_raw(start, end)
606     }
607 
608     /// Returns an iterator over all occurrences of the needle bytes in the
609     /// given haystack.
610     ///
611     /// The iterator returned implements `DoubleEndedIterator`. This means it
612     /// can also be used to find occurrences in reverse order.
613     #[inline]
iter<'a, 'h>(&'a self, haystack: &'h [u8]) -> TwoIter<'a, 'h>614     pub fn iter<'a, 'h>(&'a self, haystack: &'h [u8]) -> TwoIter<'a, 'h> {
615         TwoIter { searcher: self, it: generic::Iter::new(haystack) }
616     }
617 }
618 
619 /// An iterator over all occurrences of two possible bytes in a haystack.
620 ///
621 /// This iterator implements `DoubleEndedIterator`, which means it can also be
622 /// used to find occurrences in reverse order.
623 ///
624 /// This iterator is created by the [`Two::iter`] method.
625 ///
626 /// The lifetime parameters are as follows:
627 ///
628 /// * `'a` refers to the lifetime of the underlying [`Two`] searcher.
629 /// * `'h` refers to the lifetime of the haystack being searched.
630 #[derive(Clone, Debug)]
631 pub struct TwoIter<'a, 'h> {
632     searcher: &'a Two,
633     it: generic::Iter<'h>,
634 }
635 
636 impl<'a, 'h> Iterator for TwoIter<'a, 'h> {
637     type Item = usize;
638 
639     #[inline]
next(&mut self) -> Option<usize>640     fn next(&mut self) -> Option<usize> {
641         // SAFETY: We rely on the generic iterator to provide valid start
642         // and end pointers, but we guarantee that any pointer returned by
643         // 'find_raw' falls within the bounds of the start and end pointer.
644         unsafe { self.it.next(|s, e| self.searcher.find_raw(s, e)) }
645     }
646 
647     #[inline]
size_hint(&self) -> (usize, Option<usize>)648     fn size_hint(&self) -> (usize, Option<usize>) {
649         self.it.size_hint()
650     }
651 }
652 
653 impl<'a, 'h> DoubleEndedIterator for TwoIter<'a, 'h> {
654     #[inline]
next_back(&mut self) -> Option<usize>655     fn next_back(&mut self) -> Option<usize> {
656         // SAFETY: We rely on the generic iterator to provide valid start
657         // and end pointers, but we guarantee that any pointer returned by
658         // 'rfind_raw' falls within the bounds of the start and end pointer.
659         unsafe { self.it.next_back(|s, e| self.searcher.rfind_raw(s, e)) }
660     }
661 }
662 
663 impl<'a, 'h> core::iter::FusedIterator for TwoIter<'a, 'h> {}
664 
665 /// Finds all occurrences of three bytes in a haystack.
666 ///
667 /// That is, this reports matches of one of three possible bytes. For example,
668 /// searching for `a`, `b` or `o` in `afoobar` would report matches at offsets
669 /// `0`, `2`, `3`, `4` and `5`.
670 #[derive(Clone, Copy, Debug)]
671 pub struct Three(generic::Three<uint8x16_t>);
672 
673 impl Three {
674     /// Create a new searcher that finds occurrences of the needle bytes given.
675     ///
676     /// This particular searcher is specialized to use neon vector instructions
677     /// that typically make it quite fast.
678     ///
679     /// If neon is unavailable in the current environment, then `None` is
680     /// returned.
681     #[inline]
new(needle1: u8, needle2: u8, needle3: u8) -> Option<Three>682     pub fn new(needle1: u8, needle2: u8, needle3: u8) -> Option<Three> {
683         if Three::is_available() {
684             // SAFETY: we check that neon is available above.
685             unsafe { Some(Three::new_unchecked(needle1, needle2, needle3)) }
686         } else {
687             None
688         }
689     }
690 
691     /// Create a new finder specific to neon vectors and routines without
692     /// checking that neon is available.
693     ///
694     /// # Safety
695     ///
696     /// Callers must guarantee that it is safe to execute `neon` instructions
697     /// in the current environment.
698     ///
699     /// Note that it is a common misconception that if one compiles for an
700     /// `x86_64` target, then they therefore automatically have access to neon
701     /// instructions. While this is almost always the case, it isn't true in
702     /// 100% of cases.
703     #[target_feature(enable = "neon")]
704     #[inline]
new_unchecked( needle1: u8, needle2: u8, needle3: u8, ) -> Three705     pub unsafe fn new_unchecked(
706         needle1: u8,
707         needle2: u8,
708         needle3: u8,
709     ) -> Three {
710         Three(generic::Three::new(needle1, needle2, needle3))
711     }
712 
713     /// Returns true when this implementation is available in the current
714     /// environment.
715     ///
716     /// When this is true, it is guaranteed that [`Three::new`] will return
717     /// a `Some` value. Similarly, when it is false, it is guaranteed that
718     /// `Three::new` will return a `None` value.
719     ///
720     /// Note also that for the lifetime of a single program, if this returns
721     /// true then it will always return true.
722     #[inline]
is_available() -> bool723     pub fn is_available() -> bool {
724         #[cfg(target_feature = "neon")]
725         {
726             true
727         }
728         #[cfg(not(target_feature = "neon"))]
729         {
730             false
731         }
732     }
733 
734     /// Return the first occurrence of one of the needle bytes in the given
735     /// haystack. If no such occurrence exists, then `None` is returned.
736     ///
737     /// The occurrence is reported as an offset into `haystack`. Its maximum
738     /// value is `haystack.len() - 1`.
739     #[inline]
find(&self, haystack: &[u8]) -> Option<usize>740     pub fn find(&self, haystack: &[u8]) -> Option<usize> {
741         // SAFETY: `find_raw` guarantees that if a pointer is returned, it
742         // falls within the bounds of the start and end pointers.
743         unsafe {
744             generic::search_slice_with_raw(haystack, |s, e| {
745                 self.find_raw(s, e)
746             })
747         }
748     }
749 
750     /// Return the last occurrence of one of the needle bytes in the given
751     /// haystack. If no such occurrence exists, then `None` is returned.
752     ///
753     /// The occurrence is reported as an offset into `haystack`. Its maximum
754     /// value is `haystack.len() - 1`.
755     #[inline]
rfind(&self, haystack: &[u8]) -> Option<usize>756     pub fn rfind(&self, haystack: &[u8]) -> Option<usize> {
757         // SAFETY: `rfind_raw` guarantees that if a pointer is returned, it
758         // falls within the bounds of the start and end pointers.
759         unsafe {
760             generic::search_slice_with_raw(haystack, |s, e| {
761                 self.rfind_raw(s, e)
762             })
763         }
764     }
765 
766     /// Like `find`, but accepts and returns raw pointers.
767     ///
768     /// When a match is found, the pointer returned is guaranteed to be
769     /// `>= start` and `< end`.
770     ///
771     /// This routine is useful if you're already using raw pointers and would
772     /// like to avoid converting back to a slice before executing a search.
773     ///
774     /// # Safety
775     ///
776     /// * Both `start` and `end` must be valid for reads.
777     /// * Both `start` and `end` must point to an initialized value.
778     /// * Both `start` and `end` must point to the same allocated object and
779     /// must either be in bounds or at most one byte past the end of the
780     /// allocated object.
781     /// * Both `start` and `end` must be _derived from_ a pointer to the same
782     /// object.
783     /// * The distance between `start` and `end` must not overflow `isize`.
784     /// * The distance being in bounds must not rely on "wrapping around" the
785     /// address space.
786     ///
787     /// Note that callers may pass a pair of pointers such that `start >= end`.
788     /// In that case, `None` will always be returned.
789     #[inline]
find_raw( &self, start: *const u8, end: *const u8, ) -> Option<*const u8>790     pub unsafe fn find_raw(
791         &self,
792         start: *const u8,
793         end: *const u8,
794     ) -> Option<*const u8> {
795         if start >= end {
796             return None;
797         }
798         if end.distance(start) < uint8x16_t::BYTES {
799             // SAFETY: We require the caller to pass valid start/end pointers.
800             return generic::fwd_byte_by_byte(start, end, |b| {
801                 b == self.0.needle1()
802                     || b == self.0.needle2()
803                     || b == self.0.needle3()
804             });
805         }
806         // SAFETY: Building a `Three` means it's safe to call 'neon' routines.
807         // Also, we've checked that our haystack is big enough to run on the
808         // vector routine. Pointer validity is caller's responsibility.
809         self.find_raw_impl(start, end)
810     }
811 
812     /// Like `rfind`, but accepts and returns raw pointers.
813     ///
814     /// When a match is found, the pointer returned is guaranteed to be
815     /// `>= start` and `< end`.
816     ///
817     /// This routine is useful if you're already using raw pointers and would
818     /// like to avoid converting back to a slice before executing a search.
819     ///
820     /// # Safety
821     ///
822     /// * Both `start` and `end` must be valid for reads.
823     /// * Both `start` and `end` must point to an initialized value.
824     /// * Both `start` and `end` must point to the same allocated object and
825     /// must either be in bounds or at most one byte past the end of the
826     /// allocated object.
827     /// * Both `start` and `end` must be _derived from_ a pointer to the same
828     /// object.
829     /// * The distance between `start` and `end` must not overflow `isize`.
830     /// * The distance being in bounds must not rely on "wrapping around" the
831     /// address space.
832     ///
833     /// Note that callers may pass a pair of pointers such that `start >= end`.
834     /// In that case, `None` will always be returned.
835     #[inline]
rfind_raw( &self, start: *const u8, end: *const u8, ) -> Option<*const u8>836     pub unsafe fn rfind_raw(
837         &self,
838         start: *const u8,
839         end: *const u8,
840     ) -> Option<*const u8> {
841         if start >= end {
842             return None;
843         }
844         if end.distance(start) < uint8x16_t::BYTES {
845             // SAFETY: We require the caller to pass valid start/end pointers.
846             return generic::rev_byte_by_byte(start, end, |b| {
847                 b == self.0.needle1()
848                     || b == self.0.needle2()
849                     || b == self.0.needle3()
850             });
851         }
852         // SAFETY: Building a `Three` means it's safe to call 'neon' routines.
853         // Also, we've checked that our haystack is big enough to run on the
854         // vector routine. Pointer validity is caller's responsibility.
855         self.rfind_raw_impl(start, end)
856     }
857 
858     /// Execute a search using neon vectors and routines.
859     ///
860     /// # Safety
861     ///
862     /// Same as [`Three::find_raw`], except the distance between `start` and
863     /// `end` must be at least the size of a neon vector (in bytes).
864     ///
865     /// (The target feature safety obligation is automatically fulfilled by
866     /// virtue of being a method on `Three`, which can only be constructed
867     /// when it is safe to call `neon` routines.)
868     #[target_feature(enable = "neon")]
869     #[inline]
find_raw_impl( &self, start: *const u8, end: *const u8, ) -> Option<*const u8>870     unsafe fn find_raw_impl(
871         &self,
872         start: *const u8,
873         end: *const u8,
874     ) -> Option<*const u8> {
875         self.0.find_raw(start, end)
876     }
877 
878     /// Execute a search using neon vectors and routines.
879     ///
880     /// # Safety
881     ///
882     /// Same as [`Three::rfind_raw`], except the distance between `start` and
883     /// `end` must be at least the size of a neon vector (in bytes).
884     ///
885     /// (The target feature safety obligation is automatically fulfilled by
886     /// virtue of being a method on `Three`, which can only be constructed
887     /// when it is safe to call `neon` routines.)
888     #[target_feature(enable = "neon")]
889     #[inline]
rfind_raw_impl( &self, start: *const u8, end: *const u8, ) -> Option<*const u8>890     unsafe fn rfind_raw_impl(
891         &self,
892         start: *const u8,
893         end: *const u8,
894     ) -> Option<*const u8> {
895         self.0.rfind_raw(start, end)
896     }
897 
898     /// Returns an iterator over all occurrences of the needle byte in the
899     /// given haystack.
900     ///
901     /// The iterator returned implements `DoubleEndedIterator`. This means it
902     /// can also be used to find occurrences in reverse order.
903     #[inline]
iter<'a, 'h>(&'a self, haystack: &'h [u8]) -> ThreeIter<'a, 'h>904     pub fn iter<'a, 'h>(&'a self, haystack: &'h [u8]) -> ThreeIter<'a, 'h> {
905         ThreeIter { searcher: self, it: generic::Iter::new(haystack) }
906     }
907 }
908 
909 /// An iterator over all occurrences of three possible bytes in a haystack.
910 ///
911 /// This iterator implements `DoubleEndedIterator`, which means it can also be
912 /// used to find occurrences in reverse order.
913 ///
914 /// This iterator is created by the [`Three::iter`] method.
915 ///
916 /// The lifetime parameters are as follows:
917 ///
918 /// * `'a` refers to the lifetime of the underlying [`Three`] searcher.
919 /// * `'h` refers to the lifetime of the haystack being searched.
920 #[derive(Clone, Debug)]
921 pub struct ThreeIter<'a, 'h> {
922     searcher: &'a Three,
923     it: generic::Iter<'h>,
924 }
925 
926 impl<'a, 'h> Iterator for ThreeIter<'a, 'h> {
927     type Item = usize;
928 
929     #[inline]
next(&mut self) -> Option<usize>930     fn next(&mut self) -> Option<usize> {
931         // SAFETY: We rely on the generic iterator to provide valid start
932         // and end pointers, but we guarantee that any pointer returned by
933         // 'find_raw' falls within the bounds of the start and end pointer.
934         unsafe { self.it.next(|s, e| self.searcher.find_raw(s, e)) }
935     }
936 
937     #[inline]
size_hint(&self) -> (usize, Option<usize>)938     fn size_hint(&self) -> (usize, Option<usize>) {
939         self.it.size_hint()
940     }
941 }
942 
943 impl<'a, 'h> DoubleEndedIterator for ThreeIter<'a, 'h> {
944     #[inline]
next_back(&mut self) -> Option<usize>945     fn next_back(&mut self) -> Option<usize> {
946         // SAFETY: We rely on the generic iterator to provide valid start
947         // and end pointers, but we guarantee that any pointer returned by
948         // 'rfind_raw' falls within the bounds of the start and end pointer.
949         unsafe { self.it.next_back(|s, e| self.searcher.rfind_raw(s, e)) }
950     }
951 }
952 
953 impl<'a, 'h> core::iter::FusedIterator for ThreeIter<'a, 'h> {}
954 
955 #[cfg(test)]
956 mod tests {
957     use super::*;
958 
959     define_memchr_quickcheck!(super);
960 
961     #[test]
forward_one()962     fn forward_one() {
963         crate::tests::memchr::Runner::new(1).forward_iter(
964             |haystack, needles| {
965                 Some(One::new(needles[0])?.iter(haystack).collect())
966             },
967         )
968     }
969 
970     #[test]
reverse_one()971     fn reverse_one() {
972         crate::tests::memchr::Runner::new(1).reverse_iter(
973             |haystack, needles| {
974                 Some(One::new(needles[0])?.iter(haystack).rev().collect())
975             },
976         )
977     }
978 
979     #[test]
count_one()980     fn count_one() {
981         crate::tests::memchr::Runner::new(1).count_iter(|haystack, needles| {
982             Some(One::new(needles[0])?.iter(haystack).count())
983         })
984     }
985 
986     #[test]
forward_two()987     fn forward_two() {
988         crate::tests::memchr::Runner::new(2).forward_iter(
989             |haystack, needles| {
990                 let n1 = needles.get(0).copied()?;
991                 let n2 = needles.get(1).copied()?;
992                 Some(Two::new(n1, n2)?.iter(haystack).collect())
993             },
994         )
995     }
996 
997     #[test]
reverse_two()998     fn reverse_two() {
999         crate::tests::memchr::Runner::new(2).reverse_iter(
1000             |haystack, needles| {
1001                 let n1 = needles.get(0).copied()?;
1002                 let n2 = needles.get(1).copied()?;
1003                 Some(Two::new(n1, n2)?.iter(haystack).rev().collect())
1004             },
1005         )
1006     }
1007 
1008     #[test]
forward_three()1009     fn forward_three() {
1010         crate::tests::memchr::Runner::new(3).forward_iter(
1011             |haystack, needles| {
1012                 let n1 = needles.get(0).copied()?;
1013                 let n2 = needles.get(1).copied()?;
1014                 let n3 = needles.get(2).copied()?;
1015                 Some(Three::new(n1, n2, n3)?.iter(haystack).collect())
1016             },
1017         )
1018     }
1019 
1020     #[test]
reverse_three()1021     fn reverse_three() {
1022         crate::tests::memchr::Runner::new(3).reverse_iter(
1023             |haystack, needles| {
1024                 let n1 = needles.get(0).copied()?;
1025                 let n2 = needles.get(1).copied()?;
1026                 let n3 = needles.get(2).copied()?;
1027                 Some(Three::new(n1, n2, n3)?.iter(haystack).rev().collect())
1028             },
1029         )
1030     }
1031 }
1032