Transmute integer to fieldless enum

kmir/src/kmir/kdist/mir-semantics/rt/data.md

@@ -1560,23 +1560,58 @@ Casting an integer to a `[u8; N]` array materialises its little-endian bytes.
   rule #staticArrayLenBits(_OTHER) => 0 [owise]
 ```
 
-Another specialisation is getting the discriminant of `enum`s without fields after converting some integer data to it
-(see `#discriminant` and `rvalueDiscriminant`).
-If none of the `enum` variants has any fields, the `Transmute` of a number to the `enum` data is necessarily just the discriminant itself., and can be returned as the integer value afgter adjusting to the byte length of the discriminant:
+A transmutation from an integer to an enum is wellformed if:
+- The bit width of the incoming integer is the same as the discriminant type of the enum
+    (e.g. `u8 -> i8` fine, `u8 -> u16` not fine) - this is guaranteed by the compiler;
+- The incoming integer has a bit pattern that matches a discriminant of the enum
+    (e.g. `255_u8` and `-1_i8` fine iff `0b1111_1111` is a discriminant of the enum);
+
+Note that discriminants are stored as `u128` in the type data even if they are signed
+or unsigned at the source level. This means that our approach to soundly transmute an
+integer into a enum is to treat the incoming integer as unsigned (converting if signed),
+and check if the value is in the discriminants. If yes, find the corresponding variant
+index; if not, return `#UBErrorInvalidDiscriminantsInEnumCast`.
 
 ```k
-  rule <k> #discriminant(
-              thunk(#cast (Integer(DATA, _, false), castKindTransmute, _, TY)),
-              TY
-            ) => Integer(DATA, 0, false) // HACK: bit width 0 means "flexible"
-          ...
-        </k>
-    requires #isEnumWithoutFields(lookupTy(TY))
-
   syntax Bool ::= #isEnumWithoutFields ( TypeInfo ) [function, total]
   // ----------------------------------------------------------------
   rule #isEnumWithoutFields(typeInfoEnumType(_, _, _, FIELDSS, _)) => #noFields(FIELDSS)
   rule #isEnumWithoutFields(_OTHER) => false [owise]
+
+  // TODO: Connect this with MirError
+  syntax Evaulation ::= "#UBErrorInvalidDiscriminantsInEnumCast"
+  rule <k>
+           #cast( Integer ( VAL , WIDTH , _SIGNED ) , castKindTransmute , _TY_FROM , TY_TO ) ~> _REST
+        =>
+           #UBErrorInvalidDiscriminantsInEnumCast
+      </k>
+      requires #isEnumWithoutFields(lookupTy(TY_TO))
+        andBool notBool #validDiscriminant( truncate(VAL, WIDTH, Unsigned) , lookupTy(TY_TO) )
+
+  rule <k>
+           #cast( Integer ( VAL , WIDTH , _SIGNED ) , castKindTransmute , _TY_FROM , TY_TO )
+        =>
+           Aggregate( #findVariantIdxFromTy( truncate(VAL, WIDTH, Unsigned), lookupTy(TY_TO) ) , .List )
+       ...
+      </k>
+      requires #isEnumWithoutFields(lookupTy(TY_TO))
+        andBool #validDiscriminant( truncate(VAL, WIDTH, Unsigned) , lookupTy(TY_TO))
+
+  syntax VariantIdx ::= #findVariantIdxFromTy ( Int , TypeInfo ) [function, total]
+  //------------------------------------------------------------------------------
+  rule #findVariantIdxFromTy( VAL , typeInfoEnumType(_, _, DISCRIMINANTS, _, _) ) => #findVariantIdx( VAL, DISCRIMINANTS)
+  rule #findVariantIdxFromTy( _ , _ ) => err("NotAnEnum") [owise]
+
+  syntax Bool ::= #validDiscriminant    ( Int , TypeInfo )      [function, total]
+  // ----------------------------------------------------------------------------
+  rule #validDiscriminant( VAL , typeInfoEnumType(_, _, DISCRIMINANTS, _, _) ) => #validDiscriminantAux( VAL , DISCRIMINANTS )
+  rule #validDiscriminant( _ , _ ) => false [owise]
+
+  syntax Bool ::= #validDiscriminantAux ( Int , Discriminants ) [function, total]
+  // ----------------------------------------------------------------------------
+  rule #validDiscriminantAux( VAL, discriminant(mirInt(DISCRIMINANT)) REST ) => VAL ==Int DISCRIMINANT orBool #validDiscriminantAux( VAL, REST )
+  rule #validDiscriminantAux( VAL, discriminant(    DISCRIMINANT    ) REST ) => VAL ==Int DISCRIMINANT orBool #validDiscriminantAux( VAL, REST )
+  rule #validDiscriminantAux( _VAL, .Discriminants ) => false
 ```
 
 

kmir/src/tests/integration/data/prove-rs/show/transmute-u8-to-enum-changed-discriminant-signed-fail.main.expected

@@ -0,0 +1,15 @@
+
+┌─ 1 (root, init)
+│   #execTerminator ( terminator ( ... kind: terminatorKindCall ( ... func: operandC
+│   span: 0
+│
+│  (100 steps)
+└─ 3 (stuck, leaf)
+    #traverseProjection ( toLocal ( 2 ) , thunk ( #decodeConstant ( constantKindAllo
+    span: 153
+
+
+┌─ 2 (root, leaf, target, terminal)
+│   #EndProgram ~> .K
+
+

kmir/src/tests/integration/data/prove-rs/show/transmute-u8-to-enum-fail.main.expected

@@ -0,0 +1,15 @@
+
+┌─ 1 (root, init)
+│   #execTerminator ( terminator ( ... kind: terminatorKindCall ( ... func: operandC
+│   span: 0
+│
+│  (13 steps)
+└─ 3 (stuck, leaf)
+    #UBErrorInvalidDiscriminantsInEnumCast ~> .K
+    function: main
+
+
+┌─ 2 (root, leaf, target, terminal)
+│   #EndProgram ~> .K
+
+

kmir/src/tests/integration/data/prove-rs/transmute-try-from-symbolic.rs

@@ -0,0 +1,40 @@
+use std::convert::TryFrom;
+
+pub enum ProgramError {
+    InvalidAccountData
+}
+
+#[repr(u8)]
+#[derive(Clone, Copy, Debug, PartialEq)]
+pub enum AccountState {
+    Uninitialized,
+    Initialized,
+    Frozen,
+}
+
+impl TryFrom<u8> for AccountState {
+    type Error = ProgramError;
+
+    #[inline(always)]
+    fn try_from(value: u8) -> Result<Self, Self::Error> {
+        match value {
+            // SAFETY: `value` is guaranteed to be in the range of the enum variants.
+            0..=2 => Ok(unsafe { core::mem::transmute::<u8, AccountState>(value) }),
+            _ => Err(ProgramError::InvalidAccountData),
+        }
+    }
+}
+
+fn main() {
+    num_into_state(0);
+}
+
+fn num_into_state(arg: u8) {
+    let state: Result<AccountState, ProgramError> = AccountState::try_from(arg);
+
+    if arg <= 2 {
+        assert!(state.is_ok());
+    } else {
+        assert!(state.is_err());
+    }
+}

kmir/src/tests/integration/data/prove-rs/transmute-u8-to-enum-changed-discriminant-signed-fail.rs

@@ -0,0 +1,20 @@
+#[repr(i8)]
+#[derive(Clone, Copy, Debug, PartialEq)]
+pub enum AccountState {
+    Uninitialized = 0,
+    Initialized = 1,
+    Frozen = -1,
+}
+
+fn main() {
+    unsafe {
+        assert_eq!(core::mem::transmute::<u8, AccountState>(0), AccountState::Uninitialized);
+        assert_eq!(core::mem::transmute::<i8, AccountState>(0), AccountState::Uninitialized);
+
+        assert_eq!(core::mem::transmute::<u8, AccountState>(1), AccountState::Initialized);
+        assert_eq!(core::mem::transmute::<i8, AccountState>(1), AccountState::Initialized);
+
+        assert_eq!(core::mem::transmute::<u8, AccountState>(255), AccountState::Frozen);
+        assert_eq!(core::mem::transmute::<i8, AccountState>(-1), AccountState::Frozen);
+    }
+}

kmir/src/tests/integration/data/prove-rs/transmute-u8-to-enum-changed-discriminant.rs

@@ -0,0 +1,22 @@
+#[repr(u8)]
+#[derive(Clone, Copy, Debug, PartialEq)]
+pub enum AccountState {
+    Uninitialized = 11,
+    Initialized = 22,
+    Frozen = 255,
+}
+
+fn main() {
+    unsafe {
+        assert_eq!(core::mem::transmute::<u8, AccountState>(11), AccountState::Uninitialized);
+        assert_eq!(core::mem::transmute::<u8, AccountState>(22), AccountState::Initialized);
+        assert_eq!(core::mem::transmute::<u8, AccountState>(255), AccountState::Frozen);
+
+        // Compiler rejects transmuting between different width types
+        // assert_eq!(core::mem::transmute::<u64, AccountState>(-1), AccountState::Frozen);
+
+        // The within the same width the type can be different, here an `-1_i8` is used instead
+        // of a `255_u8`, but the bit pattern matches it creates the correct type
+        assert_eq!(core::mem::transmute::<i8, AccountState>(-1), AccountState::Frozen);
+    }
+}

kmir/src/tests/integration/data/prove-rs/transmute-u8-to-enum-fail.rs

@@ -0,0 +1,15 @@
+#[repr(u8)]
+#[derive(Clone, Copy, Debug, PartialEq)]
+pub enum AccountState {
+    Uninitialized,
+    Initialized,
+    Frozen,
+}
+
+fn main() {
+    unsafe {
+        std::hint::black_box(
+            core::mem::transmute::<u8, AccountState>(4)
+        );
+    }
+}

kmir/src/tests/integration/data/prove-rs/transmute-u8-to-enum.rs

@@ -0,0 +1,15 @@
+#[repr(u8)]
+#[derive(Clone, Copy, Debug, PartialEq)]
+pub enum AccountState {
+    Uninitialized,
+    Initialized,
+    Frozen,
+}
+
+fn main() {
+    unsafe {
+        assert_eq!(core::mem::transmute::<u8, AccountState>(0), AccountState::Uninitialized);
+        assert_eq!(core::mem::transmute::<u8, AccountState>(1), AccountState::Initialized);
+        assert_eq!(core::mem::transmute::<u8, AccountState>(2), AccountState::Frozen);
+    }
+}

kmir/src/tests/integration/test_integration.py

@@ -52,6 +52,8 @@
     'niche-enum',
     'assume-cheatcode-conflict-fail',
     'raw-ptr-cast-fail',
+    'transmute-u8-to-enum-fail',
+    'transmute-u8-to-enum-changed-discriminant-signed-fail',
 ]