Add more methods on Array and Dictionary (#212)

crates/types/src/array.rs

@@ -1,3 +1,5 @@
+use core::ops::{Deref, DerefMut};
+
 use luajit as lua;
 
 use crate::kvec::{self, KVec};
@@ -56,6 +58,35 @@ impl Array {
     {
         self.0.push(value.into());
     }
+
+    /// Removes an `Object` from the `Array` and returns it.
+    ///
+    /// The removed object is replaced by the last element of the array.
+    ///
+    /// # Panics
+    ///
+    /// Panics if `index` is out of bounds.
+    #[track_caller]
+    #[inline]
+    pub fn swap_remove(&mut self, index: usize) -> Object {
+        self.0.swap_remove(index)
+    }
+}
+
+impl Deref for Array {
+    type Target = [Object];
+
+    #[inline]
+    fn deref(&self) -> &Self::Target {
+        self.0.as_slice()
+    }
+}
+
+impl DerefMut for Array {
+    #[inline]
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        self.0.as_mut_slice()
+    }
 }
 
 impl<T: Into<Object>> FromIterator<T> for Array {

crates/types/src/dictionary.rs

@@ -10,12 +10,12 @@ use crate::Object;
 #[repr(transparent)]
 pub struct Dictionary(pub(super) KVec<KeyValuePair>);
 
-/// A key-value pair mapping a Neovim [`String`] to a Neovim [`Object`].
+/// A key-value pair mapping a [`String`] to an [`Object`].
 //
 // https://github.com/neovim/neovim/blob/v0.9.0/src/nvim/api/private/defs.h#L122-L125
 #[derive(Clone, PartialEq)]
 #[repr(C)]
-pub(super) struct KeyValuePair {
+pub struct KeyValuePair {
     key: crate::String,
     value: Object,
 }
@@ -42,13 +42,36 @@ impl core::fmt::Debug for Dictionary {
 }
 
 impl Dictionary {
+    /// Returns a slice of all key-value pairs in the dictionary.
+    #[inline]
+    pub fn as_slice(&self) -> &[KeyValuePair] {
+        &self.0
+    }
+
+    /// Returns a mutable slice of all key-value pairs in the dictionary.
+    #[inline]
+    pub fn as_mut_slice(&mut self) -> &mut [KeyValuePair] {
+        &mut self.0
+    }
+
     /// Returns a reference to the value corresponding to the key.
     #[inline]
     pub fn get<Q>(&self, query: &Q) -> Option<&Object>
     where
         Q: ?Sized + PartialEq<crate::String>,
     {
-        self.iter().find_map(|(key, value)| (query == key).then_some(value))
+        self.get_index(query).map(|idx| self.as_slice()[idx].value())
+    }
+
+    /// Returns the index of the key-value pair corresponding to the key.
+    #[inline]
+    pub fn get_index<Q>(&self, query: &Q) -> Option<usize>
+    where
+        Q: ?Sized + PartialEq<crate::String>,
+    {
+        self.keys()
+            .enumerate()
+            .find_map(|(idx, key)| (query == key).then_some(idx))
     }
 
     /// Returns a mutable reference to the value corresponding to the key.
@@ -57,8 +80,7 @@ impl Dictionary {
     where
         Q: ?Sized + PartialEq<crate::String>,
     {
-        self.iter_mut()
-            .find_map(|(key, value)| (query == key).then_some(value))
+        self.get_index(query).map(|idx| self.as_mut_slice()[idx].value_mut())
     }
 
     /// Inserts a key-value pair into the dictionary.
@@ -98,6 +120,12 @@ impl Dictionary {
         self.0.len()
     }
 
+    /// Returns an iterator over the keys of the dictionary.
+    #[inline]
+    pub fn keys(&self) -> impl Iterator<Item = &crate::String> + '_ {
+        self.iter().map(|(key, _)| key)
+    }
+
     /// Creates a new, empty `Dictionary`.
     #[inline]
     pub fn new() -> Self {
@@ -110,6 +138,75 @@ impl Dictionary {
         #[allow(clippy::unnecessary_struct_initialization)]
         NonOwning::new(Self(KVec { ..self.0 }))
     }
+
+    /// Removes a `KeyValuePair` from the `Dictionary` and returns it.
+    ///
+    /// The removed pair is replaced by the last element of the dictionary.
+    ///
+    /// # Panics
+    ///
+    /// Panics if `index` is out of bounds.
+    #[track_caller]
+    #[inline]
+    pub fn swap_remove(&mut self, index: usize) -> KeyValuePair {
+        self.0.swap_remove(index)
+    }
+}
+
+impl KeyValuePair {
+    /// Consumes the `KeyValuePair` and returns the key.
+    #[inline]
+    pub fn into_key(self) -> crate::String {
+        self.key
+    }
+
+    /// Consumes the `KeyValuePair` and returns a `(key, value)` tuple.
+    #[inline]
+    pub fn into_tuple(self) -> (crate::String, Object) {
+        (self.key, self.value)
+    }
+
+    /// Consumes the `KeyValuePair` and returns the value.
+    #[inline]
+    pub fn into_value(self) -> Object {
+        self.value
+    }
+
+    /// Returns a shared reference to the key of the `KeyValuePair`.
+    #[inline]
+    pub fn key(&self) -> &crate::String {
+        &self.key
+    }
+
+    /// Returns an exclusive reference to the key of the `KeyValuePair`.
+    #[inline]
+    pub fn key_mut(&mut self) -> &mut crate::String {
+        &mut self.key
+    }
+
+    /// Returns references to both the key and value as a tuple.
+    #[inline]
+    pub fn tuple(&self) -> (&crate::String, &Object) {
+        (&self.key, &self.value)
+    }
+
+    /// Returns exclusive references to both the key and value as a tuple.
+    #[inline]
+    pub fn tuple_mut(&mut self) -> (&mut crate::String, &mut Object) {
+        (&mut self.key, &mut self.value)
+    }
+
+    /// Returns a shared reference to the value of the `KeyValuePair`.
+    #[inline]
+    pub fn value(&self) -> &Object {
+        &self.value
+    }
+
+    /// Returns an exclusive reference to the value of the `KeyValuePair`.
+    #[inline]
+    pub fn value_mut(&mut self) -> &mut Object {
+        &mut self.value
+    }
 }
 
 impl<S> core::ops::Index<S> for Dictionary
@@ -173,7 +270,7 @@ impl Iterator for DictIterator {
 
     #[inline]
     fn next(&mut self) -> Option<Self::Item> {
-        self.0.next().map(|pair| (pair.key, pair.value))
+        self.0.next().map(KeyValuePair::into_tuple)
     }
 
     #[inline]
@@ -192,7 +289,7 @@ impl ExactSizeIterator for DictIterator {
 impl DoubleEndedIterator for DictIterator {
     #[inline]
     fn next_back(&mut self) -> Option<Self::Item> {
-        self.0.next_back().map(|pair| (pair.key, pair.value))
+        self.0.next_back().map(KeyValuePair::into_tuple)
     }
 }
 
@@ -207,7 +304,7 @@ impl<'a> Iterator for DictIter<'a> {
 
     #[inline]
     fn next(&mut self) -> Option<Self::Item> {
-        self.0.next().map(|pair| (&pair.key, &pair.value))
+        self.0.next().map(KeyValuePair::tuple)
     }
 
     #[inline]
@@ -226,7 +323,7 @@ impl ExactSizeIterator for DictIter<'_> {
 impl DoubleEndedIterator for DictIter<'_> {
     #[inline]
     fn next_back(&mut self) -> Option<Self::Item> {
-        self.0.next_back().map(|pair| (&pair.key, &pair.value))
+        self.0.next_back().map(KeyValuePair::tuple)
     }
 }
 
@@ -240,7 +337,7 @@ impl<'a> Iterator for DictIterMut<'a> {
 
     #[inline]
     fn next(&mut self) -> Option<Self::Item> {
-        self.0.next().map(|pair| (&mut pair.key, &mut pair.value))
+        self.0.next().map(KeyValuePair::tuple_mut)
     }
 
     #[inline]
@@ -259,7 +356,7 @@ impl ExactSizeIterator for DictIterMut<'_> {
 impl DoubleEndedIterator for DictIterMut<'_> {
     #[inline]
     fn next_back(&mut self) -> Option<Self::Item> {
-        self.0.next_back().map(|pair| (&mut pair.key, &mut pair.value))
+        self.0.next_back().map(KeyValuePair::tuple_mut)
     }
 }
 

crates/types/src/kvec.rs

@@ -115,6 +115,28 @@ impl<T> KVec<T> {
         self.size += 1;
     }
 
+    /// Removes an element from the `KVec` and returns it.
+    ///
+    /// The removed element is replaced by the last element of the vector.
+    ///
+    /// # Panics
+    ///
+    /// Panics if `index` is out of bounds.
+    #[track_caller]
+    #[inline]
+    pub(crate) fn swap_remove(&mut self, index: usize) -> T {
+        let len = self.len();
+        if index >= len {
+            panic!("swap_remove index is {index}, but len is {len}");
+        }
+        unsafe {
+            let item = ptr::read(self.items.add(index));
+            ptr::copy(self.items.add(len - 1), self.items.add(index), 1);
+            self.size -= 1;
+            item
+        }
+    }
+
     /// Creates a new, empty `KVec<T>` with the specified capacity.
     #[inline]
     pub(crate) fn with_capacity(capacity: usize) -> Self {
@@ -150,6 +172,34 @@ impl<T: PartialEq> PartialEq for KVec<T> {
     }
 }
 
+impl<T: PartialEq> PartialEq<[T]> for KVec<T> {
+    #[inline]
+    fn eq(&self, other: &[T]) -> bool {
+        self.as_slice() == other
+    }
+}
+
+impl<T: PartialEq> PartialEq<&[T]> for KVec<T> {
+    #[inline]
+    fn eq(&self, other: &&[T]) -> bool {
+        self.as_slice() == *other
+    }
+}
+
+impl<const N: usize, T: PartialEq> PartialEq<[T; N]> for KVec<T> {
+    #[inline]
+    fn eq(&self, other: &[T; N]) -> bool {
+        self == other.as_slice()
+    }
+}
+
+impl<const N: usize, T: PartialEq> PartialEq<&[T; N]> for KVec<T> {
+    #[inline]
+    fn eq(&self, other: &&[T; N]) -> bool {
+        self == *other
+    }
+}
+
 impl<T> Deref for KVec<T> {
     type Target = [T];
 
@@ -427,4 +477,18 @@ mod tests {
         assert_eq!(Some("baz"), iter2.next().as_deref());
         assert_eq!(None, iter2.next());
     }
+
+    #[test]
+    fn swap_remove() {
+        let mut kvec = KVec::from_iter([1, 2, 3, 4]);
+        assert_eq!(kvec.swap_remove(1), 2);
+        assert_eq!(kvec, &[1, 4, 3]);
+    }
+
+    #[should_panic]
+    #[test]
+    fn swap_remove_oob() {
+        let mut kvec = KVec::from_iter([1, 2, 3, 4]);
+        kvec.swap_remove(kvec.len());
+    }
 }

crates/types/src/lib.rs

@@ -18,7 +18,7 @@ mod string;
 
 pub use arena::{arena, arena_init, Arena};
 pub use array::Array;
-pub use dictionary::Dictionary;
+pub use dictionary::{Dictionary, KeyValuePair};
 pub use error::Error;
 pub use function::Function;
 pub use non_owning::NonOwning;