diff --git a/components/salsa-macros/src/database_storage.rs b/components/salsa-macros/src/database_storage.rs index d167b2d9..cf00c379 100644 --- a/components/salsa-macros/src/database_storage.rs +++ b/components/salsa-macros/src/database_storage.rs @@ -128,11 +128,11 @@ pub(crate) fn database_storage(input: TokenStream) -> TokenStream { // let mut for_each_ops = proc_macro2::TokenStream::new(); - for (query_group, Query { query_name, .. }) in each_query() { + for query_group in &query_groups { let group_storage = query_group.group_storage(); for_each_ops.extend(quote! { let storage: &#group_storage<#database_name> = ::salsa::plumbing::GetQueryGroupStorage::from(self); - op(&storage.#query_name); + storage.for_each_query(self, &mut op); }); } output.extend(quote! { diff --git a/components/salsa-macros/src/query_group.rs b/components/salsa-macros/src/query_group.rs index 79621a6a..fff0ec32 100644 --- a/components/salsa-macros/src/query_group.rs +++ b/components/salsa-macros/src/query_group.rs @@ -213,7 +213,7 @@ pub(crate) fn query_group(args: TokenStream, input: TokenStream) -> TokenStream }); // Emit the query types. - for query in queries { + for query in &queries { let qt = &query.query_type; let storage = Ident::new( match query.storage { @@ -255,7 +255,7 @@ pub(crate) fn query_group(args: TokenStream, input: TokenStream) -> TokenStream }; let invoke = match &query.invoke { Some(i) => i.into_token_stream(), - None => query.fn_name.into_token_stream(), + None => query.fn_name.clone().into_token_stream(), }; output.extend(quote_spanned! {span=> impl salsa::plumbing::QueryFunction for #qt @@ -296,14 +296,41 @@ pub(crate) fn query_group(args: TokenStream, input: TokenStream) -> TokenStream } }); + let mut for_each_ops = proc_macro2::TokenStream::new(); + for Query { fn_name, .. } in &queries { + for_each_ops.extend(quote! { + op(&self.#fn_name); + }); + } + // Emit query group storage struct output.extend(quote! { #[derive(Default)] #trait_vis struct #group_storage { #storage_fields } + + impl #group_storage + where + DB__: #trait_name, + DB__: ::salsa::plumbing::GetQueryGroupStorage<#group_storage>, + { + #trait_vis fn for_each_query( + &self, + db: &DB__, + mut op: &mut dyn FnMut(&dyn ::salsa::plumbing::QueryStorageMassOps), + ) { + #for_each_ops + } + } }); + if std::env::var("SALSA_DUMP").is_ok() { + println!("~~~ query_group"); + println!("{}", output.to_string()); + println!("~~~ query_group"); + } + output.into() } diff --git a/src/plumbing.rs b/src/plumbing.rs index fa8480da..ad671323 100644 --- a/src/plumbing.rs +++ b/src/plumbing.rs @@ -87,6 +87,7 @@ pub trait GetQueryGroupStorage: Database { pub trait QueryStorageOps: Default where + Self: QueryStorageMassOps, DB: Database, Q: Query, {