Pack QueryEdge memory layout (#886)

* pack `QueryEdge` into 12 bytes

* document internals of `QueryEdge`

src/active_query.rs

@@ -90,7 +90,7 @@ impl ActiveQuery {
     ) {
         self.durability = self.durability.min(durability);
         self.changed_at = self.changed_at.max(changed_at);
-        self.input_outputs.insert(QueryEdge::Input(input));
+        self.input_outputs.insert(QueryEdge::input(input));
         self.accumulated_inputs = self.accumulated_inputs.or_else(|| match has_accumulated {
             true => InputAccumulatedValues::Any,
             false => accumulated_inputs.load(),
@@ -106,7 +106,7 @@ impl ActiveQuery {
     ) {
         self.durability = self.durability.min(durability);
         self.changed_at = self.changed_at.max(revision);
-        self.input_outputs.insert(QueryEdge::Input(input));
+        self.input_outputs.insert(QueryEdge::input(input));
     }
 
     pub(super) fn add_untracked_read(&mut self, changed_at: Revision) {
@@ -127,12 +127,12 @@ impl ActiveQuery {
 
     /// Adds a key to our list of outputs.
     pub(super) fn add_output(&mut self, key: DatabaseKeyIndex) {
-        self.input_outputs.insert(QueryEdge::Output(key));
+        self.input_outputs.insert(QueryEdge::output(key));
     }
 
     /// True if the given key was output by this query.
     pub(super) fn is_output(&self, key: DatabaseKeyIndex) -> bool {
-        self.input_outputs.contains(&QueryEdge::Output(key))
+        self.input_outputs.contains(&QueryEdge::output(key))
     }
 
     pub(super) fn disambiguate(&mut self, key: IdentityHash) -> Disambiguator {

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, QueryOriginRef};
+use crate::zalsa_local::{QueryEdgeKind, QueryOriginRef};
 use crate::{AsDynDatabase as _, Id, Revision};
 
 /// Result of memo validation.
@@ -399,8 +399,8 @@ where
                 // 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 {
-                    match edge {
-                        QueryEdge::Input(dependency_index) => {
+                    match edge.kind() {
+                        QueryEdgeKind::Input(dependency_index) => {
                             match dependency_index.maybe_changed_after(
                                 dyn_db,
                                 zalsa,
@@ -413,7 +413,7 @@ where
                                 }
                             }
                         }
-                        QueryEdge::Output(dependency_index) => {
+                        QueryEdgeKind::Output(dependency_index) => {
                             // Subtle: Mark outputs as validated now, even though we may
                             // later find an input that requires us to re-execute the function.
                             // Even if it re-execute, the function will wind up writing the same value,

src/id.rs

@@ -1,6 +1,6 @@
 use std::fmt::Debug;
 use std::hash::Hash;
-use std::num::NonZeroU64;
+use std::num::NonZeroU32;
 
 use crate::zalsa::Zalsa;
 
@@ -19,7 +19,8 @@ use crate::zalsa::Zalsa;
 /// it is wrapped in new types.
 #[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
 pub struct Id {
-    value: NonZeroU64,
+    index: NonZeroU32,
+    generation: u32,
 }
 
 impl Id {
@@ -39,16 +40,17 @@ impl Id {
     #[doc(hidden)]
     #[track_caller]
     #[inline]
-    pub const unsafe fn from_index(v: u32) -> Self {
-        debug_assert!(v < Self::MAX_U32);
+    pub const unsafe fn from_index(index: u32) -> Self {
+        debug_assert!(index < Self::MAX_U32);
 
         Id {
             // SAFETY: Caller obligation.
-            value: unsafe { NonZeroU64::new_unchecked((v + 1) as u64) },
+            index: unsafe { NonZeroU32::new_unchecked(index + 1) },
+            generation: 0,
         }
     }
 
-    /// Create a `salsa::Id` from a u64 value.
+    /// Create a `salsa::Id` from a `u64` value.
     ///
     /// This should only be used to recreate an `Id` together with `Id::as_u64`.
     ///
@@ -59,11 +61,18 @@ impl Id {
     #[doc(hidden)]
     #[track_caller]
     #[inline]
-    pub const unsafe fn from_bits(v: u64) -> Self {
-        Id {
-            // SAFETY: Caller obligation.
-            value: unsafe { NonZeroU64::new_unchecked(v) },
-        }
+    pub const unsafe fn from_bits(bits: u64) -> Self {
+        // SAFETY: Caller obligation.
+        let index = unsafe { NonZeroU32::new_unchecked(bits as u32) };
+        let generation = (bits >> 32) as u32;
+
+        Id { index, generation }
+    }
+
+    /// Return a `u64` representation of this `Id`.
+    #[inline]
+    pub fn as_bits(self) -> u64 {
+        u64::from(self.index.get()) | (u64::from(self.generation) << 32)
     }
 
     /// Returns a new `Id` with same index, but the generation incremented by one.
@@ -84,35 +93,22 @@ impl Id {
     /// provided generation.
     #[inline]
     pub const fn with_generation(self, generation: u32) -> Id {
-        let mut value = self.value.get();
-
-        value &= 0xFFFFFFFF;
-        value |= (generation as u64) << 32;
-
         Id {
-            // SAFETY: The niche of `value` is in the lower bits, which we did not touch.
-            value: unsafe { NonZeroU64::new_unchecked(value) },
+            index: self.index,
+            generation,
         }
     }
 
     /// Return the index portion of this `Id`.
     #[inline]
     pub const fn index(self) -> u32 {
-        // Truncate the high-order bits.
-        (self.value.get() as u32) - 1
+        self.index.get() - 1
     }
 
     /// Return the generation of this `Id`.
     #[inline]
     pub const fn generation(self) -> u32 {
-        // Shift away the low-order bits.
-        (self.value.get() >> 32) as u32
-    }
-
-    /// Return the internal `u64` representation of this `Id`.
-    #[inline]
-    pub const fn as_bits(self) -> u64 {
-        self.value.get()
+        self.generation
     }
 }
 

src/key.rs

@@ -12,8 +12,7 @@ use crate::{Database, Id};
 /// only for inserting into maps and the like.
 #[derive(Copy, Clone, PartialEq, Eq, Hash)]
 pub struct DatabaseKeyIndex {
-    key_index: u32,
-    key_generation: u32,
+    key_index: Id,
     ingredient_index: IngredientIndex,
 }
 // ANCHOR_END: DatabaseKeyIndex
@@ -22,8 +21,7 @@ impl DatabaseKeyIndex {
     #[inline]
     pub(crate) fn new(ingredient_index: IngredientIndex, key_index: Id) -> Self {
         Self {
-            key_index: key_index.index(),
-            key_generation: key_index.generation(),
+            key_index,
             ingredient_index,
         }
     }
@@ -33,8 +31,7 @@ impl DatabaseKeyIndex {
     }
 
     pub const fn key_index(self) -> Id {
-        // SAFETY: `self.key_index` was returned by `Id::data`.
-        unsafe { Id::from_index(self.key_index) }.with_generation(self.key_generation)
+        self.key_index
     }
 
     pub(crate) fn maybe_changed_after(

src/zalsa.rs

@@ -77,9 +77,14 @@ static NONCE: NonceGenerator<StorageNonce> = NonceGenerator::new();
 pub struct IngredientIndex(u32);
 
 impl IngredientIndex {
-    /// Create an ingredient index from a usize.
+    /// The maximum supported ingredient index.
+    ///
+    /// This reserves one bit for an optional tag.
+    const MAX_INDEX: u32 = 0x7FFF_FFFF;
+
+    /// Create an ingredient index from a `usize`.
     pub(crate) fn from(v: usize) -> Self {
-        assert!(v <= u32::MAX as usize);
+        assert!(v <= Self::MAX_INDEX as usize);
         Self(v as u32)
     }
 
@@ -91,6 +96,18 @@ impl IngredientIndex {
     pub fn successor(self, index: usize) -> Self {
         IngredientIndex(self.0 + 1 + index as u32)
     }
+
+    /// Returns a new `IngredientIndex` with the tag bit set to the provided value.
+    pub(crate) fn with_tag(mut self, tag: bool) -> IngredientIndex {
+        self.0 &= Self::MAX_INDEX;
+        self.0 |= (tag as u32) << 31;
+        self
+    }
+
+    /// Returns the value of the tag bit.
+    pub(crate) fn tag(self) -> bool {
+        self.0 & !Self::MAX_INDEX != 0
+    }
 }
 
 /// A special secondary index *just* for ingredients that attach
@@ -309,7 +326,7 @@ impl Zalsa {
                 actual_index,
                 "ingredient `{:?}` was predicted to have index `{:?}` but actually has index `{:?}`",
                 self.ingredients_vec[actual_index],
-                expected_index,
+                expected_index.as_u32(),
                 actual_index,
             );
         }
@@ -466,7 +483,7 @@ where
             ) -> IngredientIndex {
                 let index = create_index();
                 let nonce = zalsa.nonce().into_u32().get() as u64;
-                let packed = (nonce << u32::BITS) | (index.0 as u64);
+                let packed = (nonce << u32::BITS) | (index.as_u32() as u64);
                 debug_assert_ne!(packed, IngredientCache::<I>::UNINITIALIZED);
 
                 // Discard the result, whether we won over the cache or not does not matter

src/zalsa_local.rs

@@ -745,8 +745,56 @@ impl std::fmt::Debug for QueryOrigin {
     }
 }
 
+/// An input or output query edge.
+///
+/// This type is a packed version of `QueryEdgeKind`, tagging the `IngredientIndex`
+/// in `key` with a discriminator for the input and output variants without increasing
+/// the size of the type. Notably, this type is 12 bytes as opposed to the 16 byte
+/// `QueryEdgeKind`, which is meaningful as inputs and outputs are stored contiguously.
+#[derive(Copy, Clone, PartialEq, Eq, Hash)]
+pub struct QueryEdge {
+    key: DatabaseKeyIndex,
+}
+
+impl QueryEdge {
+    /// Create an input query edge with the given index.
+    pub fn input(key: DatabaseKeyIndex) -> QueryEdge {
+        Self { key }
+    }
+
+    /// Create an output query edge with the given index.
+    pub fn output(key: DatabaseKeyIndex) -> QueryEdge {
+        let ingredient_index = key.ingredient_index().with_tag(true);
+
+        Self {
+            key: DatabaseKeyIndex::new(ingredient_index, key.key_index()),
+        }
+    }
+
+    /// Returns the kind of this query edge.
+    pub fn kind(self) -> QueryEdgeKind {
+        // Clear the tag to restore the original index.
+        let untagged = DatabaseKeyIndex::new(
+            self.key.ingredient_index().with_tag(false),
+            self.key.key_index(),
+        );
+
+        if self.key.ingredient_index().tag() {
+            QueryEdgeKind::Output(untagged)
+        } else {
+            QueryEdgeKind::Input(untagged)
+        }
+    }
+}
+
+impl std::fmt::Debug for QueryEdge {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        self.kind().fmt(f)
+    }
+}
+
 #[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
-pub enum QueryEdge {
+pub enum QueryEdgeKind {
     Input(DatabaseKeyIndex),
     Output(DatabaseKeyIndex),
 }
@@ -757,9 +805,9 @@ pub enum QueryEdge {
 pub(crate) fn input_edges(
     input_outputs: &[QueryEdge],
 ) -> impl DoubleEndedIterator<Item = DatabaseKeyIndex> + '_ {
-    input_outputs.iter().filter_map(|&edge| match edge {
-        QueryEdge::Input(dependency_index) => Some(dependency_index),
-        QueryEdge::Output(_) => None,
+    input_outputs.iter().filter_map(|&edge| match edge.kind() {
+        QueryEdgeKind::Input(dependency_index) => Some(dependency_index),
+        QueryEdgeKind::Output(_) => None,
     })
 }
 
@@ -769,9 +817,9 @@ pub(crate) fn input_edges(
 pub(crate) fn output_edges(
     input_outputs: &[QueryEdge],
 ) -> impl DoubleEndedIterator<Item = DatabaseKeyIndex> + '_ {
-    input_outputs.iter().filter_map(|&edge| match edge {
-        QueryEdge::Output(dependency_index) => Some(dependency_index),
-        QueryEdge::Input(_) => None,
+    input_outputs.iter().filter_map(|&edge| match edge.kind() {
+        QueryEdgeKind::Output(dependency_index) => Some(dependency_index),
+        QueryEdgeKind::Input(_) => None,
     })
 }