Extend Pblock value kind to handle variants (#703)

Author: mshinwell

backend/cmmgen.ml

@@ -339,7 +339,7 @@ let join_unboxed_number_kind ~strict k1 k2 =
 
 let is_strict = function
   | Pfloatval | Pboxedintval _ -> false
-  | Pintval | Pgenval | Pblock _ | Parrayval _ -> true
+  | Pintval | Pgenval | Pvariant _ | Parrayval _ -> true
 
 let rec is_unboxed_number_cmm = function
     | Cop(Calloc mode, [Cconst_natint (hdr, _); _], dbg)
@@ -386,10 +386,10 @@ let rec is_unboxed_number_cmm = function
     | Cassign _
     | Ctuple _
     | Cop _ -> No_unboxing
-    | Cifthenelse (_, _, _, _, _, _, Vval (Pintval | Pblock _))
-    | Cswitch (_, _,  _, _, Vval (Pintval | Pblock _))
-    | Ctrywith (_, _, _, _, _, Vval (Pintval | Pblock _))
-    | Ccatch (_, _, _, Vval (Pintval | Pblock _)) ->
+    | Cifthenelse (_, _, _, _, _, _, Vval (Pintval | Pvariant _))
+    | Cswitch (_, _,  _, _, Vval (Pintval | Pvariant _))
+    | Ctrywith (_, _, _, _, _, Vval (Pintval | Pvariant _))
+    | Ccatch (_, _, _, Vval (Pintval | Pvariant _)) ->
       No_unboxing
     | Cifthenelse (_, _, a, _, b, _, Vval kind) ->
       join_unboxed_number_kind ~strict:(is_strict kind)

middle_end/clambda_primitives.ml

@@ -135,7 +135,10 @@ and array_kind = Lambda.array_kind =
 and value_kind = Lambda.value_kind =
   (* CR mshinwell: Pfloatval should be renamed to Pboxedfloatval *)
     Pgenval | Pfloatval | Pboxedintval of boxed_integer | Pintval
-  | Pblock of { tag : int; fields : value_kind list }
+  | Pvariant of {
+      consts : int list;
+      non_consts : (int * value_kind list) list;
+    }
   | Parrayval of array_kind
 
 and block_shape = Lambda.block_shape

middle_end/clambda_primitives.mli

@@ -138,7 +138,10 @@ and array_kind = Lambda.array_kind =
 and value_kind = Lambda.value_kind =
   (* CR mshinwell: Pfloatval should be renamed to Pboxedfloatval *)
     Pgenval | Pfloatval | Pboxedintval of boxed_integer | Pintval
-  | Pblock of { tag : int; fields : value_kind list }
+  | Pvariant of {
+      consts : int list;
+      non_consts : (int * value_kind list) list;
+    }
   | Parrayval of array_kind
 
 and block_shape = Lambda.block_shape

middle_end/flambda2/from_lambda/lambda_to_flambda_primitives.ml

@@ -30,7 +30,7 @@ let convert_block_of_values_field (value_kind : L.value_kind) :
   | Pboxedintval Pint64 -> Boxed_int64
   | Pboxedintval Pnativeint -> Boxed_nativeint
   | Pintval -> Immediate
-  | Pblock _ | Parrayval _ -> Any_value
+  | Pvariant _ | Parrayval _ -> Any_value
 
 let convert_integer_comparison_prim (comp : L.integer_comparison) :
     P.binary_primitive =

middle_end/flambda2/kinds/flambda_kind.ml

@@ -276,9 +276,9 @@ module With_subkind = struct
       | Boxed_int64
       | Boxed_nativeint
       | Tagged_immediate
-      | Block of
-          { tag : Tag.t;
-            fields : t list
+      | Variant of
+          { consts : Targetint_31_63.Set.t;
+            non_consts : t list Tag.Scannable.Map.t
           }
       | Float_block of { num_fields : int }
       | Float_array
@@ -300,19 +300,34 @@ module With_subkind = struct
       | Value_array, Value_array
       | Generic_array, Generic_array ->
         true
-      | ( Block { tag = t1; fields = fields1 },
-          Block { tag = t2; fields = fields2 } ) ->
-        Tag.equal t1 t2
-        && List.length fields1 = List.length fields2
-        && List.for_all2
-             (fun d when_used_at -> compatible d ~when_used_at)
-             fields1 fields2
+      | ( Variant { consts = consts1; non_consts = non_consts1 },
+          Variant { consts = consts2; non_consts = non_consts2 } ) ->
+        if not (Targetint_31_63.Set.equal consts1 consts2)
+        then false
+        else
+          let tags1 = Tag.Scannable.Map.keys non_consts1 in
+          let tags2 = Tag.Scannable.Map.keys non_consts2 in
+          if not (Tag.Scannable.Set.equal tags1 tags2)
+          then false
+          else
+            let field_lists1 = Tag.Scannable.Map.data non_consts1 in
+            let field_lists2 = Tag.Scannable.Map.data non_consts2 in
+            assert (List.compare_lengths field_lists1 field_lists2 = 0);
+            List.for_all2
+              (fun fields1 fields2 ->
+                if List.compare_lengths fields1 fields2 <> 0
+                then false
+                else
+                  List.for_all2
+                    (fun d when_used_at -> compatible d ~when_used_at)
+                    fields1 fields2)
+              field_lists1 field_lists2
       | ( Float_block { num_fields = num_fields1 },
           Float_block { num_fields = num_fields2 } ) ->
         num_fields1 = num_fields2
       (* Subkinds of [Value] may always be used at [Value] (but not the
          converse): *)
-      | ( ( Block _ | Float_block _ | Float_array | Immediate_array
+      | ( ( Variant _ | Float_block _ | Float_array | Immediate_array
           | Value_array | Generic_array | Boxed_float | Boxed_int32
           | Boxed_int64 | Boxed_nativeint | Tagged_immediate ),
           Anything ) ->
@@ -324,7 +339,7 @@ module With_subkind = struct
         true
       (* All other combinations are incompatible: *)
       | ( ( Anything | Boxed_float | Boxed_int32 | Boxed_int64 | Boxed_nativeint
-          | Tagged_immediate | Block _ | Float_block _ | Float_array
+          | Tagged_immediate | Variant _ | Float_block _ | Float_array
           | Immediate_array | Value_array | Generic_array ),
           _ ) ->
         false
@@ -335,7 +350,7 @@ module With_subkind = struct
       let rec print ppf t =
         let colour = Flambda_colours.subkind () in
         match t with
-        | Anything -> ()
+        | Anything -> Format.fprintf ppf "*"
         | Tagged_immediate ->
           Format.fprintf ppf "@<0>%s=tagged_@<1>\u{2115}@<1>\u{1d55a}@<0>%s"
             colour
@@ -356,10 +371,15 @@ module With_subkind = struct
           Format.fprintf ppf "@<0>%s=boxed_@<1>\u{2115}@<1>\u{2115}@<0>%s"
             colour
             (Flambda_colours.normal ())
-        | Block { tag; fields } ->
-          Format.fprintf ppf "@<0>%s=Block{%a: %a}@<0>%s" colour Tag.print tag
-            (Format.pp_print_list ~pp_sep:Format.pp_print_space print)
-            fields
+        | Variant { consts; non_consts } ->
+          Format.fprintf ppf
+            "@<0>%s=Variant((consts (%a))@ (non_consts (%a)))@<0>%s" colour
+            Targetint_31_63.Set.print consts
+            (Tag.Scannable.Map.print (fun ppf fields ->
+                 Format.fprintf ppf "[%a]"
+                   (Format.pp_print_list ~pp_sep:Format.pp_print_space print)
+                   fields))
+            non_consts
             (Flambda_colours.normal ())
         | Float_block { num_fields } ->
           Format.fprintf ppf "@<0>%s=Float_block(%d)@<0>%s" colour num_fields
@@ -399,7 +419,7 @@ module With_subkind = struct
       match subkind with
       | Anything -> ()
       | Boxed_float | Boxed_int32 | Boxed_int64 | Boxed_nativeint
-      | Tagged_immediate | Block _ | Float_block _ | Float_array
+      | Tagged_immediate | Variant _ | Float_block _ | Float_array
       | Immediate_array | Value_array | Generic_array ->
         Misc.fatal_errorf "Subkind %a is not valid for kind %a" Subkind.print
           subkind print kind));
@@ -448,7 +468,14 @@ module With_subkind = struct
         "Block with fields of non-Value kind (use \
          [Flambda_kind.With_subkind.float_block] for float records)";
     let fields = List.map (fun t -> t.subkind) fields in
-    create value (Block { tag; fields })
+    match Tag.Scannable.of_tag tag with
+    | Some tag ->
+      create value
+        (Variant
+           { consts = Targetint_31_63.Set.empty;
+             non_consts = Tag.Scannable.Map.singleton tag fields
+           })
+    | None -> Misc.fatal_errorf "Tag %a is not scannable" Tag.print tag
 
   let float_block ~num_fields = create value (Float_block { num_fields })
 
@@ -468,12 +495,34 @@ module With_subkind = struct
     | Pboxedintval Pint64 -> boxed_int64
     | Pboxedintval Pnativeint -> boxed_nativeint
     | Pintval -> tagged_immediate
-    | Pblock { tag; fields } ->
-      (* If we have [Obj.double_array_tag] here, this is always an all-float
-         block, not an array. *)
-      if tag = Obj.double_array_tag
-      then float_block ~num_fields:(List.length fields)
-      else block (Tag.create_exn tag) (List.map from_lambda fields)
+    | Pvariant { consts; non_consts } -> (
+      match consts, non_consts with
+      | [], [] -> Misc.fatal_error "[Pvariant] with no constructors at all"
+      | [], [(tag, fields)] when tag = Obj.double_array_tag ->
+        (* If we have [Obj.double_array_tag] here, this is always an all-float
+           block, not an array. *)
+        float_block ~num_fields:(List.length fields)
+      | [], _ :: _ | _ :: _, [] | _ :: _, _ :: _ ->
+        let consts =
+          Targetint_31_63.Set.of_list
+            (List.map
+               (fun const ->
+                 Targetint_31_63.int (Targetint_31_63.Imm.of_int const))
+               consts)
+        in
+        let non_consts =
+          List.fold_left
+            (fun non_consts (tag, fields) ->
+              match Tag.Scannable.create tag with
+              | Some tag ->
+                Tag.Scannable.Map.add tag
+                  (List.map (fun vk -> subkind (from_lambda vk)) fields)
+                  non_consts
+              | None ->
+                Misc.fatal_errorf "Non-scannable tag %d in [Pvariant]" tag)
+            Tag.Scannable.Map.empty non_consts
+        in
+        create value (Variant { consts; non_consts }))
     | Parrayval Pfloatarray -> float_array
     | Parrayval Pintarray -> immediate_array
     | Parrayval Paddrarray -> value_array
@@ -489,7 +538,7 @@ module With_subkind = struct
         Format.fprintf ppf "@[%a%a@]" print kind Subkind.print subkind
       | ( (Naked_number _ | Region | Rec_info),
           ( Boxed_float | Boxed_int32 | Boxed_int64 | Boxed_nativeint
-          | Tagged_immediate | Block _ | Float_block _ | Float_array
+          | Tagged_immediate | Variant _ | Float_block _ | Float_array
           | Immediate_array | Value_array | Generic_array ) ) ->
         assert false
     (* see [create] *)
@@ -511,7 +560,9 @@ module With_subkind = struct
     match t.subkind with
     | Anything -> false
     | Boxed_float | Boxed_int32 | Boxed_int64 | Boxed_nativeint
-    | Tagged_immediate | Block _ | Float_block _ | Float_array | Immediate_array
-    | Value_array | Generic_array ->
+    | Tagged_immediate | Variant _ | Float_block _ | Float_array
+    | Immediate_array | Value_array | Generic_array ->
       true
+
+  let erase_subkind t = { t with subkind = Anything }
 end

middle_end/flambda2/kinds/flambda_kind.mli

@@ -131,9 +131,9 @@ module With_subkind : sig
       | Boxed_int64
       | Boxed_nativeint
       | Tagged_immediate
-      | Block of
-          { tag : Tag.t;
-            fields : t list
+      | Variant of
+          { consts : Targetint_31_63.Set.t;
+            non_consts : t list Tag.Scannable.Map.t
           }
       | Float_block of { num_fields : int }
       | Float_array
@@ -198,5 +198,7 @@ module With_subkind : sig
 
   val compatible : t -> when_used_at:t -> bool
 
+  val erase_subkind : t -> t
+
   include Container_types.S with type t := t
 end

middle_end/flambda2/parser/print_fexpr.ml

@@ -142,7 +142,7 @@ let kind_with_subkind ppf (k : kind_with_subkind) =
   | Value -> (
     match Flambda_kind.With_subkind.subkind k with
     | Anything -> str "val"
-    | Block _ -> str "block" (* CR mshinwell: improve this *)
+    | Variant _ -> str "variant" (* CR mshinwell: improve this *)
     | Float_block _ -> str "float_block"
     | Boxed_float -> str "float boxed"
     | Boxed_int32 -> str "int32 boxed"

middle_end/flambda2/types/grammar/more_type_creators.ml

@@ -315,16 +315,19 @@ let rec unknown_with_subkind (kind : Flambda_kind.With_subkind.t) =
   | Boxed_int64 -> any_boxed_int64
   | Boxed_nativeint -> any_boxed_nativeint
   | Tagged_immediate -> any_tagged_immediate
-  | Block { tag; fields } ->
-    assert (not (Tag.equal tag Tag.double_array_tag));
-    immutable_block ~is_unique:false tag ~field_kind:Flambda_kind.value
-      ~fields:
-        (List.map
-           (fun subkind ->
-             unknown_with_subkind
-               (Flambda_kind.With_subkind.create Flambda_kind.value subkind))
-           fields)
-      Unknown
+  | Variant { consts; non_consts } ->
+    let const_ctors = these_naked_immediates consts in
+    let non_const_ctors =
+      Tag.Scannable.Map.map
+        (fun fields ->
+          List.map
+            (fun subkind ->
+              unknown_with_subkind
+                (Flambda_kind.With_subkind.create Flambda_kind.value subkind))
+            fields)
+        non_consts
+    in
+    variant ~const_ctors ~non_const_ctors Unknown
   | Float_block { num_fields } ->
     immutable_block ~is_unique:false Tag.double_array_tag
       ~field_kind:Flambda_kind.naked_float

middle_end/printclambda.ml

@@ -24,23 +24,33 @@ let mutable_flag = function
   | Lambda.Mutable-> "[mut]"
   | Lambda.Immutable | Lambda.Immutable_unique -> ""
 
-let value_kind =
+let rec value_kind0 ppf kind =
   let open Lambda in
-  function
-  | Pgenval -> ""
-  | Pintval -> ":int"
-  | Pfloatval -> ":float"
-  | Parrayval Pgenarray -> ":genarray"
-  | Parrayval Pintarray -> ":intarray"
-  | Parrayval Pfloatarray -> ":floatarray"
-  | Parrayval Paddrarray -> ":addrarray"
-  | Pboxedintval Pnativeint -> ":nativeint"
-  | Pboxedintval Pint32 -> ":int32"
-  | Pboxedintval Pint64 -> ":int64"
-  | Pblock { tag; fields } ->
-    asprintf ":[%d: %a]" tag
-      (Format.pp_print_list ~pp_sep:(fun ppf () -> fprintf ppf ",@ ")
-         Printlambda.value_kind') fields
+  match kind with
+  | Pgenval -> Format.pp_print_string ppf ""
+  | Pintval -> Format.pp_print_string ppf ":int"
+  | Pfloatval -> Format.pp_print_string ppf ":float"
+  | Parrayval Pgenarray -> Format.pp_print_string ppf ":genarray"
+  | Parrayval Pintarray -> Format.pp_print_string ppf ":intarray"
+  | Parrayval Pfloatarray -> Format.pp_print_string ppf ":floatarray"
+  | Parrayval Paddrarray -> Format.pp_print_string ppf ":addrarray"
+  | Pboxedintval Pnativeint -> Format.pp_print_string ppf ":nativeint"
+  | Pboxedintval Pint32 -> Format.pp_print_string ppf ":int32"
+  | Pboxedintval Pint64 -> Format.pp_print_string ppf ":int64"
+  | Pvariant { consts; non_consts } ->
+    Format.fprintf ppf "@[<hov 1>[(consts (%a))@ (non_consts (%a))]@]"
+      (Format.pp_print_list ~pp_sep:Format.pp_print_space Format.pp_print_int)
+      consts
+      (Format.pp_print_list ~pp_sep:Format.pp_print_space
+          (fun ppf (tag, fields) ->
+            fprintf ppf "@[<hov 1>[%d:@ %a]@]" tag
+              (Format.pp_print_list
+                ~pp_sep:(fun ppf () -> fprintf ppf ",@ ")
+                value_kind0)
+              fields))
+      non_consts
+
+let value_kind kind = Format.asprintf "%a" value_kind0 kind
 
 let rec structured_constant ppf = function
   | Uconst_float x -> fprintf ppf "%F" x

ocaml/asmcomp/cmmgen.ml

@@ -711,7 +711,7 @@ and transl_catch env nfail ids body handler dbg =
          let strict =
            match kind with
            | Pfloatval | Pboxedintval _ -> false
-           | Pintval | Pgenval | Pblock _ | Parrayval _ -> true
+           | Pintval | Pgenval | Pvariant _ | Parrayval _ -> true
          in
          u := join_unboxed_number_kind ~strict !u
              (is_unboxed_number_cmm ~strict c)
@@ -1179,7 +1179,7 @@ and transl_let env str kind id exp transl_body =
            we do it only if this indeed allows us to get rid of
            some allocations in the bound expression. *)
         is_unboxed_number_cmm ~strict:false cexp
-    | _, (Pgenval | Pblock _ | Parrayval _) ->
+    | _, (Pgenval | Pvariant _ | Parrayval _) ->
         (* Here we don't know statically that the bound expression
            evaluates to an unboxable number type.  We need to be stricter
            and ensure that all possible branches in the expression