Implement SliceIndex for ByteStr

Author: thaliaarchi

library/core/src/bstr/mod.rs

@@ -92,6 +92,13 @@ impl ByteStr {
     pub fn as_bytes(&self) -> &[u8] {
         &self.0
     }
+
+    #[doc(hidden)]
+    #[unstable(feature = "bstr_internals", issue = "none")]
+    #[inline]
+    pub fn as_bytes_mut(&mut self) -> &mut [u8] {
+        &mut self.0
+    }
 }
 
 #[unstable(feature = "bstr", issue = "134915")]

library/core/src/bstr/traits.rs

@@ -2,10 +2,8 @@
 
 use crate::bstr::ByteStr;
 use crate::cmp::Ordering;
-use crate::hash;
-use crate::ops::{
-    Index, IndexMut, Range, RangeFrom, RangeFull, RangeInclusive, RangeTo, RangeToInclusive,
-};
+use crate::slice::SliceIndex;
+use crate::{hash, ops, range};
 
 #[unstable(feature = "bstr", issue = "134915")]
 impl Ord for ByteStr {
@@ -149,127 +147,131 @@ impl_partial_eq_n!(ByteStr, [u8; N]);
 impl_partial_eq_n!(ByteStr, &[u8; N]);
 
 #[unstable(feature = "bstr", issue = "134915")]
-impl Index<usize> for ByteStr {
-    type Output = u8;
+impl<I> ops::Index<I> for ByteStr
+where
+    I: SliceIndex<ByteStr>,
+{
+    type Output = I::Output;
 
     #[inline]
-    fn index(&self, idx: usize) -> &u8 {
-        &self.0[idx]
+    fn index(&self, index: I) -> &I::Output {
+        index.index(self)
     }
 }
 
 #[unstable(feature = "bstr", issue = "134915")]
-impl Index<RangeFull> for ByteStr {
-    type Output = ByteStr;
-
+impl<I> ops::IndexMut<I> for ByteStr
+where
+    I: SliceIndex<ByteStr>,
+{
     #[inline]
-    fn index(&self, _: RangeFull) -> &ByteStr {
-        self
+    fn index_mut(&mut self, index: I) -> &mut I::Output {
+        index.index_mut(self)
     }
 }
 
 #[unstable(feature = "bstr", issue = "134915")]
-impl Index<Range<usize>> for ByteStr {
+unsafe impl SliceIndex<ByteStr> for ops::RangeFull {
     type Output = ByteStr;
-
     #[inline]
-    fn index(&self, r: Range<usize>) -> &ByteStr {
-        ByteStr::from_bytes(&self.0[r])
+    fn get(self, slice: &ByteStr) -> Option<&Self::Output> {
+        Some(slice)
     }
-}
-
-#[unstable(feature = "bstr", issue = "134915")]
-impl Index<RangeInclusive<usize>> for ByteStr {
-    type Output = ByteStr;
-
     #[inline]
-    fn index(&self, r: RangeInclusive<usize>) -> &ByteStr {
-        ByteStr::from_bytes(&self.0[r])
+    fn get_mut(self, slice: &mut ByteStr) -> Option<&mut Self::Output> {
+        Some(slice)
     }
-}
-
-#[unstable(feature = "bstr", issue = "134915")]
-impl Index<RangeFrom<usize>> for ByteStr {
-    type Output = ByteStr;
-
     #[inline]
-    fn index(&self, r: RangeFrom<usize>) -> &ByteStr {
-        ByteStr::from_bytes(&self.0[r])
+    unsafe fn get_unchecked(self, slice: *const ByteStr) -> *const Self::Output {
+        slice
     }
-}
-
-#[unstable(feature = "bstr", issue = "134915")]
-impl Index<RangeTo<usize>> for ByteStr {
-    type Output = ByteStr;
-
     #[inline]
-    fn index(&self, r: RangeTo<usize>) -> &ByteStr {
-        ByteStr::from_bytes(&self.0[r])
+    unsafe fn get_unchecked_mut(self, slice: *mut ByteStr) -> *mut Self::Output {
+        slice
     }
-}
-
-#[unstable(feature = "bstr", issue = "134915")]
-impl Index<RangeToInclusive<usize>> for ByteStr {
-    type Output = ByteStr;
-
     #[inline]
-    fn index(&self, r: RangeToInclusive<usize>) -> &ByteStr {
-        ByteStr::from_bytes(&self.0[r])
+    fn index(self, slice: &ByteStr) -> &Self::Output {
+        slice
     }
-}
-
-#[unstable(feature = "bstr", issue = "134915")]
-impl IndexMut<usize> for ByteStr {
     #[inline]
-    fn index_mut(&mut self, idx: usize) -> &mut u8 {
-        &mut self.0[idx]
+    fn index_mut(self, slice: &mut ByteStr) -> &mut Self::Output {
+        slice
     }
 }
 
 #[unstable(feature = "bstr", issue = "134915")]
-impl IndexMut<RangeFull> for ByteStr {
+unsafe impl SliceIndex<ByteStr> for usize {
+    type Output = u8;
     #[inline]
-    fn index_mut(&mut self, _: RangeFull) -> &mut ByteStr {
-        self
+    fn get(self, slice: &ByteStr) -> Option<&Self::Output> {
+        self.get(slice.as_bytes())
     }
-}
-
-#[unstable(feature = "bstr", issue = "134915")]
-impl IndexMut<Range<usize>> for ByteStr {
     #[inline]
-    fn index_mut(&mut self, r: Range<usize>) -> &mut ByteStr {
-        ByteStr::from_bytes_mut(&mut self.0[r])
+    fn get_mut(self, slice: &mut ByteStr) -> Option<&mut Self::Output> {
+        self.get_mut(slice.as_bytes_mut())
     }
-}
-
-#[unstable(feature = "bstr", issue = "134915")]
-impl IndexMut<RangeInclusive<usize>> for ByteStr {
     #[inline]
-    fn index_mut(&mut self, r: RangeInclusive<usize>) -> &mut ByteStr {
-        ByteStr::from_bytes_mut(&mut self.0[r])
+    unsafe fn get_unchecked(self, slice: *const ByteStr) -> *const Self::Output {
+        // SAFETY: the caller has to uphold the safety contract for `get_unchecked`.
+        unsafe { self.get_unchecked(slice as *const [u8]) }
     }
-}
-
-#[unstable(feature = "bstr", issue = "134915")]
-impl IndexMut<RangeFrom<usize>> for ByteStr {
     #[inline]
-    fn index_mut(&mut self, r: RangeFrom<usize>) -> &mut ByteStr {
-        ByteStr::from_bytes_mut(&mut self.0[r])
+    unsafe fn get_unchecked_mut(self, slice: *mut ByteStr) -> *mut Self::Output {
+        // SAFETY: the caller has to uphold the safety contract for `get_unchecked_mut`.
+        unsafe { self.get_unchecked_mut(slice as *mut [u8]) }
     }
-}
-
-#[unstable(feature = "bstr", issue = "134915")]
-impl IndexMut<RangeTo<usize>> for ByteStr {
     #[inline]
-    fn index_mut(&mut self, r: RangeTo<usize>) -> &mut ByteStr {
-        ByteStr::from_bytes_mut(&mut self.0[r])
+    fn index(self, slice: &ByteStr) -> &Self::Output {
+        self.index(slice.as_bytes())
     }
-}
-
-#[unstable(feature = "bstr", issue = "134915")]
-impl IndexMut<RangeToInclusive<usize>> for ByteStr {
     #[inline]
-    fn index_mut(&mut self, r: RangeToInclusive<usize>) -> &mut ByteStr {
-        ByteStr::from_bytes_mut(&mut self.0[r])
+    fn index_mut(self, slice: &mut ByteStr) -> &mut Self::Output {
+        self.index_mut(slice.as_bytes_mut())
     }
 }
+
+macro_rules! impl_slice_index {
+    ($index:ty) => {
+        #[unstable(feature = "bstr", issue = "134915")]
+        unsafe impl SliceIndex<ByteStr> for $index {
+            type Output = ByteStr;
+            #[inline]
+            fn get(self, slice: &ByteStr) -> Option<&Self::Output> {
+                self.get(slice.as_bytes()).map(ByteStr::from_bytes)
+            }
+            #[inline]
+            fn get_mut(self, slice: &mut ByteStr) -> Option<&mut Self::Output> {
+                self.get_mut(slice.as_bytes_mut()).map(ByteStr::from_bytes_mut)
+            }
+            #[inline]
+            unsafe fn get_unchecked(self, slice: *const ByteStr) -> *const Self::Output {
+                // SAFETY: the caller has to uphold the safety contract for `get_unchecked`.
+                unsafe { self.get_unchecked(slice as *const [u8]) as *const ByteStr }
+            }
+            #[inline]
+            unsafe fn get_unchecked_mut(self, slice: *mut ByteStr) -> *mut Self::Output {
+                // SAFETY: the caller has to uphold the safety contract for `get_unchecked_mut`.
+                unsafe { self.get_unchecked_mut(slice as *mut [u8]) as *mut ByteStr }
+            }
+            #[inline]
+            fn index(self, slice: &ByteStr) -> &Self::Output {
+                ByteStr::from_bytes(self.index(slice.as_bytes()))
+            }
+            #[inline]
+            fn index_mut(self, slice: &mut ByteStr) -> &mut Self::Output {
+                ByteStr::from_bytes_mut(self.index_mut(slice.as_bytes_mut()))
+            }
+        }
+    };
+}
+
+impl_slice_index!(ops::IndexRange);
+impl_slice_index!(ops::Range<usize>);
+impl_slice_index!(range::Range<usize>);
+impl_slice_index!(ops::RangeTo<usize>);
+impl_slice_index!(ops::RangeFrom<usize>);
+impl_slice_index!(range::RangeFrom<usize>);
+impl_slice_index!(ops::RangeInclusive<usize>);
+impl_slice_index!(range::RangeInclusive<usize>);
+impl_slice_index!(ops::RangeToInclusive<usize>);
+impl_slice_index!((ops::Bound<usize>, ops::Bound<usize>));