From ed1632d151c8d7e30761a0d7cccb04a3d37a48e5 Mon Sep 17 00:00:00 2001
From: Niko Matsakis <niko@alum.mit.edu>
Date: Tue, 30 Jun 2020 15:12:53 +0000
Subject: [PATCH] compute database-key-index for derived slots

---
 src/derived.rs      | 26 +++++++++++++++++---------
 src/derived/slot.rs |  8 ++++++--
 src/runtime.rs      |  1 +
 3 files changed, 24 insertions(+), 11 deletions(-)

diff --git a/src/derived.rs b/src/derived.rs
index 3f7107bc..9eafe9c7 100644
--- a/src/derived.rs
+++ b/src/derived.rs
@@ -7,10 +7,10 @@ use crate::plumbing::LruQueryStorageOps;
 use crate::plumbing::QueryFunction;
 use crate::plumbing::QueryStorageMassOps;
 use crate::plumbing::QueryStorageOps;
-use crate::runtime::StampedValue;
-use crate::{CycleError, Database, SweepStrategy};
+use crate::runtime::{FxIndexMap, StampedValue};
+use crate::{CycleError, Database, DatabaseKeyIndex, SweepStrategy};
 use parking_lot::RwLock;
-use rustc_hash::FxHashMap;
+use std::convert::TryFrom;
 use std::marker::PhantomData;
 use std::sync::Arc;
 
@@ -35,8 +35,9 @@ where
     DB: Database + HasQueryGroup<Q::Group>,
     MP: MemoizationPolicy<DB, Q>,
 {
+    group_index: u16,
     lru_list: Lru<Slot<DB, Q, MP>>,
-    slot_map: RwLock<FxHashMap<Q::Key, Arc<Slot<DB, Q, MP>>>>,
+    slot_map: RwLock<FxIndexMap<Q::Key, Arc<Slot<DB, Q, MP>>>>,
     policy: PhantomData<MP>,
 }
 
@@ -103,9 +104,15 @@ where
         }
 
         let mut write = self.slot_map.write();
-        write
-            .entry(key.clone())
-            .or_insert_with(|| Arc::new(Slot::new(key.clone())))
+        let entry = write.entry(key.clone());
+        let key_index = u32::try_from(entry.index()).unwrap();
+        let database_key_index = DatabaseKeyIndex {
+            group_index: self.group_index,
+            query_index: Q::QUERY_INDEX,
+            key_index: key_index,
+        };
+        entry
+            .or_insert_with(|| Arc::new(Slot::new(key.clone(), database_key_index)))
             .clone()
     }
 }
@@ -116,9 +123,10 @@ where
     DB: Database + HasQueryGroup<Q::Group>,
     MP: MemoizationPolicy<DB, Q>,
 {
-    fn new(_group_index: u16) -> Self {
+    fn new(group_index: u16) -> Self {
         DerivedStorage {
-            slot_map: RwLock::new(FxHashMap::default()),
+            group_index,
+            slot_map: RwLock::new(FxIndexMap::default()),
             lru_list: Default::default(),
             policy: PhantomData,
         }
diff --git a/src/derived/slot.rs b/src/derived/slot.rs
index b91e5a7b..402055de 100644
--- a/src/derived/slot.rs
+++ b/src/derived/slot.rs
@@ -14,7 +14,9 @@ use crate::revision::Revision;
 use crate::runtime::Runtime;
 use crate::runtime::RuntimeId;
 use crate::runtime::StampedValue;
-use crate::{CycleError, Database, DiscardIf, DiscardWhat, Event, EventKind, SweepStrategy};
+use crate::{
+    CycleError, Database, DatabaseKeyIndex, DiscardIf, DiscardWhat, Event, EventKind, SweepStrategy,
+};
 use log::{debug, info};
 use parking_lot::Mutex;
 use parking_lot::{RawRwLock, RwLock};
@@ -30,6 +32,7 @@ where
     MP: MemoizationPolicy<DB, Q>,
 {
     key: Q::Key,
+    database_key_index: DatabaseKeyIndex,
     state: RwLock<QueryState<DB, Q>>,
     policy: PhantomData<MP>,
     lru_index: LruIndex,
@@ -109,9 +112,10 @@ where
     DB: Database + HasQueryGroup<Q::Group>,
     MP: MemoizationPolicy<DB, Q>,
 {
-    pub(super) fn new(key: Q::Key) -> Self {
+    pub(super) fn new(key: Q::Key, database_key_index: DatabaseKeyIndex) -> Self {
         Self {
             key,
+            database_key_index,
             state: RwLock::new(QueryState::NotComputed),
             lru_index: LruIndex::default(),
             policy: PhantomData,
diff --git a/src/runtime.rs b/src/runtime.rs
index 68ec675c..aa0a7700 100644
--- a/src/runtime.rs
+++ b/src/runtime.rs
@@ -14,6 +14,7 @@ use std::sync::atomic::{AtomicUsize, Ordering};
 use std::sync::Arc;
 
 pub(crate) type FxIndexSet<K> = indexmap::IndexSet<K, BuildHasherDefault<FxHasher>>;
+pub(crate) type FxIndexMap<K, V> = indexmap::IndexMap<K, V, BuildHasherDefault<FxHasher>>;
 
 mod local_state;
 use local_state::LocalState;