fix(rt): generalize direct-tag enum decoding to any variant count and discriminant (#955)

## Summary

Replace the specialized \`#decodeOptionTag01\` rule (which only handled
exactly 2 variants with discriminants 0/1 and the Option-like 0+1 field
pattern) with a general \`#decodeEnumDirect\` decoder that handles:

- Any number of variants (not just two)
- Any discriminant values (not just 0 and 1)
- Any unsigned integer tag width (u8, u16, etc.)
- Any number of fields per variant

New K helper functions (all total with `owise` fallbacks):
- `#decodeEnumDirectFields`: given variant index, decode its fields
using per-variant offsets
- `#decodeEnumDirectTag`: read tag bytes as unsigned little-endian
integer
- `#nthTys`: index into per-variant field type lists
- `#nthVariantOffsets`: extract field offsets from nth variant layout

## Context

Before this change, direct-tag enum decoding was specialized to
`discriminants: 0, 1` and single-byte tag checks (modeled after
`Option`). Any enum with different discriminant values or more than two
variants would produce `UnableToDecode`.

## Changes

- `kmir/src/kmir/kdist/mir-semantics/rt/decoding.md`: replace
`#decodeOptionTag01` with general `#decodeEnumDirect` rules
- `kmir/src/tests/integration/data/prove-rs/enum-direct-tag-decode.rs`:
add Rust test covering repr(u8) 0/1, repr(u8) 2/5, repr(u16) 0/256, and
3-variant enums
- `kmir/src/tests/integration/test_integration.py`: register the new
prove-rs test

## Testing

- `make build` ✅
- `uv --directory kmir run pytest
kmir/src/tests/integration/test_integration.py -k
'enum-direct-tag-decode' -vv` ✅

---------

Co-authored-by: Daniel Cumming <124537596+dkcumming@users.noreply.github.com>

Author: Stevengre

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

@@ -297,27 +297,26 @@ If there are no fields, the enum can be decoded by using their data as the discr
    requires TAG =/=Int DISCRIMINANT
 ```
 
-#### Enums with two variants
+#### Enums with direct tag encoding and fields
 
-Having two variants with possibly zero or one field each is a very common case,
-it includes a number of standard library types such as `Option` and `Result`.
-
-The following rules are specialised to the case of encoding an `Option`.
-An important distinction here is whether or not the tag is niche-encoded.
-If the option holds data that has all-zero as a possible value, a separate tag is used, usually as the first field.
-In both cases we expect the tag to be in the single shared field, and the discriminant to be just 0 and 1.
+Decodes any enum with `tagEncodingDirect` and a `primitiveInt` tag,
+with arbitrary variant counts, discriminant values, and field counts per variant.
+The tag is read as an unsigned little-endian integer, matched against the discriminant
+list to find the variant index, and then the variant's fields are decoded using the
+per-variant layout offsets.
 
 ```k
+  // General entry rule: direct-tag enum with at least one field somewhere.
   rule #decodeValue(
          BYTES
        , typeInfoEnumType(...
            name: _
          , adtDef: _
-         , discriminants: discriminant(0) discriminant(1) .Discriminants
-         , fields: (.Tys : (FIELD_TYPE .Tys) : .Tyss)
+         , discriminants: DISCRIMINANTS
+         , fields: FIELD_TYPESS
          , layout:
             someLayoutShape(layoutShape(...
-                fields: fieldsShapeArbitrary(mk(... offsets: machineSize(0) .MachineSizes))
+                fields: _FIELDS
               , variants:
                   variantsShapeMultiple(
                     mk(...
@@ -329,7 +328,7 @@ In both cases we expect the tag to be in the single shared field, and the discri
                         )
                       , tagEncoding: tagEncodingDirect
                       , tagField: 0
-                      , variants: _VARIANTS
+                      , variants: VARIANT_LAYOUTS
                       )
                     )
               , abi: _ABI
@@ -338,22 +337,59 @@ In both cases we expect the tag to be in the single shared field, and the discri
             ))
          ) #as ENUM_TYPE
        )
-    => #decodeOptionTag01(BYTES, TAG_WIDTH, FIELD_TYPE, ENUM_TYPE)
-
-  syntax Evaluation ::= #decodeOptionTag01 ( Bytes , IntegerLength , Ty , TypeInfo ) [function, total]
-  // --------------------------------------------------------------------------------------
-  rule #decodeOptionTag01(BYTES, _LEN, _TY, _ENUM_TYPE)
-    => Aggregate(variantIdx(0), .List)
-    requires 0 ==Int BYTES[0] // expect only 0 or 1 as tags, so higher bytes do not matter
-    [preserves-definedness]
-  rule #decodeOptionTag01(BYTES,  LEN,  TY, _ENUM_TYPE)
-    => Aggregate(variantIdx(1), ListItem(#decodeValue(substrBytes(BYTES, #byteLength(LEN), lengthBytes(BYTES)), lookupTy(TY))))
-    requires 1 ==Int BYTES[0] // expect only 0 or 1 as tags, so higher bytes do not matter
+    => #decodeEnumDirectFields(
+         BYTES,
+         #findVariantIdx(#decodeEnumDirectTag(BYTES, TAG_WIDTH), DISCRIMINANTS),
+         FIELD_TYPESS,
+         VARIANT_LAYOUTS,
+         ENUM_TYPE
+       )
+    requires notBool #noFields(FIELD_TYPESS)
+
+  // ---------------------------------------------------------------------------
+  // #decodeEnumDirectFields: given the variant index, decode its fields
+  // ---------------------------------------------------------------------------
+  syntax Evaluation ::= #decodeEnumDirectFields ( Bytes , VariantIdx , Tyss , LayoutShapes , TypeInfo ) [function, total]
+  // --------------------------------------------------------------------------------------------------------------------------
+  rule #decodeEnumDirectFields(BYTES, variantIdx(IDX), FIELD_TYPESS, VARIANT_LAYOUTS, _ENUM_TYPE)
+    => Aggregate(
+         variantIdx(IDX),
+         #decodeFieldsWithOffsets(BYTES, #nthTys(FIELD_TYPESS, IDX), #nthVariantOffsets(VARIANT_LAYOUTS, IDX))
+       )
+    requires 0 <=Int IDX
     [preserves-definedness]
-  rule #decodeOptionTag01(BYTES, _LEN, _TY,  ENUM_TYPE)
+
+  // Error cases: variant not found or other failure
+  rule #decodeEnumDirectFields(BYTES, _, _FIELD_TYPESS, _VARIANT_LAYOUTS, ENUM_TYPE)
     => UnableToDecode(BYTES, ENUM_TYPE)
     [owise]
 
+  // ---------------------------------------------------------------------------
+  // #decodeEnumDirectTag: read the tag bytes as an unsigned little-endian int
+  // ---------------------------------------------------------------------------
+  syntax Int ::= #decodeEnumDirectTag ( Bytes , IntegerLength ) [function, total]
+  rule #decodeEnumDirectTag(BYTES, TAG_WIDTH)
+    => Bytes2Int(substrBytes(BYTES, 0, #byteLength(TAG_WIDTH)), LE, Unsigned)
+    requires lengthBytes(BYTES) >=Int #byteLength(TAG_WIDTH)
+    [preserves-definedness]
+  rule #decodeEnumDirectTag(_, _) => -1 [owise]
+
+  // ---------------------------------------------------------------------------
+  // List-indexing helpers
+  // ---------------------------------------------------------------------------
+
+  // Index into a Tyss (list of per-variant field type lists)
+  syntax Tys ::= #nthTys ( Tyss , Int ) [function, total]
+  rule #nthTys(TYS : _REST, 0) => TYS
+  rule #nthTys(_TYS : REST, N) => #nthTys(REST, N -Int 1) requires N >Int 0
+  rule #nthTys(_, _) => .Tys [owise]
+
+  // Index into variant layouts and extract field offsets in one step (total)
+  syntax MachineSizes ::= #nthVariantOffsets ( LayoutShapes , Int ) [function, total]
+  rule #nthVariantOffsets(layoutShape(fieldsShapeArbitrary(mk(OFFSETS)), _, _, _, _) _REST, 0) => OFFSETS
+  rule #nthVariantOffsets(_L REST, N) => #nthVariantOffsets(REST, N -Int 1) requires N >Int 0
+  rule #nthVariantOffsets(_, _) => .MachineSizes [owise]
+
   syntax Int ::= #byteLength ( IntegerLength ) [function, total]
   // -----------------------------------------------------------
   rule #byteLength(integerLengthI8  ) => 1

kmir/src/tests/integration/data/prove-rs/enum-direct-tag-decode.rs

@@ -0,0 +1,72 @@
+/// Comprehensive test for direct-tag enum decoding.
+///
+/// Covers:
+///   - `#[repr(u8)]` with discriminants 0/1 (baseline, like Option)
+///   - `#[repr(u8)]` with discriminants 2/5 (non-zero start)
+///   - `#[repr(u16)]` with discriminants 0/256 (wide tag exceeding u8 range)
+///   - `#[repr(u8)]` with 3 variants (beyond two-variant restriction)
+
+// --- Baseline: repr(u8), disc 0 / 1 -------------------------------------------
+
+#[repr(u8)]
+enum TwoU8 {
+    A(bool) = 0,
+    B(bool) = 1,
+}
+
+const TWO_U8_A: TwoU8 = TwoU8::A(true);
+const TWO_U8_B: TwoU8 = TwoU8::B(false);
+
+// --- Non-zero discriminants: repr(u8), disc 2 / 5 -----------------------------
+
+#[repr(u8)]
+enum NonZero {
+    X(bool) = 2,
+    Y(bool) = 5,
+}
+
+const NZ_X: NonZero = NonZero::X(false);
+const NZ_Y: NonZero = NonZero::Y(true);
+
+// --- Wide tag: repr(u16), disc 0 / 256 ----------------------------------------
+
+#[repr(u16)]
+enum WideTag {
+    Lo(bool) = 0,
+    Hi(bool) = 256,
+}
+
+const WIDE_LO: WideTag = WideTag::Lo(true);
+const WIDE_HI: WideTag = WideTag::Hi(false);
+
+// --- Three variants: repr(u8), disc 0 / 1 / 2 ---------------------------------
+
+#[repr(u8)]
+enum Three {
+    First(bool)  = 0,
+    Second(bool) = 1,
+    Third(bool)  = 2,
+}
+
+const THREE_A: Three = Three::First(true);
+const THREE_B: Three = Three::Second(false);
+const THREE_C: Three = Three::Third(true);
+
+fn main() {
+    // Baseline
+    assert!(matches!(TWO_U8_A, TwoU8::A(true)));
+    assert!(matches!(TWO_U8_B, TwoU8::B(false)));
+
+    // Non-zero discriminants
+    assert!(matches!(NZ_X, NonZero::X(false)));
+    assert!(matches!(NZ_Y, NonZero::Y(true)));
+
+    // Wide tag
+    assert!(matches!(WIDE_LO, WideTag::Lo(true)));
+    assert!(matches!(WIDE_HI, WideTag::Hi(false)));
+
+    // Three variants
+    assert!(matches!(THREE_A, Three::First(true)));
+    assert!(matches!(THREE_B, Three::Second(false)));
+    assert!(matches!(THREE_C, Three::Third(true)));
+}