feat: add lamport snapshot helpers

Author: kamiyo-ai

library/kani_solana_agent/.gitignore

@@ -0,0 +1 @@
+target/

library/kani_solana_agent/Cargo.toml

@@ -11,4 +11,7 @@ publish = false
 [dependencies]
 solana-program = "=3.0.0"
 
+[lints.rust]
+unexpected_cfgs = { level = "warn", check-cfg = ["cfg(kani)"] }
+
 [workspace]

library/kani_solana_agent/README.md

@@ -7,22 +7,25 @@ Non-goals:
 - Anchor-specific modeling
 - Token program semantics
 
-The primary entry point is `any_account_info::<DATA_LEN>(...)`, which constructs an `AccountInfo<'static>` backed by leaked allocations so it can be passed into program logic during verification.
+The primary entry point is `any_account_info::<DATA_LEN>(...)` (alias: `any_agent_account::<DATA_LEN>(...)`), which constructs an `AccountInfo<'static>` backed by leaked allocations so it can be passed into program logic during verification.
 
 ## Example
 
 ```rust
 #[cfg(kani)]
 mod proofs {
-    use kani_solana_agent::{AccountConfig, any_account_info};
+    use kani_solana_agent::{AccountConfig, LamportSnapshot, any_agent_account};
 
     #[kani::proof]
     fn can_build_accounts() {
-        let payer = any_account_info::<0>(AccountConfig::new().payer());
-        let escrow = any_account_info::<128>(AccountConfig::new().writable());
+        let payer = any_agent_account::<0>(AccountConfig::new().payer());
+        let escrow = any_agent_account::<128>(AccountConfig::new().writable());
 
         kani::assert(payer.is_signer, "payer must be a signer");
         kani::assert(escrow.is_writable, "escrow must be writable");
+
+        let before = LamportSnapshot::new(&[&payer, &escrow]);
+        before.assert_unchanged(&[&payer, &escrow], "snapshot must be stable without mutation");
     }
 }
 ```

library/kani_solana_agent/src/account.rs

@@ -3,9 +3,13 @@
 
 use solana_program::account_info::AccountInfo;
 use solana_program::pubkey::Pubkey;
+use std::ops::RangeInclusive;
+
+#[cfg(kani)]
 use solana_program::rent::Rent;
+#[cfg(kani)]
 use std::cell::RefCell;
-use std::ops::RangeInclusive;
+#[cfg(kani)]
 use std::rc::Rc;
 
 #[derive(Clone, Debug)]
@@ -100,14 +104,13 @@ impl AccountConfig {
     }
 }
 
-#[cfg(kani)]
-fn any_pubkey() -> Pubkey {
-    Pubkey::new_from_array(kani::any())
+pub fn any_agent_account<const DATA_LEN: usize>(cfg: AccountConfig) -> AccountInfo<'static> {
+    any_account_info::<DATA_LEN>(cfg)
 }
 
-#[cfg(not(kani))]
+#[cfg(kani)]
 fn any_pubkey() -> Pubkey {
-    Pubkey::default()
+    Pubkey::new_from_array(kani::any())
 }
 
 #[cfg(kani)]
@@ -129,12 +132,6 @@ fn pick_lamports<const DATA_LEN: usize>(cfg: &AccountConfig) -> u64 {
     lamports
 }
 
-#[cfg(not(kani))]
-fn pick_lamports<const DATA_LEN: usize>(_cfg: &AccountConfig) -> u64 {
-    let _ = DATA_LEN;
-    0
-}
-
 #[cfg(kani)]
 pub fn any_account_info<const DATA_LEN: usize>(cfg: AccountConfig) -> AccountInfo<'static> {
     let key = cfg.key.unwrap_or_else(any_pubkey);
@@ -161,7 +158,16 @@ pub fn any_account_info<const DATA_LEN: usize>(cfg: AccountConfig) -> AccountInf
 }
 
 #[cfg(not(kani))]
-pub fn any_account_info<const DATA_LEN: usize>(_cfg: AccountConfig) -> AccountInfo<'static> {
+pub fn any_account_info<const DATA_LEN: usize>(cfg: AccountConfig) -> AccountInfo<'static> {
     let _ = DATA_LEN;
+    match cfg.lamports {
+        Lamports::Any => {}
+        Lamports::Exact(v) => {
+            let _ = v;
+        }
+        Lamports::Range(r) => {
+            let _ = r;
+        }
+    }
     panic!("any_account_info is only available under `cargo kani` (cfg(kani))");
 }

library/kani_solana_agent/src/invariants.rs

@@ -0,0 +1,31 @@
+// Copyright Kani Contributors
+// SPDX-License-Identifier: Apache-2.0 OR MIT
+
+use solana_program::account_info::AccountInfo;
+
+pub fn lamports(account: &AccountInfo<'_>) -> u64 {
+    **account.lamports.borrow()
+}
+
+pub fn sum_lamports(accounts: &[&AccountInfo<'_>]) -> u128 {
+    accounts.iter().map(|a| lamports(a) as u128).sum()
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub struct LamportSnapshot(pub u128);
+
+impl LamportSnapshot {
+    pub fn new(accounts: &[&AccountInfo<'_>]) -> Self {
+        Self(sum_lamports(accounts))
+    }
+
+    pub fn assert_unchanged(self, accounts: &[&AccountInfo<'_>], msg: &'static str) {
+        let after = sum_lamports(accounts);
+
+        #[cfg(kani)]
+        kani::assert(after == self.0, msg);
+
+        #[cfg(not(kani))]
+        assert!(after == self.0, "{msg}");
+    }
+}

library/kani_solana_agent/src/lib.rs

@@ -2,7 +2,9 @@
 // SPDX-License-Identifier: Apache-2.0 OR MIT
 
 pub mod account;
+pub mod invariants;
 
-pub use account::{AccountConfig, any_account_info};
+pub use account::{AccountConfig, any_account_info, any_agent_account};
+pub use invariants::{LamportSnapshot, lamports, sum_lamports};
 pub use solana_program::account_info::AccountInfo;
 pub use solana_program::pubkey::Pubkey;

tests/cargo-kani/solana-accountinfo/Cargo.toml

@@ -10,4 +10,7 @@ edition = "2018"
 kani-solana-agent = { path = "../../../library/kani_solana_agent" }
 solana-program = "=3.0.0"
 
+[lints.rust]
+unexpected_cfgs = { level = "warn", check-cfg = ["cfg(kani)"] }
+
 [workspace]

tests/cargo-kani/solana-accountinfo/src/lib.rs

@@ -1,16 +1,15 @@
 // Copyright Kani Contributors
 // SPDX-License-Identifier: Apache-2.0 OR MIT
 
-use kani_solana_agent::{AccountConfig, AccountInfo, any_account_info};
+#[cfg(kani)]
+use kani_solana_agent::AccountInfo;
 
+#[cfg(kani)]
 fn authorized(now: i64, expires_at: i64, agent: &AccountInfo<'_>, api: &AccountInfo<'_>) -> bool {
     agent.is_signer || (api.is_signer && now >= expires_at)
 }
 
-fn lamports(account: &AccountInfo<'_>) -> u64 {
-    **account.lamports.borrow()
-}
-
+#[cfg(kani)]
 fn transfer(from: &AccountInfo<'_>, to: &AccountInfo<'_>, amount: u64) -> Result<(), ()> {
     let mut from_lamports = from.lamports.borrow_mut();
     let Some(new_from) = (**from_lamports).checked_sub(amount) else {
@@ -27,6 +26,7 @@ fn transfer(from: &AccountInfo<'_>, to: &AccountInfo<'_>, amount: u64) -> Result
     Ok(())
 }
 
+#[cfg(kani)]
 fn release_funds(
     now: i64,
     expires_at: i64,
@@ -46,14 +46,15 @@ fn release_funds(
 #[cfg(kani)]
 mod proofs {
     use super::*;
+    use kani_solana_agent::{AccountConfig, LamportSnapshot, any_agent_account, lamports};
 
     #[kani::proof]
     fn timelock_policy_matches_release_rule() {
         let now: i64 = kani::any();
         let expires_at: i64 = kani::any();
 
-        let mut agent = any_account_info::<0>(AccountConfig::new());
-        let mut api = any_account_info::<0>(AccountConfig::new());
+        let mut agent = any_agent_account::<0>(AccountConfig::new());
+        let mut api = any_agent_account::<0>(AccountConfig::new());
         agent.is_signer = kani::any();
         api.is_signer = kani::any();
 
@@ -75,31 +76,30 @@ mod proofs {
         let expires_at: i64 = kani::any();
         let amount: u64 = kani::any::<u32>() as u64;
 
-        let mut agent = any_account_info::<0>(AccountConfig::new());
-        let mut api = any_account_info::<0>(AccountConfig::new());
+        let mut agent = any_agent_account::<0>(AccountConfig::new());
+        let mut api = any_agent_account::<0>(AccountConfig::new());
         agent.is_signer = kani::any();
         api.is_signer = kani::any();
 
-        let escrow = any_account_info::<0>(
+        let escrow = any_agent_account::<0>(
             AccountConfig::new().writable().lamports_range(0..=(u32::MAX as u64)),
         );
-        let payee = any_account_info::<0>(
+        let payee = any_agent_account::<0>(
             AccountConfig::new().writable().lamports_range(0..=(u32::MAX as u64)),
         );
 
         let escrow_before = lamports(&escrow);
         let payee_before = lamports(&payee);
         kani::assume(escrow_before >= amount);
 
-        let total_before = escrow_before as u128 + payee_before as u128;
+        let total_before = LamportSnapshot::new(&[&escrow, &payee]);
 
         let res = release_funds(now, expires_at, &agent, &api, &escrow, &payee, amount);
         let escrow_after = lamports(&escrow);
         let payee_after = lamports(&payee);
 
         if res.is_ok() {
-            let total_after = escrow_after as u128 + payee_after as u128;
-            kani::assert(total_after == total_before, "release must conserve lamports");
+            total_before.assert_unchanged(&[&escrow, &payee], "release must conserve lamports");
             kani::assert(escrow_after + amount == escrow_before, "escrow must decrease by amount");
             kani::assert(payee_before + amount == payee_after, "payee must increase by amount");
         } else {