Pack QueryOrigin memory layout (#885)

* pack `QueryOrigin` into 13 bytes

* nits

Co-authored-by: David Barsky <[email protected]>

---------

Co-authored-by: David Barsky <[email protected]>

Author: ibraheemdev

src/active_query.rs

@@ -11,7 +11,7 @@ use crate::key::DatabaseKeyIndex;
 use crate::runtime::Stamp;
 use crate::sync::atomic::AtomicBool;
 use crate::tracked_struct::{Disambiguator, DisambiguatorMap, IdentityHash, IdentityMap};
-use crate::zalsa_local::{QueryEdge, QueryEdges, QueryOrigin, QueryRevisions};
+use crate::zalsa_local::{QueryEdge, QueryOrigin, QueryRevisions};
 use crate::{Accumulator, IngredientIndex, Revision};
 
 #[derive(Debug)]
@@ -193,11 +193,10 @@ impl ActiveQuery {
             iteration_count: _,
         } = self;
 
-        let edges = QueryEdges::new(input_outputs.drain(..));
         let origin = if untracked_read {
-            QueryOrigin::DerivedUntracked(edges)
+            QueryOrigin::derived_untracked(input_outputs.drain(..))
         } else {
-            QueryOrigin::Derived(edges)
+            QueryOrigin::derived(input_outputs.drain(..))
         };
         disambiguator_map.clear();
         let accumulated = accumulated

src/function.rs

@@ -17,7 +17,7 @@ use crate::table::memo::MemoTableTypes;
 use crate::table::Table;
 use crate::views::DatabaseDownCaster;
 use crate::zalsa::{IngredientIndex, MemoIngredientIndex, Zalsa};
-use crate::zalsa_local::QueryOrigin;
+use crate::zalsa_local::QueryOriginRef;
 use crate::{Database, Id, Revision};
 
 mod accumulated;
@@ -268,7 +268,7 @@ where
         }
     }
 
-    fn origin(&self, zalsa: &Zalsa, key: Id) -> Option<QueryOrigin> {
+    fn origin<'db>(&self, zalsa: &'db Zalsa, key: Id) -> Option<QueryOriginRef<'db>> {
         self.origin(zalsa, key)
     }
 

src/function/accumulated.rs

@@ -3,7 +3,7 @@ use crate::accumulator::{self};
 use crate::function::{Configuration, IngredientImpl};
 use crate::hash::FxHashSet;
 use crate::zalsa::ZalsaDatabase;
-use crate::zalsa_local::QueryOrigin;
+use crate::zalsa_local::QueryOriginRef;
 use crate::{AsDynDatabase, DatabaseKeyIndex, Id};
 
 impl<C> IngredientImpl<C>
@@ -73,8 +73,9 @@ where
                 continue;
             };
 
-            if let QueryOrigin::Derived(edges) | QueryOrigin::DerivedUntracked(edges) = &origin {
-                stack.reserve(edges.input_outputs.len());
+            if let QueryOriginRef::Derived(edges) | QueryOriginRef::DerivedUntracked(edges) = origin
+            {
+                stack.reserve(edges.len());
             }
 
             stack.extend(

src/function/diff_outputs.rs

@@ -28,6 +28,7 @@ where
         let mut old_outputs: FxIndexSet<_> = old_memo
             .revisions
             .origin
+            .as_ref()
             .outputs()
             .map(|a| (a.ingredient_index(), a.key_index().index()))
             .collect();
@@ -38,7 +39,7 @@ where
 
         // Iterate over the outputs of the current query
         // and remove elements from `old_outputs` when we find them
-        for new_output in revisions.origin.outputs() {
+        for new_output in revisions.origin.as_ref().outputs() {
             old_outputs.swap_remove(&(
                 new_output.ingredient_index(),
                 new_output.key_index().index(),

src/function/inputs.rs

@@ -1,15 +1,15 @@
 use crate::function::{Configuration, IngredientImpl};
 use crate::zalsa::Zalsa;
-use crate::zalsa_local::QueryOrigin;
+use crate::zalsa_local::QueryOriginRef;
 use crate::Id;
 
 impl<C> IngredientImpl<C>
 where
     C: Configuration,
 {
-    pub(super) fn origin(&self, zalsa: &Zalsa, key: Id) -> Option<QueryOrigin> {
+    pub(super) fn origin<'db>(&self, zalsa: &'db Zalsa, key: Id) -> Option<QueryOriginRef<'db>> {
         let memo_ingredient_index = self.memo_ingredient_index(zalsa, key);
         self.get_memo_from_table_for(zalsa, key, memo_ingredient_index)
-            .map(|m| m.revisions.origin.clone())
+            .map(|m| m.revisions.origin.as_ref())
     }
 }

src/function/maybe_changed_after.rs

@@ -7,7 +7,7 @@ use crate::key::DatabaseKeyIndex;
 use crate::plumbing::ZalsaLocal;
 use crate::sync::atomic::Ordering;
 use crate::zalsa::{MemoIngredientIndex, Zalsa, ZalsaDatabase};
-use crate::zalsa_local::{QueryEdge, QueryOrigin};
+use crate::zalsa_local::{QueryEdge, QueryOriginRef};
 use crate::{AsDynDatabase as _, Id, Revision};
 
 /// Result of memo validation.
@@ -353,8 +353,8 @@ where
             return VerifyResult::unchanged();
         }
 
-        match &old_memo.revisions.origin {
-            QueryOrigin::Assigned(_) => {
+        match old_memo.revisions.origin.as_ref() {
+            QueryOriginRef::Assigned(_) => {
                 // If the value was assigned by another query,
                 // and that query were up-to-date,
                 // then we would have updated the `verified_at` field already.
@@ -372,15 +372,15 @@ where
             // when we hit the cycle. Any dependencies accessed when creating the fixpoint initial
             // are tracked by the outer query. Nothing should have changed assuming that the
             // fixpoint initial function is deterministic.
-            QueryOrigin::FixpointInitial => {
+            QueryOriginRef::FixpointInitial => {
                 cycle_heads.push_initial(database_key_index);
                 VerifyResult::unchanged()
             }
-            QueryOrigin::DerivedUntracked(_) => {
+            QueryOriginRef::DerivedUntracked(_) => {
                 // Untracked inputs? Have to assume that it changed.
                 VerifyResult::Changed
             }
-            QueryOrigin::Derived(edges) => {
+            QueryOriginRef::Derived(edges) => {
                 let is_provisional = old_memo.may_be_provisional();
 
                 // If the value is from the same revision but is still provisional, consider it changed
@@ -398,7 +398,7 @@ where
                 // they executed. It's possible that if the value of some input I0 is no longer
                 // valid, then some later input I1 might never have executed at all, so verifying
                 // it is still up to date is meaningless.
-                for &edge in edges.input_outputs.iter() {
+                for &edge in edges {
                     match edge {
                         QueryEdge::Input(dependency_index) => {
                             match dependency_index.maybe_changed_after(

src/function/memo.rs

@@ -10,7 +10,7 @@ use crate::revision::AtomicRevision;
 use crate::sync::atomic::Ordering;
 use crate::table::memo::MemoTableWithTypesMut;
 use crate::zalsa::{MemoIngredientIndex, Zalsa};
-use crate::zalsa_local::{QueryOrigin, QueryRevisions};
+use crate::zalsa_local::{QueryOriginRef, QueryRevisions};
 use crate::{Event, EventKind, Id, Revision};
 
 impl<C: Configuration> IngredientImpl<C> {
@@ -64,16 +64,16 @@ impl<C: Configuration> IngredientImpl<C> {
         memo_ingredient_index: MemoIngredientIndex,
     ) {
         let map = |memo: &mut Memo<C::Output<'static>>| {
-            match &memo.revisions.origin {
-                QueryOrigin::Assigned(_)
-                | QueryOrigin::DerivedUntracked(_)
-                | QueryOrigin::FixpointInitial => {
+            match memo.revisions.origin.as_ref() {
+                QueryOriginRef::Assigned(_)
+                | QueryOriginRef::DerivedUntracked(_)
+                | QueryOriginRef::FixpointInitial => {
                     // Careful: Cannot evict memos whose values were
                     // assigned as output of another query
                     // or those with untracked inputs
                     // as their values cannot be reconstructed.
                 }
-                QueryOrigin::Derived(_) => {
+                QueryOriginRef::Derived(_) => {
                     // Set the memo value to `None`.
                     memo.value = None;
                 }
@@ -101,7 +101,7 @@ pub struct Memo<V> {
 #[cfg(not(feature = "shuttle"))]
 #[cfg(target_pointer_width = "64")]
 const _: [(); std::mem::size_of::<Memo<std::num::NonZeroUsize>>()] =
-    [(); std::mem::size_of::<[usize; 13]>()];
+    [(); std::mem::size_of::<[usize; 11]>()];
 
 impl<V> Memo<V> {
     pub(super) fn new(value: Option<V>, revision_now: Revision, revisions: QueryRevisions) -> Self {
@@ -230,7 +230,7 @@ impl<V> Memo<V> {
         zalsa: &Zalsa,
         database_key_index: DatabaseKeyIndex,
     ) {
-        for output in self.revisions.origin.outputs() {
+        for output in self.revisions.origin.as_ref().outputs() {
             output.mark_validated_output(zalsa, database_key_index);
         }
     }
@@ -262,7 +262,7 @@ impl<V> Memo<V> {
 }
 
 impl<V: Send + Sync + Any> crate::table::memo::Memo for Memo<V> {
-    fn origin(&self) -> &QueryOrigin {
-        &self.revisions.origin
+    fn origin(&self) -> QueryOriginRef<'_> {
+        self.revisions.origin.as_ref()
     }
 }

src/function/specify.rs

@@ -5,7 +5,7 @@ use crate::revision::AtomicRevision;
 use crate::sync::atomic::AtomicBool;
 use crate::tracked_struct::TrackedStructInDb;
 use crate::zalsa::{Zalsa, ZalsaDatabase};
-use crate::zalsa_local::{QueryOrigin, QueryRevisions};
+use crate::zalsa_local::{QueryOrigin, QueryOriginRef, QueryRevisions};
 use crate::{DatabaseKeyIndex, Id};
 
 impl<C> IngredientImpl<C>
@@ -65,7 +65,7 @@ where
         let mut revisions = QueryRevisions {
             changed_at: current_deps.changed_at,
             durability: current_deps.durability,
-            origin: QueryOrigin::Assigned(active_query_key),
+            origin: QueryOrigin::assigned(active_query_key),
             tracked_struct_ids: Default::default(),
             accumulated: Default::default(),
             accumulated_inputs: Default::default(),
@@ -116,8 +116,8 @@ where
 
         // If we are marking this as validated, it must be a value that was
         // assigned by `executor`.
-        match memo.revisions.origin {
-            QueryOrigin::Assigned(by_query) => assert_eq!(by_query, executor),
+        match memo.revisions.origin.as_ref() {
+            QueryOriginRef::Assigned(by_query) => assert_eq!(by_query, executor),
             _ => panic!(
                 "expected a query assigned by `{:?}`, not `{:?}`",
                 executor, memo.revisions.origin,

src/ingredient.rs

@@ -9,7 +9,7 @@ use crate::sync::Arc;
 use crate::table::memo::MemoTableTypes;
 use crate::table::Table;
 use crate::zalsa::{transmute_data_mut_ptr, transmute_data_ptr, IngredientIndex, Zalsa};
-use crate::zalsa_local::QueryOrigin;
+use crate::zalsa_local::QueryOriginRef;
 use crate::{Database, DatabaseKeyIndex, Id, Revision};
 
 /// A "jar" is a group of ingredients that are added atomically.
@@ -150,7 +150,7 @@ pub trait Ingredient: Any + std::fmt::Debug + Send + Sync {
     }
 
     /// What were the inputs (if any) that were used to create the value at `key_index`.
-    fn origin(&self, zalsa: &Zalsa, key_index: Id) -> Option<QueryOrigin> {
+    fn origin<'db>(&self, zalsa: &'db Zalsa, key_index: Id) -> Option<QueryOriginRef<'db>> {
         let _ = (zalsa, key_index);
         unreachable!("only function ingredients have origins")
     }

src/memo_ingredient_indices.rs

@@ -77,7 +77,7 @@ impl NewMemoIngredientIndices for MemoIngredientIndices {
         };
         let mut indices = Vec::new();
         indices.resize(
-            last.as_usize() + 1,
+            (last.as_u32() as usize) + 1,
             MemoIngredientIndex::from_usize((u32::MAX - 1) as usize),
         );
         for &struct_ingredient in &struct_indices.indices {
@@ -88,7 +88,7 @@ impl NewMemoIngredientIndices for MemoIngredientIndices {
             let mi = zalsa.next_memo_ingredient_index(struct_ingredient, ingredient);
             memo_types.set(mi, &memo_type);
 
-            indices[struct_ingredient.as_usize()] = mi;
+            indices[struct_ingredient.as_u32() as usize] = mi;
         }
         MemoIngredientIndices {
             indices: indices.into_boxed_slice(),
@@ -125,7 +125,7 @@ impl MemoIngredientMap for MemoIngredientIndices {
     }
     #[inline(always)]
     fn get(&self, index: IngredientIndex) -> MemoIngredientIndex {
-        self.indices[index.as_usize()]
+        self.indices[index.as_u32() as usize]
     }
 }
 

src/table/memo.rs

@@ -9,7 +9,7 @@ use thin_vec::ThinVec;
 
 use crate::sync::atomic::{AtomicPtr, Ordering};
 use crate::sync::{OnceLock, RwLock};
-use crate::{zalsa::MemoIngredientIndex, zalsa_local::QueryOrigin};
+use crate::{zalsa::MemoIngredientIndex, zalsa_local::QueryOriginRef};
 
 /// The "memo table" stores the memoized results of tracked function calls.
 /// Every tracked function must take a salsa struct as its first argument
@@ -21,7 +21,7 @@ pub(crate) struct MemoTable {
 
 pub trait Memo: Any + Send + Sync {
     /// Returns the `origin` of this memo
-    fn origin(&self) -> &QueryOrigin;
+    fn origin(&self) -> QueryOriginRef<'_>;
 }
 
 /// Data for a memoized entry.
@@ -103,7 +103,7 @@ impl MemoEntryType {
 struct DummyMemo {}
 
 impl Memo for DummyMemo {
-    fn origin(&self) -> &QueryOrigin {
+    fn origin(&self) -> QueryOriginRef<'_> {
         unreachable!("should not get here")
     }
 }

src/tracked_struct.rs

@@ -447,7 +447,7 @@ where
             // shouldn't be a problem in general as sufficient bits are reserved for the generation.
             let Some(id) = id.next_generation() else {
                 tracing::info!(
-                    "leaking tracked struct {:?} due to generation oveflow",
+                    "leaking tracked struct {:?} due to generation overflow",
                     self.database_key_index(id)
                 );
 
@@ -547,7 +547,7 @@ where
             // the previous slot and allocate a new value.
             if id.generation() == u32::MAX {
                 tracing::info!(
-                    "leaking tracked struct {:?} due to generation oveflow",
+                    "leaking tracked struct {:?} due to generation overflow",
                     self.database_key_index(id)
                 );
 

src/zalsa.rs

@@ -83,9 +83,9 @@ impl IngredientIndex {
         Self(v as u32)
     }
 
-    /// Convert the ingredient index back into a usize.
-    pub(crate) fn as_usize(self) -> usize {
-        self.0 as usize
+    /// Convert the ingredient index back into a `u32`.
+    pub(crate) fn as_u32(self) -> u32 {
+        self.0
     }
 
     pub fn successor(self, index: usize) -> Self {
@@ -202,7 +202,7 @@ impl Zalsa {
 
     #[inline]
     pub(crate) fn lookup_ingredient(&self, index: IngredientIndex) -> &dyn Ingredient {
-        let index = index.as_usize();
+        let index = index.as_u32() as usize;
         self.ingredients_vec
             .get(index)
             .unwrap_or_else(|| panic!("index `{index}` is uninitialized"))
@@ -214,7 +214,7 @@ impl Zalsa {
         struct_ingredient_index: IngredientIndex,
         memo_ingredient_index: MemoIngredientIndex,
     ) -> IngredientIndex {
-        self.memo_ingredient_indices.read()[struct_ingredient_index.as_usize()]
+        self.memo_ingredient_indices.read()[struct_ingredient_index.as_u32() as usize]
             [memo_ingredient_index.as_usize()]
     }
 
@@ -240,7 +240,7 @@ impl Zalsa {
         ingredient_index: IngredientIndex,
     ) -> MemoIngredientIndex {
         let mut memo_ingredients = self.memo_ingredient_indices.write();
-        let idx = struct_ingredient_index.as_usize();
+        let idx = struct_ingredient_index.as_u32() as usize;
         let memo_ingredients = if let Some(memo_ingredients) = memo_ingredients.get_mut(idx) {
             memo_ingredients
         } else {
@@ -305,7 +305,7 @@ impl Zalsa {
 
             let actual_index = self.ingredients_vec.push(ingredient);
             assert_eq!(
-                expected_index.as_usize(),
+                expected_index.as_u32() as usize,
                 actual_index,
                 "ingredient `{:?}` was predicted to have index `{:?}` but actually has index `{:?}`",
                 self.ingredients_vec[actual_index],
@@ -328,7 +328,7 @@ impl Zalsa {
         &mut self,
         index: IngredientIndex,
     ) -> (&mut dyn Ingredient, &mut Runtime) {
-        let index = index.as_usize();
+        let index = index.as_u32() as usize;
         let ingredient = self
             .ingredients_vec
             .get_mut(index)
@@ -358,7 +358,7 @@ impl Zalsa {
         let _span = tracing::debug_span!("new_revision", ?new_revision).entered();
 
         for (_, index) in self.ingredients_requiring_reset.iter() {
-            let index = index.as_usize();
+            let index = index.as_u32() as usize;
             let ingredient = self
                 .ingredients_vec
                 .get_mut(index)
@@ -375,7 +375,7 @@ impl Zalsa {
     pub fn evict_lru(&mut self) {
         let _span = tracing::debug_span!("evict_lru").entered();
         for (_, index) in self.ingredients_requiring_reset.iter() {
-            let index = index.as_usize();
+            let index = index.as_u32() as usize;
             self.ingredients_vec
                 .get_mut(index)
                 .unwrap_or_else(|| panic!("index `{index}` is uninitialized"))