review suggestions

examples/calculator/src/calculator.rs

@@ -2,7 +2,7 @@ use p3_air::{Air, AirBuilder, AirBuilderWithPublicValues, BaseAir};
 use p3_field::Field;
 use p3_matrix::{dense::RowMajorMatrix, Matrix};
 use rayon::iter::{IntoParallelRefIterator, ParallelIterator};
-use std::{mem::size_of, slice::from_raw_parts};
+use std::mem::size_of;
 
 use loam::pointers::GPtr;
 
@@ -61,11 +61,72 @@ impl<T> VarFrame<T> {
         &shorts[0]
     }
 
-    /// Unsafely reads `&VarFrame<T>` as `&[T]`, similarly to `from_slice`
     #[inline]
-    fn to_slice(&self) -> &[T] {
-        let pointer = self as *const Self as *const T;
-        unsafe { from_raw_parts(pointer, NUM_COLS) }
+    fn into_array(self) -> [T; NUM_COLS]
+    where
+        T: Copy + Default,
+    {
+        let Self {
+            entry,
+            stack,
+            o_entry,
+            o_stack,
+            slot0,
+            slot1,
+            slot2,
+            slot3,
+            entry_tag_is_nil,
+            entry_tag_is_cons,
+            entry_head,
+            entry_tail,
+            entry_head_tag_is_num,
+            new_stack_push,
+            entry_head_tag_is_op_add,
+            first_num,
+            stack_tail_1,
+            second_num,
+            stack_tail_2,
+            sum,
+            new_stack_op_add,
+        } = self;
+        let mut array = [T::default(); NUM_COLS];
+        let mut idx = 0;
+        let populate_with_atom = |a, array: &mut [T; NUM_COLS], idx: &mut usize| {
+            array[*idx] = a;
+            *idx += 1;
+        };
+        let populate_with_ptr = |ptr: GPtr<T, T>, array: &mut [T; NUM_COLS], idx: &mut usize| {
+            let (tag, val) = ptr.into_parts();
+            populate_with_atom(tag, array, idx);
+            populate_with_atom(val, array, idx);
+        };
+        let populate_with_slot = |slot: [T; 4], array: &mut [T; NUM_COLS], idx: &mut usize| {
+            for a in slot {
+                populate_with_atom(a, array, idx);
+            }
+        };
+        populate_with_ptr(entry, &mut array, &mut idx);
+        populate_with_ptr(stack, &mut array, &mut idx);
+        populate_with_ptr(o_entry, &mut array, &mut idx);
+        populate_with_ptr(o_stack, &mut array, &mut idx);
+        populate_with_slot(slot0, &mut array, &mut idx);
+        populate_with_slot(slot1, &mut array, &mut idx);
+        populate_with_slot(slot2, &mut array, &mut idx);
+        populate_with_slot(slot3, &mut array, &mut idx);
+        populate_with_atom(entry_tag_is_nil, &mut array, &mut idx);
+        populate_with_atom(entry_tag_is_cons, &mut array, &mut idx);
+        populate_with_ptr(entry_head, &mut array, &mut idx);
+        populate_with_ptr(entry_tail, &mut array, &mut idx);
+        populate_with_atom(entry_head_tag_is_num, &mut array, &mut idx);
+        populate_with_ptr(new_stack_push, &mut array, &mut idx);
+        populate_with_atom(entry_head_tag_is_op_add, &mut array, &mut idx);
+        populate_with_ptr(first_num, &mut array, &mut idx);
+        populate_with_ptr(stack_tail_1, &mut array, &mut idx);
+        populate_with_ptr(second_num, &mut array, &mut idx);
+        populate_with_ptr(stack_tail_2, &mut array, &mut idx);
+        populate_with_ptr(sum, &mut array, &mut idx);
+        populate_with_ptr(new_stack_op_add, &mut array, &mut idx);
+        array
     }
 }
 
@@ -233,12 +294,7 @@ impl<F: Field> Calculator<F> {
     pub(crate) fn compute_trace(frames: &[WitnessFrame], store: &Store<F>) -> RowMajorMatrix<F> {
         let values = frames
             .par_iter()
-            .map(|frame| {
-                Self::materialize_var_frame(frame, store)
-                    .to_slice()
-                    .to_vec()
-            })
-            .flatten()
+            .flat_map(|frame| Self::materialize_var_frame(frame, store).into_array())
             .collect();
         RowMajorMatrix::new(values, NUM_COLS)
     }