From bc11aa1ce7ad70b29fd3de85267231590b265e66 Mon Sep 17 00:00:00 2001
From: Arrowana <Arrowana@users.noreply.github.com>
Date: Wed, 18 Dec 2024 09:34:45 +1100
Subject: [PATCH] feat: declare program allow overriding address

---
 .../program/src/declare_program/mod.rs        | 28 +++++++++++++++++--
 .../programs/declare-program/src/lib.rs       |  5 +++-
 2 files changed, 30 insertions(+), 3 deletions(-)

diff --git a/lang/attribute/program/src/declare_program/mod.rs b/lang/attribute/program/src/declare_program/mod.rs
index 2b96995ee9..c677c5a6e9 100644
--- a/lang/attribute/program/src/declare_program/mod.rs
+++ b/lang/attribute/program/src/declare_program/mod.rs
@@ -4,7 +4,10 @@ mod mods;
 use anchor_lang_idl::{convert::convert_idl, types::Idl};
 use anyhow::anyhow;
 use quote::{quote, ToTokens};
-use syn::parse::{Parse, ParseStream};
+use syn::{
+    parse::{Parse, ParseStream},
+    Expr, LitStr, Token,
+};
 
 use common::gen_docs;
 use mods::{
@@ -18,10 +21,31 @@ pub struct DeclareProgram {
     idl: Idl,
 }
 
+// Custom keyword for program_id
+mod kw {
+    syn::custom_keyword!(address);
+}
+
 impl Parse for DeclareProgram {
     fn parse(input: ParseStream) -> syn::Result<Self> {
         let name = input.parse()?;
-        let idl = get_idl(&name).map_err(|e| syn::Error::new(name.span(), e))?;
+
+        let address = if input.peek(Token![,]) {
+            input.parse::<Token![,]>()?;
+            input.parse::<kw::address>()?;
+            input.parse::<Token![=]>()?;
+
+            // Parse the address expression
+            Some(input.parse::<LitStr>()?)
+        } else {
+            None
+        };
+
+        let mut idl = get_idl(&name).map_err(|e| syn::Error::new(name.span(), e))?;
+
+        if let Some(address) = address {
+            idl.address = address.value();
+        }
         Ok(Self { name, idl })
     }
 }
diff --git a/tests/declare-program/programs/declare-program/src/lib.rs b/tests/declare-program/programs/declare-program/src/lib.rs
index dc6c36f8f6..1b9343f61a 100644
--- a/tests/declare-program/programs/declare-program/src/lib.rs
+++ b/tests/declare-program/programs/declare-program/src/lib.rs
@@ -2,7 +2,10 @@ use anchor_lang::prelude::*;
 
 declare_id!("Dec1areProgram11111111111111111111111111111");
 
-declare_program!(external);
+declare_program!(
+    external,
+    address = "Externa111111111111111111111111111111111111"
+);
 use external::program::External;
 
 // Compilation check for legacy IDL (pre Anchor `0.30`)