|
19 | 19 | open! Flambda.Import |
20 | 20 |
|
21 | 21 | module BLB = Bindable_let_bound |
| 22 | +module KP = Kinded_parameter |
22 | 23 | module LC = Lifted_constant |
23 | 24 | module LCS = Lifted_constant_state |
24 | 25 | module P = Flambda_primitive |
25 | 26 | module UA = Upwards_acc |
| 27 | +module UE = Upwards_env |
26 | 28 | module VB = Var_in_binding_pos |
27 | 29 |
|
28 | 30 | (* The constructed values of this type aren't currently used, but will be |
@@ -537,3 +539,137 @@ let create_switch uacc ~scrutinee ~arms = |
537 | 539 | Expr.create_switch switch, |
538 | 540 | UA.cost_metrics_add ~added:(Cost_metrics.switch switch) uacc |
539 | 541 |
|
| 542 | +type add_wrapper_for_fixed_arity_continuation0_result = |
| 543 | + | This_continuation of Continuation.t |
| 544 | + | Apply_cont of Apply_cont.t |
| 545 | + | New_wrapper of Continuation.t * Continuation_handler.t * Cost_metrics.t |
| 546 | + |
| 547 | +type cont_or_apply_cont = |
| 548 | + | Continuation of Continuation.t |
| 549 | + | Apply_cont of Apply_cont.t |
| 550 | + |
| 551 | +let add_wrapper_for_fixed_arity_continuation0 uacc cont_or_apply_cont |
| 552 | + ~use_id arity : add_wrapper_for_fixed_arity_continuation0_result = |
| 553 | + let uenv = UA.uenv uacc in |
| 554 | + let cont = |
| 555 | + match cont_or_apply_cont with |
| 556 | + | Continuation cont -> cont |
| 557 | + | Apply_cont apply_cont -> Apply_cont.continuation apply_cont |
| 558 | + in |
| 559 | + let original_cont = cont in |
| 560 | + let cont = UE.resolve_continuation_aliases uenv cont in |
| 561 | + match UE.find_apply_cont_rewrite uenv original_cont with |
| 562 | + | None -> This_continuation cont |
| 563 | + | Some rewrite when Apply_cont_rewrite.does_nothing rewrite -> |
| 564 | + (* CR mshinwell: think more about this check w.r.t. subkinds *) |
| 565 | + let arity = Flambda_arity.With_subkinds.to_arity arity in |
| 566 | + let arity_in_rewrite = |
| 567 | + Apply_cont_rewrite.original_params_arity rewrite |
| 568 | + |> Flambda_arity.With_subkinds.to_arity |
| 569 | + in |
| 570 | + if not (Flambda_arity.equal arity arity_in_rewrite) then begin |
| 571 | + Misc.fatal_errorf "Arity %a provided to fixed-arity-wrapper \ |
| 572 | + addition function does not match arity %a in rewrite:@ %a" |
| 573 | + Flambda_arity.print arity |
| 574 | + Flambda_arity.print arity_in_rewrite |
| 575 | + Apply_cont_rewrite.print rewrite |
| 576 | + end; |
| 577 | + This_continuation cont |
| 578 | + | Some rewrite -> |
| 579 | + (* CR-someday mshinwell: This area should be improved and hence |
| 580 | + simplified. Allowing [Apply] to take extra arguments is probably the |
| 581 | + way forward. Although unboxing of variants requires untagging |
| 582 | + expressions to be inserted, so wrappers cannot always be avoided. *) |
| 583 | + let params = List.map (fun _kind -> Variable.create "param") arity in |
| 584 | + let kinded_params = List.map2 KP.create params arity in |
| 585 | + let new_wrapper expr ~free_names ~cost_metrics = |
| 586 | + let new_cont = Continuation.create () in |
| 587 | + let new_handler = |
| 588 | + Continuation_handler.create kinded_params ~handler:expr |
| 589 | + ~free_names_of_handler:free_names |
| 590 | + ~is_exn_handler:false |
| 591 | + in |
| 592 | + New_wrapper (new_cont, new_handler, cost_metrics) |
| 593 | + in |
| 594 | + match cont_or_apply_cont with |
| 595 | + | Continuation cont -> |
| 596 | + (* In this case, any generated [Apply_cont] will sit inside a wrapper |
| 597 | + that binds [kinded_params]. *) |
| 598 | + let args = List.map KP.simple kinded_params in |
| 599 | + let apply_cont = Apply_cont.create cont ~args ~dbg:Debuginfo.none in |
| 600 | + begin match Apply_cont_rewrite.rewrite_use rewrite use_id apply_cont with |
| 601 | + | Apply_cont apply_cont -> |
| 602 | + new_wrapper (Expr.create_apply_cont apply_cont) |
| 603 | + ~free_names:(Known (Apply_cont.free_names apply_cont)) |
| 604 | + ~cost_metrics:(Cost_metrics.apply_cont apply_cont) |
| 605 | + | Expr build_expr -> |
| 606 | + let expr, cost_metrics, free_names = |
| 607 | + build_expr ~apply_cont_to_expr:(fun apply_cont -> |
| 608 | + Expr.create_apply_cont apply_cont, |
| 609 | + Cost_metrics.apply_cont apply_cont, |
| 610 | + Apply_cont.free_names apply_cont) |
| 611 | + in |
| 612 | + new_wrapper expr ~free_names:(Known free_names) ~cost_metrics |
| 613 | + end |
| 614 | + | Apply_cont apply_cont -> |
| 615 | + let apply_cont = Apply_cont.update_continuation apply_cont cont in |
| 616 | + match Apply_cont_rewrite.rewrite_use rewrite use_id apply_cont with |
| 617 | + | Apply_cont apply_cont -> Apply_cont apply_cont |
| 618 | + | Expr build_expr -> |
| 619 | + let expr, cost_metrics, free_names = |
| 620 | + build_expr ~apply_cont_to_expr:(fun apply_cont -> |
| 621 | + Expr.create_apply_cont apply_cont, |
| 622 | + Cost_metrics.apply_cont apply_cont, |
| 623 | + Apply_cont.free_names apply_cont) |
| 624 | + in |
| 625 | + new_wrapper expr ~free_names:(Known free_names) ~cost_metrics |
| 626 | + |
| 627 | +type add_wrapper_for_switch_arm_result = |
| 628 | + | Apply_cont of Apply_cont.t |
| 629 | + | New_wrapper of Continuation.t * Continuation_handler.t * Cost_metrics.t |
| 630 | + |
| 631 | +let add_wrapper_for_switch_arm uacc apply_cont ~use_id arity |
| 632 | + : add_wrapper_for_switch_arm_result = |
| 633 | + match |
| 634 | + add_wrapper_for_fixed_arity_continuation0 uacc (Apply_cont apply_cont) |
| 635 | + ~use_id arity |
| 636 | + with |
| 637 | + | This_continuation cont -> |
| 638 | + Apply_cont (Apply_cont.update_continuation apply_cont cont) |
| 639 | + | Apply_cont apply_cont -> Apply_cont apply_cont |
| 640 | + | New_wrapper (cont, wrapper, cost_metrics) -> New_wrapper (cont, wrapper, cost_metrics) |
| 641 | + |
| 642 | +let add_wrapper_for_fixed_arity_continuation uacc cont ~use_id arity ~around = |
| 643 | + match |
| 644 | + add_wrapper_for_fixed_arity_continuation0 uacc (Continuation cont) |
| 645 | + ~use_id arity |
| 646 | + with |
| 647 | + | This_continuation cont -> around uacc cont |
| 648 | + | Apply_cont _ -> assert false |
| 649 | + | New_wrapper (new_cont, new_handler, cost_metrics_of_handler) -> |
| 650 | + let body, uacc = around uacc new_cont in |
| 651 | + let added = |
| 652 | + Cost_metrics.increase_due_to_let_cont_non_recursive ~cost_metrics_of_handler |
| 653 | + in |
| 654 | + Let_cont.create_non_recursive new_cont new_handler ~body |
| 655 | + ~free_names_of_body:(Known (Expr.free_names body)), |
| 656 | + UA.cost_metrics_add ~added uacc |
| 657 | + |
| 658 | +let add_wrapper_for_fixed_arity_apply uacc ~use_id arity apply = |
| 659 | + match Apply.continuation apply with |
| 660 | + | Never_returns -> |
| 661 | + Expr.create_apply apply, |
| 662 | + UA.cost_metrics_add ~added:(Cost_metrics.apply apply) uacc |
| 663 | + | Return cont -> |
| 664 | + add_wrapper_for_fixed_arity_continuation uacc cont |
| 665 | + ~use_id arity |
| 666 | + ~around:(fun uacc return_cont -> |
| 667 | + let exn_cont = |
| 668 | + UE.resolve_exn_continuation_aliases (UA.uenv uacc) |
| 669 | + (Apply.exn_continuation apply) |
| 670 | + in |
| 671 | + let apply = |
| 672 | + Apply.with_continuations apply (Return return_cont) exn_cont |
| 673 | + in |
| 674 | + Expr.create_apply apply, |
| 675 | + UA.cost_metrics_add ~added:(Cost_metrics.apply apply) uacc) |
0 commit comments