Mappings

pragma language_version >= 0.23;

import CompactStandardLibrary;

// Simple balance tracker
export ledger balances: Map<Either<ZswapCoinPublicKey, ContractAddress>, Uint<128>>;

// Get balance for an account
export circuit getBalance(account: Either<ZswapCoinPublicKey, ContractAddress>): Uint<128> {
    // Check if account exists in map
    if (!balances.member(disclose(account))) {
        return 0; // Return 0 for accounts with no balance
    }
    return balances.lookup(disclose(account));
}

// Set balance for an account
export circuit setBalance(
    account: Either<ZswapCoinPublicKey, ContractAddress>,
    amount: Uint<128>
): [] {
    balances.insert(disclose(account), disclose(amount));
}

// Add to existing balance
export circuit addBalance(
    account: Either<ZswapCoinPublicKey, ContractAddress>,
    amount: Uint<128>
): [] {
    const MAX_UINT128: Uint<128> = 340282366920938463463374607431768211455 as Uint<128>;
    const currentBalance = getBalance(account);

    // Check for overflow before casting back
    assert(currentBalance <= MAX_UINT128 - amount, "Balance overflow");

    const newBalance: Uint<128> = (currentBalance + amount) as Uint<128>;
    balances.insert(disclose(account), disclose(newBalance));
}

// Subtract from balance
export circuit subtractBalance(
    account: Either<ZswapCoinPublicKey, ContractAddress>,
    amount: Uint<128>
): [] {
    const currentBalance = getBalance(account);

    // Check sufficient balance
    assert(currentBalance >= amount, "Insufficient balance");

    const newBalance = currentBalance - amount;
    balances.insert(disclose(account), disclose(newBalance));
}