More precise prove_variant function (#408)

Author: Gbury

middle_end/flambda/types/flambda_type.mli

@@ -514,15 +514,15 @@ val prove_naked_int64s : Typing_env.t -> t -> Numbers.Int64.Set.t proof
 
 val prove_naked_nativeints : Typing_env.t -> t -> Targetint.Set.t proof
 
-type variant_proof = private {
-  const_ctors : Target_imm.Set.t;
+type variant_like_proof = private {
+  const_ctors : Target_imm.Set.t Or_unknown.t;
   non_const_ctors_with_sizes : Targetint.OCaml.t Tag.Scannable.Map.t;
 }
 
-val prove_variant
+val prove_variant_like
    : Typing_env.t
   -> t
-  -> variant_proof proof_allowing_kind_mismatch
+  -> variant_like_proof proof_allowing_kind_mismatch
 
 (** If [ty] is known to represent a boxed number or a tagged integer,
     [prove_is_a_boxed_number env ty] is [Proved kind]. [kind] is the kind of

middle_end/flambda/types/template/flambda_type.templ.ml

@@ -504,60 +504,57 @@ let prove_unique_tag_and_size env t
       | None -> Unknown
       | Some (tag, size) -> Proved (tag, size)
 
-type variant_proof = {
-  const_ctors : Target_imm.Set.t;
+type variant_like_proof = {
+  const_ctors : Target_imm.Set.t Or_unknown.t;
   non_const_ctors_with_sizes : Targetint.OCaml.t Tag.Scannable.Map.t;
 }
 
-let prove_variant env t : variant_proof proof_allowing_kind_mismatch =
+let prove_variant_like env t : variant_like_proof proof_allowing_kind_mismatch =
   (* Format.eprintf "prove_variant:@ %a\n%!" print t; *)
   match expand_head t env with
-  | Const (Tagged_immediate _) -> Unknown
+  | Const (Tagged_immediate imm) ->
+    Proved {
+      const_ctors = Known (Target_imm.Set.singleton imm);
+      non_const_ctors_with_sizes = Tag.Scannable.Map.empty;
+    }
   | Const _ -> Wrong_kind
   | Value (Ok (Variant blocks_imms)) ->
-    begin match blocks_imms.immediates with
+    begin match blocks_imms.blocks with
     | Unknown -> Unknown
-    | Known imms ->
-      let const_ctors : _ Or_unknown.t =
-        match prove_naked_immediates env imms with
-        | Unknown -> Unknown
-        | Invalid -> Known Target_imm.Set.empty
-        | Proved const_ctors -> Known const_ctors
-      in
-      match const_ctors with
+    | Known blocks ->
+      match Row_like.For_blocks.all_tags_and_sizes blocks with
       | Unknown -> Unknown
-      | Known const_ctors ->
-        let valid =
-          Target_imm.Set.for_all Target_imm.is_non_negative const_ctors
+      | Known non_const_ctors_with_sizes ->
+        let non_const_ctors_with_sizes =
+          Tag.Map.fold
+            (fun tag size (result : _ Or_unknown.t) : _ Or_unknown.t ->
+               match result with
+               | Unknown -> Unknown
+               | Known result ->
+                 match Tag.Scannable.of_tag tag with
+                 | None -> Unknown
+                 | Some tag ->
+                   Known (Tag.Scannable.Map.add tag size result))
+            non_const_ctors_with_sizes
+            (Or_unknown.Known Tag.Scannable.Map.empty)
         in
-        if not valid then Invalid
-        else
-          match blocks_imms.blocks with
-          | Unknown -> Unknown
-          | Known blocks ->
-            match Row_like.For_blocks.all_tags_and_sizes blocks with
+        match non_const_ctors_with_sizes with
+        | Unknown -> Unknown
+        | Known non_const_ctors_with_sizes ->
+          let const_ctors : _ Or_unknown.t =
+            match blocks_imms.immediates with
             | Unknown -> Unknown
-            | Known non_const_ctors_with_sizes ->
-              let non_const_ctors_with_sizes =
-                Tag.Map.fold
-                  (fun tag size (result : _ Or_bottom.t) : _ Or_bottom.t ->
-                    match result with
-                    | Bottom -> Bottom
-                    | Ok result ->
-                      match Tag.Scannable.of_tag tag with
-                      | None -> Bottom
-                      | Some tag ->
-                        Ok (Tag.Scannable.Map.add tag size result))
-                  non_const_ctors_with_sizes
-                  (Or_bottom.Ok Tag.Scannable.Map.empty)
-              in
-              match non_const_ctors_with_sizes with
-              | Bottom -> Invalid
-              | Ok non_const_ctors_with_sizes ->
-                Proved {
-                  const_ctors;
-                  non_const_ctors_with_sizes;
-                }
+            | Known imms ->
+              begin match prove_naked_immediates env imms with
+              | Unknown -> Unknown
+              | Invalid -> Known Target_imm.Set.empty
+              | Proved const_ctors -> Known const_ctors
+              end
+          in
+          Proved {
+            const_ctors;
+            non_const_ctors_with_sizes;
+          }
     end
   | Value (Ok _) -> Invalid
   | Value Unknown -> Unknown

middle_end/flambda/unboxing/unbox_continuation_params.ml

@@ -460,15 +460,15 @@ end
 module Variant : sig
   type t
 
-  val create : T.variant_proof -> t
+  val create : T.variant_like_proof -> t
 
   val max_size : t -> Targetint.OCaml.t
 
-  val const_ctors : t -> Target_imm.Set.t
+  val const_ctors : t -> Target_imm.Set.t Or_unknown.t
 
   val non_const_ctors_with_sizes : t -> Targetint.OCaml.t Tag.Scannable.Map.t
 end = struct
-  type t = T.variant_proof
+  type t = T.variant_like_proof
 
   let create variant = variant
 
@@ -1059,7 +1059,7 @@ let rec make_unboxing_decision denv ~depth ~arg_types_by_use_id
             DE.print denv
         end
     | Wrong_kind | Invalid | Unknown ->
-      match T.prove_variant (DE.typing_env denv) param_type with
+      match T.prove_variant_like (DE.typing_env denv) param_type with
       | Proved variant ->
         (*
         Format.eprintf "Starting variant unboxing\n%!";