refactor: cleanup runtime/memory code

crates/mun_memory/src/gc/mark_sweep.rs

@@ -1,14 +1,13 @@
 use crate::{
     cast,
-    diff::{Diff, FieldDiff},
     gc::{Event, GcPtr, GcRuntime, Observer, RawGcPtr, Stats, TypeTrace},
-    mapping::{self, field_mapping, FieldMappingDesc, MemoryMapper},
-    TypeDesc, TypeFields, TypeLayout,
+    mapping::{self, FieldMapping, MemoryMapper},
+    TypeLayout,
 };
-use once_cell::unsync::OnceCell;
+use mapping::{Conversion, Mapping};
 use parking_lot::RwLock;
 use std::{
-    collections::{HashMap, HashSet, VecDeque},
+    collections::{HashMap, VecDeque},
     hash::Hash,
     ops::Deref,
     pin::Pin,
@@ -200,125 +199,79 @@ where
 
 impl<T, O> MemoryMapper<T> for MarkSweep<T, O>
 where
-    T: TypeDesc + TypeLayout + TypeFields<T> + TypeTrace + Clone + Eq + Hash,
+    T: TypeLayout + TypeTrace + Clone + Eq + Hash,
     O: Observer<Event = Event>,
 {
-    fn map_memory(&self, old: &[T], new: &[T], diff: &[Diff]) -> Vec<GcPtr> {
-        // Collect all deleted types.
-        let deleted: HashSet<T> = diff
-            .iter()
-            .filter_map(|diff| {
-                if let Diff::Delete { index } = diff {
-                    Some(unsafe { old.get_unchecked(*index) }.clone())
-                } else {
-                    None
-                }
-            })
-            .collect();
-
+    fn map_memory(&self, mapping: Mapping<T, T>) -> Vec<GcPtr> {
         let mut objects = self.objects.write();
 
         // Determine which types are still allocated with deleted types
         let deleted = objects
             .iter()
             .filter_map(|(ptr, object_info)| {
-                if deleted.contains(&object_info.ty) {
+                if mapping.deletions.contains(&object_info.ty) {
                     Some(*ptr)
                 } else {
                     None
                 }
             })
             .collect();
 
-        for diff in diff.iter() {
-            match diff {
-                Diff::Delete { .. } => (), // Already handled
-                Diff::Edit {
-                    diff,
-                    old_index,
-                    new_index,
-                } => {
-                    let old_ty = unsafe { old.get_unchecked(*old_index) };
-                    let new_ty = unsafe { new.get_unchecked(*new_index) };
-
-                    // Use `OnceCell` to lazily construct `TypeData` for `old_ty` and `new_ty`.
-                    let mut type_data = OnceCell::new();
-
-                    for object_info in objects.values_mut() {
-                        if object_info.ty == *old_ty {
-                            map_fields(object_info, old_ty, new_ty, diff, &mut type_data);
-                        }
-                    }
+        for (old_ty, conversion) in mapping.conversions {
+            for object_info in objects.values_mut() {
+                if object_info.ty == old_ty {
+                    map_fields(object_info, &conversion);
                 }
-                Diff::Insert { .. } | Diff::Move { .. } => (),
             }
         }
 
         return deleted;
 
-        struct TypeData<'a, T> {
-            old_fields: Vec<(&'a str, T)>,
-            new_fields: Vec<(&'a str, T)>,
-            old_offsets: &'a [u16],
-            new_offsets: &'a [u16],
-        }
-
-        fn map_fields<'a, T: TypeDesc + TypeLayout + TypeFields<T> + TypeTrace + Clone + Eq>(
+        fn map_fields<T: Clone + TypeLayout + TypeTrace>(
             object_info: &mut Pin<Box<ObjectInfo<T>>>,
-            old_ty: &'a T,
-            new_ty: &'a T,
-            diff: &[FieldDiff],
-            type_data: &mut OnceCell<TypeData<'a, T>>,
+            conversion: &Conversion<T>,
         ) {
-            let TypeData {
-                old_fields,
-                new_fields,
-                old_offsets,
-                new_offsets,
-            } = type_data.get_or_init(|| TypeData {
-                old_fields: old_ty.fields(),
-                new_fields: new_ty.fields(),
-                old_offsets: old_ty.offsets(),
-                new_offsets: new_ty.offsets(),
-            });
-            let ptr = unsafe { std::alloc::alloc_zeroed(new_ty.layout()) };
-
-            let mapping = field_mapping(&old_fields, &diff);
-            for (new_index, map) in mapping.into_iter().enumerate() {
-                if let Some(FieldMappingDesc { old_index, action }) = map {
+            let ptr = unsafe { std::alloc::alloc_zeroed(conversion.new_ty.layout()) };
+
+            for map in conversion.field_mapping.iter() {
+                if let Some(FieldMapping {
+                    old_offset,
+                    new_offset,
+                    action,
+                }) = map
+                {
                     let src = {
                         let mut src = object_info.ptr as usize;
-                        src += usize::from(unsafe { *old_offsets.get_unchecked(old_index) });
+                        src += old_offset;
                         src as *mut u8
                     };
                     let dest = {
                         let mut dest = ptr as usize;
-                        dest += usize::from(unsafe { *new_offsets.get_unchecked(new_index) });
+                        dest += new_offset;
                         dest as *mut u8
                     };
-                    let old_field = unsafe { old_fields.get_unchecked(old_index) };
-                    if action == mapping::Action::Cast {
-                        let new_field = unsafe { new_fields.get_unchecked(new_index) };
-                        if !cast::try_cast_from_to(
-                            *old_field.1.guid(),
-                            *new_field.1.guid(),
-                            unsafe { NonNull::new_unchecked(src) },
-                            unsafe { NonNull::new_unchecked(dest) },
-                        ) {
-                            // Failed to cast. Use the previously zero-initialized value instead
-                        }
-                    } else {
-                        unsafe {
-                            std::ptr::copy_nonoverlapping(src, dest, old_field.1.layout().size())
+                    match action {
+                        mapping::Action::Cast { old, new } => {
+                            if !cast::try_cast_from_to(
+                                *old,
+                                *new,
+                                unsafe { NonNull::new_unchecked(src) },
+                                unsafe { NonNull::new_unchecked(dest) },
+                            ) {
+                                // Failed to cast. Use the previously zero-initialized value instead
+                            }
                         }
+                        mapping::Action::Copy { size } => unsafe {
+                            std::ptr::copy_nonoverlapping(src, dest, *size)
+                        },
                     }
                 }
             }
             object_info.set(ObjectInfo {
                 ptr,
                 roots: object_info.roots,
                 color: object_info.color,
-                ty: new_ty.clone(),
+                ty: conversion.new_ty.clone(),
             });
         }
     }

crates/mun_memory/src/lib.rs

@@ -7,7 +7,7 @@ pub mod mapping;
 
 pub mod prelude {
     pub use crate::diff::{diff, Diff, FieldDiff, FieldEditKind};
-    pub use crate::mapping::{Action, FieldMappingDesc};
+    pub use crate::mapping::{Action, FieldMapping};
 }
 
 /// A trait used to obtain a type's description.

crates/mun_memory/src/mapping.rs

@@ -1,37 +1,89 @@
 use crate::{
-    diff::{Diff, FieldDiff, FieldEditKind},
+    diff::{diff, Diff, FieldDiff, FieldEditKind},
     gc::GcPtr,
+    TypeDesc, TypeFields, TypeLayout,
+};
+use std::{
+    collections::{HashMap, HashSet},
+    hash::Hash,
 };
-use std::collections::HashSet;
 
-/// A trait used to map allocated memory using type differences.
-pub trait MemoryMapper<T> {
-    /// Maps the values memory from `old` to `new` using `diff`.
-    ///
-    /// A `Vec<GcPtr>` is returned containing all objects of types that were deleted. The
-    /// corresponding types have to remain in-memory until the objects have been deallocated.
-    fn map_memory(&self, old: &[T], new: &[T], diff: &[Diff]) -> Vec<GcPtr>;
+pub struct Mapping<T: Eq + Hash, U: TypeLayout> {
+    pub deletions: HashSet<T>,
+    pub conversions: HashMap<T, Conversion<U>>,
 }
 
-/// The `Action` to take when mapping memory from A to B.
-#[derive(Eq, PartialEq)]
-pub enum Action {
-    Cast,
-    Copy,
+pub struct Conversion<T: TypeLayout> {
+    pub field_mapping: Vec<Option<FieldMapping>>,
+    pub new_ty: T,
 }
 
 /// Description of the mapping of a single field. When stored together with the new index, this
 /// provides all information necessary for a mapping function.
-pub struct FieldMappingDesc {
-    pub old_index: usize,
+pub struct FieldMapping {
+    pub old_offset: usize,
+    pub new_offset: usize,
     pub action: Action,
 }
 
+/// The `Action` to take when mapping memory from A to B.
+#[derive(Eq, PartialEq)]
+pub enum Action {
+    Cast { old: abi::Guid, new: abi::Guid },
+    Copy { size: usize },
+}
+
+impl<T> Mapping<T, T>
+where
+    T: TypeDesc + TypeFields<T> + TypeLayout + Copy + Eq + Hash,
+{
+    ///
+    pub fn new(old: &[T], new: &[T]) -> Self {
+        let diff = diff(old, new);
+
+        let mut deletions = HashSet::new();
+        let mut conversions = HashMap::new();
+
+        for diff in diff.iter() {
+            match diff {
+                Diff::Delete { index } => {
+                    deletions.insert(unsafe { *old.get_unchecked(*index) });
+                }
+                Diff::Edit {
+                    diff,
+                    old_index,
+                    new_index,
+                } => {
+                    let old_ty = unsafe { *old.get_unchecked(*old_index) };
+                    let new_ty = unsafe { *new.get_unchecked(*new_index) };
+                    conversions.insert(old_ty, unsafe { field_mapping(old_ty, new_ty, diff) });
+                }
+                Diff::Insert { .. } | Diff::Move { .. } => (),
+            }
+        }
+
+        Self {
+            deletions,
+            conversions,
+        }
+    }
+}
+
 /// Given a set of `old_fields` of type `T` and their corresponding `diff`, calculates the mapping
 /// `new_index -> Option<FieldMappingDesc>` for each new field.
 ///
 /// The indices of the returned `Vec`'s elements should be used as indices for the new fields.
-pub fn field_mapping<T>(old_fields: &[T], diff: &[FieldDiff]) -> Vec<Option<FieldMappingDesc>> {
+///
+/// # Safety
+///
+/// Expects the `diff` to be based on `old_ty` and `new_ty`. If not, it causes undefined behavior.
+pub unsafe fn field_mapping<T: TypeDesc + TypeFields<T> + TypeLayout>(
+    old_ty: T,
+    new_ty: T,
+    diff: &[FieldDiff],
+) -> Conversion<T> {
+    let old_fields = old_ty.fields();
+
     let deletions: HashSet<usize> = diff
         .iter()
         .filter_map(|diff| match diff {
@@ -41,6 +93,17 @@ pub fn field_mapping<T>(old_fields: &[T], diff: &[FieldDiff]) -> Vec<Option<Fiel
         })
         .collect();
 
+    struct FieldMappingDesc {
+        old_index: usize,
+        action: ActionDesc,
+    }
+
+    #[derive(PartialEq)]
+    enum ActionDesc {
+        Cast,
+        Copy,
+    }
+
     // Add mappings for all `old_fields`, unless they were deleted or moved.
     let mut mapping: Vec<Option<FieldMappingDesc>> = (0..old_fields.len())
         .filter_map(|idx| {
@@ -49,7 +112,7 @@ pub fn field_mapping<T>(old_fields: &[T], diff: &[FieldDiff]) -> Vec<Option<Fiel
             } else {
                 Some(Some(FieldMappingDesc {
                     old_index: idx,
-                    action: Action::Copy,
+                    action: ActionDesc::Copy,
                 }))
             }
         })
@@ -69,11 +132,11 @@ pub fn field_mapping<T>(old_fields: &[T], diff: &[FieldDiff]) -> Vec<Option<Fiel
                 *new_index,
                 Some(FieldMappingDesc {
                     old_index: *old_index,
-                    action: edit.as_ref().map_or(Action::Copy, |kind| {
+                    action: edit.as_ref().map_or(ActionDesc::Copy, |kind| {
                         if *kind == FieldEditKind::ConvertType {
-                            Action::Cast
+                            ActionDesc::Cast
                         } else {
-                            Action::Copy
+                            ActionDesc::Copy
                         }
                     }),
                 }),
@@ -93,12 +156,50 @@ pub fn field_mapping<T>(old_fields: &[T], diff: &[FieldDiff]) -> Vec<Option<Fiel
         if let FieldDiff::Edit { index, kind } = diff {
             if let Some(map) = mapping.get_mut(*index).unwrap() {
                 map.action = if *kind == FieldEditKind::ConvertType {
-                    Action::Cast
+                    ActionDesc::Cast
                 } else {
-                    Action::Copy
+                    ActionDesc::Copy
                 };
             }
         }
     }
-    mapping
+
+    let new_fields = new_ty.fields();
+    let old_offsets = old_ty.offsets();
+    let new_offsets = new_ty.offsets();
+    Conversion {
+        field_mapping: mapping
+            .into_iter()
+            .enumerate()
+            .map(|(new_index, desc)| {
+                desc.map(|desc| {
+                    let old_field = old_fields.get_unchecked(desc.old_index);
+                    FieldMapping {
+                        old_offset: usize::from(*old_offsets.get_unchecked(desc.old_index)),
+                        new_offset: usize::from(*new_offsets.get_unchecked(new_index)),
+                        action: if desc.action == ActionDesc::Cast {
+                            Action::Cast {
+                                old: *old_field.1.guid(),
+                                new: *new_fields.get_unchecked(new_index).1.guid(),
+                            }
+                        } else {
+                            Action::Copy {
+                                size: old_field.1.layout().size(),
+                            }
+                        },
+                    }
+                })
+            })
+            .collect(),
+        new_ty,
+    }
+}
+
+/// A trait used to map allocated memory using type differences.
+pub trait MemoryMapper<T: Eq + Hash + TypeLayout> {
+    /// Maps its allocated memory using the provided `mapping`.
+    ///
+    /// A `Vec<GcPtr>` is returned containing all objects of types that were deleted. The
+    /// corresponding types have to remain in-memory until the objects have been deallocated.
+    fn map_memory(&self, mapping: Mapping<T, T>) -> Vec<GcPtr>;
 }