From 1dfaa0e4a0a0eddc8df5d0d1bfc7f564fb3dfa21 Mon Sep 17 00:00:00 2001 From: Sayan Nandan Date: Sun, 22 Aug 2021 21:41:46 -0700 Subject: [PATCH] Switch to using `next_unchecked` --- server/src/actions/flushdb.rs | 2 +- server/src/actions/lskeys.rs | 2 +- server/src/actions/update.rs | 4 ++-- server/src/protocol/iter.rs | 19 ++++++++++++++++++- server/src/queryengine/mod.rs | 8 +++----- 5 files changed, 25 insertions(+), 10 deletions(-) diff --git a/server/src/actions/flushdb.rs b/server/src/actions/flushdb.rs index c3cb66b0..1eedab89 100644 --- a/server/src/actions/flushdb.rs +++ b/server/src/actions/flushdb.rs @@ -37,7 +37,7 @@ action!( get_tbl!(handle, con).truncate_table(); } else { // flush the entity - let raw_entity = unsafe { act.next().unsafe_unwrap() }; + let raw_entity = unsafe { act.next_unchecked() }; let entity = handle_entity!(con, raw_entity); get_tbl!(entity, handle, con).truncate_table(); } diff --git a/server/src/actions/lskeys.rs b/server/src/actions/lskeys.rs index 3cd0b183..b14daa4d 100644 --- a/server/src/actions/lskeys.rs +++ b/server/src/actions/lskeys.rs @@ -39,7 +39,7 @@ action!( (get_tbl!(handle, con), DEFAULT_COUNT) } else if act.len() == 1 { // two args, could either be count or an entity - let nextret = unsafe { act.next().unsafe_unwrap() }; + let nextret = unsafe { act.next_unchecked() }; if unsafe { nextret.get_unchecked(0) }.is_ascii_digit() { // noice, this is a number; let's try to parse it let count = if let Ok(cnt) = String::from_utf8_lossy(nextret).parse::() { diff --git a/server/src/actions/update.rs b/server/src/actions/update.rs index 6f882510..a934b556 100644 --- a/server/src/actions/update.rs +++ b/server/src/actions/update.rs @@ -43,8 +43,8 @@ action!( // UNSAFE(@ohsayan): This is completely safe as we've already checked // that there are exactly 2 arguments writer.update( - Data::copy_from_slice(act.next().unsafe_unwrap()), - Data::copy_from_slice(act.next().unsafe_unwrap()), + Data::copy_from_slice(act.next_unchecked()), + Data::copy_from_slice(act.next_unchecked()), ) } { Ok(true) => Some(true), diff --git a/server/src/protocol/iter.rs b/server/src/protocol/iter.rs index 86fe8a41..39fdb6df 100644 --- a/server/src/protocol/iter.rs +++ b/server/src/protocol/iter.rs @@ -34,10 +34,16 @@ use core::ops::Deref; use core::slice::ChunksExact; use core::slice::Iter; +/// An iterator over an [`AnyArray`] (an [`UnsafeSlice`]). The validity of the iterator is +/// left to the caller who has to guarantee: +/// - Source pointers for the unsafe slice are valid +/// - Source pointers exist as long as this iterator is used pub struct AnyArrayIter<'a> { iter: Iter<'a, UnsafeSlice>, } +/// Same as [`AnyArrayIter`] with the exception that it directly dereferences to the actual +/// slice iterator pub struct BorrowedAnyArrayIter<'a> { iter: Iter<'a, UnsafeSlice>, } @@ -50,29 +56,40 @@ impl<'a> Deref for BorrowedAnyArrayIter<'a> { } impl<'a> AnyArrayIter<'a> { + /// Create a new `AnyArrayIter`. + /// + /// ## Safety + /// - Valid source pointers + /// - Source pointers exist as long as the iterator is used pub const unsafe fn new(iter: Iter<'a, UnsafeSlice>) -> AnyArrayIter<'a> { Self { iter } } + /// Returns a [`ChunksExact`] (similar to [`ChunksExact` provided by core::slice](core::slice::ChunksExact)) pub fn chunks_exact(&'a self, chunks_exact: usize) -> ChunksExact<'a, UnsafeSlice> { self.iter.as_ref().chunks_exact(chunks_exact) } + /// Returns a borrowed iterator => simply put, advancing the returned iterator does not + /// affect the base iterator owned by this object pub fn as_ref(&'a self) -> BorrowedAnyArrayIter<'a> { BorrowedAnyArrayIter { iter: self.iter.as_ref().iter(), } } + /// Returns the next value in uppercase pub fn next_uppercase(&mut self) -> Option> { self.iter.next().map(|v| unsafe { // SAFETY: Only construction is unsafe, forwarding is not v.as_slice().to_ascii_uppercase().into_boxed_slice() }) } + /// Returns the next value without any checks pub unsafe fn next_unchecked(&mut self) -> &'a [u8] { match self.next() { Some(s) => s, None => unreachable_unchecked(), } } + /// Returns the next value without any checks as an owned copy of [`Bytes`] pub unsafe fn next_unchecked_bytes(&mut self) -> Bytes { Bytes::copy_from_slice(self.next_unchecked()) } @@ -130,7 +147,7 @@ fn test_iter() { } }; let it = arr.iter(); - let mut iter =unsafe { AnyArrayIter::new(it)}; + let mut iter = unsafe { AnyArrayIter::new(it) }; assert_eq!(iter.next_uppercase().unwrap().as_ref(), "SET".as_bytes()); assert_eq!(iter.next().unwrap(), "x".as_bytes()); assert_eq!(iter.next().unwrap(), "100".as_bytes()); diff --git a/server/src/queryengine/mod.rs b/server/src/queryengine/mod.rs index 9e082144..f71213f3 100644 --- a/server/src/queryengine/mod.rs +++ b/server/src/queryengine/mod.rs @@ -33,7 +33,6 @@ use crate::protocol::element::UnsafeElement; use crate::protocol::iter::AnyArrayIter; use crate::protocol::responses; use crate::protocol::SimpleQuery; -use crate::protocol::UnsafeSlice; use crate::{actions, admin}; use core::hint::unreachable_unchecked; mod ddl; @@ -53,11 +52,10 @@ macro_rules! gen_constants_and_matches { pub const $action: &[u8] = stringify!($action).as_bytes(); )* } - let mut first = match $buf.next_uppercase() { + let first = match $buf.next_uppercase() { Some(frst) => frst, None => return $con.write_response(responses::groups::PACKET_ERR).await, }; - first.make_ascii_uppercase(); match first.as_ref() { $( tags::$action => $fns($db, $con, $buf).await?, @@ -105,7 +103,7 @@ where if !buf.is_any_array() { return con.write_response(responses::groups::WRONGTYPE_ERR).await; } - let bufref: Box<[UnsafeSlice]>; + let bufref; let _rawiter; let mut iter; unsafe { @@ -165,7 +163,7 @@ action! { err_if_len_is!(act, con, not 1); let entity = unsafe { // SAFETY: Already checked len - act.next().unsafe_unwrap() + act.next_unchecked() }; swap_entity!(con, handle, entity); Ok(())