Add mutable accessors to TypedArray and TypedMut

engine/src/lhs_types/array.rs

@@ -17,11 +17,13 @@ use std::{
     ops::Deref,
 };
 
+use super::{map::InnerMap, TypedMap};
+
 // Ideally, we would want to use Cow<'a, LhsValue<'a>> here
 // but it doesnt work for unknown reasons
 // See https://github.com/rust-lang/rust/issues/23707#issuecomment-557312736
 #[derive(Debug, Clone)]
-enum InnerArray<'a> {
+pub(crate) enum InnerArray<'a> {
     Owned(Vec<LhsValue<'a>>),
     Borrowed(&'a [LhsValue<'a>]),
 }
@@ -89,7 +91,7 @@ impl Default for InnerArray<'_> {
 #[derive(Debug, Clone)]
 pub struct Array<'a> {
     val_type: CompoundType,
-    data: InnerArray<'a>,
+    pub(crate) data: InnerArray<'a>,
 }
 
 impl<'a> Array<'a> {
@@ -486,6 +488,7 @@ impl<'a, 'b> From<&'a mut Array<'b>> for ArrayMut<'a, 'b> {
 /// Typed wrapper over an `Array` which provides
 /// infaillible operations.
 #[derive(Debug)]
+#[repr(transparent)]
 pub struct TypedArray<'a, V>
 where
     V: IntoValue<'a>,
@@ -594,6 +597,58 @@ impl<'a, V: IntoValue<'a>> IntoValue<'a> for TypedArray<'a, V> {
     }
 }
 
+impl<'a, V: IntoValue<'a>> TypedArray<'a, TypedArray<'a, V>> {
+    /// Returns a mutable reference to an element or None if the index is out of bounds.
+    pub fn get(&self, index: usize) -> Option<&TypedArray<'a, V>> {
+        self.array.get(index).map(|val| match val {
+            LhsValue::Array(array) => {
+                // Safety: this is safe because `TypedArray` is a repr(transparent)
+                // newtype over `InnerArray`.
+                unsafe { std::mem::transmute::<&InnerArray<'a>, &TypedArray<'a, V>>(&array.data) }
+            }
+            _ => unreachable!(),
+        })
+    }
+
+    /// Returns a mutable reference to an element or None if the index is out of bounds.
+    pub fn get_mut(&mut self, index: usize) -> Option<&mut TypedArray<'a, V>> {
+        self.array.get_mut(index).map(|val| match val {
+            LhsValue::Array(array) => {
+                // Safety: this is safe because `TypedArray` is a repr(transparent)
+                // newtype over `InnerArray`.
+                unsafe { std::mem::transmute::<&mut InnerArray<'a>, &mut TypedArray<'a, V>>(&mut array.data) }
+            }
+            _ => unreachable!(),
+        })
+    }
+}
+
+impl<'a, V: IntoValue<'a>> TypedArray<'a, TypedMap<'a, V>> {
+    /// Returns a mutable reference to an element or None if the index is out of bounds.
+    pub fn get(&self, index: usize) -> Option<&TypedMap<'a, V>> {
+        self.array.get(index).map(|val| match val {
+            LhsValue::Map(map) => {
+                // Safety: this is safe because `TypedMap` is a repr(transparent)
+                // newtype over `InnerMap`.
+                unsafe { std::mem::transmute::<&InnerMap<'a>, &TypedMap<'a, V>>(&map.data) }
+            }
+            _ => unreachable!(),
+        })
+    }
+
+    /// Returns a mutable reference to an element or None if the index is out of bounds.
+    pub fn get_mut(&mut self, index: usize) -> Option<&mut TypedMap<'a, V>> {
+        self.array.get_mut(index).map(|val| match val {
+            LhsValue::Map(map) => {
+                // Safety: this is safe because `TypedMap` is a repr(transparent)
+                // newtype over `InnerMap`.
+                unsafe { std::mem::transmute::<&mut InnerMap<'a>, &mut TypedMap<'a, V>>(&mut map.data) }
+            }
+            _ => unreachable!(),
+        })
+    }
+}
+
 #[test]
 fn test_size_of_array() {
     assert_eq!(std::mem::size_of::<Array<'_>>(), 32);

engine/src/lhs_types/map.rs

@@ -19,8 +19,10 @@ use std::{
     ops::Deref,
 };
 
+use super::{array::InnerArray, TypedArray};
+
 #[derive(Debug, Clone)]
-enum InnerMap<'a> {
+pub(crate) enum InnerMap<'a> {
     Owned(BTreeMap<Box<[u8]>, LhsValue<'a>>),
     Borrowed(&'a BTreeMap<Box<[u8]>, LhsValue<'a>>),
 }
@@ -78,7 +80,7 @@ impl Default for InnerMap<'_> {
 #[derive(Debug, Clone)]
 pub struct Map<'a> {
     val_type: CompoundType,
-    data: InnerMap<'a>,
+    pub(crate) data: InnerMap<'a>,
 }
 
 impl<'a> Map<'a> {
@@ -480,6 +482,7 @@ impl<'a, 'b> From<&'a mut Map<'b>> for MapMut<'a, 'b> {
 /// Typed wrapper over a `Map` which provides
 /// infaillible operations.
 #[derive(Debug)]
+#[repr(transparent)]
 pub struct TypedMap<'a, V>
 where
     V: IntoValue<'a>,
@@ -558,6 +561,58 @@ impl<'a, V: IntoValue<'a>> IntoValue<'a> for TypedMap<'a, V> {
     }
 }
 
+impl<'a, V: IntoValue<'a>> TypedMap<'a, TypedMap<'a, V>> {
+    /// Returns a mutable reference to the value corresponding to the key.
+    pub fn get(&self, key: &[u8]) -> Option<&TypedMap<'a, V>> {
+        self.map.get(key).map(|val| match val {
+            LhsValue::Map(map) => {
+                // Safety: this is safe because `TypedMap` is a repr(transparent)
+                // newtype over `InnerMap`.
+                unsafe { std::mem::transmute::<&InnerMap<'a>, &TypedMap<'a, V>>(&map.data) }
+            }
+            _ => unreachable!(),
+        })
+    }
+
+    /// Returns a mutable reference to the value corresponding to the key.
+    pub fn get_mut(&mut self, key: &[u8]) -> Option<&mut TypedMap<'a, V>> {
+        self.map.get_mut(key).map(|val| match val {
+            LhsValue::Map(map) => {
+                // Safety: this is safe because `TypedMap` is a repr(transparent)
+                // newtype over `InnerMap`.
+                unsafe { std::mem::transmute::<&mut InnerMap<'a>, &mut TypedMap<'a, V>>(&mut map.data) }
+            }
+            _ => unreachable!(),
+        })
+    }
+}
+
+impl<'a, V: IntoValue<'a>> TypedMap<'a, TypedArray<'a, V>> {
+    /// Returns a mutable reference to the value corresponding to the key.
+    pub fn get(&self, key: &[u8]) -> Option<&TypedArray<'a, V>> {
+        self.map.get(key).map(|val| match val {
+            LhsValue::Array(array) => {
+                // Safety: this is safe because `TypedArray` is a repr(transparent)
+                // newtype over `InnerArray`.
+                unsafe { std::mem::transmute::<&InnerArray<'a>, &TypedArray<'a, V>>(&array.data) }
+            }
+            _ => unreachable!(),
+        })
+    }
+
+    /// Returns a mutable reference to the value corresponding to the key.
+    pub fn get_mut(&mut self, key: &[u8]) -> Option<&mut TypedArray<'a, V>> {
+        self.map.get_mut(key).map(|val| match val {
+            LhsValue::Array(array) => {
+                // Safety: this is safe because `TypedArray` is a repr(transparent)
+                // newtype over `InnerArray`.
+                unsafe { std::mem::transmute::<&mut InnerArray<'a>, &mut TypedArray<'a, V>>(&mut array.data) }
+            }
+            _ => unreachable!(),
+        })
+    }
+}
+
 #[test]
 fn test_size_of_map() {
     assert_eq!(std::mem::size_of::<Map<'_>>(), 40);

engine/src/lib.rs

@@ -102,7 +102,7 @@ pub use self::{
         SimpleFunctionParam,
     },
     lex::LexErrorKind,
-    lhs_types::{Array, ArrayMut, Map, MapIter, MapMut, TypedArray},
+    lhs_types::{Array, ArrayMut, Map, MapIter, MapMut, TypedArray, TypedMap},
     list_matcher::{
         AlwaysList, AlwaysListMatcher, ListDefinition, ListMatcher, NeverList, NeverListMatcher,
     },

engine/src/types.rs

@@ -453,8 +453,7 @@ impl<'a> BytesOrString<'a> {
 
 mod private {
     use super::IntoValue;
-    use crate::lhs_types::TypedMap;
-    use crate::TypedArray;
+    use crate::{TypedArray, TypedMap};
     use std::borrow::Cow;
     use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};