Add data_layout_element_at()

2 files changed, 102 insertions, 1 deletions

diff --git a/patio-to-cranelift/src/lib.rs b/patio-to-cranelift/src/lib.rs
index f5108617..687687f0 100644
--- a/patio-to-cranelift/src/lib.rs
+++ b/patio-to-cranelift/src/lib.rs
@@ -817,3 +817,104 @@ mod data_layout_byte_offset_test {
         assert_eq!(byte_offset, 13);
     }
 }
+
+fn data_layout_element_at<'a>(
+    layout: &'a PatioDataLayout,
+    element_offsets: &[u32],
+) -> &'a PatioDataLayout {
+    let mut current_element = layout;
+
+    for element_offset in element_offsets {
+        match current_element {
+            PatioDataLayout::Array { element, size: _ } => {
+                current_element = element;
+            }
+
+            PatioDataLayout::Bool => panic!("offset not allowed for bool"),
+            PatioDataLayout::Byte => panic!("offset not allowed for byte"),
+            PatioDataLayout::DataPtr => panic!("offset not allowed for data ptr"),
+            PatioDataLayout::FunctionPtr => panic!("offset not allowed for function ptr"),
+            PatioDataLayout::Int => panic!("offset not allowed for int"),
+            PatioDataLayout::Int32 => panic!("offset not allowed for int32"),
+
+            PatioDataLayout::Struct { elements } => {
+                current_element = &elements[usize::try_from(*element_offset).unwrap()];
+            }
+        }
+    }
+
+    current_element
+}
+
+#[cfg(test)]
+mod data_layout_element_at_test {
+    use crate::patio::types::PatioDataLayout;
+
+    fn data_layout_element_at(
+        layout: &PatioDataLayout,
+        element_offsets: Vec<u32>,
+    ) -> &PatioDataLayout {
+        super::data_layout_element_at(layout, &element_offsets)
+    }
+
+    #[test]
+    fn element_of_array_is_element_layout() {
+        let layout = PatioDataLayout::array(PatioDataLayout::Int32, 2);
+
+        let element = data_layout_element_at(&layout, vec![0]);
+
+        assert_eq!(element, &PatioDataLayout::Int32);
+    }
+
+    #[test]
+    fn element_of_element_of_array() {
+        let layout = PatioDataLayout::array(PatioDataLayout::array(PatioDataLayout::Byte, 7), 5);
+
+        let element = data_layout_element_at(&layout, vec![2, 3]);
+
+        assert_eq!(element, &PatioDataLayout::Byte);
+    }
+
+    #[test]
+    fn first_element_of_struct() {
+        let layout = PatioDataLayout::struct_(vec![
+            PatioDataLayout::Int32,
+            PatioDataLayout::Byte,
+            PatioDataLayout::DataPtr,
+        ]);
+
+        let element = data_layout_element_at(&layout, vec![0]);
+
+        assert_eq!(element, &PatioDataLayout::Int32);
+    }
+
+    #[test]
+    fn non_first_element_of_struct() {
+        let layout = PatioDataLayout::struct_(vec![
+            PatioDataLayout::Int32,
+            PatioDataLayout::Byte,
+            PatioDataLayout::DataPtr,
+        ]);
+
+        let element = data_layout_element_at(&layout, vec![2]);
+
+        assert_eq!(element, &PatioDataLayout::DataPtr);
+    }
+
+    #[test]
+    fn element_of_element_of_struct() {
+        let layout = PatioDataLayout::struct_(vec![
+            PatioDataLayout::Int32,
+            PatioDataLayout::struct_(vec![
+                PatioDataLayout::Byte,
+                PatioDataLayout::DataPtr,
+                PatioDataLayout::Byte,
+            ]),
+            PatioDataLayout::Int32,
+        ]);
+
+        let element = data_layout_element_at(&layout, vec![1, 2]);
+
+        assert_eq!(element, &PatioDataLayout::Byte);
+    }
+}
diff --git a/patio-to-cranelift/src/patio/types.rs b/patio-to-cranelift/src/patio/types.rs
index 73736ff9..524edd52 100644
--- a/patio-to-cranelift/src/patio/types.rs
+++ b/patio-to-cranelift/src/patio/types.rs
@@ -9,7 +9,7 @@ pub(crate) enum PatioType {
     Void,
 }
 
-#[derive(Debug)]
+#[derive(Debug, PartialEq)]
 pub(crate) enum PatioDataLayout {
     Array {
         element: Box<PatioDataLayout>,