fix(symbolic-spl): discharge AccountState discriminant branches

Author: Stevengre

kmir/src/kmir/kdist/mir-semantics/lemmas/kmir-lemmas.md

@@ -94,6 +94,10 @@ For symbolic enum values, the variant index remains unevaluated but the original
     requires isOneOf(DISCR, DISCRS)
     [simplification, symbolic(DISCR)]
 
+  // Normalize enum-discriminant comparisons into the same shape emitted by branch splits.
+  rule #lookupDiscrAux(DISCRS, IDX) ==Int DISCR => DISCR ==Int #lookupDiscrAux(DISCRS, IDX)
+    [simplification]
+
   rule 0 <=Int asInt(#findVariantIdxAux(DISCR, DISCRS, _)) => true
     requires isOneOf(DISCR, DISCRS)
     [simplification, symbolic(DISCR)]

kmir/src/kmir/kdist/mir-semantics/symbolic/spl-token.md

@@ -362,6 +362,9 @@ The `#initBorrow` helper resets borrow counters to 0 and sets the correct dynami
       andBool #isSplPubkey(?SplDelegateKey)
       andBool 0 <=Int ?SplAmount andBool ?SplAmount <Int (1 <<Int 64)
       andBool 0 <=Int ?SplAccountState andBool ?SplAccountState <=Int 2
+      andBool (0 ==Int #lookupDiscrAux(discriminant(0) discriminant(1) discriminant(2) .Discriminants, variantIdx(?SplAccountState))
+               orBool 1 ==Int #lookupDiscrAux(discriminant(0) discriminant(1) discriminant(2) .Discriminants, variantIdx(?SplAccountState))
+               orBool 2 ==Int #lookupDiscrAux(discriminant(0) discriminant(1) discriminant(2) .Discriminants, variantIdx(?SplAccountState)))
       andBool 0 <=Int ?SplDelegatedAmount andBool ?SplDelegatedAmount <Int (1 <<Int 64)
       andBool 0 <=Int ?SplIsNativeLamportsVariant andBool ?SplIsNativeLamportsVariant <=Int 1
       andBool (0 ==Int #lookupDiscrAux(discriminant(0) discriminant(1) .Discriminants, variantIdx(?SplIsNativeLamportsVariant)) orBool 1 ==Int #lookupDiscrAux(discriminant(0) discriminant(1) .Discriminants, variantIdx(?SplIsNativeLamportsVariant)))

kmir/src/tests/integration/data/prove-rs/spl-accountstate-discriminant-branch.rs

@@ -0,0 +1,116 @@
+use std::cell::RefCell;
+use std::convert::TryInto;
+use std::rc::Rc;
+
+type ProgramResult = Result<(), ProgramError>;
+
+fn main() {
+    let _keep: fn(&AccountInfo<'_>) -> ProgramResult = repro;
+}
+
+#[no_mangle]
+pub fn repro(source_account_info: &AccountInfo<'_>) -> ProgramResult {
+    cheatcode_is_spl_account(source_account_info);
+    let account = get_account(source_account_info)?;
+
+    if account.state == AccountState::Frozen {
+        Err(ProgramError::Custom(7))
+    } else {
+        Ok(())
+    }
+}
+
+#[inline(never)]
+fn cheatcode_is_spl_account(_: &AccountInfo<'_>) {}
+
+#[derive(Clone)]
+struct AccountInfo<'a> {
+    key: &'a Pubkey,
+    lamports: Rc<RefCell<&'a mut u64>>,
+    data: Rc<RefCell<&'a mut [u8]>>,
+    owner: &'a Pubkey,
+    rent_epoch: u64,
+    is_signer: bool,
+    is_writable: bool,
+    executable: bool,
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+struct Pubkey([u8; 32]);
+
+impl Pubkey {
+    fn new(bytes: [u8; 32]) -> Self {
+        Self(bytes)
+    }
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+enum ProgramError {
+    InvalidAccountData,
+    Custom(u32),
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+enum AccountState {
+    Uninitialized,
+    Initialized,
+    Frozen,
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+enum COption<T> {
+    None,
+    Some(T),
+}
+
+#[derive(Clone, Copy, Debug)]
+struct Account {
+    mint: Pubkey,
+    owner: Pubkey,
+    amount: u64,
+    delegate: COption<Pubkey>,
+    state: AccountState,
+    is_native: COption<u64>,
+    delegated_amount: u64,
+    close_authority: COption<Pubkey>,
+}
+
+impl Account {
+    const LEN: usize = 165;
+
+    fn unpack_from_slice(data: &[u8]) -> Result<Self, ProgramError> {
+        if data.len() < Self::LEN {
+            return Err(ProgramError::InvalidAccountData);
+        }
+
+        let mint = Pubkey::new(data[0..32].try_into().unwrap());
+        let owner = Pubkey::new(data[32..64].try_into().unwrap());
+        let amount = u64::from_le_bytes(data[64..72].try_into().unwrap());
+        let state = match data[108] {
+            0 => AccountState::Uninitialized,
+            1 => AccountState::Initialized,
+            2 => AccountState::Frozen,
+            _ => return Err(ProgramError::InvalidAccountData),
+        };
+
+        Ok(Self {
+            mint,
+            owner,
+            amount,
+            delegate: COption::None,
+            state,
+            is_native: COption::None,
+            delegated_amount: 0,
+            close_authority: COption::None,
+        })
+    }
+
+    fn unpack_unchecked(data: &[u8]) -> Result<Self, ProgramError> {
+        Self::unpack_from_slice(data)
+    }
+}
+
+#[inline(never)]
+fn get_account(acc: &AccountInfo<'_>) -> Result<Account, ProgramError> {
+    Account::unpack_unchecked(&acc.data.borrow())
+}

kmir/src/tests/integration/test_integration.py

@@ -39,6 +39,7 @@
     'transmute-bytes': ['bytes_to_u64', 'u64_to_bytes'],
     'test_offset_from-fail': ['testing'],
     'iter-eq-copied-take-dereftruncate': ['repro'],
+    'spl-accountstate-discriminant-branch': ['repro'],
     'spl-multisig-iter-eq-copied-next-fail': ['repro'],
     'spl-multisig-signer-index': ['repro'],
 }