Repository: cryspen/hax (language: rust)

Issue:
Backend/Engine: `cast_to_repr` not counting correctly
This:
```rust
enum SimpleEnum {
    A,
    B,
    C,
    D,
    E,
}
```
produces this f* code:
```fstar
type t_SimpleEnum =
  | SimpleEnum_A : t_SimpleEnum
  | SimpleEnum_B : t_SimpleEnum
  | SimpleEnum_C : t_SimpleEnum
  | SimpleEnum_D : t_SimpleEnum
  | SimpleEnum_E : t_SimpleEnum

let t_SimpleEnum_cast_to_repr (x: t_SimpleEnum) : isize =
  match x with
  | SimpleEnum_A  -> isz 0
  | SimpleEnum_B  -> isz 1
  | SimpleEnum_C  -> isz 3
  | SimpleEnum_D  -> isz 6
  | SimpleEnum_E  -> isz 10
```

The `t_SimpleEnum_cast_to_repr` should be counting from 0 to 4 but instead counts in this pattern:
A -> 0
B -> A + 1
C -> B + 2
D -> C + 3
E -> D + 4

The bug is that it's adding the running count to the previous variant instead of the first variant.

The behavior seems correct for custom counters:
```rust
enum CustomEnum {
    A = 10,
    B,
    C,
    D = 20,
    E,
}
```
```fstar
let discriminant_CustomEnum_A: isize = isz 10

let discriminant_CustomEnum_D: isize = isz 20

type t_CustomEnum =
  | CustomEnum_A : t_CustomEnum
  | CustomEnum_B : t_CustomEnum
  | CustomEnum_C : t_CustomEnum
  | CustomEnum_D : t_CustomEnum
  | CustomEnum_E : t_CustomEnum

let t_CustomEnum_cast_to_repr (x: t_CustomEnum) : isize =
  match x with
  | CustomEnum_A  -> discriminant_CustomEnum_A
  | CustomEnum_B  -> discriminant_CustomEnum_A +! isz 1
  | CustomEnum_C  -> discriminant_CustomEnum_A +! isz 2
  | CustomEnum_D  -> discriminant_CustomEnum_D
  | CustomEnum_E  -> discriminant_CustomEnum_D +! isz 1
```
But more comprehensive testing might be needed in general.

[Open this code snippet in the playground](https://hax-playground.cryspen.com/#fstar/4291b195f4/gist=e716d98e0208fa954e1d76756b37a3cd)



Discussion / hints:
This issue has been marked as stale due to a lack of activity for 60 days. If you believe this issue is still relevant, please provide an update or comment to keep it open. Otherwise, it will be closed in 7 days.
This issue has been marked as stale due to a lack of activity for 60 days. If you believe this issue is still relevant, please provide an update or comment to keep it open. Otherwise, it will be closed in 7 days.
I had a quick look, the problem is [here](https://github.com/cryspen/hax/blob/321a6137ee93ddc547c3ca7357b46188d430ad19/engine/lib/import_thir.ml#L1409), we add the relative path to the sum of all relative indices, not sure about the right way to fix this (what is the reference of relative indices?)

Fix patch:
diff --git a/engine/lib/import_thir.ml b/engine/lib/import_thir.ml
index df9ff0aa5..ad0c272c5 100644
--- a/engine/lib/import_thir.ml
+++ b/engine/lib/import_thir.ml
@@ -37,6 +37,7 @@ end
 
 module U = Ast_utils.Make (Features.Rust)
 module W = Features.On
+module Ast_builder = Ast_builder.Make (Features.Rust)
 open Ast
 
 let def_id ~value (def_id : Thir.def_id) : global_ident =
@@ -1358,95 +1359,60 @@ let generic_param_to_value ({ ident; kind; span; _ } : generic_param) :
   | GPType -> GType (TParam ident)
   | GPConst { typ } -> GConst { e = LocalVar ident; typ; span }
 
-type discriminant_expr =
-  | Lit of Int64.t
-  | Exp of expr  (** Helper type for [cast_of_enum]. *)
-
 (** Generate a cast function from an inductive to its represantant type. *)
 let cast_of_enum typ_name generics typ thir_span
     (variants : (variant * Types.variant_for__decorated_for__expr_kind) list) :
     item =
+  let span = Span.of_thir thir_span in
+  let (module M) = Ast_builder.make span in
   let self =
-    TApp
-      {
-        ident = `Concrete typ_name;
-        args = List.map ~f:generic_param_to_value generics.params;
-      }
+    let args = List.map ~f:generic_param_to_value generics.params in
+    TApp { ident = `Concrete typ_name; args }
   in
-  let span = Span.of_thir thir_span in
-  let init = Lit (Int64.of_int 0) in
-  let to_expr (n : Int64.t) : expr =
-    match typ with
-    | TInt kind ->
-        let value = Int64.to_string n in
-        {
-          e = Literal (Int { value; negative = Int64.is_negative n; kind });
-          span;
-          typ;
-        }
-    | typ ->
-        assertion_failure [ thir_span ]
-        @@ "disc_literal_to_expr: got repr type "
-        ^ [%show: ty] typ
+  let expr_of_int (n : Int64.t) : expr =
+    let kind =
+      match typ with
+      | TInt kind -> kind
+      | typ ->
+          assertion_failure [ thir_span ]
+            ("cast_of_enum: expected in type, got " ^ [%show: ty] typ)
+    in
+    let value = Int64.to_string n in
+    M.expr_Literal ~typ (Int { value; negative = Int64.is_negative n; kind })
   in
   let arms =
-    List.folding_map variants ~init ~f:(fun acc (variant, thir_variant) ->
+    (* Each variant comes with a [rustc_middle::ty::VariantDiscr]. Some variant have [Explicit] discr (i.e. an expression)
+       while other have [Relative] discr (the distance to the previous last explicit discr). *)
+    List.folding_map variants ~init:None
+      ~f:(fun previous_explicit_discriminator (variant, { discr; _ }) ->
         let pat =
-          PConstruct
-            {
-              is_record = variant.is_record;
-              is_struct = false;
-              fields =
-                List.map
-                  ~f:(fun (cid, typ, _) ->
-                    { field = `Concrete cid; pat = { p = PWild; typ; span } })
-                  variant.arguments;
-              constructor = `Concrete variant.name;
-            }
+          let mk_wild_field (cid, typ, _) =
+            { field = `Concrete cid; pat = M.pat_PWild ~typ }
+          in
+          M.pat_PConstruct ~constructor:(`Concrete variant.name)
+            ~is_struct:false ~typ ~is_record:variant.is_record
+            ~fields:(List.map ~f:mk_wild_field variant.arguments)
         in
-        let pat = { p = pat; typ = self; span } in
-        match (acc, thir_variant.discr) with
-        | Lit n, Relative m ->
-            let acc = Lit Int64.(n + m) in
-            (acc, (pat, acc))
+        match (previous_explicit_discriminator, discr) with
+        | None, Relative m -> (None, (pat, expr_of_int m))
         | _, Explicit did ->
-            let acc =
-              Exp { e = GlobalVar (def_id ~value:true did); span; typ }
-            in
-            (acc, (pat, acc))
-        | Exp e, Relative n ->
-            let acc =
-              Exp (U.call Core__ops__arith__Add__add [ e; to_expr n ] span typ)
-            in
-            (Exp e, (pat, acc)))
-    |> List.map ~f:(Fn.id *** function Exp e -> e | Lit n -> to_expr n)
-    |> List.map ~f:(fun (arm_pat, body) ->
-           { arm = { arm_pat; body; guard = None }; span })
-  in
-  let scrutinee_var =
-    Local_ident.{ name = "x"; id = Local_ident.mk_id Expr (-1) }
+            let e = M.expr_GlobalVar ~typ (def_id ~value:true did) in
+            (Some e, (pat, e))
+        | Some e, Relative n ->
+            let n = expr_of_int n in
+            let e = U.call Core__ops__arith__Add__add [ e; n ] span typ in
+            (previous_explicit_discriminator, (pat, e)))
+    |> List.map ~f:(fun (p, e) -> M.arm p e)
   in
-  let scrutinee = { e = LocalVar scrutinee_var; typ = self; span } in
+  let scrutinee_var = Local_ident.{ name = "x"; id = mk_id Expr (-1) } in
+  let scrutinee = M.expr_LocalVar ~typ:self scrutinee_var in
   let ident = cast_name_for_type typ_name in
-  let v =
-    Fn
-      {
-        name = ident;
-        generics;
-        body = { e = Match { scrutinee; arms }; typ; span };
-        params =
-          [
-            {
-              pat = U.make_var_pat scrutinee_var self span;
-              typ = self;
-              typ_span = None;
-              attrs = [];
-            };
-          ];
-        safety = Safe;
-      }
+  let params =
+    let pat = U.make_var_pat scrutinee_var self span in
+    [ { pat; typ = self; typ_span = None; attrs = [] } ]
   in
-  { v; span; ident; attrs = [] }
+  let body = M.expr_Match ~typ ~scrutinee ~arms in
+  M.item_Fn ~ident ~attrs:[] ~name:ident ~generics ~params ~safety:Safe ~body
 
 let rec c_item ~ident ~type_only (item : Thir.item) : item list =
   try


Test patch:
diff --git a/test-harness/src/snapshots/toolchain__enum-repr into-coq.snap b/test-harness/src/snapshots/toolchain__enum-repr into-coq.snap
index b72ba94ec..ee26f94a2 100644
--- a/test-harness/src/snapshots/toolchain__enum-repr into-coq.snap
+++ b/test-harness/src/snapshots/toolchain__enum-repr into-coq.snap
@@ -70,6 +70,44 @@ Definition t_EnumWithRepr_cast_to_repr (x : t_EnumWithRepr) : t_u16 :=
     f_add (anon_const_EnumWithRepr_ExplicitDiscr2__anon_const_0) (2)
   end.
 
+Inductive t_ImplicitReprs : Type :=
+| ImplicitReprs_A
+| ImplicitReprs_B
+| ImplicitReprs_C
+| ImplicitReprs_D
+| ImplicitReprs_E
+| ImplicitReprs_F
+| ImplicitReprs_G
+| ImplicitReprs_H
+| ImplicitReprs_I.
+Arguments t_ImplicitReprs:clear implicits.
+Arguments t_ImplicitReprs.
+
+Definition anon_const_ImplicitReprs_E__anon_const_0 : t_u64 :=
+  30.
+
+Definition t_ImplicitReprs_cast_to_repr (x : t_ImplicitReprs) : t_u64 :=
+  match x with
+  | ImplicitReprs_A =>
+    0
+  | ImplicitReprs_B =>
+    1
+  | ImplicitReprs_C =>
+    2
+  | ImplicitReprs_D =>
+    3
+  | ImplicitReprs_E =>
+    anon_const_ImplicitReprs_E__anon_const_0
+  | ImplicitReprs_F =>
+    f_add (anon_const_ImplicitReprs_E__anon_const_0) (1)
+  | ImplicitReprs_G =>
+    f_add (anon_const_ImplicitReprs_E__anon_const_0) (2)
+  | ImplicitReprs_H =>
+    f_add (anon_const_ImplicitReprs_E__anon_const_0) (3)
+  | ImplicitReprs_I =>
+    f_add (anon_const_ImplicitReprs_E__anon_const_0) (4)
+  end.
+
 Definition f (_ : unit) : t_u32 :=
   let e_x := cast (f_add (anon_const_EnumWithRepr_ExplicitDiscr2__anon_const_0) (0)) in
   f_add (cast (t_EnumWithRepr_cast_to_repr (EnumWithRepr_ImplicitDiscrEmptyTuple))) (cast (t_EnumWithRepr_cast_to_repr (EnumWithRepr_ImplicitDiscrEmptyStruct))).
diff --git a/test-harness/src/snapshots/toolchain__enum-repr into-fstar.snap b/test-harness/src/snapshots/toolchain__enum-repr into-fstar.snap
index f81ab79a1..a80c36176 100644
--- a/test-harness/src/snapshots/toolchain__enum-repr into-fstar.snap
+++ b/test-harness/src/snapshots/toolchain__enum-repr into-fstar.snap
@@ -52,6 +52,31 @@ let t_EnumWithRepr_cast_to_repr (x: t_EnumWithRepr) : u16 =
   | EnumWithRepr_ImplicitDiscrEmptyStruct  ->
     anon_const_EnumWithRepr_ExplicitDiscr2__anon_const_0 +! mk_u16 2
 
+type t_ImplicitReprs =
+  | ImplicitReprs_A : t_ImplicitReprs
+  | ImplicitReprs_B : t_ImplicitReprs
+  | ImplicitReprs_C : t_ImplicitReprs
+  | ImplicitReprs_D : t_ImplicitReprs
+  | ImplicitReprs_E : t_ImplicitReprs
+  | ImplicitReprs_F : t_ImplicitReprs
+  | ImplicitReprs_G : t_ImplicitReprs
+  | ImplicitReprs_H : t_ImplicitReprs
+  | ImplicitReprs_I : t_ImplicitReprs
+
+let anon_const_ImplicitReprs_E__anon_const_0: u64 = mk_u64 30
+
+let t_ImplicitReprs_cast_to_repr (x: t_ImplicitReprs) : u64 =
+  match x <: t_ImplicitReprs with
+  | ImplicitReprs_A  -> mk_u64 0
+  | ImplicitReprs_B  -> mk_u64 1
+  | ImplicitReprs_C  -> mk_u64 2
+  | ImplicitReprs_D  -> mk_u64 3
+  | ImplicitReprs_E  -> anon_const_ImplicitReprs_E__anon_const_0
+  | ImplicitReprs_F  -> anon_const_ImplicitReprs_E__anon_const_0 +! mk_u64 1
+  | ImplicitReprs_G  -> anon_const_ImplicitReprs_E__anon_const_0 +! mk_u64 2
+  | ImplicitReprs_H  -> anon_const_ImplicitReprs_E__anon_const_0 +! mk_u64 3
+  | ImplicitReprs_I  -> anon_const_ImplicitReprs_E__anon_const_0 +! mk_u64 4
+
 let f (_: Prims.unit) : u32 =
   let e_x:u16 =
     cast (anon_const_EnumWithRepr_ExplicitDiscr2__anon_const_0 +! mk_u16 0 <: u16) <: u16
diff --git a/test-harness/src/snapshots/toolchain__enum-repr into-ssprove.snap b/test-harness/src/snapshots/toolchain__enum-repr into-ssprove.snap
index 3df6ab5b2..031e56753 100644
--- a/test-harness/src/snapshots/toolchain__enum-repr into-ssprove.snap
+++ b/test-harness/src/snapshots/toolchain__enum-repr into-ssprove.snap
@@ -103,6 +103,83 @@ Equations t_EnumWithRepr_cast_to_repr {L1 : {fset Location}} {I1 : Interface} (x
     end : both L1 I1 int16.
 Fail Next Obligation.
 
+Definition t_ImplicitReprs : choice_type :=
+  ('unit ∐ 'unit ∐ 'unit ∐ 'unit ∐ 'unit ∐ 'unit ∐ 'unit ∐ 'unit ∐ 'unit).
+Notation "'ImplicitReprs_A_case'" := (inl (inl (inl (inl (inl (inl (inl (inl tt)))))))) (at level 100).
+Equations ImplicitReprs_A {L : {fset Location}} {I : Interface} : both L I t_ImplicitReprs :=
+  ImplicitReprs_A  :=
+    solve_lift (ret_both (inl (inl (inl (inl (inl (inl (inl (inl (tt : 'unit)))))))) : t_ImplicitReprs)) : both L I t_ImplicitReprs.
+Fail Next Obligation.
+Notation "'ImplicitReprs_B_case'" := (inl (inl (inl (inl (inl (inl (inl (inr tt)))))))) (at level 100).
+Equations ImplicitReprs_B {L : {fset Location}} {I : Interface} : both L I t_ImplicitReprs :=
+  ImplicitReprs_B  :=
+    solve_lift (ret_both (inl (inl (inl (inl (inl (inl (inl (inr (tt : 'unit)))))))) : t_ImplicitReprs)) : both L I t_ImplicitReprs.
+Fail Next Obligation.
+Notation "'ImplicitReprs_C_case'" := (inl (inl (inl (inl (inl (inl (inr tt))))))) (at level 100).
+Equations ImplicitReprs_C {L : {fset Location}} {I : Interface} : both L I t_ImplicitReprs :=
+  ImplicitReprs_C  :=
+    solve_lift (ret_both (inl (inl (inl (inl (inl (inl (inr (tt : 'unit))))))) : t_ImplicitReprs)) : both L I t_ImplicitReprs.
+Fail Next Obligation.
+Notation "'ImplicitReprs_D_case'" := (inl (inl (inl (inl (inl (inr tt)))))) (at level 100).
+Equations ImplicitReprs_D {L : {fset Location}} {I : Interface} : both L I t_ImplicitReprs :=
+  ImplicitReprs_D  :=
+    solve_lift (ret_both (inl (inl (inl (inl (inl (inr (tt : 'unit)))))) : t_ImplicitReprs)) : both L I t_ImplicitReprs.
+Fail Next Obligation.
+Notation "'ImplicitReprs_E_case'" := (inl (inl (inl (inl (inr tt))))) (at level 100).
+Equations ImplicitReprs_E {L : {fset Location}} {I : Interface} : both L I t_ImplicitReprs :=
+  ImplicitReprs_E  :=
+    solve_lift (ret_both (inl (inl (inl (inl (inr (tt : 'unit))))) : t_ImplicitReprs)) : both L I t_ImplicitReprs.
+Fail Next Obligation.
+Notation "'ImplicitReprs_F_case'" := (inl (inl (inl (inr tt)))) (at level 100).
+Equations ImplicitReprs_F {L : {fset Location}} {I : Interface} : both L I t_ImplicitReprs :=
+  ImplicitReprs_F  :=
+    solve_lift (ret_both (inl (inl (inl (inr (tt : 'unit)))) : t_ImplicitReprs)) : both L I t_ImplicitReprs.
+Fail Next Obligation.
+Notation "'ImplicitReprs_G_case'" := (inl (inl (inr tt))) (at level 100).
+Equations ImplicitReprs_G {L : {fset Location}} {I : Interface} : both L I t_ImplicitReprs :=
+  ImplicitReprs_G  :=
+    solve_lift (ret_both (inl (inl (inr (tt : 'unit))) : t_ImplicitReprs)) : both L I t_ImplicitReprs.
+Fail Next Obligation.
+Notation "'ImplicitReprs_H_case'" := (inl (inr tt)) (at level 100).
+Equations ImplicitReprs_H {L : {fset Location}} {I : Interface} : both L I t_ImplicitReprs :=
+  ImplicitReprs_H  :=
+    solve_lift (ret_both (inl (inr (tt : 'unit)) : t_ImplicitReprs)) : both L I t_ImplicitReprs.
+Fail Next Obligation.
+Notation "'ImplicitReprs_I_case'" := (inr tt) (at level 100).
+Equations ImplicitReprs_I {L : {fset Location}} {I : Interface} : both L I t_ImplicitReprs :=
+  ImplicitReprs_I  :=
+    solve_lift (ret_both (inr (tt : 'unit) : t_ImplicitReprs)) : both L I t_ImplicitReprs.
+Fail Next Obligation.
+
+Equations anon_const_ImplicitReprs_E__anon_const_0 {L : {fset Location}} {I : Interface} : both L I int64 :=
+  anon_const_ImplicitReprs_E__anon_const_0  :=
+    solve_lift (ret_both (30 : int64)) : both L I int64.
+Fail Next Obligation.
+
+Equations t_ImplicitReprs_cast_to_repr {L1 : {fset Location}} {I1 : Interface} (x : both L1 I1 t_ImplicitReprs) : both L1 I1 int64 :=
+  t_ImplicitReprs_cast_to_repr x  :=
+    matchb x with
+    | ImplicitReprs_A_case  =>
+      solve_lift (ret_both (0 : int64))
+    | ImplicitReprs_B_case  =>
+      solve_lift (ret_both (1 : int64))
+    | ImplicitReprs_C_case  =>
+      solve_lift (ret_both (2 : int64))
+    | ImplicitReprs_D_case  =>
+      solve_lift (ret_both (3 : int64))
+    | ImplicitReprs_E_case  =>
+      solve_lift anon_const_ImplicitReprs_E__anon_const_0
+    | ImplicitReprs_F_case  =>
+      solve_lift (anon_const_ImplicitReprs_E__anon_const_0 .+ (ret_both (1 : int64)))
+    | ImplicitReprs_G_case  =>
+      solve_lift (anon_const_ImplicitReprs_E__anon_const_0 .+ (ret_both (2 : int64)))
+    | ImplicitReprs_H_case  =>
+      solve_lift (anon_const_ImplicitReprs_E__anon_const_0 .+ (ret_both (3 : int64)))
+    | ImplicitReprs_I_case  =>
+      solve_lift (anon_const_ImplicitReprs_E__anon_const_0 .+ (ret_both (4 : int64)))
+    end : both L1 I1 int64.
+Fail Next Obligation.
+
 Equations f {L1 : {fset Location}} {I1 : Interface} (_ : both L1 I1 'unit) : both L1 I1 int32 :=
   f _  :=
     letb e_x := cast_int (WS2 := _) (anon_const_EnumWithRepr_ExplicitDiscr2__anon_const_0 .+ (ret_both (0 : int16))) in
diff --git a/tests/enum-repr/src/lib.rs b/tests/enum-repr/src/lib.rs
index 9d15e05fe..b76ebbcd6 100644
--- a/tests/enum-repr/src/lib.rs
+++ b/tests/enum-repr/src/lib.rs
@@ -8,6 +8,19 @@ enum EnumWithRepr {
     ImplicitDiscrEmptyStruct {},
 }
 
+#[repr(u64)]
+enum ImplicitReprs {
+    A,
+    B(),
+    C {},
+    D,
+    E = 30,
+    F,
+    G,
+    H {},
+    I(),
+}
+
 fn f() -> u32 {
     const CONST: u16 = EnumWithRepr::ExplicitDiscr1 as u16;
     let _x = EnumWithRepr::ExplicitDiscr2 as u16;