Direct conversion of lambda switches to ilambda (#429)

Keryan-dev · web-flow · commit 5fdd91cbb3b1 · 2021-05-07T15:47:40.000+02:00
diff --git a/lambda/lambda.ml b/lambda/lambda.ml
@@ -467,7 +467,6 @@ and lambda_switch =
     sw_numblocks: int;
     sw_blocks: (lambda_switch_block_key * lambda) list;
     sw_failaction : lambda option;
-    sw_tags_to_sizes : Targetint.OCaml.t Tag.Scannable.Map.t;
   }
 
 and lambda_switch_block_key =
@@ -948,7 +947,6 @@ let shallow_map f = function
                  sw_numblocks = sw.sw_numblocks;
                  sw_blocks = List.map (fun (n, e) -> (n, f e)) sw.sw_blocks;
                  sw_failaction = Option.map f sw.sw_failaction;
-                 sw_tags_to_sizes = sw.sw_tags_to_sizes;
                },
                loc)
   | Lstringswitch (e, sw, default, loc) ->
diff --git a/lambda/lambda.mli b/lambda/lambda.mli
@@ -339,7 +339,6 @@ and lambda_switch =
     sw_numblocks: int;                  (* Number of tag block cases *)
     sw_blocks: (lambda_switch_block_key * lambda) list;  (* Tag block cases *)
     sw_failaction : lambda option;      (* Action to take if failure *)
-    sw_tags_to_sizes : Targetint.OCaml.t Tag.Scannable.Map.t;
   }
 
 and lambda_switch_block_key =
diff --git a/lambda/matching.ml b/lambda/matching.ml
@@ -1999,7 +1999,6 @@ let inline_lazy_force_switch arg loc =
                         } )
                   ];
                 sw_failaction = Some varg;
-                sw_tags_to_sizes = Tag.Scannable.Map.empty;
               },
               loc ) ) )
 
@@ -2457,7 +2456,6 @@ module SArg = struct
           sw_numblocks = 0;
           sw_blocks = [];
           sw_failaction = None;
-          sw_tags_to_sizes = Tag.Scannable.Map.empty;
         },
         loc )
 
@@ -2959,7 +2957,6 @@ let combine_constructor loc arg pat_env cstr partial ctx def
                         sw_numblocks = cstr.cstr_nonconsts;
                         sw_blocks = nonconsts;
                         sw_failaction = fail_opt;
-                        sw_tags_to_sizes = Tag.Scannable.Map.empty;
                       }
                     in
                     let hs, sw = share_actions_sw sw in
diff --git a/lambda/simplif.ml b/lambda/simplif.ml
@@ -59,7 +59,6 @@ let rec eliminate_ref id = function
             List.map (fun (n, e) -> (n, eliminate_ref id e)) sw.sw_blocks;
          sw_failaction =
             Option.map (eliminate_ref id) sw.sw_failaction;
-         sw_tags_to_sizes = sw.sw_tags_to_sizes;
         },
         loc)
   | Lstringswitch(e, sw, default, loc) ->
diff --git a/middle_end/flambda/from_lambda/cps_conversion.ml b/middle_end/flambda/from_lambda/cps_conversion.ml
@@ -24,12 +24,6 @@ module I = Ilambda
 module L = Lambda
 module C = Lambda_conversions
 
-type proto_switch = {
-  numconsts : int;
-  consts : (int * L.lambda) list;
-  failaction : L.lambda option;
-}
-
 type primitive_transform_result =
   | Primitive of L.primitive * L.lambda list * L.scoped_location
   | Transformed of L.lambda
@@ -133,7 +127,6 @@ let switch_for_if_then_else ~cond ~ifso ~ifnot =
       sw_numblocks = 0;
       sw_blocks = [];
       sw_failaction = None;
-      sw_tags_to_sizes = Tag.Scannable.Map.empty;
     }
   in
   L.Lswitch (cond, switch, Loc_unknown)
@@ -449,23 +442,11 @@ let rec cps_non_tail (lam : L.lambda) (k : Ident.t -> Ilambda.t)
         k_exn
     | Transformed lam -> cps_non_tail lam k k_exn
     end
-  | Lswitch (scrutinee,
-      { sw_numconsts; sw_consts; sw_numblocks = _; sw_blocks; sw_failaction;
-        sw_tags_to_sizes = _; }, _loc) ->
-    begin match sw_blocks with
-    | [] -> ()
-    | _ -> Misc.fatal_error "Lswitch `block' cases are forbidden"
-    end;
+  | Lswitch (scrutinee, switch, _loc) ->
     let after_switch = Continuation.create () in
     let result_var = Ident.create_local "switch_result" in
     let after = k result_var in
-    let proto_switch : proto_switch =
-      { numconsts = sw_numconsts;
-        consts = sw_consts;
-        failaction = sw_failaction;
-      }
-    in
-    let body = cps_switch proto_switch ~scrutinee after_switch k_exn in
+    let body = cps_switch switch ~scrutinee after_switch k_exn in
     Let_cont {
       name = after_switch;
       is_exn_handler = false;
@@ -474,8 +455,9 @@ let rec cps_non_tail (lam : L.lambda) (k : Ident.t -> Ilambda.t)
       body;
       handler = after;
     }
-  | Lstringswitch _ ->
-    Misc.fatal_error "Lstringswitch must be expanded prior to CPS conversion"
+  | Lstringswitch (scrutinee, cases, default, loc) ->
+    cps_non_tail (Matching.expand_stringswitch loc scrutinee cases default)
+      k k_exn
   | Lstaticraise (static_exn, args) ->
     let continuation =
       match Numbers.Int.Map.find static_exn !static_exn_env with
@@ -783,22 +765,10 @@ and cps_tail (lam : L.lambda) (k : Continuation.t) (k_exn : Continuation.t)
             Apply_cont (k, None, [Ilambda.Var result_var]))) k_exn
     | Transformed lam -> cps_tail lam k k_exn
     end
-  | Lswitch (scrutinee,
-      { sw_numconsts; sw_consts; sw_numblocks = _; sw_blocks; sw_failaction;
-        sw_tags_to_sizes = _; }, _loc) ->
-    begin match sw_blocks with
-    | [] -> ()
-    | _ -> Misc.fatal_error "Lswitch `block' cases are forbidden"
-    end;
-    let proto_switch : proto_switch =
-      { numconsts = sw_numconsts;
-        consts = sw_consts;
-        failaction = sw_failaction;
-      }
-    in
-    cps_switch proto_switch ~scrutinee k k_exn
-  | Lstringswitch _ ->
-    Misc.fatal_error "Lstringswitch must be expanded prior to CPS conversion"
+  | Lswitch (scrutinee,switch, _loc) ->
+    cps_switch switch ~scrutinee k k_exn
+  | Lstringswitch (scrutinee, cases, default, loc) ->
+    cps_tail (Matching.expand_stringswitch loc scrutinee cases default) k k_exn
   | Lstaticraise (static_exn, args) ->
     let continuation =
       match Numbers.Int.Map.find static_exn !static_exn_env with
@@ -958,9 +928,35 @@ and cps_function ({ kind; params; return; body; attr; loc; } : L.lfunction)
     stub = false;
   }
 
-and cps_switch (switch : proto_switch) ~scrutinee (k : Continuation.t)
+and cps_switch (switch : L.lambda_switch) ~scrutinee (k : Continuation.t)
       (k_exn : Continuation.t) : Ilambda.t =
-  let consts_rev, wrappers =
+  let block_nums, sw_blocks = List.split switch.sw_blocks in
+  let block_nums =
+    List.map (fun ({ sw_tag; _ } : L.lambda_switch_block_key) ->
+        begin match Tag.Scannable.create sw_tag with
+        | Some tag ->
+          let tag' = Tag.Scannable.to_tag tag in
+          if Tag.is_structured_block_but_not_a_variant tag' then
+            Misc.fatal_errorf "Bad tag %a in [Lswitch] (tag is that of a \
+              scannable block, but not one treated like a variant; \
+              [Lswitch] can only be used for variant matching)"
+              Tag.print tag'
+        | None ->
+           Misc.fatal_errorf "Bad tag %d in [Lswitch] (not the tag \
+             of a GC-scannable block)"
+             sw_tag
+        end;
+        sw_tag)
+      block_nums
+  in
+  if switch.sw_numblocks > Obj.last_non_constant_constructor_tag + 1
+  then begin
+      Misc.fatal_errorf "Too many blocks (%d) in [Lswitch], would \
+        overlap into tag space for blocks that are not treated like variants; \
+        [Lswitch] can only be used for variant matching"
+        switch.sw_numblocks
+    end;
+  let convert_arms_rev cases wrappers =
     List.fold_left (fun (consts_rev, wrappers) (arm, (action : L.lambda)) ->
         match action with
         | Lvar var when not (Ident.Set.mem var !mutable_variables) ->
@@ -999,26 +995,76 @@ and cps_switch (switch : proto_switch) ~scrutinee (k : Continuation.t)
             let consts_rev = (arm, cont, None, []) :: consts_rev in
             let wrappers = (cont, action) :: wrappers in
             consts_rev, wrappers)
-      ([], [])
-      switch.consts
+      ([], wrappers)
+      cases
+  in
+  let consts_rev, wrappers = convert_arms_rev switch.sw_consts [] in
+  let blocks_rev, wrappers =
+    convert_arms_rev (List.combine block_nums sw_blocks) wrappers
   in
   let consts = List.rev consts_rev in
+  let blocks = List.rev blocks_rev in
   let failaction, wrappers =
-    match switch.failaction with
+    match switch.sw_failaction with
     | None -> None, wrappers
     | Some action ->
       let cont = Continuation.create () in
       let action = cps_tail action k k_exn in
       let wrappers = (cont, action) :: wrappers in
       Some (cont, None, []), wrappers
   in
-  let switch : I.switch =
-    { numconsts = switch.numconsts;
+  let const_switch : I.switch =
+    { numconsts = switch.sw_numconsts;
       consts;
       failaction;
     }
   in
+  let block_switch : I.switch =
+    { numconsts = switch.sw_numblocks;
+      consts = blocks;
+      failaction;
+    }
+  in
+  let build_switch scrutinee wrappers =
+    let const_switch = I.Switch (scrutinee, const_switch) in
+    let scrutinee_tag = Ident.create_local "scrutinee_tag" in
+    let block_switch =
+      I.Let (scrutinee_tag,
+        Not_user_visible,
+        Pintval,
+        Prim { prim = Pgettag;
+          args = [Var scrutinee];
+          loc = Loc_unknown;
+          exn_continuation = None; },
+        I.Switch(scrutinee_tag, block_switch))
+    in
+    if switch.sw_numblocks = 0 then const_switch, wrappers
+    else if switch.sw_numconsts = 0 then block_switch, wrappers
+    else
+      let const_cont = Continuation.create () in
+      let block_cont = Continuation.create () in
+      let isint_switch : I.switch =
+        { numconsts = 2;
+          consts = [ (0, block_cont, None, []); (1, const_cont, None, []) ];
+          failaction = None;
+        }
+      in
+      let is_scrutinee_int = Ident.create_local "is_scrutinee_int" in
+      let isint_switch =
+        I.Let (is_scrutinee_int,
+          Not_user_visible,
+          Pintval,
+          Prim { prim = Pflambda_isint;
+            args = [Var scrutinee];
+            loc = Loc_unknown;
+            exn_continuation = None; },
+          I.Switch(is_scrutinee_int, isint_switch))
+      in
+      isint_switch,
+      ((const_cont, const_switch)::(block_cont, block_switch)::wrappers)
+  in
   cps_non_tail scrutinee (fun scrutinee ->
+      let switch, wrappers = build_switch scrutinee wrappers in
       List.fold_left (fun body (cont, action) ->
           I.Let_cont {
             name = cont;
@@ -1028,8 +1074,7 @@ and cps_switch (switch : proto_switch) ~scrutinee (k : Continuation.t)
             body;
             handler = action;
           })
-        (I.Switch (scrutinee, switch))
-        wrappers)
+        switch wrappers)
     k_exn
 
 let lambda_to_ilambda lam : Ilambda.program =
diff --git a/middle_end/flambda/from_lambda/prepare_lambda.ml b/middle_end/flambda/from_lambda/prepare_lambda.ml
@@ -215,99 +215,22 @@ let rec prepare env (lam : L.lambda) (k : L.lambda -> L.lambda) =
       prepare_option env switch.sw_failaction (fun sw_failaction ->
         prepare_list env sw_consts (fun sw_consts ->
           prepare_list env sw_blocks (fun sw_blocks ->
-            let sw_failaction, wrap_switch =
-              match sw_failaction with
-              | None -> None, (fun lam -> lam)
-              | Some failaction ->
-                let failaction_cont = L.next_raise_count () in
-                let wrap_switch lam : L.lambda =
-                  Lstaticcatch (lam, (failaction_cont, []), failaction)
-                in
-                Some (L.Lstaticraise (failaction_cont, [])), wrap_switch
-            in
-            let consts_switch : L.lambda_switch =
+            let switch : L.lambda_switch =
               { sw_numconsts = switch.sw_numconsts;
                 sw_consts = List.combine const_nums sw_consts;
-                sw_numblocks = 0;
-                sw_blocks = [];
-                sw_failaction;
-                sw_tags_to_sizes = Tag.Scannable.Map.empty;
-              }
-            in
-            (* CR mshinwell: Merge this file into Cps_conversion then delete
-               [sw_tags_to_sizes]. *)
-            let tags_to_sizes =
-              List.fold_left (fun tags_to_sizes
-                        ({ sw_tag; sw_size; } : L.lambda_switch_block_key) ->
-                  match Tag.Scannable.create sw_tag with
-                  | Some tag ->
-                    let tag' = Tag.Scannable.to_tag tag in
-                    if Tag.is_structured_block_but_not_a_variant tag' then begin
-                      Misc.fatal_errorf "Bad tag %a in [Lswitch] (tag is that \
-                          of a scannable block, but not one treated like a \
-                          variant; [Lswitch] can only be used for variant \
-                          matching)"
-                        Tag.print tag'
-                    end;
-                    let size = Targetint.OCaml.of_int sw_size in
-                    Tag.Scannable.Map.add tag size tags_to_sizes
-                  | None ->
-                    Misc.fatal_errorf "Bad tag %d in [Lswitch] (not the tag \
-                        of a GC-scannable block)"
-                      sw_tag)
-                Tag.Scannable.Map.empty
-                block_nums
-            in
-            let block_nums =
-              List.map (fun ({ sw_tag; _} : L.lambda_switch_block_key) ->
-                  sw_tag)
-                block_nums
-            in
-            if switch.sw_numblocks > Obj.last_non_constant_constructor_tag + 1
-            then begin
-              Misc.fatal_errorf "Too many blocks (%d) in [Lswitch], would \
-                  overlap into tag space for blocks that are not treated \
-                  like variants; [Lswitch] can only be used for variant \
-                  matching"
-                switch.sw_numblocks
-            end;
-            let blocks_switch : L.lambda_switch =
-              { sw_numconsts = switch.sw_numblocks;
-                sw_consts = List.combine block_nums sw_blocks;
-                sw_numblocks = 0;
-                sw_blocks = [];
+                sw_numblocks = switch.sw_numblocks;
+                sw_blocks = List.combine block_nums sw_blocks;
                 sw_failaction;
-                (* XXX What about the size for the failaction? ... *)
-                sw_tags_to_sizes = tags_to_sizes;
               }
             in
-            let consts_switch : L.lambda =
-              L.Lswitch (scrutinee, consts_switch, loc)
-            in
-            let blocks_switch : L.lambda =
-              L.Lswitch (
-               Lprim (Pgettag, [scrutinee], Loc_unknown),
-               blocks_switch, loc)
-            in
-            let isint_switch : L.lambda_switch =
-              { sw_numconsts = 2;
-                sw_consts = [0, blocks_switch; 1, consts_switch];
-                sw_numblocks = 0;
-                sw_blocks = [];
-                sw_failaction = None;
-                sw_tags_to_sizes = Tag.Scannable.Map.empty;
-              }
-            in
-            let switch =
-              if switch.sw_numconsts = 0 then blocks_switch
-              else if switch.sw_numblocks = 0 then consts_switch
-              else
-                L.Lswitch (L.Lprim (Pflambda_isint, [scrutinee], Loc_unknown),
-                  isint_switch, loc)
-            in
-            k (wrap_switch switch)))))
+            k (Lswitch (scrutinee, switch, loc))))))
   | Lstringswitch (scrutinee, cases, default, loc) ->
-    prepare env (Matching.expand_stringswitch loc scrutinee cases default) k
+    prepare env scrutinee (fun scrutinee ->
+      let patterns, actions = List.split cases in
+      prepare_list env actions (fun actions ->
+        prepare_option env default (fun default ->
+          let cases = List.combine patterns actions in
+          k (L.Lstringswitch (scrutinee, cases, default, loc)))))
   | Lstaticraise (cont, args) ->
     prepare_list env args (fun args ->
       k (L.Lstaticraise (cont, args)))

Original file line number	Diff line number	Diff line change
`@@ -339,7 +339,6 @@ and lambda_switch =`
`339`	`339`	`sw_numblocks: int; (* Number of tag block cases *)`
`340`	`340`	`sw_blocks: (lambda_switch_block_key * lambda) list; (* Tag block cases *)`
`341`	`341`	`sw_failaction : lambda option; (* Action to take if failure *)`
`342`		`- sw_tags_to_sizes : Targetint.OCaml.t Tag.Scannable.Map.t;`
`343`	`342`	`}`
`344`	`343`
`345`	`344`	`and lambda_switch_block_key =`