lang: Fix cpi feature instructions not accounting for discriminator…

… overrides
coral-xyz · Nov 20, 2024 · f3c3ea3 · f3c3ea3
1 parent 8ec7ec8
commit f3c3ea3
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 15 deletions.
diff --git a/lang/syn/src/codegen/program/common.rs b/lang/syn/src/codegen/program/common.rs
@@ -25,10 +25,7 @@ pub fn gen_discriminator(namespace: &str, name: impl ToString) -> proc_macro2::T
 
 pub fn generate_ix_variant(name: String, args: &[IxArg]) -> proc_macro2::TokenStream {
     let ix_arg_names: Vec<&syn::Ident> = args.iter().map(|arg| &arg.name).collect();
-    let ix_name_camel: proc_macro2::TokenStream = {
-        let n = name.to_camel_case();
-        n.parse().unwrap()
-    };
+    let ix_name_camel = generate_ix_variant_name(name);
 
     if args.is_empty() {
         quote! {
@@ -42,3 +39,8 @@ pub fn generate_ix_variant(name: String, args: &[IxArg]) -> proc_macro2::TokenSt
         }
     }
 }
+
+pub fn generate_ix_variant_name(name: String) -> proc_macro2::TokenStream {
+    let n = name.to_camel_case();
+    n.parse().unwrap()
+}
diff --git a/lang/syn/src/codegen/program/cpi.rs b/lang/syn/src/codegen/program/cpi.rs
@@ -1,6 +1,4 @@
-use crate::codegen::program::common::{
-    gen_discriminator, generate_ix_variant, SIGHASH_GLOBAL_NAMESPACE,
-};
+use crate::codegen::program::common::{generate_ix_variant, generate_ix_variant_name};
 use crate::Program;
 use heck::SnakeCase;
 use quote::{quote, ToTokens};
@@ -14,10 +12,13 @@ pub fn generate(program: &Program) -> proc_macro2::TokenStream {
             let accounts_ident: proc_macro2::TokenStream = format!("crate::cpi::accounts::{}", &ix.anchor_ident.to_string()).parse().unwrap();
             let cpi_method = {
                 let name = &ix.raw_method.sig.ident;
-                let ix_variant = generate_ix_variant(name.to_string(), &ix.args);
+                let ix_variant = generate_ix_variant(&name, &ix.args);
                 let method_name = &ix.ident;
                 let args: Vec<&syn::PatType> = ix.args.iter().map(|arg| &arg.raw_arg).collect();
-                let discriminator = gen_discriminator(SIGHASH_GLOBAL_NAMESPACE, name);
+                let discriminator = {
+                    let name = generate_ix_variant_name(name.to_string());
+                    quote! { <instruction::#name as anchor_lang::Discriminator>::DISCRIMINATOR }
+                };
                 let ret_type = &ix.returns.ty.to_token_stream();
                 let ix_cfgs = &ix.cfgs;
                 let (method_ret, maybe_return) = match ret_type.to_string().as_str() {

diff --git a/lang/syn/src/codegen/program/handlers.rs b/lang/syn/src/codegen/program/handlers.rs
@@ -1,7 +1,6 @@
 use crate::codegen::program::common::*;
 use crate::program_codegen::idl::idl_accounts_and_functions;
 use crate::Program;
-use heck::CamelCase;
 use quote::{quote, ToTokens};
 
 // Generate non-inlined wrappers for each instruction handler, since Solana's
@@ -191,11 +190,6 @@ pub fn generate(program: &Program) -> proc_macro2::TokenStream {
     }
 }
 
-fn generate_ix_variant_name(name: String) -> proc_macro2::TokenStream {
-    let n = name.to_camel_case();
-    n.parse().unwrap()
-}
-
 /// Generate the event module based on whether the `event-cpi` feature is enabled.
 fn generate_event_cpi_mod() -> proc_macro2::TokenStream {
     #[cfg(feature = "event-cpi")]