From 90eaa8291e18369f016b061cd95607dac8dca57a Mon Sep 17 00:00:00 2001
From: Jonathan Protzenko <protz@microsoft.com>
Date: Mon, 28 Oct 2024 15:55:13 -0700
Subject: [PATCH] Fix pretty-printing and variant analysis (with Aymeric)

---
 lib/AstToMiniRust.ml | 39 +++++++++++++++++++++++++--------------
 lib/PrintMiniRust.ml | 25 ++++++++++++++-----------
 2 files changed, 39 insertions(+), 25 deletions(-)

diff --git a/lib/AstToMiniRust.ml b/lib/AstToMiniRust.ml
index e349f7b5..c9c75db5 100644
--- a/lib/AstToMiniRust.ml
+++ b/lib/AstToMiniRust.ml
@@ -460,7 +460,7 @@ let rec translate_type_with_config (env: env) (config: config) (t: Ast.typ): Min
       begin try
         Name (lookup_type env lid, generic_params)
       with Not_found ->
-        KPrint.bprintf "Type name not found: %a\n" PrintAst.plid lid;
+        (* KPrint.bprintf "Type name not found: %a\n" PrintAst.plid lid; *)
         Name (translate_unknown_lid lid, generic_params)
       end
   | TArrow _ ->
@@ -1214,11 +1214,13 @@ let is_handled_primitively = function
   | _ ->
       false
 
-let compute_derives heap_structs _pointer_holding_structs files =
+let compute_derives heap_structs pointer_holding_structs files =
+  (* A map from lid to Ast definition, of type decl *)
   let definitions = List.fold_left (fun map (_, decls) ->
     List.fold_left (fun map decl -> LidMap.add (Ast.lid_of_decl decl) decl map) map decls
   ) LidMap.empty files in
 
+  (* The bottom element of our lattice *)
   let everything = TraitSet.of_list [ MiniRust.PartialEq; Clone; Copy ] in
 
   let module F = Fix.Fix.ForOrderedType(struct
@@ -1232,19 +1234,28 @@ let compute_derives heap_structs _pointer_holding_structs files =
   end) in
 
   let equations lid valuation =
-    let traits = object
-        inherit [_] Ast.reduce
-        method zero = everything
-        method plus = TraitSet.inter
-        method! visit_TQualified _ lid =
-          valuation lid
-      end#visit_decl () (LidMap.find lid definitions)
-    in
-    if LidSet.mem lid heap_structs then
-      (* If this type will contain a Box<...> then it cannot have trait copy. *)
-      TraitSet.diff traits (TraitSet.of_list [ MiniRust.Copy ])
+    if not (LidMap.mem lid definitions) then
+      everything
     else
-      traits
+      let traits = object
+          inherit [_] Ast.reduce
+          method zero = everything
+          method plus = TraitSet.inter
+          method! visit_TQualified _ lid =
+            valuation lid
+        end#visit_decl () (LidMap.find lid definitions)
+      in
+      let traits = 
+        if LidSet.mem lid heap_structs || LidSet.mem lid pointer_holding_structs then
+          (* If this type will contain a Box<...> then it cannot have trait copy. *)
+          TraitSet.diff traits (TraitSet.of_list [ MiniRust.Copy ])
+        else
+          traits
+      in
+      if LidSet.mem lid pointer_holding_structs then
+        TraitSet.diff traits (TraitSet.of_list [ MiniRust.Clone ])
+      else
+        traits
   in
 
   F.lfp equations
diff --git a/lib/PrintMiniRust.ml b/lib/PrintMiniRust.ml
index 017a0c82..68179879 100644
--- a/lib/PrintMiniRust.ml
+++ b/lib/PrintMiniRust.ml
@@ -452,17 +452,20 @@ and print_expr env (context: int) (e: expr): document =
 
   | Struct (cons, fields) ->
       group @@
-      print_data_type_name env cons ^/^ braces_with_nesting (
-        separate_map (comma ^^ break1) (fun (f, e) ->
-          group @@
-          if string f = print_expr env max_int e then
-            (* If the field name is the same as the expression assigned to it
-               (typically, a variable name), we do not need to duplicate it *)
-            string f
-          else
-            string f ^^ colon ^/^ group (print_expr env max_int e)
-        ) fields
-      )
+      print_data_type_name env cons ^/^ (
+        match cons with
+        | `Variant _ when fields = [] -> empty
+        | _ -> braces_with_nesting (
+            separate_map (comma ^^ break1) (fun (f, e) ->
+              group @@
+              if string f = print_expr env max_int e then
+                (* If the field name is the same as the expression assigned to it
+                   (typically, a variable name), we do not need to duplicate it *)
+                string f
+              else
+                string f ^^ colon ^/^ group (print_expr env max_int e)
+            ) fields
+      ))
 
   | Var v ->
       begin match lookup env v with