Add IIConst and TIConst

cryspen · Nov 15, 2023 · 5d398fe · 5d398fe
1 parent 4b2756e
commit 5d398fe
Show file tree

Hide file tree

Showing 5 changed files with 27 additions and 6 deletions.
diff --git a/engine/lib/ast.ml b/engine/lib/ast.ml
@@ -75,6 +75,7 @@ type global_ident = (Global_ident.t[@visitors.opaque])
     yojson,
     hash,
     eq,
+    compare,
     visitors { variety = "reduce"; name = "global_ident_reduce" },
     visitors { variety = "mapreduce"; name = "global_ident_mapreduce" },
     visitors { variety = "map"; name = "global_ident_map" }]
@@ -577,6 +578,7 @@ functor
     and impl_item' =
       | IIType of ty
       | IIFn of { body : expr; params : param list }
+      | IIConst of { body : expr }
 
     and impl_item = {
       ii_span : span;
@@ -586,7 +588,7 @@ functor
       ii_attrs : attrs;
     }
 
-    and trait_item' = TIType of trait_ref list | TIFn of ty
+    and trait_item' = TIType of trait_ref list | TIFn of ty | TIConst of ty
 
     and trait_item = {
       (* TODO: why do I need to prefix by `ti_` here? I guess visitors fail or something *)

diff --git a/engine/lib/import_thir.ml b/engine/lib/import_thir.ml
@@ -1010,7 +1010,7 @@ end) : EXPR = struct
     | Const (_, Some _) ->
         unimplemented [ span ]
           "TODO: traits: no support for defaults in traits for now"
-    | Const (ty, None) -> TIFn (c_ty span ty)
+    | Const (ty, None) -> TIConst (c_ty span ty)
     | ProvidedFn (sg, _) | RequiredFn (sg, _) ->
         let (Thir.{ inputs; output; _ } : Thir.fn_decl) = sg.decl in
         let output =
@@ -1038,7 +1038,11 @@ let c_trait_item (item : Thir.trait_item) : trait_item =
   let open (val make ~krate:item.owner_id.krate : EXPR) in
   let { params; constraints } = c_generics item.generics in
   (* TODO: see TODO in impl items *)
-  let ti_ident = Concrete_ident.of_def_id Field item.owner_id in
+  let ti_ident =
+    Concrete_ident.of_def_id
+      (match (item.kind : _) with Const _ -> Value | _ -> Field)
+      item.owner_id
+  in
   {
     ti_span = Span.of_thir item.span;
     ti_generics = { params; constraints };
@@ -1272,7 +1276,11 @@ and c_item_unwrapped ~ident (item : Thir.item) : item list =
                         backend will see traits and impls as
                         records. See https://github.com/hacspec/hacspec-v2/issues/271. *)
                      let ii_ident =
-                       Concrete_ident.of_def_id Field item.owner_id
+                       Concrete_ident.of_def_id
+                         (match (item.kind : _) with
+                         | Const _ -> Value
+                         | _ -> Field)
+                         item.owner_id
                      in
                      {
                        ii_span = Span.of_thir item.span;
@@ -1285,8 +1293,7 @@ and c_item_unwrapped ~ident (item : Thir.item) : item list =
                                  body = c_expr body;
                                  params = List.map ~f:(c_param item.span) params;
                                }
-                         | Const (_ty, e) ->
-                             IIFn { body = c_expr e; params = [] }
+                         | Const (_ty, e) -> IIConst { body = c_expr e }
                          | Type ty -> IIType (c_ty item.span ty));
                        ii_ident;
                        ii_attrs = c_item_attrs item.attributes;

diff --git a/engine/lib/phases/phase_and_mut_defsite.ml b/engine/lib/phases/phase_and_mut_defsite.ml
@@ -235,6 +235,9 @@ struct
             | IIFn { params; body } ->
                 let* params, body = rewrite_function params body in
                 Some (B.IIFn { body; params })
+            | IIConst { body } ->
+                let* _, body = rewrite_function [] body in
+                Some (B.IIConst { body })
             | _ -> None)
             |> Option.value_or_thunk
                  ~default:(Fn.flip super#visit_impl_item' item')

diff --git a/engine/lib/print_rust.ml b/engine/lib/print_rust.ml
@@ -427,6 +427,9 @@ module Raw = struct
         in
         !"fn " & ident & generics & !"(" & params & !") -> " & return_type
         & bounds & !";"
+    | TIConst ty ->
+        let return_type = pty ti.ti_span ty in
+        !"fn " & ident & generics & !" -> " & return_type & bounds & !";"
 
   let pparam span ({ pat; typ; typ_span; attrs } : param) =
     let ( ! ) = pure span in
@@ -449,6 +452,10 @@ module Raw = struct
         let return_type = pty span body.typ in
         !"fn " & ident & generics & pparams span params & !" -> " & return_type
         & bounds & !"{" & pexpr body & !"}"
+    | IIConst { body } ->
+        let return_type = pty span body.typ in
+        !"fn " & ident & generics & !" -> " & return_type & bounds & !"{"
+        & pexpr body & !"}"
 
   let pitem (e : item) =
     let exception NotImplemented in

diff --git a/engine/lib/subtype.ml b/engine/lib/subtype.ml
@@ -355,6 +355,7 @@ struct
       match ti with
       | TIType g -> TIType (List.map ~f:(dtrait_ref span) g)
       | TIFn t -> TIFn (dty span t)
+      | TIConst t -> TIConst (dty span t)
 
     and dtrait_item (ti : A.trait_item) : B.trait_item =
       {
@@ -370,6 +371,7 @@ struct
       | IIType g -> IIType (dty span g)
       | IIFn { body; params } ->
           IIFn { body = dexpr body; params = List.map ~f:(dparam span) params }
+      | IIConst { body } -> IIConst { body = dexpr body }
 
     and dimpl_item (ii : A.impl_item) : B.impl_item =
       {