1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
extern crate proc_macro;
extern crate quote;
extern crate syn;

use proc_macro::TokenStream;

use quote::quote;
use syn::{parse_macro_input, ItemFn};

/// Generates a discriminant for a given function name within a global namespace.
///
/// This macro is used to create a unique identifier (discriminant) for a function
/// within the Anchor framework. It's typically used for generating unique
/// instruction identifiers in Solana programs.
///
/// Anchor typically requires adds the namespace to the function name to generate
/// a unique identifier. In this macro, we default to the global namespace if no
/// namespace is provided.
///
/// # Arguments
///
/// * `input` - A `TokenStream` containing the function name to generate a discriminant for.
///
/// # Returns
///
/// A `TokenStream` representing an array of bytes, which is the generated discriminant.
///
/// # Example
///
/// ```rust
/// use sol_dev_proc_macros::anchor_discriminant;
/// const DISCRIMINANT: [u8; 8] = anchor_discriminant!(initialize);
/// const DISCRIMINANT_WITH_NAMESPACE: [u8; 8] = anchor_discriminant!(global:initialize);
/// assert_eq!(
///    DISCRIMINANT,
///    [175, 175, 109, 31, 13, 152, 155, 237]
/// );
/// assert_eq!(
///     DISCRIMINANT_WITH_NAMESPACE,
///     DISCRIMINANT
/// );
/// ```
#[proc_macro]
pub fn anchor_discriminant(input: TokenStream) -> TokenStream {
    const NAMESPACE: &str = "global";
    let function_name = input.to_string();
    // If the function does not contain a namespace, we add the global namespace.
    let full_name = if function_name.contains(':') {
        function_name
    } else {
        format!("{}:{}", NAMESPACE, function_name)
    };
    let arr = sol_dev_utils::anchor_discriminant(&full_name);
    let expanded = quote::quote! {
        [#(#arr),*]
    };
    TokenStream::from(expanded)
}

/// Attribute macro for instrumenting functions with compute unit logging.
///
/// This macro wraps the decorated function with additional logging statements
/// that print the function name and the number of compute units used before and after
/// the function execution.
///
/// # Usage
///
/// ```rust,ignore
/// #[compute_fn]
/// fn my_function() {
///     // Function body
/// }
/// ```
///
/// # Effects
///
/// - Adds a log message with the function name at the start of execution.
/// - Logs the number of compute units before and after the function execution.
/// - Adds a closing log message with the function name at the end of execution.
///
/// # Note on Compute Units Used by `compute_fn!`
///
/// ## Testing Results (as of 2024-09-01)
///
///  TOTAL COST OF LOGGING: 445 - 36 = 409
///  EXTRA COST INSIDE THE FUNCTION: 101
///
///  EMPTY_WITH_LOG where nothing happens inside the log.
///  TOTAL COMPUTE UNITS USED: 445
///  INNER LOG COST: 101
///
///  "Program EMPTY_WITH_LOG invoke [1]",
///  "Program log: Program log: process_instruction {{",
///  "Program consumption: 199762 units remaining",
///  "Program consumption: 199661 units remaining", // 199762 - 199661 = 101
///  "Program log: Program log: }} // process_instruction",
///  "Program EMPTY_WITH_LOG consumed 445 of 200000 compute units",
///  "Program EMPTY_WITH_LOG success"
///
///  EMPTY where nothing happens at all.
///  TOTAL COMPUTE UNITS USED: 36
///  
///  "Program EMPTY invoke [1]",
///  "Program EMPTY consumed 36 of 200000 compute units",
///  "Program EMPTY success"
///  
///  Total extra compute units used per `compute_fn!` call: 409 CU
///  For more details, see:
///  - https://github.com/anza-xyz/agave/blob/d88050cda335f87e872eddbdf8506bc063f039d3/programs/bpf_loader/src/syscalls/logging.rs#L70
///  - https://github.com/anza-xyz/agave/blob/d88050cda335f87e872eddbdf8506bc063f039d3/program-runtime/src/compute_budget.rs#L150
#[proc_macro_attribute]
pub fn compute_fn(_attr: TokenStream, item: TokenStream) -> TokenStream {
    let mut input = parse_macro_input!(item as ItemFn);
    let fn_name = &input.sig.ident;
    let block = &input.block;

    input.block = syn::parse_quote!({
        ::solana_program::msg!(concat!(stringify!(#fn_name), " {{"));
        ::solana_program::log::sol_log_compute_units();

        let __result = (|| #block)();

        ::solana_program::log::sol_log_compute_units();
        ::solana_program::msg!(concat!("}} // ", stringify!(#fn_name)));

        __result
    });

    quote!(#input).into()
}