import the code

Cargo.toml

@@ -6,3 +6,5 @@ edition = "2018"
 license = "Apache-2.0/MIT"
 
 [dependencies]
+derive-new = "0.5.5"
+rustc-hash = "1.0"

src/dyn_descriptor.rs

@@ -0,0 +1,46 @@
+use crate::BaseQueryContext;
+use crate::Query;
+use crate::QueryTable;
+use rustc_hash::FxHashMap;
+use std::any::{Any, TypeId};
+use std::cell::RefCell;
+use std::collections::hash_map::Entry;
+use std::fmt::Debug;
+use std::fmt::Display;
+use std::fmt::Write;
+use std::hash::Hash;
+
+// Total hack for now: assume that the Debug string
+// for the key, combined with the type-id of the query,
+// is sufficient for an equality comparison.
+
+/// A simple-to-use query descriptor that is meant only for dumping
+/// out cycle stack errors and not for any real recovery; also, not
+/// especially efficient.
+#[derive(PartialEq, Eq)]
+crate struct DynDescriptor {
+    type_id: TypeId,
+    debug_string: String,
+}
+
+impl DynDescriptor {
+    crate fn from_key<QC, Q>(_query: &QC, key: &Q::Key) -> DynDescriptor
+    where
+        QC: BaseQueryContext,
+        Q: Query<QC>,
+    {
+        let type_id = TypeId::of::<Q>();
+        let query = Q::default();
+        let debug_string = format!("Query `{:?}` applied to `{:?}`", query, key);
+        DynDescriptor {
+            type_id,
+            debug_string,
+        }
+    }
+}
+
+impl std::fmt::Debug for DynDescriptor {
+    fn fmt(&self, fmt: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(fmt, "{}", self.debug_string)
+    }
+}

src/lib.rs

@@ -1,7 +1,159 @@
-#[cfg(test)]
+#![deny(rust_2018_idioms)]
-mod tests {
+#![feature(in_band_lifetimes)]
-    #[test]
+#![feature(box_patterns)]
-    fn it_works() {
+#![feature(crate_visibility_modifier)]
-        assert_eq!(2 + 2, 4);
+#![feature(nll)]
+#![feature(min_const_fn)]
+#![feature(const_fn)]
+#![feature(const_let)]
+#![feature(try_from)]
+#![feature(macro_at_most_once_rep)]
+#![allow(dead_code)]
+#![allow(unused_imports)]
+
+use derive_new::new;
+use rustc_hash::FxHashMap;
+use std::any::Any;
+use std::cell::RefCell;
+use std::collections::hash_map::Entry;
+use std::fmt::Debug;
+use std::fmt::Display;
+use std::fmt::Write;
+use std::hash::Hash;
+
+pub mod dyn_descriptor;
+pub mod storage;
+
+pub trait BaseQueryContext: Sized {
+    /// A "query descriptor" packages up all the possible queries and a key.
+    /// It is used to store information about (e.g.) the stack.
+    ///
+    /// At runtime, it can be implemented in various ways: a monster enum
+    /// works for a fixed set of queries, but a boxed trait object is good
+    /// for a more open-ended option.
+    type QueryDescriptor: Debug + Eq;
+
+    fn execute_query_implementation<Q>(
+        &self,
+        descriptor: Self::QueryDescriptor,
+        key: &Q::Key,
+    ) -> Q::Value
+    where
+        Q: Query<Self>;
+
+    /// Reports an unexpected cycle attempting to access the query Q with the given key.
+    fn report_unexpected_cycle(&self, descriptor: Self::QueryDescriptor) -> !;
+}
+
+pub trait Query<QC: BaseQueryContext>: Debug + Default + Sized + 'static {
+    type Key: Clone + Debug + Hash + Eq;
+    type Value: Clone + Debug + Hash + Eq;
+    type Storage: QueryStorageOps<QC, Self>;
+
+    fn execute(query: &QC, key: Self::Key) -> Self::Value;
+}
+
+pub trait QueryStorageOps<QC, Q>
+where
+    QC: BaseQueryContext,
+    Q: Query<QC>,
+{
+    fn try_fetch<'q>(
+        &self,
+        query: &'q QC,
+        key: &Q::Key,
+        descriptor: impl FnOnce() -> QC::QueryDescriptor,
+    ) -> Result<Q::Value, CycleDetected>;
+}
+
+#[derive(new)]
+pub struct QueryTable<'me, QC, Q>
+where
+    QC: BaseQueryContext,
+    Q: Query<QC>,
+{
+    pub query: &'me QC,
+    pub storage: &'me Q::Storage,
+    pub descriptor_fn: fn(&QC, &Q::Key) -> QC::QueryDescriptor,
+}
+
+#[derive(Debug)]
+pub enum QueryState<V> {
+    InProgress,
+    Memoized(V),
+}
+
+pub struct CycleDetected;
+
+impl<QC, Q> QueryTable<'me, QC, Q>
+where
+    QC: BaseQueryContext,
+    Q: Query<QC>,
+{
+    pub fn of(&self, key: Q::Key) -> Q::Value {
+        self.storage
+            .try_fetch(self.query, &key, || self.descriptor(&key))
+            .unwrap_or_else(|CycleDetected| {
+                self.query.report_unexpected_cycle(self.descriptor(&key))
+            })
+    }
+
+    fn descriptor(&self, key: &Q::Key) -> QC::QueryDescriptor {
+        (self.descriptor_fn)(self.query, key)
+    }
+}
+
+/// A macro helper for writing the query contexts in traits that helps
+/// you avoid repeating information.
+///
+/// Example:
+///
+/// ```
+/// trait TypeckQueryContext {
+///     query_prototype!(fn <method>() for <type>);
+/// }
+/// ```
+#[macro_export]
+macro_rules! query_prototype {
+    (
+        $(#[$attr:meta])*
+        fn $method_name:ident() for $query_type:ty
+    ) => {
+        $(#[$attr])*
+        fn $method_name(&self) -> $crate::query::QueryTable<'_, Self, $query_type>;
+    }
+}
+
+/// Example:
+///
+/// ```
+/// query_definition! {
+///     QueryName(query: &impl TypeckQueryContext, key: DefId) -> Arc<Vec<DefId>> {
+///         ...       
+///     }
+/// }
+#[macro_export]
+macro_rules! query_definition {
+    (
+        $(#[$attr:meta])*
+        $v:vis $name:ident($query:tt : &impl $query_trait:path, $key:tt : $key_ty:ty) -> $value_ty:ty {
+            $($body:tt)*
+        }
+    ) => {
+        #[derive(Default, Debug)]
+        $v struct $name;
+
+        impl<QC> $crate::query::Query<QC> for $name
+        where
+            QC: $query_trait,
+        {
+            type Key = $key_ty;
+            type Value = $value_ty;
+            type Storage = $crate::query::storage::MemoizedStorage<QC, Self>;
+
+            fn execute($query: &QC, $key: $key_ty) -> $value_ty {
+                $($body)*
+            }
+        }
     }
 }

src/storage.rs

@@ -0,0 +1,74 @@
+use crate::BaseQueryContext;
+use crate::CycleDetected;
+use crate::Query;
+use crate::QueryState;
+use crate::QueryStorageOps;
+use crate::QueryTable;
+use rustc_hash::FxHashMap;
+use std::any::Any;
+use std::cell::RefCell;
+use std::collections::hash_map::Entry;
+use std::fmt::Debug;
+use std::fmt::Display;
+use std::fmt::Write;
+use std::hash::Hash;
+
+// The master implementation that knits together all the queries
+// contains a certain amount of boilerplate. This file aims to
+// reduce that.
+
+crate struct MemoizedStorage<QC, Q>
+where
+    Q: Query<QC>,
+    QC: BaseQueryContext,
+{
+    map: RefCell<FxHashMap<Q::Key, QueryState<Q::Value>>>,
+}
+
+impl<QC, Q> QueryStorageOps<QC, Q> for MemoizedStorage<QC, Q>
+where
+    Q: Query<QC>,
+    QC: BaseQueryContext,
+{
+    fn try_fetch<'q>(
+        &self,
+        query: &'q QC,
+        key: &Q::Key,
+        descriptor: impl FnOnce() -> QC::QueryDescriptor,
+    ) -> Result<Q::Value, CycleDetected> {
+        {
+            let mut map = self.map.borrow_mut();
+            match map.entry(key.clone()) {
+                Entry::Occupied(entry) => {
+                    return match entry.get() {
+                        QueryState::InProgress => Err(CycleDetected),
+                        QueryState::Memoized(value) => Ok(value.clone()),
+                    };
+                }
+                Entry::Vacant(entry) => {
+                    entry.insert(QueryState::InProgress);
+                }
+            }
+        }
+
+        // If we get here, the query is in progress, and we are the
+        // ones tasked with finding its final value.
+        let descriptor = descriptor();
+        let value = query.execute_query_implementation::<Q>(descriptor, key);
+
+        {
+            let mut map = self.map.borrow_mut();
+            let old_value = map.insert(key.clone(), QueryState::Memoized(value.clone()));
+            assert!(
+                match old_value {
+                    Some(QueryState::InProgress) => true,
+                    _ => false,
+                },
+                "expected in-progress state, not {:?}",
+                old_value
+            );
+        }
+
+        Ok(value)
+    }
+}