program-error: Add option to specify solana_program crate

libraries/program-error/derive/src/macro_impl.rs

@@ -1,7 +1,7 @@
 //! The actual token generator for the macro
 
 use {
-    crate::parser::SplProgramErrorArgs,
+    crate::parser::{SolanaProgram, SplProgramErrorArgs},
     proc_macro2::Span,
     quote::quote,
     sha2::{Digest, Sha256},
@@ -36,10 +36,13 @@ pub enum MacroType {
 impl MacroType {
     /// Generates the corresponding tokens based on variant selection
     pub fn generate_tokens(&mut self) -> proc_macro2::TokenStream {
+        let default_solana_program = SolanaProgram::default();
         match self {
-            Self::IntoProgramError { ident } => into_program_error(ident),
-            Self::DecodeError { ident } => decode_error(ident),
-            Self::PrintProgramError { ident, variants } => print_program_error(ident, variants),
+            Self::IntoProgramError { ident } => into_program_error(ident, &default_solana_program),
+            Self::DecodeError { ident } => decode_error(ident, &default_solana_program),
+            Self::PrintProgramError { ident, variants } => {
+                print_program_error(ident, variants, &default_solana_program)
+            }
             Self::SplProgramError { args, item_enum } => spl_program_error(args, item_enum),
         }
     }
@@ -48,59 +51,63 @@ impl MacroType {
 /// Builds the implementation of
 /// `Into<solana_program::program_error::ProgramError>` More specifically,
 /// implements `From<Self> for solana_program::program_error::ProgramError`
-pub fn into_program_error(ident: &Ident) -> proc_macro2::TokenStream {
-    quote! {
-        impl From<#ident> for ::solana_program::program_error::ProgramError {
+pub fn into_program_error(ident: &Ident, import: &SolanaProgram) -> proc_macro2::TokenStream {
+    let this_impl = quote! {
+        impl From<#ident> for #import::program_error::ProgramError {
             fn from(e: #ident) -> Self {
-                ::solana_program::program_error::ProgramError::Custom(e as u32)
+                #import::program_error::ProgramError::Custom(e as u32)
             }
         }
-    }
+    };
+    import.wrap(this_impl)
 }
 
 /// Builds the implementation of `solana_program::decode_error::DecodeError<T>`
-pub fn decode_error(ident: &Ident) -> proc_macro2::TokenStream {
-    quote! {
-        impl<T> ::solana_program::decode_error::DecodeError<T> for #ident {
+pub fn decode_error(ident: &Ident, import: &SolanaProgram) -> proc_macro2::TokenStream {
+    let this_impl = quote! {
+        impl<T> #import::decode_error::DecodeError<T> for #ident {
             fn type_of() -> &'static str {
                 stringify!(#ident)
             }
         }
-    }
+    };
+    import.wrap(this_impl)
 }
 
 /// Builds the implementation of
 /// `solana_program::program_error::PrintProgramError`
 pub fn print_program_error(
     ident: &Ident,
     variants: &Punctuated<Variant, Comma>,
+    import: &SolanaProgram,
 ) -> proc_macro2::TokenStream {
     let ppe_match_arms = variants.iter().map(|variant| {
         let variant_ident = &variant.ident;
         let error_msg = get_error_message(variant)
             .unwrap_or_else(|| String::from("Unknown custom program error"));
         quote! {
             #ident::#variant_ident => {
-                ::solana_program::msg!(#error_msg)
+                #import::msg!(#error_msg)
             }
         }
     });
-    quote! {
-        impl ::solana_program::program_error::PrintProgramError for #ident {
+    let this_impl = quote! {
+        impl #import::program_error::PrintProgramError for #ident {
             fn print<E>(&self)
             where
                 E: 'static
                     + std::error::Error
-                    + ::solana_program::decode_error::DecodeError<E>
-                    + ::solana_program::program_error::PrintProgramError
+                    + #import::decode_error::DecodeError<E>
+                    + #import::program_error::PrintProgramError
                     + num_traits::FromPrimitive,
             {
                 match self {
                     #(#ppe_match_arms),*
                 }
             }
         }
-    }
+    };
+    import.wrap(this_impl)
 }
 
 /// Helper to parse out the string literal from the `#[error(..)]` attribute
@@ -128,9 +135,9 @@ pub fn spl_program_error(
 
     let ident = &item_enum.ident;
     let variants = &item_enum.variants;
-    let into_program_error = into_program_error(ident);
-    let decode_error = decode_error(ident);
-    let print_program_error = print_program_error(ident, variants);
+    let into_program_error = into_program_error(ident, &args.import);
+    let decode_error = decode_error(ident, &args.import);
+    let print_program_error = print_program_error(ident, variants, &args.import);
 
     quote! {
         #[repr(u32)]

libraries/program-error/derive/src/parser.rs

@@ -1,11 +1,12 @@
 //! Token parsing
 
 use {
-    proc_macro2::Ident,
+    proc_macro2::{Ident, Span, TokenStream},
+    quote::quote,
     syn::{
         parse::{Parse, ParseStream},
         token::Comma,
-        LitInt, Token,
+        LitInt, LitStr, Token,
     },
 };
 
@@ -14,51 +15,131 @@ pub struct SplProgramErrorArgs {
     /// Whether to hash the error codes using `solana_program::hash`
     /// or to use the default error code assigned by `num_traits`.
     pub hash_error_code_start: Option<u32>,
+    /// Crate to use for solana_program
+    pub import: SolanaProgram,
+}
+
+/// Struct representing the path to a `solana_program` crate, which may be
+/// renamed or otherwise.
+pub struct SolanaProgram {
+    import: Ident,
+    explicit: bool,
+}
+impl quote::ToTokens for SolanaProgram {
+    fn to_tokens(&self, tokens: &mut TokenStream) {
+        self.import.to_tokens(tokens);
+    }
+}
+impl SolanaProgram {
+    pub fn wrap(&self, output: TokenStream) -> TokenStream {
+        if self.explicit {
+            output
+        } else {
+            anon_const_trick(output)
+        }
+    }
+}
+impl Default for SolanaProgram {
+    fn default() -> Self {
+        Self {
+            import: Ident::new("_solana_program", Span::call_site()),
+            explicit: false,
+        }
+    }
 }
 
 impl Parse for SplProgramErrorArgs {
     fn parse(input: ParseStream) -> syn::Result<Self> {
-        if input.is_empty() {
-            return Ok(Self {
-                hash_error_code_start: None,
-            });
-        }
-        match SplProgramErrorArgParser::parse(input)? {
-            SplProgramErrorArgParser::HashErrorCodes { value, .. } => Ok(Self {
-                hash_error_code_start: Some(value.base10_parse::<u32>()?),
-            }),
+        let mut hash_error_code_start = None;
+        let mut import = None;
+        while !input.is_empty() {
+            match SplProgramErrorArgParser::parse(input)? {
+                SplProgramErrorArgParser::HashErrorCodes { value, .. } => {
+                    hash_error_code_start = Some(value.base10_parse::<u32>()?);
+                }
+                SplProgramErrorArgParser::SolanaProgramCrate { value, .. } => {
+                    import = Some(SolanaProgram {
+                        import: value.parse()?,
+                        explicit: true,
+                    });
+                }
+            }
         }
+        Ok(Self {
+            hash_error_code_start,
+            import: import.unwrap_or(SolanaProgram::default()),
+        })
     }
 }
 
 /// Parser for args to the `#[spl_program_error]` attribute
 /// ie. `#[spl_program_error(hash_error_code_start = 1275525928)]`
 enum SplProgramErrorArgParser {
     HashErrorCodes {
-        _ident: Ident,
         _equals_sign: Token![=],
         value: LitInt,
         _comma: Option<Comma>,
     },
+    SolanaProgramCrate {
+        _equals_sign: Token![=],
+        value: LitStr,
+        _comma: Option<Comma>,
+    },
 }
 
 impl Parse for SplProgramErrorArgParser {
     fn parse(input: ParseStream) -> syn::Result<Self> {
-        let _ident = {
-            let ident = input.parse::<Ident>()?;
-            if ident != "hash_error_code_start" {
-                return Err(input.error("Expected argument 'hash_error_code_start'"));
+        let ident = input.parse::<Ident>()?;
+        match ident.to_string().as_str() {
+            "hash_error_code_start" => {
+                let _equals_sign = input.parse::<Token![=]>()?;
+                let value = input.parse::<LitInt>()?;
+                let _comma: Option<Comma> = input.parse().unwrap_or(None);
+                Ok(Self::HashErrorCodes {
+                    _equals_sign,
+                    value,
+                    _comma,
+                })
+            }
+            "solana_program" => {
+                let _equals_sign = input.parse::<Token![=]>()?;
+                let value = input.parse::<LitStr>()?;
+                let _comma: Option<Comma> = input.parse().unwrap_or(None);
+                Ok(Self::SolanaProgramCrate {
+                    _equals_sign,
+                    value,
+                    _comma,
+                })
             }
-            ident
+            _ => Err(input.error("Expected argument 'hash_error_code_start' or 'solana_program'")),
+        }
+    }
+}
+
+// Within `exp`, you can bring things into scope with `extern crate`.
+//
+// We don't want to assume that `solana_program::` is in scope - the user may
+// have imported it under a different name, or may have imported it in a
+// non-toplevel module (common when putting impls behind a feature gate).
+//
+// Solution: let's just generate `extern crate solana_program as
+// _solana_program` and then refer to `_solana_program` in the derived code.
+// However, macros are not allowed to produce `extern crate` statements at the
+// toplevel.
+//
+// Solution: let's generate `mod _impl_foo` and import solana_program within
+// that.  However, now we lose access to private members of the surrounding
+// module.  This is a problem if, for example, we're deriving for a newtype,
+// where the inner type is defined in the same module, but not exported.
+//
+// Solution: use the anonymous const trick.  For some reason, `extern crate`
+// statements are allowed here, but everything from the surrounding module is in
+// scope.  This trick is taken from serde and num_traits.
+fn anon_const_trick(exp: TokenStream) -> TokenStream {
+    quote! {
+        const _: () = {
+            extern crate solana_program as _solana_program;
+            #exp
         };
-        let _equals_sign = input.parse::<Token![=]>()?;
-        let value = input.parse::<LitInt>()?;
-        let _comma: Option<Comma> = input.parse().unwrap_or(None);
-        Ok(Self::HashErrorCodes {
-            _ident,
-            _equals_sign,
-            value,
-            _comma,
-        })
     }
 }

libraries/program-error/tests/spl.rs

@@ -72,3 +72,20 @@ fn test_library_error_codes() {
         first_error_as_u32 + 3,
     );
 }
+
+/// Example error with solana_program crate set
+#[spl_program_error(solana_program = "solana_program")]
+enum ExampleSolanaProgramCrateError {
+    /// This is a very informative error
+    #[error("This is a very informative error")]
+    VeryInformativeError,
+    /// This is a super important error
+    #[error("This is a super important error")]
+    SuperImportantError,
+}
+
+/// Tests that all macros compile
+#[test]
+fn test_macros_compile_with_solana_program_crate() {
+    let _ = ExampleSolanaProgramCrateError::VeryInformativeError;
+}

libraries/type-length-value/src/error.rs

@@ -3,7 +3,10 @@
 use spl_program_error::*;
 
 /// Errors that may be returned by the Token program.
-#[spl_program_error(hash_error_code_start = 1_202_666_432)]
+#[spl_program_error(
+    hash_error_code_start = 1_202_666_432,
+    solana_program = "solana_program"
+)]
 pub enum TlvError {
     /// Type not found in TLV data
     #[error("Type not found in TLV data")]