feat: add active heap management (Lightprotocol#393)

ananas-block · web-flow · commit 1ab435f80b67 · 2024-01-22T06:57:27.000Z
- add heap-neutral proc macro which frees all heap memory allocated during the use of a method at the end of the method
- the macro has a check that it can only be used with methods which do not modify its struct
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/macros/light/src/lib.rs b/macros/light/src/lib.rs
@@ -1,5 +1,6 @@
 use proc_macro::TokenStream;
-use syn::{parse_macro_input, ItemStruct};
+use quote::quote;
+use syn::{parse_macro_input, parse_quote, FnArg, ItemFn, ItemStruct, Receiver, Signature};
 
 mod expand;
 
@@ -23,3 +24,48 @@ pub fn light_verifier_accounts(attr: TokenStream, item: TokenStream) -> TokenStr
         .unwrap_or_else(|err| err.to_compile_error())
         .into()
 }
+
+#[proc_macro_attribute]
+pub fn heap_neutral(_: TokenStream, input: TokenStream) -> TokenStream {
+    let mut function = parse_macro_input!(input as ItemFn);
+
+    // Check if the function signature uses `&self` and not `&mut self`
+    if !is_immutable_self(&function.sig) {
+        return syn::Error::new_spanned(
+            &function.sig,
+            "This macro requires the function to use `&self` and not `&mut self`",
+        )
+        .to_compile_error()
+        .into();
+    }
+
+    // Insert memory management code at the beginning of the function
+    let init_code: syn::Stmt = parse_quote! {
+        #[cfg(target_os = "solana")]
+        let pos = custom_heap::get_heap_pos();
+    };
+    function.block.stmts.insert(0, init_code);
+
+    // Insert memory management code at the end of the function
+    let cleanup_code: syn::Stmt = parse_quote! {
+        #[cfg(target_os = "solana")]
+        custom_heap::free_heap(pos);
+    };
+    let len = function.block.stmts.len();
+    function.block.stmts.insert(len - 1, cleanup_code);
+
+    TokenStream::from(quote! { #function })
+}
+
+fn is_immutable_self(signature: &Signature) -> bool {
+    signature.inputs.iter().any(|arg| {
+        matches!(
+            arg,
+            FnArg::Receiver(Receiver {
+                reference: Some(_),
+                mutability: None,
+                ..
+            })
+        )
+    })
+}
diff --git a/programs/psp10in2out/Cargo.toml b/programs/psp10in2out/Cargo.toml
@@ -14,7 +14,9 @@ no-entrypoint = []
 no-idl = []
 no-log-ix-name = []
 cpi = ["no-entrypoint"]
-default = []
+default = ["custom-heap", "memory-test"]
+custom-heap = []
+memory-test = []
 
 [dependencies]
 anchor-lang = "0.28.0"
@@ -26,7 +28,7 @@ solana-security-txt = "1.1.0"
 groth16-solana = { git= "https://github.com/Lightprotocol/groth16-solana", branch="master"}
 
 light-macros = { version = "0.3.1", path = "../../macros/light" }
-light-verifier-sdk = { version = "0.3.1", path = "../../verifier-sdk" }
+light-verifier-sdk = { version = "0.3.1", path = "../../verifier-sdk", features = ["custom-heap", "mem-profiling"]  }
 
 # TODO: Remove once https://github.com/solana-labs/solana/issues/33504 is resolved.
 ahash = "=0.8.6"
diff --git a/programs/psp10in2out/src/lib.rs b/programs/psp10in2out/src/lib.rs
@@ -136,7 +136,17 @@ pub mod light_psp10in2out {
         };
         let mut tx =
             Transaction::<0, 2, 10, 17, LightInstructionSecond<'info, 0, 2, 10>>::new(input);
-        tx.transact()
+        tx.transact()?;
+
+        #[cfg(all(feature = "memory-test", target_os = "solana"))]
+        assert!(
+            light_verifier_sdk::light_transaction::custom_heap::log_total_heap("memory_check")
+                < 7000u64,
+            "memory degression detected {} {}",
+            light_verifier_sdk::light_transaction::custom_heap::log_total_heap("memory_check"),
+            7000u64
+        );
+        Ok(())
     }
 
     /// Close the verifier state to reclaim rent in case the proofdata is wrong and does not verify.
diff --git a/programs/psp2in2out-storage/Cargo.toml b/programs/psp2in2out-storage/Cargo.toml
@@ -14,7 +14,9 @@ no-entrypoint = []
 no-idl = []
 no-log-ix-name = []
 cpi = ["no-entrypoint"]
-default = []
+default = ["custom-heap", "memory-test"]
+custom-heap = []
+memory-test = []
 
 [dependencies]
 anchor-lang = { version = "0.28.0", features = ["init-if-needed"] }
diff --git a/programs/psp2in2out-storage/src/lib.rs b/programs/psp2in2out-storage/src/lib.rs
@@ -115,7 +115,17 @@ pub mod light_psp2in2out_storage {
         };
         let mut transaction = Transaction::<0, 2, 2, 9, LightInstructionSecond<'info>>::new(input);
 
-        transaction.transact()
+        transaction.transact()?;
+
+        #[cfg(all(feature = "memory-test", target_os = "solana"))]
+        assert!(
+            light_verifier_sdk::light_transaction::custom_heap::log_total_heap("memory_check")
+                < 5000u64,
+            "memory degression detected {} {}",
+            light_verifier_sdk::light_transaction::custom_heap::log_total_heap("memory_check"),
+            5000u64
+        );
+        Ok(())
     }
 }
 
diff --git a/programs/psp2in2out/Cargo.toml b/programs/psp2in2out/Cargo.toml
@@ -10,11 +10,13 @@ crate-type = ["cdylib", "lib"]
 name = "light_psp2in2out"
 
 [features]
+default = ["custom-heap", "memory-test"]
+custom-heap = []
+memory-test = []
 no-entrypoint = []
 no-idl = []
 no-log-ix-name = []
 cpi = ["no-entrypoint"]
-default = []
 
 [dependencies]
 anchor-lang = "0.28.0"
@@ -25,7 +27,7 @@ solana-security-txt = "1.1.0"
 # Light Deps
 groth16-solana = { git= "https://github.com/Lightprotocol/groth16-solana", branch="master"}
 light-macros = { version = "0.3.1", path = "../../macros/light" }
-light-verifier-sdk = { version = "0.3.1", path = "../../verifier-sdk" }
+light-verifier-sdk = { version = "0.3.1", path = "../../verifier-sdk" , features = ["custom-heap", "mem-profiling"] }
 
 # TODO: Remove once https://github.com/solana-labs/solana/issues/33504 is resolved.
 ahash = "=0.8.6"
diff --git a/programs/psp2in2out/src/lib.rs b/programs/psp2in2out/src/lib.rs
@@ -3,9 +3,9 @@ use light_macros::light_verifier_accounts;
 use light_verifier_sdk::light_transaction::{
     Amounts, ProofCompressed, Transaction, TransactionInput,
 };
+use verifying_key::VERIFYINGKEY_TRANSACTION_MASP2_MAIN;
 
 pub mod verifying_key;
-use verifying_key::VERIFYINGKEY_TRANSACTION_MASP2_MAIN;
 
 declare_id!("J1RRetZ4ujphU75LP8RadjXMf3sA12yC2R44CF7PmU7i");
 
@@ -39,6 +39,7 @@ pub mod light_psp2in2out {
             InstructionDataCompressedTransferFirst::try_deserialize_unchecked(
                 &mut [vec![0u8; 8], inputs].concat().as_slice(),
             )?;
+
         let len_missing_bytes = 256 - inputs.encrypted_utxos.len();
         let mut enc_utxos = inputs.encrypted_utxos;
         enc_utxos.append(&mut vec![0u8; len_missing_bytes]);
@@ -68,7 +69,17 @@ pub mod light_psp2in2out {
         };
         let mut transaction = Transaction::<0, 2, 2, 9, LightInstruction<'info>>::new(input);
 
-        transaction.transact()
+        transaction.transact()?;
+
+        #[cfg(all(feature = "memory-test", target_os = "solana"))]
+        assert!(
+            light_verifier_sdk::light_transaction::custom_heap::log_total_heap("memory_check")
+                < 6000u64,
+            "memory degression detected {} {}",
+            light_verifier_sdk::light_transaction::custom_heap::log_total_heap("memory_check"),
+            6000u64
+        );
+        Ok(())
     }
 }
 
diff --git a/programs/psp4in4out-app-storage/Cargo.toml b/programs/psp4in4out-app-storage/Cargo.toml
@@ -14,7 +14,9 @@ no-entrypoint = []
 no-idl = []
 no-log-ix-name = []
 cpi = ["no-entrypoint"]
-default = []
+default = ["custom-heap", "memory-test"]
+custom-heap = []
+memory-test = []
 
 [dependencies]
 anchor-lang = "0.28.0"
diff --git a/programs/psp4in4out-app-storage/src/lib.rs b/programs/psp4in4out-app-storage/src/lib.rs
@@ -92,7 +92,17 @@ pub mod light_psp4in4out_app_storage {
         };
         let mut tx = Transaction::<2, 4, 4, 15, LightInstruction<'info>>::new(input);
 
-        tx.transact()
+        tx.transact()?;
+
+        #[cfg(all(feature = "memory-test", target_os = "solana"))]
+        assert!(
+            light_verifier_sdk::light_transaction::custom_heap::log_total_heap("memory_check")
+                < 7000u64,
+            "memory degression detected {} {}",
+            light_verifier_sdk::light_transaction::custom_heap::log_total_heap("memory_check"),
+            7000u64
+        );
+        Ok(())
     }
 }
 
diff --git a/psp-examples/streaming-payments/Cargo.lock b/psp-examples/streaming-payments/Cargo.lock
diff --git a/psp-examples/streaming-payments/programs/streaming-payments/Cargo.toml b/psp-examples/streaming-payments/programs/streaming-payments/Cargo.toml
@@ -13,7 +13,8 @@ no-entrypoint = []
 no-idl = []
 no-log-ix-name = []
 cpi = ["no-entrypoint"]
-default = []
+default = ["custom-heap"]
+custom-heap = []
 
 [dependencies]
 anchor-lang = "0.28.0"
diff --git a/psp-examples/streaming-payments/programs/streaming-payments/src/lib.rs b/psp-examples/streaming-payments/programs/streaming-payments/src/lib.rs
@@ -18,6 +18,8 @@ pub const PROGRAM_ID: &str = "Fg6PaFpoGXkYsidMpWTK6W2BeZ7FEfcYkg476zPFsLnS";
 #[program]
 pub mod streaming_payments {
     use super::*;
+    #[cfg(target_os = "solana")]
+    use light_verifier_sdk::light_transaction::custom_heap::log_total_heap;
     use solana_program::sysvar;
     /// This instruction is the first step of a compressed transaction.
     /// It creates and initializes a verifier state account to save state of a verification during
@@ -120,8 +122,15 @@ pub mod streaming_payments {
             panic!("Escrow still locked");
         }
         msg!("checked inputs {:?}", verifier_state.checked_public_inputs);
+        #[cfg(target_os = "solana")]
+        log_total_heap("pre verify_program_proof");
         verify_program_proof(&ctx, &inputs)?;
-        cpi_verifier_two(&ctx, &inputs)
+        #[cfg(target_os = "solana")]
+        log_total_heap("pre cpi_verifier_two");
+        cpi_verifier_two(&ctx, &inputs)?;
+        #[cfg(target_os = "solana")]
+        log_total_heap("post cpi_verifier_two");
+        Ok(())
     }
 
     /// Close the verifier state to reclaim rent in case the proofdata is wrong and does not verify.
diff --git a/psp-examples/streaming-payments/tests/streaming-payments.ts b/psp-examples/streaming-payments/tests/streaming-payments.ts
@@ -236,11 +236,10 @@ describe("Streaming Payments tests", () => {
       rootIndex,
     };
 
-    const res = await sendAndConfirmCompressedTransaction({
+    await sendAndConfirmCompressedTransaction({
       solanaTransactionInputs,
       provider: lightProvider,
     });
-    console.log("tx Hash : ", res.txHash);
   }
 
   async function paymentStreaming(
diff --git a/psp-examples/swap/Cargo.lock b/psp-examples/swap/Cargo.lock
diff --git a/psp-examples/swap/programs/swaps/Cargo.toml b/psp-examples/swap/programs/swaps/Cargo.toml
@@ -13,7 +13,8 @@ no-entrypoint = []
 no-idl = []
 no-log-ix-name = []
 cpi = ["no-entrypoint"]
-default = []
+default = ["custom-heap"]
+custom-heap = []
 
 [dependencies]
 anchor-lang = "0.28.0"
diff --git a/psp-template/psp-template/programs_psp/program_name/Cargo.toml b/psp-template/psp-template/programs_psp/program_name/Cargo.toml
@@ -13,7 +13,8 @@ no-entrypoint = []
 no-idl = []
 no-log-ix-name = []
 cpi = ["no-entrypoint"]
-default = []
+default = ["custom-heap"]
+custom-heap = []
 
 [dependencies]
 anchor-lang = "0.28.0"
diff --git a/verifier-sdk/Cargo.toml b/verifier-sdk/Cargo.toml
@@ -5,6 +5,11 @@ description = "SDK for Private Solana Programs"
 license = "GPL-3.0"
 edition = "2021"
 
+[features]
+custom-heap = []
+mem-profiling = ["custom-heap"]
+default = ["custom-heap"]
+
 [lib]
 crate-type = ["cdylib", "lib"]
 name = "light_verifier_sdk"
@@ -19,8 +24,12 @@ ark-ff = { version = "^0.3.0", default-features = false }
 ark-ec = { version = "0.3.0" }
 ark-bn254 = "0.3.0"
 ark-std = { version = "^0.3.0", default-features = false }
-spl-token = "3.3.0"
 borsh = "0.9.3"
 groth16-solana = { git= "https://github.com/Lightprotocol/groth16-solana", branch="master"}
+light-macros = { version = "0.3.1", path = "../macros/light" }
+spl-token = "3.3.0"
 light-utils = { path = "../utils", version = "0.1.0" }
 light-sparse-merkle-tree = { path = "../merkle-tree/sparse", version = "0.1.1" }
+
+[package.metadata.docs.rs]
+targets = ["x86_64-unknown-linux-gnu"]
diff --git a/verifier-sdk/src/light_transaction.rs b/verifier-sdk/src/light_transaction.rs
