chore(heap): Use SoAVec for Set storage

nova_vm/src/ecmascript/builtins/control_abstraction_objects/promise_objects/promise_constructor.rs

@@ -948,7 +948,7 @@ fn perform_promise_group<'gc>(
         #[cfg(feature = "set")]
         Object::Map(map) => agent[map].size(),
         #[cfg(feature = "set")]
-        Object::Set(set) => agent[set].size(),
+        Object::Set(set) => set.get(agent).set_data.borrow().len() as u32,
         _ => 0,
     };
 

nova_vm/src/ecmascript/builtins/keyed_collections/set_objects/set_constructor.rs

@@ -142,7 +142,7 @@ impl SetConstructor {
                 let array_heap = ArrayHeap::new(elements, arrays);
                 let primitive_heap = PrimitiveHeap::new(bigints, numbers, strings);
 
-                let set_heap_data = &mut sets[set].borrow_mut(&primitive_heap);
+                let mut set_heap_data = set.get_direct_mut(sets);
                 let values = &mut set_heap_data.values;
                 let set_data = set_heap_data.set_data.get_mut();
 
@@ -181,7 +181,7 @@ impl SetConstructor {
                             }
                             hashbrown::hash_table::Entry::Vacant(vacant) => {
                                 vacant.insert(next_index);
-                                values.push(Some(value.unbind()));
+                                values.push(Some(value));
                             }
                         }
                     });

nova_vm/src/ecmascript/builtins/keyed_collections/set_objects/set_iterator_objects/set_iterator_prototype.rs

@@ -60,14 +60,14 @@ impl SetIteratorPrototype {
         // b. Let entries be set.[[SetData]].
         // c. Let numEntries be the number of elements in entries.
         // d. Repeat, while index < numEntries,
-        while agent[iterator].next_index < agent[set].values(gc).len() {
+        while agent[iterator].next_index < set.get(agent).values.len() {
             // i. Let e be entries[index].
             // ii. Set index to index + 1.
             let index = agent[iterator].next_index;
             agent[iterator].next_index += 1;
 
             // iii. if e is not EMPTY, then
-            let Some(e) = agent[set].values(gc)[index] else {
+            let Some(e) = set.get(agent).values[index] else {
                 continue;
             };
 
@@ -91,7 +91,7 @@ impl SetIteratorPrototype {
                 .map(|o| o.into_value());
         }
 
-        debug_assert_eq!(agent[iterator].next_index, agent[set].values(gc).len());
+        debug_assert_eq!(agent[iterator].next_index, set.get(agent).values.len());
 
         // e. Return undefined.
         agent[iterator].set = None;

nova_vm/src/ecmascript/builtins/keyed_collections/set_objects/set_prototype.rs

@@ -115,8 +115,8 @@ impl SetPrototype {
         // 3. Set value to CanonicalizeKeyedCollectionKey(value).
         let value = canonicalize_keyed_collection_key(numbers, value);
 
-        let set_heap_data = &mut sets[s].borrow_mut(&primitive_heap);
-        let values = &mut set_heap_data.values;
+        let set_heap_data = s.get_direct_mut(sets);
+        let values = set_heap_data.values;
         let set_data = set_heap_data.set_data.get_mut();
         let hasher = |value: Value| {
             let mut hasher = AHasher::default();
@@ -140,7 +140,7 @@ impl SetPrototype {
             // 5. Append value to S.[[SetData]].
             let index = u32::try_from(values.len()).unwrap();
             entry.insert(index);
-            values.push(Some(value.unbind()));
+            values.push(Some(value));
         }
         // i. Return S.
         // 6. Return S.
@@ -166,7 +166,9 @@ impl SetPrototype {
         // 3. For each element e of S.[[SetData]], do
         // a. Replace the element of S.[[SetData]] whose value is e with an
         // element whose value is EMPTY.
-        agent[s].clear();
+        let data = s.get_mut(agent);
+        data.set_data.borrow_mut().clear();
+        data.values.clear();
         // 4. Return undefined.
         Ok(Value::Undefined)
     }
@@ -206,8 +208,8 @@ impl SetPrototype {
             value.hash(&primitive_heap, &mut hasher);
             hasher.finish()
         };
-        let set_heap_data = &mut sets[s].borrow_mut(&primitive_heap);
-        let values = &mut set_heap_data.values;
+        let set_heap_data = s.get_direct_mut(sets);
+        let values = set_heap_data.values;
         let set_data = set_heap_data.set_data.get_mut();
         // 4. For each element e of S.[[SetData]], do
         if let Ok(entry) = set_data.find_entry(value_hash, |hash_equal_index| {
@@ -309,7 +311,7 @@ impl SetPrototype {
         // 5. Let numEntries be the number of elements in entries.
         // Note: We must use the values vector length, not the size. The size
         // does not contain empty slots.
-        let mut num_entries = agent[s].values(gc.nogc()).len() as u32;
+        let mut num_entries = s.get(agent).values.len() as u32;
 
         let callback_fn = callback_fn.scope(agent, nogc);
         let scoped_s = s.scope(agent, nogc);
@@ -320,7 +322,7 @@ impl SetPrototype {
         // 7. Repeat, while index < numEntries,
         while index < num_entries {
             // a. Let e be entries[index].
-            let e = agent[s].values(gc.nogc())[index as usize];
+            let e = s.get(agent).values[index as usize];
             // b. Set index to index + 1.
             index += 1;
             // c. If e is not EMPTY, then
@@ -341,7 +343,7 @@ impl SetPrototype {
                 // ii. NOTE: The number of elements in entries may have increased during execution of callbackfn.
                 // iii. Set numEntries to the number of elements in entries.
                 s = scoped_s.get(agent).bind(gc.nogc());
-                num_entries = agent[s].values(gc.nogc()).len() as u32;
+                num_entries = s.get(agent).values.len() as u32;
             }
         }
         // 8. Return undefined.
@@ -370,8 +372,8 @@ impl SetPrototype {
             ..
         } = &agent.heap;
         let primitive_heap = PrimitiveHeap::new(bigints, numbers, strings);
-        let set_heap_data = &sets[s].borrow(&primitive_heap);
-        let values = &set_heap_data.values;
+        let set_heap_data = s.get_direct(sets);
+        let values = set_heap_data.values;
         let set_data = set_heap_data.set_data.borrow();
 
         // 3. Set value to CanonicalizeKeyedCollectionKey(value).
@@ -410,7 +412,7 @@ impl SetPrototype {
         // 2. Perform ? RequireInternalSlot(S, [[SetData]]).
         let s = require_set_data_internal_slot(agent, this_value, gc)?;
         // 3. Let size be SetDataSize(S.[[SetData]]).
-        let size = agent[s].size();
+        let size = s.get(agent).set_data.borrow().len() as u32;
         // 4. Return 𝔽(size).
         Ok(Number::from(size).into_value())
     }

nova_vm/src/ecmascript/builtins/set.rs

@@ -2,8 +2,6 @@
 // License, v. 2.0. If a copy of the MPL was not distributed with this
 // file, You can obtain one at https://mozilla.org/MPL/2.0/.
 
-use core::ops::{Index, IndexMut};
-
 use crate::{
     Heap,
     ecmascript::{
@@ -20,22 +18,57 @@ use crate::{
     },
 };
 
-use self::data::SetHeapData;
+use self::data::{SetHeapData, SetHeapDataMut, SetHeapDataRef};
+use soavec::SoAVec;
 
 pub mod data;
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
 #[repr(transparent)]
 pub struct Set<'a>(BaseIndex<'a, SetHeapData<'static>>);
 
-impl Set<'_> {
+impl<'gc> Set<'gc> {
     pub(crate) const fn _def() -> Self {
         Self(BaseIndex::from_u32_index(0))
     }
 
     pub(crate) const fn get_index(self) -> usize {
         self.0.into_index()
     }
+
+    #[inline(always)]
+    pub(crate) fn get<'a>(self, agent: &'a Agent) -> SetHeapDataRef<'a, 'gc> {
+        self.get_direct(&agent.heap.sets)
+    }
+
+    #[inline(always)]
+    pub(crate) fn get_mut<'a>(self, agent: &'a mut Agent) -> SetHeapDataMut<'a, 'gc> {
+        self.get_direct_mut(&mut agent.heap.sets)
+    }
+
+    #[inline(always)]
+    pub(crate) fn get_direct<'a>(
+        self,
+        sets: &'a SoAVec<SetHeapData<'static>>,
+    ) -> SetHeapDataRef<'a, 'gc> {
+        sets.get(self.0.into_u32_index())
+            .expect("Invalid Set reference")
+    }
+
+    #[inline(always)]
+    pub(crate) fn get_direct_mut<'a>(
+        self,
+        sets: &'a mut SoAVec<SetHeapData<'static>>,
+    ) -> SetHeapDataMut<'a, 'gc> {
+        // SAFETY: Lifetime transmute to thread GC lifetime to temporary heap
+        // reference.
+        unsafe {
+            core::mem::transmute::<SetHeapDataMut<'a, 'static>, SetHeapDataMut<'a, 'gc>>(
+                sets.get_mut(self.0.into_u32_index())
+                    .expect("Invalid Set reference"),
+            )
+        }
+    }
 }
 
 bindable_handle!(Set);
@@ -81,12 +114,12 @@ impl<'a> InternalSlots<'a> for Set<'a> {
 
     #[inline(always)]
     fn get_backing_object(self, agent: &Agent) -> Option<OrdinaryObject<'static>> {
-        agent[self].object_index
+        self.get(agent).object_index.unbind()
     }
 
     fn set_backing_object(self, agent: &mut Agent, backing_object: OrdinaryObject<'static>) {
         assert!(
-            agent[self]
+            self.get_mut(agent)
                 .object_index
                 .replace(backing_object.unbind())
                 .is_none()
@@ -112,34 +145,6 @@ impl HeapSweepWeakReference for Set<'static> {
     }
 }
 
-impl Index<Set<'_>> for Agent {
-    type Output = SetHeapData<'static>;
-
-    fn index(&self, index: Set) -> &Self::Output {
-        &self.heap.sets[index]
-    }
-}
-
-impl IndexMut<Set<'_>> for Agent {
-    fn index_mut(&mut self, index: Set) -> &mut Self::Output {
-        &mut self.heap.sets[index]
-    }
-}
-
-impl Index<Set<'_>> for Vec<SetHeapData<'static>> {
-    type Output = SetHeapData<'static>;
-
-    fn index(&self, index: Set) -> &Self::Output {
-        self.get(index.get_index()).expect("Set out of bounds")
-    }
-}
-
-impl IndexMut<Set<'_>> for Vec<SetHeapData<'static>> {
-    fn index_mut(&mut self, index: Set) -> &mut Self::Output {
-        self.get_mut(index.get_index()).expect("Set out of bounds")
-    }
-}
-
 impl TryFrom<HeapRootData> for Set<'_> {
     type Error = ();
 
@@ -155,8 +160,52 @@ impl TryFrom<HeapRootData> for Set<'_> {
 
 impl<'a> CreateHeapData<SetHeapData<'a>, Set<'a>> for Heap {
     fn create(&mut self, data: SetHeapData<'a>) -> Set<'a> {
-        self.sets.push(data.unbind());
+        let i = self.sets.len();
+        self.sets
+            .push(data.unbind())
+            .expect("Failed to allocate Set");
         self.alloc_counter += core::mem::size_of::<SetHeapData<'static>>();
-        Set(BaseIndex::last(&self.sets))
+        Set(BaseIndex::from_u32_index(i))
+    }
+}
+
+impl HeapMarkAndSweep for SetHeapDataRef<'_, 'static> {
+    fn mark_values(&self, queues: &mut WorkQueues) {
+        let Self {
+            set_data: _,
+            values,
+            object_index,
+            needs_primitive_rehashing: _,
+        } = self;
+        values.mark_values(queues);
+        object_index.mark_values(queues);
+    }
+
+    fn sweep_values(&mut self, _: &CompactionLists) {
+        unreachable!()
+    }
+}
+
+impl HeapMarkAndSweep for SetHeapDataMut<'_, 'static> {
+    fn mark_values(&self, queues: &mut WorkQueues) {
+        let Self {
+            set_data: _,
+            values,
+            object_index,
+            needs_primitive_rehashing: _,
+        } = self;
+        values.mark_values(queues);
+        object_index.mark_values(queues);
+    }
+
+    fn sweep_values(&mut self, compactions: &CompactionLists) {
+        let Self {
+            set_data: _,
+            values,
+            object_index,
+            needs_primitive_rehashing: _,
+        } = self;
+        values.sweep_values(compactions);
+        object_index.sweep_values(compactions);
     }
 }