Fold untagged matches by evaluating the dispatch plan

changes/untagged-fold-by-plan.polish.md

@@ -0,0 +1 @@
+- Fold a match on an untagged constant by evaluating the dispatch the code is printed from, so folding and the generated code agree, also for values outside every case.

compiler/core/lam_compile.ml

@@ -738,11 +738,6 @@ let compile output_prefix =
       default_action ~saturated:sw_blocks_full sw_failaction
     in
     let tag_name, literal_cases = dispatch_info sw_dispatch in
-    let untagged_facts =
-      match sw_dispatch with
-      | Switch_variant ({block_types = _ :: _} as facts) -> Some facts
-      | Switch_variant {block_types = []} | Switch_direct -> None
-    in
     let payload_only = sw_consts_full && sw_consts = [] in
     let compile_whole (cxt : Lam_compile_context.t) =
       match
@@ -760,38 +755,14 @@ let compile output_prefix =
         in
         block
         @
-        match untagged_facts with
-        | Some facts when payload_only || not (sw_blocks_full && sw_blocks = [])
-          ->
-          let plan =
-            Untagged_dispatch.plan ~same:Lambda.eq_approx
-              {
-                facts;
-                literal_arms =
-                  List.map
-                    (fun (key, lam) ->
-                      match tag_of_switch_key key with
-                      | Some (Literal literal) -> (literal, lam)
-                      | Some (Payload_shape _) | None -> assert false)
-                    sw_consts;
-                payload_arms =
-                  List.map
-                    (fun (key, lam) ->
-                      match tag_of_switch_key key with
-                      | Some (Payload_shape block_type) -> (block_type, lam)
-                      | Some (Literal _) | None -> assert false)
-                    sw_blocks;
-                literals_complete = sw_consts_full;
-                payloads_complete = sw_blocks_full;
-                fallback = sw_failaction;
-              }
-          in
+        match Lambda.untagged_plan sw with
+        | Some plan ->
           let compile e =
             compile_untagged_plan ~cxt ~switch_exp:e ~fallback:sw_failaction
               plan
           in
           if payload_only then compile e else with_variable compile
-        | Some _ | None ->
+        | None ->
           if payload_only then
             compile_cases ~cxt ~switch_exp:(E.tag ~name:tag_name e)
               ~default:sw_blocks_default sw_blocks

compiler/ml/UntaggedVariants.md

@@ -480,11 +480,16 @@ primitives of the variant never reach it.
 A `switch` on an untagged value the compiler knows, a literal or a block of
 constants, folds to the arm that the run-time dispatch selects for its
 JavaScript value, so `S("a")` matched against `@as("a") A | S(string)` folds
-to `A`. A value built at run time, such as `A([1])` with its array, is
+to `A`. Folding evaluates the plan the code is printed from, so it agrees
+with the generated code by construction, also for a value outside every case
+(rule 10). A value built at run time, such as `A([1])` with its array, is
 matched at run time.
 
-- Implementation: `lambda.ml:switch` (`untagged_action`, `literal_denotes_same`).
-- Tests: `tests/tests/src/unboxed_variant_fold_test.res`.
+- Implementation: `lambda.ml:switch` (`untagged_value`, `untagged_plan`),
+  `untagged_dispatch.ml:select`.
+- Tests: `tests/tests/src/unboxed_variant_fold_test.res`; `spec_test` rule 10
+  (a folded value outside every case); `ounit_untagged_dispatch_tests.ml`
+  compares the plans of random variants and matches with rules 7 to 9.
 - Shaped by #8631.
 
 ### 12. `option` of an untagged variant

compiler/ml/lambda.ml

@@ -796,15 +796,6 @@ and eq_option l1 l2 =
 
 and eq_approx_list ls ls1 = Ext_list.for_all2_no_exn ls ls1 eq_approx
 
-(* Classify constants by the JavaScript value emitted by [Lam_compile_const]
-   and [Js_dump], not by their source constructor. Unknown representations
-   must retain runtime dispatch. *)
-type value_kind =
-  | Is_literal of Variant_runtime.literal_tag
-  | Is_object
-  | Is_array
-  | Unknown_value
-
 let record_fields_are_array fields =
   let len = Array.length fields in
   let rec loop i =
@@ -815,53 +806,77 @@ let record_fields_are_array fields =
   in
   len <> 0 && loop 0
 
-let runtime_value_kind (c : structured_constant) =
+(* The JavaScript value a constant is emitted as by [Lam_compile_const] and
+   [Js_dump], not its source constructor. A constant of another
+   representation keeps the dispatch at run time. *)
+let untagged_value (c : structured_constant) : Untagged_dispatch.value option =
   match c with
-  | Const_string s -> Is_literal (String s)
-  | Const_int i -> Is_literal (Int (Int32.to_int i))
-  | Const_float f -> Is_literal (Float f)
-  | Const_js_true -> Is_literal (Bool true)
-  | Const_js_false -> Is_literal (Bool false)
-  | Const_js_null -> Is_literal Null
-  | Const_js_undefined _ -> Is_literal Undefined
-  | Const_polyvar name -> Is_literal (String name)
-  | Const_constructor {literal} -> Is_literal literal
+  | Const_string s -> Some (Literal (String s))
+  | Const_int i -> Some (Literal (Int (Int32.to_int i)))
+  | Const_float f -> Some (Literal (Float f))
+  | Const_js_true -> Some (Literal (Bool true))
+  | Const_js_false -> Some (Literal (Bool false))
+  | Const_js_null -> Some (Literal Null)
+  | Const_js_undefined _ -> Some (Literal Undefined)
+  | Const_polyvar name -> Some (Literal (String name))
+  | Const_constructor {literal} -> Some (Literal literal)
   | Const_bigint (sign, digits) ->
-    Is_literal (BigInt (Bigint_utils.to_string sign digits))
-  | Const_block (Blk_tuple, _) -> Is_array
+    Some (Literal (BigInt (Bigint_utils.to_string sign digits)))
+  | Const_block (Blk_tuple, _) -> Some Array
   | Const_block (Blk_record {fields}, _) ->
-    if record_fields_are_array fields then Is_array else Is_object
+    if record_fields_are_array fields then Some Array else Some Object
   | Const_block
       ( ( Blk_constructor _ | Blk_record_inlined _ | Blk_poly_var
         | Blk_record_ext _ | Blk_module _ | Blk_module_export _ | Blk_extension
           ),
         _ ) ->
-    Is_object
-  | Const_char _ | Const_some _ | Const_module_alias | Const_assertfalse ->
-    Unknown_value
+    Some Object
+  | Const_char _ | Const_some _ | Const_module_alias | Const_assertfalse -> None
 
 (* Runtime equality, rather than equality of tags: [@as(1)] and a payload
    [1.0] are the same JavaScript number, and [1_0n] and [10n] the same bigint.
    Integer tags are emitted through [Int32.of_int], just like [Const_int]. *)
 let literal_denotes_same = Variant_runtime.same_literal_value
 
-(* Whether the payload dispatch sends a value of this kind to a case of this
-   block type. Literals are checked as the generated code checks them; the
-   object case excludes arrays when an array case exists. *)
-let value_has_block_type ~block_types (kind : value_kind)
-    (block_type : Variant_runtime.block_type) =
-  match (block_type, kind) with
-  | _, Is_literal literal -> Variant_runtime.literal_passes literal block_type
-  | ObjectType, Is_object -> true
-  | InstanceType Array, Is_array -> true
-  | ObjectType, Is_array ->
-    not (List.mem (Variant_runtime.InstanceType Array) block_types)
-  | UnknownType, (Is_object | Is_array) -> true
-  | UnknownType, Unknown_value -> false
-  | ( ( NumberType | StringType | BooleanType | ObjectType | BigintType
-      | SymbolType | FunctionType | InstanceType _ ),
-      _ ) ->
-    false
+(* The dispatch of a switch on an untagged variant, which the code generator
+   prints and constant folding evaluates. A switch with literal arms and no
+   payload arm left is a plain switch on the literals. *)
+let untagged_plan (sw : lambda_switch) =
+  match sw.sw_dispatch with
+  | Switch_variant ({block_types = _ :: _} as facts)
+    when (sw.sw_consts_full && sw.sw_consts = [])
+         || not (sw.sw_blocks_full && sw.sw_blocks = []) -> (
+    let exception Not_untagged in
+    match
+      ( List.map
+          (fun (key, action) ->
+            match key with
+            | Switch_constructor (Constant {literal}) -> (literal, action)
+            | Switch_constructor (Block _) | Switch_int _ -> raise Not_untagged)
+          sw.sw_consts,
+        List.map
+          (fun (key, action) ->
+            match key with
+            | Switch_constructor (Block (Untagged {block_type})) ->
+              (block_type, action)
+            | Switch_constructor (Constant _ | Block (Tagged _)) | Switch_int _
+              ->
+              raise Not_untagged)
+          sw.sw_blocks )
+    with
+    | literal_arms, payload_arms ->
+      Some
+        (Untagged_dispatch.plan ~same:eq_approx
+           {
+             facts;
+             literal_arms;
+             payload_arms;
+             literals_complete = sw.sw_consts_full;
+             payloads_complete = sw.sw_blocks_full;
+             fallback = sw.sw_failaction;
+           })
+    | exception Not_untagged -> None)
+  | Switch_variant _ | Switch_direct -> None
 
 let switch lam (lam_switch : lambda_switch) : t =
   let action_or_switch = function
@@ -871,42 +886,18 @@ let switch lam (lam_switch : lambda_switch) : t =
       | Some action -> action
       | None -> Lswitch (lam, lam_switch))
   in
-  (* An untagged variant is dispatched on the value, so a constant scrutinee is
-     decided here rather than by the constructor it was written with - which
-     has no runtime existence and may be shared with a literal constructor.
-     [`Undecided] means this layer cannot name the constant's runtime shape, so
-     the switch has to stay; it is not the same as "no case matches". *)
-  let untagged_action (facts : Variant_runtime.matching_facts) cst =
-    let find_in cases matches =
-      `Case
-        (Ext_list.find_opt cases (fun (key, action) ->
-             match key with
-             | Switch_constructor case when matches case -> Some action
-             | Switch_int _ | Switch_constructor _ -> None))
-    in
-    let matches_block = value_has_block_type ~block_types:facts.block_types in
-    let kind = runtime_value_kind cst in
-    match kind with
-    | Unknown_value -> `Undecided
-    | Is_literal literal
-      when Ext_list.exists facts.literal_tags (literal_denotes_same literal) ->
-      (* The literal side wins, exactly as it does at runtime. *)
-      find_in lam_switch.sw_consts (function
-        | Constant tag -> literal_denotes_same literal tag.literal
-        | Block _ -> false)
-    | Is_literal _ | Is_object | Is_array ->
-      (* Not a declared literal, so the payload's runtime shape decides. *)
-      if Ext_list.exists facts.block_types (matches_block kind) then
-        find_in lam_switch.sw_blocks (function
-          | Block (Untagged {block_type}) -> matches_block kind block_type
-          | Constant _ | Block (Tagged _) -> false)
-      else `Undecided
-  in
   match (lam, lam_switch.sw_dispatch) with
-  | Lconst cst, Switch_variant ({block_types = _ :: _} as facts) -> (
-    match untagged_action facts cst with
-    | `Case action -> action_or_switch action
-    | `Undecided -> Lswitch (lam, lam_switch))
+  | Lconst cst, Switch_variant {block_types = _ :: _} -> (
+    (* An untagged variant is dispatched on the value, so a constant
+       scrutinee takes the arm the generated code takes for its JavaScript
+       value, not the arm of the constructor it was written with *)
+    match (untagged_plan lam_switch, untagged_value cst) with
+    | Some plan, Some value -> (
+      match Untagged_dispatch.select plan value with
+      | Arm action -> action
+      | Default -> action_or_switch None
+      | Nothing -> Lswitch (lam, lam_switch))
+    | (Some _ | None), _ -> Lswitch (lam, lam_switch))
   | Lconst (Const_constructor constant), _ ->
     (* By runtime value: a coercion accepts [@as(1.)] for [@as(1.0)] *)
     let action =

compiler/ml/lambda.mli

@@ -486,6 +486,10 @@ val letrec : (Ident.t * t) list -> t -> t
 
 val prim : primitive:primitive -> args:t list -> Location.t -> t
 
+val untagged_plan : lambda_switch -> t Untagged_dispatch.tree option
+(** The dispatch of a switch on an untagged variant, which the code generator
+    prints and constant folding evaluates *)
+
 val switch : t -> lambda_switch -> t
 
 val stringswitch : t -> (string * t) list -> t option -> t

compiler/ml/untagged_dispatch.ml

@@ -315,3 +315,65 @@ let plan ~same
     else
       Switch_value
         (clauses ~same literal_arms ~shares:(fun _ _ -> true), payload)
+
+type value = Literal of literal_tag | Object | Array
+
+type 'a outcome = Arm of 'a | Default | Nothing
+
+(* What [typeof] returns for a value, as a kind *)
+let kind_of = function
+  | Literal (String _) -> Some StringType
+  | Literal (Int _ | Float _) -> Some NumberType
+  | Literal (BigInt _) -> Some BigintType
+  | Literal (Bool _) -> Some BooleanType
+  | Literal Null | Object | Array -> Some ObjectType
+  | Literal Undefined -> None
+
+let select ?(on_test = ignore) tree value =
+  (* JavaScript semantics, including the short circuits of [&&] and [||] *)
+  let rec holds test =
+    on_test test;
+    match (test, value) with
+    | Equals literal, Literal v -> same_literal_value v literal
+    | Equals _, (Object | Array) -> false
+    | Typeof kind, _ -> kind_of value = Some kind
+    | Instance Array, _ -> value = Array
+    | Instance (Class _), _ -> false
+    | Not_null, _ -> value <> Literal Null
+    | Not test, _ -> not (holds test)
+    | And (test1, test2), _ -> holds test1 && holds test2
+    | Or (test1, test2), _ -> holds test1 || holds test2
+  in
+  let rec run = function
+    | Run body -> Arm body
+    | Fallback -> Default
+    | Skip -> Nothing
+    | If (test, then_, else_) -> if holds test then run then_ else run else_
+    | Block tree -> run tree
+    | Switch_value (clauses, default) -> (
+      match value with
+      | Literal v ->
+        from_clauses
+          (fun literal -> same_literal_value v literal)
+          clauses default
+      | Object | Array -> run default)
+    | Switch_typeof (clauses, default) -> (
+      match kind_of value with
+      | Some kind -> from_clauses (fun k -> k = kind) clauses default
+      | None -> run default)
+  (* A matching case runs its body, or falls through to the next one *)
+  and from_clauses : 'k. ('k -> bool) -> ('k * _ clause) list -> _ =
+   fun matches clauses default ->
+    let rec fall = function
+      | (_, Body tree) :: _ -> run tree
+      | (_, Falls_through _) :: rest -> fall rest
+      | [] -> run default
+    in
+    let rec find = function
+      | [] -> run default
+      | (key, _) :: _ as clauses when matches key -> fall clauses
+      | _ :: rest -> find rest
+    in
+    find clauses
+  in
+  run tree

compiler/ml/untagged_dispatch.mli

@@ -53,3 +53,20 @@ type 'a input = {
 
 val plan : same:('a -> 'a -> bool) -> 'a input -> 'a tree
 (** [same] tells two bodies that are one body apart *)
+
+(** What constant folding knows about a constant: the JavaScript value it is
+    emitted as *)
+type value =
+  | Literal of Variant_runtime.literal_tag
+  | Object  (** a record or an inline record, which is no class instance *)
+  | Array  (** a tuple, or a record whose fields are named 0, 1, … *)
+
+type 'a outcome =
+  | Arm of 'a
+  | Default  (** the match's default action *)
+  | Nothing  (** no statement runs *)
+
+val select : ?on_test:(test -> unit) -> 'a tree -> value -> 'a outcome
+(** What the code printed from a plan runs for a value: constant folding
+    evaluates the plan instead of deciding by itself, so a folded match and
+    the generated code agree. [on_test] sees each test that is evaluated. *)