import Base
import ./types.bend as Ty

# Shared Solana view. Import this file as S.
# Accounts (up to 256), instruction data, and the CPI
# queue are lists in the book. List<&2, U32> is the byte view: the checker
# sees a linked list; the compiler emits O(1) slice ops (zero-copy).

def key(a: Ty.Account) -> Ty.Pubkey:
  match a:
    case Ty.Acc{+k, own, lams, bytes, signer, writable}:
      k

def owner(a: Ty.Account) -> Ty.Pubkey:
  match a:
    case Ty.Acc{k, +own, lams, bytes, signer, writable}:
      own

def is_signer(a: Ty.Account) -> Bool:
  match a:
    case Ty.Acc{k, own, lams, bytes, +signer, writable}:
      signer

def is_writable(a: Ty.Account) -> Bool:
  match a:
    case Ty.Acc{k, own, lams, bytes, signer, +writable}:
      writable

def lamports(a: Ty.Account) -> Ty.u64:
  match a:
    case Ty.Acc{k, own, +lams, bytes, signer, writable}:
      lams

def data(a: Ty.Account) -> List<&2, U32>:
  match a:
    case Ty.Acc{k, own, lams, +bytes, signer, writable}:
      bytes

def set_data(a: Ty.Account, bytes: List<&2, U32>) -> Ty.Account:
  match a:
    case Ty.Acc{k, own, lams, old, signer, writable}:
      Ty.Acc{k, own, lams, bytes, signer, writable}

def set_lamports(a: Ty.Account, lams: Ty.u64) -> Ty.Account:
  match a:
    case Ty.Acc{k, own, old, bytes, signer, writable}:
      Ty.Acc{k, own, lams, bytes, signer, writable}

def add_lamports(+a: Ty.Account, amt: Ty.u64) -> Ty.Account:
  set_lamports(a, Ty.u64.add(lamports(a), amt))

def sub_lamports(+a: Ty.Account, amt: Ty.u64) -> Ty.Account:
  set_lamports(a, Ty.u64.sub(lamports(a), amt))

def is_owned_by(+a: Ty.Account, +program_id: Ty.Pubkey) -> Bool:
  Ty.Pubkey.eq(owner(a), program_id)

# The all-zero pubkey is the System program.
def system_id() -> Ty.Pubkey:
  Ty.Pubkey.zero()

# Key is the System program.
def is_system(+a: Ty.Account) -> Bool:
  Ty.Pubkey.eq(key(a), system_id())

# Signer flag and key comparison both hold.
def is_signer_of(is_signer: Bool, same: Bool) -> Bool:
  Bool.and(is_signer, same)

# Placeholder for an account index past the end of the list.
def empty_account() -> Ty.Account:
  Ty.Acc{
    Ty.Pubkey.zero(),
    Ty.Pubkey.zero(),
    Ty.u64.zero(),
    Nil{},
    False{},
    False{},
  }

# Ty.Account and instruction bytes are List<&2, U32> in the book.
# The compiler emits O(1) slice ops (zero-copy).

# Word i, or 0 when i is past the end.
def get_u32(xs: List<&2, U32>, i: Nat) -> U32:
  match xs i:
    case Nil{} _:
      0
    case h <> t 0n:
      h
    case h <> t 1n+p:
      get_u32(t, p)

# Replace word i. An index past the end leaves the list unchanged.
def set_u32(xs: List<&2, U32>, i: Nat, v: U32) -> List<&2, U32>:
  match xs i:
    case Nil{} _:
      Nil{}
    case h <> t 0n:
      v <> t
    case h <> t 1n+p:
      h <> set_u32(t, p, v)

def data_len(xs: List<&2, U32>) -> Nat:
  match xs:
    case Nil{}:
      0n
    case h <> t:
      1n+data_len(t)

def has_words(xs: List<&2, U32>, n: Nat) -> Bool:
  Nat.is_ge(data_len(xs), n)

# A byte slice into a word buffer. `off` and `len` are bytes.
# Ty.Account fields use this so a bump is one byte, not a u32 word.
def as_bytes(xs: List<&2, U32>) -> Ty.Slice:
  Ty.Slice{xs, 0, 0}

def slice(s: Ty.Slice, off: Nat, len: Nat) -> Ty.Slice:
  match s:
    case Ty.Slice{words, +base, +n}:
      Ty.Slice{words, U32.add(base, U32.from_nat(off)), U32.from_nat(len)}

def slice_len(s: Ty.Slice) -> U32:
  match s:
    case Ty.Slice{words, off, +n}:
      n

def byte_index(+off: U32, +i: Nat) -> U32:
  U32.add(off, U32.from_nat(i))

def word_index(+i: U32) -> Nat:
  U32.to_nat(U32.div(i, 4))

def byte_shift(+i: U32) -> U32:
  U32.mul(U32.sub(i, U32.mul(U32.div(i, 4), 4)), 8)

def get_u8(s: Ty.Slice, +i: Nat) -> U32:
  match s:
    case Ty.Slice{+words, +off, n}:
      U32.and(
        U32.shrn(
          get_u32(words, word_index(byte_index(off, i))),
          U32.to_nat(byte_shift(byte_index(off, i))),
        ),
        255,
      )

# Byte 0 equals d. An empty buffer compares as 0.
def has_discriminator(xs: List<&2, U32>, d: U32) -> Bool:
  U32.is_eq(get_u8(as_bytes(xs), 0n), d)

def set_u8(+xs: List<&2, U32>, +i: Nat, +v: U32) -> List<&2, U32>:
  set_u32(
    xs,
    word_index(U32.from_nat(i)),
    U32.or(
      U32.and(
        get_u32(xs, word_index(U32.from_nat(i))),
        U32.not(U32.shln(255, U32.to_nat(byte_shift(U32.from_nat(i))))),
      ),
      U32.shln(U32.and(v, 255), U32.to_nat(byte_shift(U32.from_nat(i)))),
    ),
  )

def le_u32(+s: Ty.Slice, +i: Nat) -> U32:
  U32.or(
    U32.or(get_u8(s, i), U32.shln(get_u8(s, 1n+i), 8n)),
    U32.or(U32.shln(get_u8(s, 2n+i), 16n), U32.shln(get_u8(s, 3n+i), 24n)),
  )

def set_le_u32(+xs: List<&2, U32>, +i: Nat, +v: U32) -> List<&2, U32>:
  set_u8(
    set_u8(
      set_u8(set_u8(xs, i, v), 1n+i, U32.shrn(v, 8n)),
      2n+i,
      U32.shrn(v, 16n),
    ),
    3n+i,
    U32.shrn(v, 24n),
  )

def le_u64(+s: Ty.Slice, +i: Nat) -> Ty.u64:
  Ty.u64.pack(le_u32(s, i), le_u32(s, 4n+i))

def as_word(v: Ty.u32) -> U32:
  U32.from_nat(Ty.u32.to_nat(v))

def set_le_u64(+xs: List<&2, U32>, +i: Nat, +v: Ty.u64) -> List<&2, U32>:
  set_le_u32(
    set_le_u32(xs, i, as_word(Ty.u64.lo(v))),
    4n+i,
    as_word(Ty.u64.hi(v)),
  )

def pubkey_at(+s: Ty.Slice, +i: Nat) -> Ty.Pubkey:
  Ty.Pk{
    le_u32(s, i),
    le_u32(s, 4n+i),
    le_u32(s, 8n+i),
    le_u32(s, 12n+i),
    le_u32(s, 16n+i),
    le_u32(s, 20n+i),
    le_u32(s, 24n+i),
    le_u32(s, 28n+i),
  }

def set_pubkey_at(+xs: List<&2, U32>, +i: Nat, +k: Ty.Pubkey) -> List<&2, U32>:
  match k:
    case Ty.Pk{+a, +b, +c, +d, +e, +f, +g, +h}:
      set_le_u32(
        set_le_u32(
          set_le_u32(
            set_le_u32(
              set_le_u32(
                set_le_u32(
                  set_le_u32(set_le_u32(xs, i, a), 4n+i, b),
                  8n+i,
                  c,
                ),
                12n+i,
                d,
              ),
              16n+i,
              e,
            ),
            20n+i,
            f,
          ),
          24n+i,
          g,
        ),
        28n+i,
        h,
      )

# Stored pubkey at byte `off` equals `k`.
def eq_pubkey(+xs: List<&2, U32>, +off: Nat, +k: Ty.Pubkey) -> Bool:
  Ty.Pubkey.eq(pubkey_at(as_bytes(xs), off), k)

# Same as get_u32.
def get_u32_word(xs: List<&2, U32>, i: Nat) -> U32:
  get_u32(xs, i)

# Same as set_u32.
def set_u32_word(xs: List<&2, U32>, i: Nat, v: U32) -> List<&2, U32>:
  set_u32(xs, i, v)

# Little-endian u64 at words i and i+1.
def get_u64(+xs: List<&2, U32>, +i: Nat) -> Ty.u64:
  Ty.u64.pack(get_u32(xs, i), get_u32(xs, 1n+i))

# Write a little-endian u64 at words i and i+1.
def set_u64(+xs: List<&2, U32>, +i: Nat, +v: Ty.u64) -> List<&2, U32>:
  set_u32(set_u32(xs, i, as_word(Ty.u64.lo(v))), 1n+i, as_word(Ty.u64.hi(v)))

# n zero words. Used to pad a short account before the first write.
def zeros(n: Nat) -> List<&2, U32>:
  match n:
    case 0n:
      Nil{}
    case 1n+p:
      0 <> zeros(p)

# Keep xs when it is already long enough; otherwise n zero words.
def ensure_if(ok: Bool, xs: List<&2, U32>, n: Nat) -> List<&2, U32>:
  match ok:
    case True{}:
      xs
    case False{}:
      zeros(n)

# Pad with zeros when xs is shorter than n words.
def ensure_words(+xs: List<&2, U32>, +n: Nat) -> List<&2, U32>:
  ensure_if(Nat.is_ge(data_len(xs), n), xs, n)

# Lengthen a word view by at most 10240 bytes. The compiler emits list_grow.
def resize(+xs: List<&2, U32>, +n: Nat) -> List<&2, U32>:
  ensure_words(xs, n)

# Ty.Pubkey at word offset off (eight words).
def get_pubkey(+xs: List<&2, U32>, +off: Nat) -> Ty.Pubkey:
  Ty.Pk{
    get_u32(xs, off),
    get_u32(xs, 1n+off),
    get_u32(xs, 2n+off),
    get_u32(xs, 3n+off),
    get_u32(xs, 4n+off),
    get_u32(xs, 5n+off),
    get_u32(xs, 6n+off),
    get_u32(xs, 7n+off),
  }

# Write eight pubkey words at off.
def set_pubkey(xs: List<&2, U32>, +off: Nat, p: Ty.Pubkey) -> List<&2, U32>:
  match p:
    case Ty.Pk{+w0, +w1, +w2, +w3, +w4, +w5, +w6, +w7}:
      set_u32(
        set_u32(
          set_u32(
            set_u32(
              set_u32(
                set_u32(
                  set_u32(set_u32(xs, off, w0), 1n+off, w1),
                  2n+off,
                  w2,
                ),
                3n+off,
                w3,
              ),
              4n+off,
              w4,
            ),
            5n+off,
            w5,
          ),
          6n+off,
          w6,
        ),
        7n+off,
        w7,
      )

# Ty.Pubkey at word offset off equals k.
def keys_eq(+xs: List<&2, U32>, +off: Nat, +k: Ty.Pubkey) -> Bool:
  Ty.Pubkey.eq(get_pubkey(xs, off), k)

# The pubkey at ao in a equals the pubkey at bo in b.
def keys_eq_words(
  +a: List<&2, U32>,
  +ao: Nat,
  +b: List<&2, U32>,
  +bo: Nat,
) -> Bool:
  Ty.Pubkey.eq(get_pubkey(a, ao), get_pubkey(b, bo))

# Byte equality. Address derivation is is_program_address, not this.
def pda_matches(+got: Ty.Pubkey, +want: Ty.Pubkey) -> Bool:
  Ty.Pubkey.eq(got, want)

# Ordered seeds for find_program_address. At most 16, each 1..=32 bytes.
# The book mark is not the hash. is_program_address / find_program_address are the natives.
def seeds_nil() -> Ty.Seeds:
  Ty.Seeds{0}

# Append `nbytes` of `bytes` (little-endian words). The chain copies those
# bytes; this body only keeps the checker from dropping the arguments.
def seeds_push(+ss: Ty.Seeds, +bytes: List<&2, U32>, nbytes: U32) -> Ty.Seeds:
  match ss bytes nbytes:
    case Ty.Seeds{mark} xs n:
      Ty.Seeds{(mark + (get_u32(xs, 0n) + n : U32) : U32)}

# Append the bytes of a constant string. The book mark ignores the text.
# The compiler native copies those bytes.
def seed_str(+ss: Ty.Seeds, +text: String) -> Ty.Seeds:
  match ss text:
    case Ty.Seeds{mark} _:
      Ty.Seeds{mark}

# Thirty-two bytes of a pubkey.
def seed_pk(+ss: Ty.Seeds, +k: Ty.Pubkey) -> Ty.Seeds:
  seeds_push(ss, set_pubkey(zeros(8n), 0n, k), 32)

# Checker dummy. Compiler native is find_program_address(seeds, program_id).
def is_program_address(
  +got: Ty.Pubkey,
  +program_id: Ty.Pubkey,
  +seeds: Ty.Seeds,
) -> Bool:
  match got program_id seeds:
    case Ty.Pk{+g0, +g1, +g2, +g3, +g4, +g5, +g6, +g7}
      Ty.Pk{+p0, +p1, +p2, +p3, +p4, +p5, +p6, +p7}
      Ty.Seeds{mark}:
      Bool.and(
        Ty.Pubkey.eq(
          Ty.Pk{g0, g1, g2, g3, g4, g5, g6, g7},
          Ty.Pk{p0, p1, p2, p3, p4, p5, p6, p7},
        ),
        U32.is_eq(mark, mark),
      )

# Address is the PDA whose only seed is `owner`. find_program_address.
def is_program_address_pk(
  +got: Ty.Pubkey,
  +program_id: Ty.Pubkey,
  +owner: Ty.Pubkey,
) -> Bool:
  is_program_address(got, program_id, seed_pk(seeds_nil(), owner))

# Checker dummy. Compiler native is create_program_address(seeds, bump).
def is_program_address_bump(
  +got: Ty.Pubkey,
  +program_id: Ty.Pubkey,
  +seeds: Ty.Seeds,
  +bump: U32,
) -> Bool:
  match got program_id seeds bump:
    case Ty.Pk{+g0, +g1, +g2, +g3, +g4, +g5, +g6, +g7}
      Ty.Pk{+p0, +p1, +p2, +p3, +p4, +p5, +p6, +p7}
      Ty.Seeds{mark}
      +b:
      Bool.and(
        Ty.Pubkey.eq(
          Ty.Pk{g0, g1, g2, g3, g4, g5, g6, g7},
          Ty.Pk{p0, p1, p2, p3, p4, p5, p6, p7},
        ),
        Bool.and(U32.is_eq(mark, mark), U32.is_eq(b, b)),
      )

# `got` is create_program_address([owner, bump]).
def is_program_address_pk_bump(
  +got: Ty.Pubkey,
  +program_id: Ty.Pubkey,
  +owner: Ty.Pubkey,
  +bump: U32,
) -> Bool:
  is_program_address_bump(got, program_id, seed_pk(seeds_nil(), owner), bump)

# Checker dummy. The compiler native reads the Rent sysvar:
# (bytes + 128) * lamports_per_byte_year * exemption_threshold.
def rent_exempt(+bytes: Ty.u64) -> Ty.u64:
  bytes

def rent_due_go(fits: Bool, +have: Ty.u64, +need: Ty.u64) -> Ty.u64:
  match fits:
    case True{}:
      Ty.u64.zero()
    case False{}:
      Ty.u64.sub(rent_exempt(need), rent_exempt(have))

# Extra lamports to raise an account from `have` bytes to `need` bytes.
def rent_due(+have: Ty.u64, +need: Ty.u64) -> Ty.u64:
  rent_due_go(Ty.u64.is_le(need, have), have, need)

# Checker dummy bump 255. Compiler native is the try_find bump.
def find_program_address(+program_id: Ty.Pubkey, +seeds: Ty.Seeds) -> U32:
  match program_id seeds:
    case _ _:
      255

# find_program_address bump for an owner-only PDA.
def find_program_address_pk(+program_id: Ty.Pubkey, +owner: Ty.Pubkey) -> U32:
  find_program_address(program_id, seed_pk(seeds_nil(), owner))

# A CPI is one queued invoke. Indices are Context account slots.
# The runtime walks the whole queue after process returns Ok.

type Cpi is Data:
  CpiNone{}
  CpiTransfer{
    from: U32, # source account index
    to: U32, # destination account index
    lamports: Ty.u64, # lamports to transfer
    sign: Bool, # true when the source PDA signs
    bump: U32, # PDA bump; unused when sign is false
    seed: U32 # unused book word; the signer list is Ty.Seeds
  }
  CpiTokenTransfer{
    from: U32, # source token-account index
    to: U32, # destination token-account index
    mint: U32, # mint account index
    auth: U32, # authority account index
    amount: Ty.u64, # token amount
    decimals: U32, # mint decimals
    sign: Bool, # true when the authority PDA signs
    bump: U32, # PDA bump; unused when sign is false
    seed: U32 # unused book word; the signer list is Ty.Seeds
  }
  # System CreateAccount. `from` pays, `to` is the new account.
  # `space` is the byte length. When `sign`, `Ty.Seeds` plus `bump` sign
  # the new account. Owner is always this program. The `seed` word is unused.
  CpiCreate{
    from: U32,
    to: U32,
    lamports: Ty.u64,
    space: Ty.u64,
    sign: Bool,
    bump: U32,
    seed: U32,
  }
  # System Assign. `from` is the account; `owner` is the new owner program.
  # When `sign`, the signer is `Ty.Seeds` plus `bump`. The `seed` word is unused.
  CpiAssign{
    from: U32,
    to: U32,
    owner: Ty.Pubkey,
    sign: Bool,
    bump: U32,
    seed: U32,
  }
  # System Allocate. `from` is the account; `space` is the byte length.
  # When `sign`, the signer is `Ty.Seeds` plus `bump`. The `seed` word is unused.
  CpiAllocate{
    from: U32,
    to: U32,
    space: Ty.u64,
    sign: Bool,
    bump: U32,
    seed: U32,
  }
  # SPL Token instruction other than TransferChecked. `kind` is the Tokenkeg
  # discriminant: 4 Approve, 5 Revoke, 7 MintTo, 10 Freeze, 11 Thaw,
  # 18 InitializeAccount3, 20 InitializeMint2.
  CpiTokenIx{
    kind: U32,
    from: U32,
    to: U32,
    auth: U32,
    amount: Ty.u64,
    decimals: U32,
    owner: Ty.Pubkey,
    sign: Bool,
    bump: U32,
    seed: U32,
  }
  # Associated Token Ty.Account program Create (0) / CreateIdempotent (1).
  CpiAta{
    kind: U32,
    payer: U32,
    ata: U32,
    wallet: U32,
    mint: U32,
    system: U32,
    token: U32,
  }
  # Invoke `program` (an account index) with `accs` (account indexes)
  # and `ix` (instruction words, little-endian on the wire).
  # `auth` is the PDA signer index when `sign` is true.
  CpiInvoke{
    program: U32,
    auth: U32,
    accs: List<&2, U32>,
    ix: List<&2, U32>,
    sign: Bool,
    bump: U32,
    seed: U32,
  }

def cpi_none() -> Cpi:
  CpiNone{}

def cpi_is_none(c: Cpi) -> Bool:
  match c:
    case CpiNone{}:
      True{}
    case CpiTransfer{_, _, _, _, _, _}:
      False{}
    case CpiTokenTransfer{
      from,
      to,
      mint,
      auth,
      amount,
      decimals,
      sign,
      bump,
      seed,
    }:
      False{}
    case CpiCreate{_, _, _, _, _, _, _}:
      False{}
    case CpiAssign{_, _, _, _, _, _}:
      False{}
    case CpiAllocate{from, to, space, sign, bump, seed}:
      False{}
    case CpiTokenIx{_, _, _, _, _, _, _, _, _, _}:
      False{}
    case CpiAta{kind, payer, ata, wallet, mint, system, token}:
      False{}
    case CpiInvoke{program, auth, accs, ix, sign, bump, seed}:
      False{}

# First queued CPI, or CpiNone when the queue is empty.
def cpi_head(xs: List<&2, Cpi>) -> Cpi:
  match xs:
    case Nil{}:
      CpiNone{}
    case h <> t:
      h

# CPI at index i, or CpiNone when i is past the end.
def cpi_at(xs: List<&2, Cpi>, i: Nat) -> Cpi:
  match xs i:
    case Nil{} _:
      CpiNone{}
    case h <> t 0n:
      h
    case h <> t 1n+p:
      cpi_at(t, p)

# Replace a leading CpiNone; otherwise append. Empty (Nil) becomes a singleton.
def cpi_put(xs: List<&2, Cpi>, c: Cpi) -> List<&2, Cpi>:
  match xs:
    case Nil{}:
      c <> Nil{}
    case h <> t:
      match h:
        case CpiNone{}:
          c <> t
        case CpiTransfer{+from, +to, +lams, +sign, +bump, +seed}:
          CpiTransfer{from, to, lams, sign, bump, seed} <> cpi_put(t, c)
        case CpiTokenTransfer{
          +from,
          +to,
          +mint,
          +auth,
          +amount,
          +decimals,
          +sign,
          +bump,
          +seed,
        }:
          CpiTokenTransfer{
            from,
            to,
            mint,
            auth,
            amount,
            decimals,
            sign,
            bump,
            seed,
          } <> cpi_put(t, c)
        case CpiCreate{+from, +to, +lams, +space, +sign, +bump, +seed}:
          CpiCreate{from, to, lams, space, sign, bump, seed} <> cpi_put(t, c)
        case CpiAssign{+from, +to, +own, +sign, +bump, +seed}:
          CpiAssign{from, to, own, sign, bump, seed} <> cpi_put(t, c)
        case CpiAllocate{+from, +to, +space, +sign, +bump, +seed}:
          CpiAllocate{from, to, space, sign, bump, seed} <> cpi_put(t, c)
        case CpiTokenIx{
          +kind,
          +from,
          +to,
          +auth,
          +amount,
          +decimals,
          +own,
          +sign,
          +bump,
          +seed,
        }:
          CpiTokenIx{
            kind,
            from,
            to,
            auth,
            amount,
            decimals,
            own,
            sign,
            bump,
            seed,
          } <> cpi_put(t, c)
        case CpiAta{+kind, +payer, +ata, +wallet, +mint, +system, +token}:
          CpiAta{kind, payer, ata, wallet, mint, system, token} <> cpi_put(t, c)
        case CpiInvoke{+program, +auth, +accs, +ix, +sign, +bump, +seed}:
          CpiInvoke{program, auth, accs, ix, sign, bump, seed} <> cpi_put(t, c)

# Context is the transaction the runtime hands the program.
# Accounts, instruction data, and CPIs are lists. The book does not cap
# them. The runtime caps accounts at 256. Instruction data is the
# bytes in the input. CPIs are limited by the heap.

# Runtime cap. The book list is unbounded; the runtime rejects more.
def MAX_ACCOUNTS() -> U32:
  256

type Context is Data:
  Context{
    n: U32, # account count
    program_id: Ty.Pubkey, # this program
    data: List<&2, U32>, # instruction data as words
    accounts: List<&2, Ty.Account>, # accounts in transaction order
    cpis: List<&2, Cpi> # cross-program invocations queued by the handler
  }

type Outcome is Data:
  Ok{context: Context}  # success; the runtime applies this context
  Err{code: Ty.u64}  # failure; a Solana program-error return

def ok(+c: Context) -> Outcome:
  Ok{c}

def fail(+code: Ty.u64) -> Outcome:
  Err{code}

# The outcome is Err{want}. An Ok outcome is False.
def err_is(outcome: Outcome, want: Ty.u64) -> Bool:
  match outcome:
    case Ok{+c}:
      False{}
    case Err{code}:
      Ty.u64.is_eq(code, want)

# The outcome is Ok. An Err outcome is False.
def is_ok(outcome: Outcome) -> Bool:
  match outcome:
    case Ok{+c}:
      True{}
    case Err{code}:
      False{}

# Ok{want}, field for field. Any other outcome is False.
# The two contexts must be the same term: same accounts, instruction, and queue.
def wrote(outcome: Outcome, want: Context) -> Bool:
  match outcome:
    case Ok{+c}:
      match c want:
        case Context{n, pid, ix, accs, queue} Context{n, pid, ix, accs, queue}:
          True{}
        case Context{n, pid, ix, accs, queue}
          Context{n2, pid2, ix2, accs2, queue2}:
          False{}
    case Err{code}:
      False{}

# One require in a `do Try<>:` block. No carries the error; Yes carries the value.
type Attempt<a, -A: Kind(a)> is Kind(a):
  No{code: Ty.u64}
  Yes{v: A}

# `do Try<>:` runs each require in order. False, or None, is Err{code}.
def Try.bind(-A: Data, m: Attempt<&2, A>, next: A -> Outcome) -> Outcome:
  match m:
    case No{code}:
      Err{code}
    case Yes{v}:
      next(v)

# `return` in a `do Try<>:` block.
def Try.pure(o: Outcome) -> Outcome:
  o

# False is No{code}. True is Yes{Unit}, so the next line runs.
def require(ok: Bool, +code: Ty.u64) -> Attempt<&2, Unit>:
  match ok:
    case False{}:
      No{code}
    case True{}:
      Yes{Unit{}}

# None is No{code}. Some is Yes of the U32, bound by `name: U32 <-`.
def require_some(m: Maybe<&2, U32>, +code: Ty.u64) -> Attempt<&2, U32>:
  match m:
    case None{}:
      No{code}
    case Some{v}:
      Yes{v}

# None is No{code}. Some is Yes of the Ty.u64, bound by `name: Ty.u64 <-`.
def require_u64(m: Maybe<&2, Ty.u64>, +code: Ty.u64) -> Attempt<&2, Ty.u64>:
  match m:
    case None{}:
      No{code}
    case Some{v}:
      Yes{v}

# Solana program-error returns. Builtin variants are `hi << 32`.
# `custom` is ProgramError::Custom: the low word, for a code with no
# standard variant. Custom(0) is not a failure; pass a non-zero code.
def Err.ok() -> Ty.u64:
  Ty.u64.zero()

def Err.not_enough_account_keys() -> Ty.u64:
  Ty.u64.pack(0, 11)

def Err.invalid_instruction_data() -> Ty.u64:
  Ty.u64.pack(0, 3)

def Err.invalid_account_data() -> Ty.u64:
  Ty.u64.pack(0, 4)

def Err.insufficient_funds() -> Ty.u64:
  Ty.u64.pack(0, 6)

def Err.account_already_initialized() -> Ty.u64:
  Ty.u64.pack(0, 9)

def Err.uninitialized_account() -> Ty.u64:
  Ty.u64.pack(0, 10)

def Err.invalid_seeds() -> Ty.u64:
  Ty.u64.pack(0, 14)

def Err.illegal_owner() -> Ty.u64:
  Ty.u64.pack(0, 18)

def Err.arithmetic_overflow() -> Ty.u64:
  Ty.u64.pack(0, 24)

def Err.incorrect_authority() -> Ty.u64:
  Ty.u64.pack(0, 26)

def Err.custom(code: U32) -> Ty.u64:
  Ty.u64.from_u32(code)

# Ty.Account i, or the empty account when i is past the end.
def account_at(xs: List<&2, Ty.Account>, i: Nat) -> Ty.Account:
  match xs i:
    case Nil{} _:
      empty_account()
    case h <> t 0n:
      h
    case h <> t 1n+p:
      account_at(t, p)

# Replace account i. An index past the end leaves the list unchanged.
def set_accounts(
  xs: List<&2, Ty.Account>,
  i: Nat,
  a: Ty.Account,
) -> List<&2, Ty.Account>:
  match xs i:
    case Nil{} _:
      Nil{}
    case h <> t 0n:
      a <> t
    case h <> t 1n+p:
      h <> set_accounts(t, p, a)

def instruction_data(c: Context) -> List<&2, U32>:
  match c:
    case Context{n, pid, +ix, accs, queue}:
      ix

def ix_u32(c: Context, i: Nat) -> U32:
  match c:
    case Context{n, pid, +ix, accs, queue}:
      get_u32(ix, i)

# Instruction discriminator. Byte `i` of the instruction data, widened to U32.
def ix_u8(c: Context, i: Nat) -> U32:
  match c:
    case Context{n, pid, +ix, accs, queue}:
      get_u8(as_bytes(ix), i)

# Little-endian u64 at byte offset `i`. Byte 0 is the discriminator.
def ix_u64(c: Context, i: Nat) -> Ty.u64:
  match c:
    case Context{n, pid, +ix, accs, queue}:
      le_u64(as_bytes(ix), i)

def accounts(c: Context) -> List<&2, Ty.Account>:
  match c:
    case Context{n, pid, ix, +accs, queue}:
      accs

def account(c: Context, i: Nat) -> Ty.Account:
  match c:
    case Context{n, pid, ix, +accs, queue}:
      account_at(accs, i)

def cpi_queue(c: Context) -> List<&2, Cpi>:
  match c:
    case Context{n, pid, ix, accs, +queue}:
      queue

def accounts_len(c: Context) -> U32:
  match c:
    case Context{+n, pid, ix, accs, queue}:
      n

def program_id(c: Context) -> Ty.Pubkey:
  match c:
    case Context{n, +pid, ix, accs, queue}:
      pid

# Head of the CPI queue.
def cpi(c: Context) -> Cpi:
  match c:
    case Context{n, pid, ix, accs, +queue}:
      cpi_head(queue)

# Second queued CPI.
def cpi1(c: Context) -> Cpi:
  match c:
    case Context{n, pid, ix, accs, +queue}:
      cpi_at(queue, 1n)

def set_account(c: Context, i: Nat, a: Ty.Account) -> Context:
  match c:
    case Context{n, pid, ix, accs, queue}:
      Context{n, pid, ix, set_accounts(accs, i, a), queue}

def context_set_cpis(c: Context, cpis: List<&2, Cpi>) -> Context:
  match c:
    case Context{n, pid, ix, accs, old}:
      Context{n, pid, ix, accs, cpis}

# Append a CPI, replacing a leading CpiNone.
def cpi_push(+c: Context, cpi: Cpi) -> Context:
  context_set_cpis(c, cpi_put(cpi_queue(c), cpi))

# Append. Does not overwrite a fixed slot.
def set_cpi(c: Context, cpi: Cpi) -> Context:
  cpi_push(c, cpi)

# Append, same as set_cpi.
def set_cpi1(c: Context, cpi1: Cpi) -> Context:
  cpi_push(c, cpi1)

# Queue an unsigned invoke. `program` and `accs` are account indexes.
# `ix` is instruction data as u32 words. Signer and writable flags are
# the ones on those accounts. The callee runs after this handler returns Ok.
def invoke(
  +c: Context,
  +program: U32,
  +accs: List<&2, U32>,
  +ix: List<&2, U32>,
) -> Context:
  cpi_push(c, CpiInvoke{program, 0, accs, ix, False{}, 0, 0})

# Queue an invoke signed by `seeds` plus `bump`. The book seed word stays 0.
def invoke_signed_seeds(
  +c: Context,
  +program: U32,
  +accs: List<&2, U32>,
  +ix: List<&2, U32>,
  +auth: U32,
  +bump: U32,
  +seeds: Ty.Seeds,
) -> Context:
  match seeds:
    case Ty.Seeds{mark}:
      cpi_push(
        c,
        CpiInvoke{program, auth, accs, ix, True{}, bump, U32.and(mark, 0)},
      )

# Queue an unsigned System transfer. from and to are account indexes.
def Sys.transfer(
  +c: Context,
  +from: U32,
  +to: U32,
  +lamports: Ty.u64,
) -> Context:
  cpi_push(c, CpiTransfer{from, to, lamports, False{}, 0, 0})

# Queue a System transfer signed by `seeds` plus `bump`.
# The chain signer is that list. The book Cpi seed word stays 0.
def Sys.transfer_signed_seeds(
  +c: Context,
  +from: U32,
  +to: U32,
  +lamports: Ty.u64,
  +bump: U32,
  +seeds: Ty.Seeds,
) -> Context:
  match seeds:
    case Ty.Seeds{mark}:
      cpi_push(
        c,
        CpiTransfer{from, to, lamports, True{}, bump, U32.and(mark, 0)},
      )

# Queue an unsigned SPL Token TransferChecked. Indexes are account slots.
def Token.transfer(
  +c: Context,
  +from: U32,
  +to: U32,
  +mint: U32,
  +auth: U32,
  +amount: Ty.u64,
  +decimals: U32,
) -> Context:
  cpi_push(
    c,
    CpiTokenTransfer{from, to, mint, auth, amount, decimals, False{}, 0, 0},
  )

# Queue a TransferChecked signed by `seeds` plus `bump`.
def Token.transfer_signed_seeds(
  +c: Context,
  +from: U32,
  +to: U32,
  +mint: U32,
  +auth: U32,
  +amount: Ty.u64,
  +decimals: U32,
  +bump: U32,
  +seeds: Ty.Seeds,
) -> Context:
  match seeds:
    case Ty.Seeds{mark}:
      cpi_push(
        c,
        CpiTokenTransfer{
          from,
          to,
          mint,
          auth,
          amount,
          decimals,
          True{},
          bump,
          U32.and(mark, 0),
        },
      )

# Queue a System CreateAccount. `payer` funds `new`.
# When `sign`, `seeds` plus `bump` sign the new account.
# Owner is this program.
def Sys.create_account(
  +c: Context,
  +payer: U32,
  +new: U32,
  +lamports: Ty.u64,
  +space: Ty.u64,
  +sign: Bool,
  +bump: U32,
  +seeds: Ty.Seeds,
) -> Context:
  match seeds:
    case Ty.Seeds{mark}:
      cpi_push(
        c,
        CpiCreate{payer, new, lamports, space, sign, bump, U32.and(mark, 0)},
      )

# Queue a System Assign. `account` is the index; `owner` is the new owner.
def Sys.assign(+c: Context, +account: U32, +owner: Ty.Pubkey) -> Context:
  cpi_push(c, CpiAssign{account, 0, owner, False{}, 0, 0})

# Queue a System Assign signed by `seeds` plus `bump`.
def Sys.assign_signed_seeds(
  +c: Context,
  +account: U32,
  +owner: Ty.Pubkey,
  +bump: U32,
  +seeds: Ty.Seeds,
) -> Context:
  match seeds:
    case Ty.Seeds{mark}:
      cpi_push(c, CpiAssign{account, 0, owner, True{}, bump, U32.and(mark, 0)})

# Queue a System Allocate. `account` is the index; `space` is the byte length.
def Sys.allocate(+c: Context, +account: U32, +space: Ty.u64) -> Context:
  cpi_push(c, CpiAllocate{account, 0, space, False{}, 0, 0})

# Queue a System Allocate signed by `seeds` plus `bump`.
def Sys.allocate_signed_seeds(
  +c: Context,
  +account: U32,
  +space: Ty.u64,
  +bump: U32,
  +seeds: Ty.Seeds,
) -> Context:
  match seeds:
    case Ty.Seeds{mark}:
      cpi_push(
        c,
        CpiAllocate{account, 0, space, True{}, bump, U32.and(mark, 0)},
      )

# SPL Token Approve (disc 4). `from` source, `to` delegate, `auth` owner.
def Token.approve(
  +c: Context,
  +from: U32,
  +to: U32,
  +auth: U32,
  +amount: Ty.u64,
) -> Context:
  cpi_push(
    c,
    CpiTokenIx{4, from, to, auth, amount, 0, Ty.Pubkey.zero(), False{}, 0, 0},
  )

# SPL Token Revoke (disc 5).
def Token.revoke(+c: Context, +from: U32, +auth: U32) -> Context:
  cpi_push(
    c,
    CpiTokenIx{
      5,
      from,
      0,
      auth,
      Ty.u64.zero(),
      0,
      Ty.Pubkey.zero(),
      False{},
      0,
      0,
    },
  )

# SPL Token MintTo (disc 7). `from` mint, `to` destination, `auth` mint authority.
def Token.mint_to(
  +c: Context,
  +from: U32,
  +to: U32,
  +auth: U32,
  +amount: Ty.u64,
) -> Context:
  cpi_push(
    c,
    CpiTokenIx{7, from, to, auth, amount, 0, Ty.Pubkey.zero(), False{}, 0, 0},
  )

# SPL Token FreezeAccount (disc 10). `from` account, `to` mint, `auth` freeze authority.
def Token.freeze(+c: Context, +from: U32, +to: U32, +auth: U32) -> Context:
  cpi_push(
    c,
    CpiTokenIx{
      10,
      from,
      to,
      auth,
      Ty.u64.zero(),
      0,
      Ty.Pubkey.zero(),
      False{},
      0,
      0,
    },
  )

# SPL Token ThawAccount (disc 11).
def Token.thaw(+c: Context, +from: U32, +to: U32, +auth: U32) -> Context:
  cpi_push(
    c,
    CpiTokenIx{
      11,
      from,
      to,
      auth,
      Ty.u64.zero(),
      0,
      Ty.Pubkey.zero(),
      False{},
      0,
      0,
    },
  )

# SPL Token InitializeMint2 (disc 20). `from` is the mint. `owner` is mint+freeze authority.
def Token.initialize_mint(
  +c: Context,
  +from: U32,
  +decimals: U32,
  +owner: Ty.Pubkey,
) -> Context:
  cpi_push(
    c,
    CpiTokenIx{20, from, 0, 0, Ty.u64.zero(), decimals, owner, False{}, 0, 0},
  )

# SPL Token InitializeAccount3 (disc 18). `from` account, `to` mint, `owner` token owner.
def Token.initialize_account(
  +c: Context,
  +from: U32,
  +to: U32,
  +owner: Ty.Pubkey,
) -> Context:
  cpi_push(
    c,
    CpiTokenIx{18, from, to, 0, Ty.u64.zero(), 0, owner, False{}, 0, 0},
  )

# Associated Token Ty.Account program id (ATokenGPvbdGVxr1b2hvZbsiqW5xWH25efTNsLJA8knL).
def ATA.id() -> Ty.Pubkey:
  Ty.Pk{
    2401605516,
    4052296782,
    688930235,
    2198703636,
    2568182283,
    2215706586,
    3631975940,
    1509485019,
  }

# Derive the ATA address for `wallet` + `mint` (Tokenkeg). Native on chain.
def ATA.address(+wallet: Ty.Pubkey, +mint: Ty.Pubkey) -> Ty.Pubkey:
  wallet

# Queue ATA Create (disc 0).
def ATA.create(
  +c: Context,
  +payer: U32,
  +ata: U32,
  +wallet: U32,
  +mint: U32,
  +system: U32,
  +token: U32,
) -> Context:
  cpi_push(c, CpiAta{0, payer, ata, wallet, mint, system, token})

# Queue ATA CreateIdempotent (disc 1).
def ATA.create_idempotent(
  +c: Context,
  +payer: U32,
  +ata: U32,
  +wallet: U32,
  +mint: U32,
  +system: U32,
  +token: U32,
) -> Context:
  cpi_push(c, CpiAta{1, payer, ata, wallet, mint, system, token})

# accounts_len is at least k.
def require_accounts(c: Context, k: U32) -> Bool:
  U32.is_ge(accounts_len(c), k)

# Checked arithmetic and small masks used by the examples.
# counter_inc wraps; the checked ones return None on failure.
def counter_inc(+value: Ty.u64) -> Ty.u64:
  Ty.u64.add(value, Ty.u64.from_u32(1))

# None when z (the value is already 0).
def dec_if(z: Bool, value: Ty.u64) -> Maybe<&2, Ty.u64>:
  match z:
    case True{}:
      None{}
    case False{}:
      Some{Ty.u64.sub(value, Ty.u64.from_u32(1))}

# value - 1, or None at 0.
def counter_dec(+value: Ty.u64) -> Maybe<&2, Ty.u64>:
  dec_if(Ty.u64.is_zero(value), value)

# None when lt (bal < amt).
def withdraw_if(lt: Bool, bal: Ty.u64, amt: Ty.u64) -> Maybe<&2, Ty.u64>:
  match lt:
    case True{}:
      None{}
    case False{}:
      Some{Ty.u64.sub(bal, amt)}

# bal - amt, or None when amt exceeds bal.
def try_withdraw(+bal: Ty.u64, +amt: Ty.u64) -> Maybe<&2, Ty.u64>:
  withdraw_if(Ty.u64.is_lt(bal, amt), bal, amt)

# None when ovf (the add would wrap).
def deposit_if(ovf: Bool, bal: Ty.u64, amt: Ty.u64) -> Maybe<&2, Ty.u64>:
  match ovf:
    case True{}:
      None{}
    case False{}:
      Some{Ty.u64.add(bal, amt)}

# bal + amt, or None on overflow.
def try_deposit(+bal: Ty.u64, +amt: Ty.u64) -> Maybe<&2, Ty.u64>:
  deposit_if(Ty.u64.would_ovf(bal, amt), bal, amt)

# Count set bits, one shift per fuel step. fuel 0 returns acc.
def popcount.go(fuel: Nat, +mask: U32, acc: U32) -> U32:
  match fuel:
    case 0n:
      acc
    case 1n+p:
      popcount.go(p, U32.shr(mask), (acc + U32.and(mask, 1) : U32))

# Number of set bits in mask.
def popcount(mask: U32) -> U32:
  popcount.go(32n, mask, 0)

# Set bit i in one U32 half of a mask.
def bit_set(mask: U32, i: U32) -> U32:
  U32.or(mask, U32.shln(1, U32.to_nat(i)))

# Set bit i of a u64 mask.
def bit_set_u64(+mask: Ty.u64, +i: U32) -> Ty.u64:
  Ty.u64.or(mask, Ty.u64.shln(Ty.u64.from_nat(1n), U32.to_nat(i)))

def popcount_u64.go(fuel: Nat, +mask: Ty.u64, acc: Ty.u64) -> Ty.u64:
  match fuel:
    case 0n:
      acc
    case 1n+p:
      popcount_u64.go(
        p,
        Ty.u64.shr(mask),
        Ty.u64.add(acc, Ty.u64.and(mask, Ty.u64.from_nat(1n))),
      )

# Number of set bits in a u64 mask.
def popcount_u64(+mask: Ty.u64) -> Ty.u64:
  popcount_u64.go(64n, mask, Ty.u64.from_nat(0n))

# At least `threshold` bits are set in the approval mask.
def can_execute(+approvals: Ty.u64, +threshold: Ty.u64) -> Bool:
  Ty.u64.is_ge(popcount_u64(approvals), threshold)

# t when b, otherwise f.
def pick_u32(b: Bool, t: U32, f: U32) -> U32:
  match b:
    case True{}:
      t
    case False{}:
      f

# Some{v} when b, otherwise None.
def maybe_u32(b: Bool, v: U32) -> Maybe<&2, U32>:
  match b:
    case True{}:
      Some{v}
    case False{}:
      None{}

# Some{0} if is0, else Some{1} if is1, else None.
def owner_pick(is0: Bool, is1: Bool) -> Maybe<&2, U32>:
  match is0:
    case True{}:
      Some{0}
    case False{}:
      maybe_u32(is1, 1)

# Index of key among the first n owners, or None.
def owner_index(
  +pk: Ty.Pubkey,
  +o0: Ty.Pubkey,
  +o1: Ty.Pubkey,
  +n: Ty.u64,
) -> Maybe<&2, U32>:
  owner_pick(
    Bool.and(Ty.u64.is_gt(n, Ty.u64.zero()), Ty.Pubkey.eq(pk, o0)),
    Bool.and(Ty.u64.is_gt(n, Ty.u64.from_u32(1)), Ty.Pubkey.eq(pk, o1)),
  )