Adapt for new cycle handling changing in Salsa

crates/query-group-macro/src/lib.rs

@@ -10,10 +10,13 @@ use queries::{
     Queries, SetterKind, TrackedQuery, Transparent,
 };
 use quote::{ToTokens, format_ident, quote};
+use syn::parse::{Parse, ParseStream};
+use syn::punctuated::Punctuated;
 use syn::spanned::Spanned;
 use syn::visit_mut::VisitMut;
 use syn::{
-    Attribute, FnArg, ItemTrait, Path, TraitItem, TraitItemFn, parse_quote, parse_quote_spanned,
+    Attribute, FnArg, ItemTrait, Path, Token, TraitItem, TraitItemFn, parse_quote,
+    parse_quote_spanned,
 };
 
 mod queries;
@@ -106,6 +109,66 @@ enum QueryKind {
     Interned,
 }
 
+#[derive(Default, Debug, Clone)]
+struct Cycle {
+    cycle_fn: Option<(syn::Ident, Path)>,
+    cycle_initial: Option<(syn::Ident, Path)>,
+    cycle_result: Option<(syn::Ident, Path)>,
+}
+
+impl Parse for Cycle {
+    fn parse(input: ParseStream<'_>) -> syn::Result<Self> {
+        let options = Punctuated::<Option, Token![,]>::parse_terminated(input)?;
+        let mut cycle_fn = None;
+        let mut cycle_initial = None;
+        let mut cycle_result = None;
+        for option in options {
+            let name = option.name.to_string();
+            match &*name {
+                "cycle_fn" => {
+                    if cycle_fn.is_some() {
+                        return Err(syn::Error::new_spanned(&option.name, "duplicate option"));
+                    }
+                    cycle_fn = Some((option.name, option.value));
+                }
+                "cycle_initial" => {
+                    if cycle_initial.is_some() {
+                        return Err(syn::Error::new_spanned(&option.name, "duplicate option"));
+                    }
+                    cycle_initial = Some((option.name, option.value));
+                }
+                "cycle_result" => {
+                    if cycle_result.is_some() {
+                        return Err(syn::Error::new_spanned(&option.name, "duplicate option"));
+                    }
+                    cycle_result = Some((option.name, option.value));
+                }
+                _ => {
+                    return Err(syn::Error::new_spanned(
+                        &option.name,
+                        "unknown cycle option. Accepted values: `cycle_result`, `cycle_fn`, `cycle_initial`",
+                    ));
+                }
+            }
+        }
+        return Ok(Self { cycle_fn, cycle_initial, cycle_result });
+
+        struct Option {
+            name: syn::Ident,
+            value: Path,
+        }
+
+        impl Parse for Option {
+            fn parse(input: ParseStream) -> syn::Result<Self> {
+                let name = input.parse()?;
+                input.parse::<Token![=]>()?;
+                let value = input.parse()?;
+                Ok(Self { name, value })
+            }
+        }
+    }
+}
+
 pub(crate) fn query_group_impl(
     _args: proc_macro::TokenStream,
     input: proc_macro::TokenStream,
@@ -155,8 +218,8 @@ pub(crate) fn query_group_impl(
             for SalsaAttr { name, tts, span } in salsa_attrs {
                 match name.as_str() {
                     "cycle" => {
-                        let path = syn::parse::<Parenthesized<Path>>(tts)?;
-                        cycle = Some(path.0.clone())
+                        let c = syn::parse::<Parenthesized<Cycle>>(tts)?;
+                        cycle = Some(c.0);
                     }
                     "input" => {
                         if !pat_and_tys.is_empty() {
@@ -415,7 +478,7 @@ impl<T> syn::parse::Parse for Parenthesized<T>
 where
     T: syn::parse::Parse,
 {
-    fn parse(input: syn::parse::ParseStream<'_>) -> syn::Result<Self> {
+    fn parse(input: ParseStream<'_>) -> syn::Result<Self> {
         let content;
         syn::parenthesized!(content in input);
         content.parse::<T>().map(Parenthesized)

crates/query-group-macro/src/queries.rs

@@ -3,13 +3,15 @@
 use quote::{ToTokens, format_ident, quote, quote_spanned};
 use syn::{FnArg, Ident, PatType, Path, Receiver, ReturnType, Type, parse_quote, spanned::Spanned};
 
+use crate::Cycle;
+
 pub(crate) struct TrackedQuery {
     pub(crate) trait_name: Ident,
     pub(crate) signature: syn::Signature,
     pub(crate) pat_and_tys: Vec<PatType>,
     pub(crate) invoke: Option<Path>,
     pub(crate) default: Option<syn::Block>,
-    pub(crate) cycle: Option<Path>,
+    pub(crate) cycle: Option<Cycle>,
     pub(crate) lru: Option<u32>,
     pub(crate) generated_struct: Option<GeneratedInputStruct>,
 }
@@ -34,12 +36,20 @@ impl ToTokens for TrackedQuery {
         let fn_ident = &sig.ident;
         let shim: Ident = format_ident!("{}_shim", fn_ident);
 
-        let annotation = match (self.cycle.clone(), self.lru) {
-            (Some(cycle), Some(lru)) => quote!(#[salsa::tracked(lru = #lru, recovery_fn = #cycle)]),
-            (Some(cycle), None) => quote!(#[salsa::tracked(recovery_fn = #cycle)]),
-            (None, Some(lru)) => quote!(#[salsa::tracked(lru = #lru)]),
-            (None, None) => quote!(#[salsa::tracked]),
-        };
+        let options = self
+            .cycle
+            .as_ref()
+            .map(|Cycle { cycle_fn, cycle_initial, cycle_result }| {
+                let cycle_fn = cycle_fn.as_ref().map(|(ident, path)| quote!(#ident=#path));
+                let cycle_initial =
+                    cycle_initial.as_ref().map(|(ident, path)| quote!(#ident=#path));
+                let cycle_result = cycle_result.as_ref().map(|(ident, path)| quote!(#ident=#path));
+                let options = cycle_fn.into_iter().chain(cycle_initial).chain(cycle_result);
+                quote!(#(#options),*)
+            })
+            .into_iter()
+            .chain(self.lru.map(|lru| quote!(lru = #lru)));
+        let annotation = quote!(#[salsa::tracked( #(#options),* )]);
 
         let pat_and_tys = &self.pat_and_tys;
         let params = self