Add CU instrumentation option

cli/src/checks.rs

@@ -186,3 +186,45 @@ in `{manifest_path:?}`."#
 
     Ok(())
 }
+
+/// Check whether the `instrument-compute-units` feature is being used correctly.
+///
+/// **Note:** The check expects the current directory to be a program directory.
+pub fn check_instrument_compute_units_build_feature() -> Result<()> {
+    let manifest_path = Path::new("Cargo.toml").canonicalize()?;
+    let manifest = Manifest::from_path(&manifest_path)?;
+
+    // Check whether the manifest has `instrument-compute-units` feature
+    let has_instrument_compute_units_feature = manifest
+        .features
+        .iter()
+        .any(|(feature, _)| feature == "instrument-compute-units");
+    if !has_instrument_compute_units_feature {
+        return Err(anyhow!(
+            r#"`instrument-compute-units` feature is missing. To solve, add
+
+[features]
+instrument-compute-units = ["anchor-lang/instrument-compute-units"]
+
+in `{manifest_path:?}`."#
+        ));
+    }
+
+    // Check if `instrument-compute-units` is enabled by default
+    manifest
+        .dependencies
+        .iter()
+        .filter(|(_, dep)| dep.req_features().contains(&"instrument-compute-units".into()))
+        .for_each(|(name, _)| {
+            eprintln!(
+                "WARNING: `instrument-compute-units` feature of crate `{name}` is enabled by default. \
+                    This is not the intended usage.\n\n\t\
+                    To solve, do not enable the `instrument-compute-units` feature and include crates that have \
+                    `instrument-compute-units` feature in the `instrument-compute-units feature list:\n\n\t\
+                    [features]\n\t\
+                    instrument-compute-units = [\"{name}/instrument-compute-units\", ...]\n"
+            )
+        });
+
+    Ok(())
+}

cli/src/lib.rs

@@ -1,3 +1,4 @@
+use crate::checks::check_instrument_compute_units_build_feature;
 use crate::config::{
     get_default_ledger_path, BootstrapMode, BuildConfig, Config, ConfigOverride, Manifest,
     PackageManager, ProgramArch, ProgramDeployment, ProgramWorkspace, ScriptsConfig, TestValidator,
@@ -134,6 +135,9 @@ pub enum Command {
         /// Architecture to use when building the program
         #[clap(value_enum, long, default_value = "sbf")]
         arch: ProgramArch,
+        /// Add log messages to each line to instrument CU usage
+        #[clap(long)]
+        instrument_compute_units: bool,
     },
     /// Expands macros (wrapper around cargo expand)
     ///
@@ -203,6 +207,9 @@ pub enum Command {
         /// to be able to check the transactions.
         #[clap(long)]
         detach: bool,
+        /// Add log messages to each line to instrument CU usage
+        #[clap(long)]
+        instrument_compute_units: bool,
         /// Run the test suites under the specified path
         #[clap(long)]
         run: Vec<String>,
@@ -750,6 +757,7 @@ fn process_command(opts: Opts) -> Result<()> {
             no_idl,
             idl,
             idl_ts,
+            instrument_compute_units,
             verifiable,
             program_name,
             solana_version,
@@ -765,6 +773,7 @@ fn process_command(opts: Opts) -> Result<()> {
             no_idl,
             idl,
             idl_ts,
+            instrument_compute_units,
             verifiable,
             skip_lint,
             program_name,
@@ -832,6 +841,7 @@ fn process_command(opts: Opts) -> Result<()> {
             skip_local_validator,
             skip_build,
             no_idl,
+            instrument_compute_units,
             detach,
             run,
             args,
@@ -847,6 +857,7 @@ fn process_command(opts: Opts) -> Result<()> {
             skip_build,
             skip_lint,
             no_idl,
+            instrument_compute_units,
             detach,
             run,
             args,
@@ -1266,6 +1277,7 @@ pub fn build(
     no_idl: bool,
     idl: Option<String>,
     idl_ts: Option<String>,
+    instrument_compute_units: bool,
     verifiable: bool,
     skip_lint: bool,
     program_name: Option<String>,
@@ -1327,6 +1339,7 @@ pub fn build(
             no_idl,
             idl_out,
             idl_ts_out,
+            instrument_compute_units,
             &build_config,
             stdout,
             stderr,
@@ -1343,6 +1356,7 @@ pub fn build(
             no_idl,
             idl_out,
             idl_ts_out,
+            instrument_compute_units,
             &build_config,
             stdout,
             stderr,
@@ -1359,6 +1373,7 @@ pub fn build(
             no_idl,
             idl_out,
             idl_ts_out,
+            instrument_compute_units,
             &build_config,
             stdout,
             stderr,
@@ -1382,6 +1397,7 @@ fn build_all(
     no_idl: bool,
     idl_out: Option<PathBuf>,
     idl_ts_out: Option<PathBuf>,
+    instrument_compute_units: bool,
     build_config: &BuildConfig,
     stdout: Option<File>, // Used for the package registry server.
     stderr: Option<File>, // Used for the package registry server.
@@ -1402,6 +1418,7 @@ fn build_all(
                     no_idl,
                     idl_out.clone(),
                     idl_ts_out.clone(),
+                    instrument_compute_units,
                     build_config,
                     stdout.as_ref().map(|f| f.try_clone()).transpose()?,
                     stderr.as_ref().map(|f| f.try_clone()).transpose()?,
@@ -1427,6 +1444,7 @@ fn build_rust_cwd(
     no_idl: bool,
     idl_out: Option<PathBuf>,
     idl_ts_out: Option<PathBuf>,
+    instrument_compute_units: bool,
     build_config: &BuildConfig,
     stdout: Option<File>,
     stderr: Option<File>,
@@ -1442,7 +1460,15 @@ fn build_rust_cwd(
     };
     match build_config.verifiable {
         false => _build_rust_cwd(
-            cfg, no_idl, idl_out, idl_ts_out, skip_lint, no_docs, arch, cargo_args,
+            cfg,
+            no_idl,
+            idl_out,
+            idl_ts_out,
+            instrument_compute_units,
+            skip_lint,
+            no_docs,
+            arch,
+            cargo_args,
         ),
         true => build_cwd_verifiable(
             cfg,
@@ -1797,11 +1823,18 @@ fn _build_rust_cwd(
     no_idl: bool,
     idl_out: Option<PathBuf>,
     idl_ts_out: Option<PathBuf>,
+    instrument_compute_units: bool,
     skip_lint: bool,
     no_docs: bool,
     arch: &ProgramArch,
-    cargo_args: Vec<String>,
+    mut cargo_args: Vec<String>,
 ) -> Result<()> {
+    if instrument_compute_units {
+        check_instrument_compute_units_build_feature()?;
+        cargo_args.insert(0, "--features".to_string());
+        cargo_args.insert(1, "instrument-compute-units".to_string());
+    }
+
     let exit = std::process::Command::new("cargo")
         .arg(arch.build_subcommand())
         .args(cargo_args.clone())
@@ -2912,6 +2945,7 @@ fn test(
     skip_build: bool,
     skip_lint: bool,
     no_idl: bool,
+    instrument_compute_units: bool,
     detach: bool,
     tests_to_run: Vec<String>,
     extra_args: Vec<String>,
@@ -2936,6 +2970,7 @@ fn test(
                 no_idl,
                 None,
                 None,
+                instrument_compute_units,
                 false,
                 skip_lint,
                 program_name.clone(),
@@ -4186,6 +4221,7 @@ fn localnet(
                 None,
                 None,
                 false,
+                false,
                 skip_lint,
                 None,
                 None,

cli/src/rust_template.rs

@@ -203,6 +203,7 @@ no-entrypoint = []
 no-idl = []
 no-log-ix-name = []
 idl-build = ["anchor-lang/idl-build"]
+instrument-compute-units = ["anchor-lang/instrument-compute-units"]
 anchor-debug = []
 custom-heap = []
 custom-panic = []

lang/Cargo.toml

@@ -35,6 +35,7 @@ idl-build = [
     "anchor-lang-idl/build",
 ]
 init-if-needed = ["anchor-derive-accounts/init-if-needed"]
+instrument-compute-units = ["anchor-attribute-program/instrument-compute-units"]
 interface-instructions = ["anchor-attribute-program/interface-instructions"]
 lazy-account = ["anchor-attribute-account/lazy-account", "anchor-derive-serde/lazy-account"]
 

lang/attribute/program/Cargo.toml

@@ -13,6 +13,7 @@ proc-macro = true
 [features]
 anchor-debug = ["anchor-syn/anchor-debug"]
 idl-build = ["anchor-syn/idl-build"]
+instrument-compute-units = ["anchor-syn/instrument-compute-units"]
 interface-instructions = ["anchor-syn/interface-instructions"]
 
 [dependencies]

lang/src/lib.rs

@@ -126,6 +126,15 @@ pub mod solana_program {
             #[cfg(not(target_os = "solana"))]
             core::hint::black_box(data);
         }
+
+        // FIXME: `solana_msg` does not currently expose `sol_log_compute_units`
+        #[inline]
+        pub fn sol_log_compute_units() {
+            #[cfg(target_os = "solana")]
+            unsafe {
+                solana_define_syscall::definitions::sol_log_compute_units_();
+            }
+        }
     }
     pub mod sysvar {
         pub use solana_sysvar_id::{declare_deprecated_sysvar_id, declare_sysvar_id, SysvarId};

lang/syn/Cargo.toml

@@ -18,6 +18,7 @@ event-cpi = []
 hash = []
 idl-build = ["cargo_toml"]
 init-if-needed = []
+instrument-compute-units = []
 interface-instructions = []
 
 [dependencies]

lang/syn/src/codegen/program/instrument.rs

@@ -0,0 +1,37 @@
+//! Instruments each line of user code with `sol_log_compute_units`
+
+use quote::ToTokens;
+use syn::{ItemMod, Stmt};
+
+pub fn generate(mut module: ItemMod) -> ItemMod {
+    let Some((_, content)) = module.content.as_mut() else {
+        return module;
+    };
+    // Insert a message and CU log after each statement
+    for insn in content.iter_mut().filter_map(|item| match item {
+        syn::Item::Fn(item_fn) => Some(item_fn),
+        _ => None,
+    }) {
+        let stmts = std::mem::take(&mut insn.block.stmts);
+        let interpsersed: Vec<Stmt> = stmts
+            .into_iter()
+            .flat_map(|stmt| match stmt {
+                // Last expression of a block - don't add code after
+                Stmt::Expr(_) => vec![stmt],
+                _ => {
+                    let as_str = stmt.to_token_stream().to_string();
+                    let log: Stmt = syn::parse_str(&format!(
+                        r#"{{
+    ::anchor_lang::solana_program::log::sol_log(concat!("Executing `", stringify!({as_str}), "`"));
+    ::anchor_lang::solana_program::log::sol_log_compute_units();
+                    }};"#,
+                    ))
+                    .unwrap();
+                    vec![stmt, log]
+                }
+            })
+            .collect();
+        insn.block.stmts = interpsersed;
+    }
+    module
+}

lang/syn/src/codegen/program/mod.rs

@@ -9,14 +9,20 @@ mod entry;
 mod handlers;
 mod idl;
 mod instruction;
+#[cfg(feature = "instrument-compute-units")]
+mod instrument;
 
 pub fn generate(program: &Program) -> proc_macro2::TokenStream {
     let mod_name = &program.name;
 
     let entry = entry::generate(program);
     let dispatch = dispatch::generate(program);
     let handlers = handlers::generate(program);
+    #[cfg(not(feature = "instrument-compute-units"))]
     let user_defined_program = &program.program_mod;
+    #[cfg(feature = "instrument-compute-units")]
+    let user_defined_program = &instrument::generate(program.program_mod.clone());
+
     let instruction = instruction::generate(program);
     let cpi = cpi::generate(program);
     let accounts = accounts::generate(program);