Merge branch 'main' into schouhy/add-block-context

programs/associated_token_account/Cargo.toml

@@ -0,0 +1,10 @@
+[package]
+name = "ata_program"
+version = "0.1.0"
+edition = "2024"
+license = { workspace = true }
+
+[dependencies]
+nssa_core.workspace = true
+token_core.workspace = true
+ata_core.workspace = true

programs/associated_token_account/core/Cargo.toml

@@ -0,0 +1,10 @@
+[package]
+name = "ata_core"
+version = "0.1.0"
+edition = "2024"
+license = { workspace = true }
+
+[dependencies]
+nssa_core.workspace = true
+serde.workspace = true
+risc0-zkvm.workspace = true

programs/associated_token_account/core/src/lib.rs

@@ -0,0 +1,82 @@
+pub use nssa_core::program::PdaSeed;
+use nssa_core::{
+    account::{AccountId, AccountWithMetadata},
+    program::ProgramId,
+};
+use serde::{Deserialize, Serialize};
+
+#[derive(Serialize, Deserialize)]
+pub enum Instruction {
+    /// Create the Associated Token Account for (owner, definition).
+    /// Idempotent: no-op if the account already exists.
+    ///
+    /// Required accounts (3):
+    /// - Owner account
+    /// - Token definition account
+    /// - Associated token account (default/uninitialized, or already initialized)
+    ///
+    /// `token_program_id` is derived from `token_definition.account.program_owner`.
+    Create { ata_program_id: ProgramId },
+
+    /// Transfer tokens FROM owner's ATA to a recipient holding account.
+    /// Uses PDA seeds to authorize the ATA in the chained Token::Transfer call.
+    ///
+    /// Required accounts (3):
+    /// - Owner account (authorized)
+    /// - Sender ATA (owner's token holding)
+    /// - Recipient token holding (any account; auto-created if default)
+    ///
+    /// `token_program_id` is derived from `sender_ata.account.program_owner`.
+    Transfer {
+        ata_program_id: ProgramId,
+        amount: u128,
+    },
+
+    /// Burn tokens FROM owner's ATA.
+    /// Uses PDA seeds to authorize the ATA in the chained Token::Burn call.
+    ///
+    /// Required accounts (3):
+    /// - Owner account (authorized)
+    /// - Owner's ATA (the holding to burn from)
+    /// - Token definition account
+    ///
+    /// `token_program_id` is derived from `holder_ata.account.program_owner`.
+    Burn {
+        ata_program_id: ProgramId,
+        amount: u128,
+    },
+}
+
+pub fn compute_ata_seed(owner_id: AccountId, definition_id: AccountId) -> PdaSeed {
+    use risc0_zkvm::sha::{Impl, Sha256};
+    let mut bytes = [0u8; 64];
+    bytes[0..32].copy_from_slice(&owner_id.to_bytes());
+    bytes[32..64].copy_from_slice(&definition_id.to_bytes());
+    PdaSeed::new(
+        Impl::hash_bytes(&bytes)
+            .as_bytes()
+            .try_into()
+            .expect("Hash output must be exactly 32 bytes long"),
+    )
+}
+
+pub fn get_associated_token_account_id(ata_program_id: &ProgramId, seed: &PdaSeed) -> AccountId {
+    AccountId::from((ata_program_id, seed))
+}
+
+/// Verify the ATA's address matches `(ata_program_id, owner, definition)` and return
+/// the [`PdaSeed`] for use in chained calls.
+pub fn verify_ata_and_get_seed(
+    ata_account: &AccountWithMetadata,
+    owner: &AccountWithMetadata,
+    definition_id: AccountId,
+    ata_program_id: ProgramId,
+) -> PdaSeed {
+    let seed = compute_ata_seed(owner.account_id, definition_id);
+    let expected_id = get_associated_token_account_id(&ata_program_id, &seed);
+    assert_eq!(
+        ata_account.account_id, expected_id,
+        "ATA account ID does not match expected derivation"
+    );
+    seed
+}

programs/associated_token_account/src/burn.rs

@@ -0,0 +1,39 @@
+use nssa_core::{
+    account::AccountWithMetadata,
+    program::{AccountPostState, ChainedCall, ProgramId},
+};
+use token_core::TokenHolding;
+
+pub fn burn_from_associated_token_account(
+    owner: AccountWithMetadata,
+    holder_ata: AccountWithMetadata,
+    token_definition: AccountWithMetadata,
+    ata_program_id: ProgramId,
+    amount: u128,
+) -> (Vec<AccountPostState>, Vec<ChainedCall>) {
+    let token_program_id = holder_ata.account.program_owner;
+    assert!(owner.is_authorized, "Owner authorization is missing");
+    let definition_id = TokenHolding::try_from(&holder_ata.account.data)
+        .expect("Holder ATA must hold a valid token")
+        .definition_id();
+    let seed =
+        ata_core::verify_ata_and_get_seed(&holder_ata, &owner, definition_id, ata_program_id);
+
+    let post_states = vec![
+        AccountPostState::new(owner.account.clone()),
+        AccountPostState::new(holder_ata.account.clone()),
+        AccountPostState::new(token_definition.account.clone()),
+    ];
+    let mut holder_ata_auth = holder_ata.clone();
+    holder_ata_auth.is_authorized = true;
+
+    let chained_call = ChainedCall::new(
+        token_program_id,
+        vec![token_definition.clone(), holder_ata_auth],
+        &token_core::Instruction::Burn {
+            amount_to_burn: amount,
+        },
+    )
+    .with_pda_seeds(vec![seed]);
+    (post_states, vec![chained_call])
+}

programs/associated_token_account/src/create.rs

@@ -0,0 +1,44 @@
+use nssa_core::{
+    account::{Account, AccountWithMetadata},
+    program::{AccountPostState, ChainedCall, ProgramId},
+};
+
+pub fn create_associated_token_account(
+    owner: AccountWithMetadata,
+    token_definition: AccountWithMetadata,
+    ata_account: AccountWithMetadata,
+    ata_program_id: ProgramId,
+) -> (Vec<AccountPostState>, Vec<ChainedCall>) {
+    // No authorization check needed: create is idempotent, so anyone can call it safely.
+    let token_program_id = token_definition.account.program_owner;
+    ata_core::verify_ata_and_get_seed(
+        &ata_account,
+        &owner,
+        token_definition.account_id,
+        ata_program_id,
+    );
+
+    // Idempotent: already initialized → no-op
+    if ata_account.account != Account::default() {
+        return (
+            vec![
+                AccountPostState::new_claimed_if_default(owner.account.clone()),
+                AccountPostState::new(token_definition.account.clone()),
+                AccountPostState::new(ata_account.account.clone()),
+            ],
+            vec![],
+        );
+    }
+
+    let post_states = vec![
+        AccountPostState::new_claimed_if_default(owner.account.clone()),
+        AccountPostState::new(token_definition.account.clone()),
+        AccountPostState::new(ata_account.account.clone()),
+    ];
+    let chained_call = ChainedCall::new(
+        token_program_id,
+        vec![token_definition.clone(), ata_account.clone()],
+        &token_core::Instruction::InitializeAccount,
+    );
+    (post_states, vec![chained_call])
+}

programs/associated_token_account/src/lib.rs

@@ -0,0 +1,10 @@
+//! The Associated Token Account Program implementation.
+
+pub use ata_core as core;
+
+pub mod burn;
+pub mod create;
+pub mod transfer;
+
+#[cfg(test)]
+mod tests;

programs/associated_token_account/src/tests.rs

@@ -0,0 +1,153 @@
+#![cfg(test)]
+
+use ata_core::{compute_ata_seed, get_associated_token_account_id};
+use nssa_core::account::{Account, AccountId, AccountWithMetadata, Data};
+use token_core::{TokenDefinition, TokenHolding};
+
+const ATA_PROGRAM_ID: nssa_core::program::ProgramId = [1u32; 8];
+const TOKEN_PROGRAM_ID: nssa_core::program::ProgramId = [2u32; 8];
+
+fn owner_id() -> AccountId {
+    AccountId::new([0x01u8; 32])
+}
+
+fn definition_id() -> AccountId {
+    AccountId::new([0x02u8; 32])
+}
+
+fn ata_id() -> AccountId {
+    get_associated_token_account_id(
+        &ATA_PROGRAM_ID,
+        &compute_ata_seed(owner_id(), definition_id()),
+    )
+}
+
+fn owner_account() -> AccountWithMetadata {
+    AccountWithMetadata {
+        account: Account::default(),
+        is_authorized: true,
+        account_id: owner_id(),
+    }
+}
+
+fn definition_account() -> AccountWithMetadata {
+    AccountWithMetadata {
+        account: Account {
+            program_owner: TOKEN_PROGRAM_ID,
+            balance: 0,
+            data: Data::from(&TokenDefinition::Fungible {
+                name: "TEST".to_string(),
+                total_supply: 1000,
+                metadata_id: None,
+            }),
+            nonce: nssa_core::account::Nonce(0),
+        },
+        is_authorized: false,
+        account_id: definition_id(),
+    }
+}
+
+fn uninitialized_ata_account() -> AccountWithMetadata {
+    AccountWithMetadata {
+        account: Account::default(),
+        is_authorized: false,
+        account_id: ata_id(),
+    }
+}
+
+fn initialized_ata_account() -> AccountWithMetadata {
+    AccountWithMetadata {
+        account: Account {
+            program_owner: TOKEN_PROGRAM_ID,
+            balance: 0,
+            data: Data::from(&TokenHolding::Fungible {
+                definition_id: definition_id(),
+                balance: 100,
+            }),
+            nonce: nssa_core::account::Nonce(0),
+        },
+        is_authorized: false,
+        account_id: ata_id(),
+    }
+}
+
+#[test]
+fn create_emits_chained_call_for_uninitialized_ata() {
+    let (post_states, chained_calls) = crate::create::create_associated_token_account(
+        owner_account(),
+        definition_account(),
+        uninitialized_ata_account(),
+        ATA_PROGRAM_ID,
+    );
+
+    assert_eq!(post_states.len(), 3);
+    assert_eq!(chained_calls.len(), 1);
+    assert_eq!(chained_calls[0].program_id, TOKEN_PROGRAM_ID);
+}
+
+#[test]
+fn create_is_idempotent_for_initialized_ata() {
+    let (post_states, chained_calls) = crate::create::create_associated_token_account(
+        owner_account(),
+        definition_account(),
+        initialized_ata_account(),
+        ATA_PROGRAM_ID,
+    );
+
+    assert_eq!(post_states.len(), 3);
+    assert!(
+        chained_calls.is_empty(),
+        "Should emit no chained call for already-initialized ATA"
+    );
+}
+
+#[test]
+#[should_panic(expected = "ATA account ID does not match expected derivation")]
+fn create_panics_on_wrong_ata_address() {
+    let wrong_ata = AccountWithMetadata {
+        account: Account::default(),
+        is_authorized: false,
+        account_id: AccountId::new([0xFFu8; 32]),
+    };
+
+    crate::create::create_associated_token_account(
+        owner_account(),
+        definition_account(),
+        wrong_ata,
+        ATA_PROGRAM_ID,
+    );
+}
+
+#[test]
+fn get_associated_token_account_id_is_deterministic() {
+    let seed = compute_ata_seed(owner_id(), definition_id());
+    let id1 = get_associated_token_account_id(&ATA_PROGRAM_ID, &seed);
+    let id2 = get_associated_token_account_id(&ATA_PROGRAM_ID, &seed);
+    assert_eq!(id1, id2);
+}
+
+#[test]
+fn get_associated_token_account_id_differs_by_owner() {
+    let other_owner = AccountId::new([0x99u8; 32]);
+    let id1 = get_associated_token_account_id(
+        &ATA_PROGRAM_ID,
+        &compute_ata_seed(owner_id(), definition_id()),
+    );
+    let id2 = get_associated_token_account_id(
+        &ATA_PROGRAM_ID,
+        &compute_ata_seed(other_owner, definition_id()),
+    );
+    assert_ne!(id1, id2);
+}
+
+#[test]
+fn get_associated_token_account_id_differs_by_definition() {
+    let other_def = AccountId::new([0x99u8; 32]);
+    let id1 = get_associated_token_account_id(
+        &ATA_PROGRAM_ID,
+        &compute_ata_seed(owner_id(), definition_id()),
+    );
+    let id2 =
+        get_associated_token_account_id(&ATA_PROGRAM_ID, &compute_ata_seed(owner_id(), other_def));
+    assert_ne!(id1, id2);
+}

programs/associated_token_account/src/transfer.rs

@@ -0,0 +1,39 @@
+use nssa_core::{
+    account::AccountWithMetadata,
+    program::{AccountPostState, ChainedCall, ProgramId},
+};
+use token_core::TokenHolding;
+
+pub fn transfer_from_associated_token_account(
+    owner: AccountWithMetadata,
+    sender_ata: AccountWithMetadata,
+    recipient: AccountWithMetadata,
+    ata_program_id: ProgramId,
+    amount: u128,
+) -> (Vec<AccountPostState>, Vec<ChainedCall>) {
+    let token_program_id = sender_ata.account.program_owner;
+    assert!(owner.is_authorized, "Owner authorization is missing");
+    let definition_id = TokenHolding::try_from(&sender_ata.account.data)
+        .expect("Sender ATA must hold a valid token")
+        .definition_id();
+    let seed =
+        ata_core::verify_ata_and_get_seed(&sender_ata, &owner, definition_id, ata_program_id);
+
+    let post_states = vec![
+        AccountPostState::new(owner.account.clone()),
+        AccountPostState::new(sender_ata.account.clone()),
+        AccountPostState::new(recipient.account.clone()),
+    ];
+    let mut sender_ata_auth = sender_ata.clone();
+    sender_ata_auth.is_authorized = true;
+
+    let chained_call = ChainedCall::new(
+        token_program_id,
+        vec![sender_ata_auth, recipient.clone()],
+        &token_core::Instruction::Transfer {
+            amount_to_transfer: amount,
+        },
+    )
+    .with_pda_seeds(vec![seed]);
+    (post_states, vec![chained_call])
+}