1(**************************************************************************) 2(* *) 3(* OCaml *) 4(* *) 5(* Pierre Chambart, OCamlPro *) 6(* Mark Shinwell and Leo White, Jane Street Europe *) 7(* *) 8(* Copyright 2013--2016 OCamlPro SAS *) 9(* Copyright 2014--2016 Jane Street Group LLC *) 10(* *) 11(* All rights reserved. This file is distributed under the terms of *) 12(* the GNU Lesser General Public License version 2.1, with the *) 13(* special exception on linking described in the file LICENSE. *) 14(* *) 15(**************************************************************************) 16 17[@@@ocaml.warning "+a-4-9-30-40-41-42"] 18 19(* CR-someday mshinwell: move to Flambda_utils *) 20let rec tail_variable : Flambda.t -> Variable.t option = function 21 | Var v -> Some v 22 | Let_rec (_, e) 23 | Let_mutable { body = e } 24 | Let { body = e; _ } -> tail_variable e 25 | _ -> None 26 27let closure_symbol ~(backend : (module Backend_intf.S)) closure_id = 28 let module Backend = (val backend) in 29 Backend.closure_symbol closure_id 30 31let make_variable_symbol prefix var = 32 Symbol.create (Compilation_unit.get_current_exn ()) 33 (Linkage_name.create 34 (prefix ^ Variable.unique_name (Variable.rename var))) 35 36(** Traverse the given expression assigning symbols to [let]- and [let rec]- 37 bound constant variables. At the same time collect the definitions of 38 such variables. *) 39let assign_symbols_and_collect_constant_definitions 40 ~(backend : (module Backend_intf.S)) 41 ~(program : Flambda.program) 42 ~(inconstants : Inconstant_idents.result) = 43 let var_to_symbol_tbl = Variable.Tbl.create 42 in 44 let var_to_definition_tbl = Variable.Tbl.create 42 in 45 let module AA = Alias_analysis in 46 let assign_symbol var (named : Flambda.named) = 47 if not (Inconstant_idents.variable var inconstants) then begin 48 let assign_symbol () = 49 let symbol = make_variable_symbol "" var in 50 Variable.Tbl.add var_to_symbol_tbl var symbol 51 in 52 let assign_existing_symbol = Variable.Tbl.add var_to_symbol_tbl var in 53 let record_definition = Variable.Tbl.add var_to_definition_tbl var in 54 match named with 55 | Symbol symbol -> 56 assign_existing_symbol symbol; 57 record_definition (AA.Symbol symbol) 58 | Const const -> record_definition (AA.Const const) 59 | Allocated_const const -> 60 assign_symbol (); 61 record_definition (AA.Allocated_const (Normal const)) 62 | Read_mutable _ -> 63 (* [Inconstant_idents] always marks these expressions as 64 inconstant, so we should never get here. *) 65 assert false 66 | Prim (Pmakeblock (tag, _, _value_kind), fields, _) -> 67 assign_symbol (); 68 record_definition (AA.Block (Tag.create_exn tag, fields)) 69 | Read_symbol_field (symbol, field) -> 70 record_definition (AA.Symbol_field (symbol, field)) 71 | Set_of_closures ( 72 { function_decls = { funs; set_of_closures_id; _ }; 73 _ } as set) -> 74 assert (not (Inconstant_idents.closure set_of_closures_id 75 inconstants)); 76 assign_symbol (); 77 record_definition (AA.Set_of_closures set); 78 Variable.Map.iter (fun fun_var _ -> 79 let closure_id = Closure_id.wrap fun_var in 80 let closure_symbol = closure_symbol ~backend closure_id in 81 Variable.Tbl.add var_to_symbol_tbl fun_var closure_symbol; 82 let project_closure = 83 Alias_analysis.Project_closure 84 { set_of_closures = var; closure_id } 85 in 86 Variable.Tbl.add var_to_definition_tbl fun_var 87 project_closure) 88 funs 89 | Move_within_set_of_closures ({ closure = _; start_from = _; move_to; } 90 as move) -> 91 assign_existing_symbol (closure_symbol ~backend move_to); 92 record_definition (AA.Move_within_set_of_closures move) 93 | Project_closure ({ closure_id } as project_closure) -> 94 assign_existing_symbol (closure_symbol ~backend closure_id); 95 record_definition (AA.Project_closure project_closure) 96 | Prim (Pfield index, [block], _) -> 97 record_definition (AA.Field (block, index)) 98 | Prim (Pfield _, _, _) -> 99 Misc.fatal_errorf "[Pfield] with the wrong number of arguments" 100 Flambda.print_named named 101 | Prim (Pmakearray (Pfloatarray as kind, mutability), args, _) -> 102 assign_symbol (); 103 record_definition (AA.Allocated_const (Array (kind, mutability, args))) 104 | Prim (Pduparray (kind, mutability), [arg], _) -> 105 assign_symbol (); 106 record_definition (AA.Allocated_const ( 107 Duplicate_array (kind, mutability, arg))) 108 | Prim _ -> 109 Misc.fatal_errorf "Primitive not expected to be constant: @.%a@." 110 Flambda.print_named named 111 | Project_var project_var -> 112 record_definition (AA.Project_var project_var) 113 | Expr e -> 114 match tail_variable e with 115 | None -> assert false (* See [Inconstant_idents]. *) 116 | Some v -> record_definition (AA.Variable v) 117 end 118 in 119 let assign_symbol_program expr = 120 Flambda_iterators.iter_all_immutable_let_and_let_rec_bindings expr 121 ~f:assign_symbol 122 in 123 Flambda_iterators.iter_exprs_at_toplevel_of_program program 124 ~f:assign_symbol_program; 125 let let_symbol_to_definition_tbl = Symbol.Tbl.create 42 in 126 let initialize_symbol_to_definition_tbl = Symbol.Tbl.create 42 in 127 let rec collect_let_and_initialize_symbols (program : Flambda.program_body) = 128 match program with 129 | Let_symbol (symbol, decl, program) -> 130 Symbol.Tbl.add let_symbol_to_definition_tbl symbol decl; 131 collect_let_and_initialize_symbols program 132 | Let_rec_symbol (decls, program) -> 133 List.iter (fun (symbol, decl) -> 134 Symbol.Tbl.add let_symbol_to_definition_tbl symbol decl) 135 decls; 136 collect_let_and_initialize_symbols program 137 | Effect (_, program) -> collect_let_and_initialize_symbols program 138 | Initialize_symbol (symbol,_tag,fields,program) -> 139 collect_let_and_initialize_symbols program; 140 let fields = List.map tail_variable fields in 141 Symbol.Tbl.add initialize_symbol_to_definition_tbl symbol fields 142 | End _ -> () 143 in 144 collect_let_and_initialize_symbols program.program_body; 145 let record_set_of_closure_equalities 146 (set_of_closures : Flambda.set_of_closures) = 147 Variable.Map.iter (fun arg (var : Flambda.specialised_to) -> 148 if not (Inconstant_idents.variable arg inconstants) then 149 Variable.Tbl.add var_to_definition_tbl arg (AA.Variable var.var)) 150 set_of_closures.free_vars; 151 Variable.Map.iter (fun arg (spec_to : Flambda.specialised_to) -> 152 if not (Inconstant_idents.variable arg inconstants) then 153 Variable.Tbl.add var_to_definition_tbl arg 154 (AA.Variable spec_to.var)) 155 set_of_closures.specialised_args 156 in 157 Flambda_iterators.iter_on_set_of_closures_of_program program 158 ~f:(fun ~constant set_of_closures -> 159 record_set_of_closure_equalities set_of_closures; 160 if constant then begin 161 Variable.Map.iter (fun fun_var _ -> 162 let closure_id = Closure_id.wrap fun_var in 163 let closure_symbol = closure_symbol ~backend closure_id in 164 Variable.Tbl.add var_to_definition_tbl fun_var 165 (AA.Symbol closure_symbol); 166 Variable.Tbl.add var_to_symbol_tbl fun_var closure_symbol) 167 set_of_closures.Flambda.function_decls.funs 168 end); 169 var_to_symbol_tbl, var_to_definition_tbl, 170 let_symbol_to_definition_tbl, initialize_symbol_to_definition_tbl 171 172let variable_field_definition 173 (var_to_symbol_tbl : Symbol.t Variable.Tbl.t) 174 (var_to_definition_tbl : 175 Alias_analysis.constant_defining_value Variable.Tbl.t) 176 (var : Variable.t) : Flambda.constant_defining_value_block_field = 177 try 178 Symbol (Variable.Tbl.find var_to_symbol_tbl var) 179 with Not_found -> 180 match Variable.Tbl.find var_to_definition_tbl var with 181 | Const c -> Const c 182 | const_defining_value -> 183 Misc.fatal_errorf "Unexpected pattern for a constant: %a: %a" 184 Variable.print var 185 Alias_analysis.print_constant_defining_value const_defining_value 186 | exception Not_found -> 187 Misc.fatal_errorf "No associated symbol for the constant %a" 188 Variable.print var 189 190let resolve_variable 191 (aliases : Alias_analysis.allocation_point Variable.Map.t) 192 (var_to_symbol_tbl : Symbol.t Variable.Tbl.t) 193 (var_to_definition_tbl : 194 Alias_analysis.constant_defining_value Variable.Tbl.t) 195 (var : Variable.t) : Flambda.constant_defining_value_block_field = 196 match Variable.Map.find var aliases with 197 | exception Not_found -> 198 variable_field_definition var_to_symbol_tbl var_to_definition_tbl var 199 | Symbol s -> Symbol s 200 | Variable aliased_variable -> 201 variable_field_definition var_to_symbol_tbl var_to_definition_tbl 202 aliased_variable 203 204let translate_set_of_closures 205 (inconstants : Inconstant_idents.result) 206 (aliases : Alias_analysis.allocation_point Variable.Map.t) 207 (var_to_symbol_tbl : Symbol.t Variable.Tbl.t) 208 (var_to_definition_tbl: 209 Alias_analysis.constant_defining_value Variable.Tbl.t) 210 (set_of_closures : Flambda.set_of_closures) = 211 let f var (named : Flambda.named) : Flambda.named = 212 if Inconstant_idents.variable var inconstants then 213 named 214 else 215 let resolved = 216 resolve_variable 217 aliases 218 var_to_symbol_tbl 219 var_to_definition_tbl 220 var 221 in 222 match resolved with 223 | Symbol s -> Symbol s 224 | Const c -> Const c 225 in 226 Flambda_iterators.map_function_bodies set_of_closures 227 ~f:(Flambda_iterators.map_all_immutable_let_and_let_rec_bindings ~f) 228 229let translate_constant_set_of_closures 230 (inconstants : Inconstant_idents.result) 231 (aliases : Alias_analysis.allocation_point Variable.Map.t) 232 (var_to_symbol_tbl : Symbol.t Variable.Tbl.t) 233 (var_to_definition_tbl: 234 Alias_analysis.constant_defining_value Variable.Tbl.t) 235 (constant_defining_values : Flambda.constant_defining_value Symbol.Map.t) = 236 Symbol.Map.map (fun (const : Flambda.constant_defining_value) -> 237 match const with 238 | Flambda.Allocated_const _ 239 | Flambda.Block _ 240 | Flambda.Project_closure _ -> 241 const 242 | Flambda.Set_of_closures set_of_closures -> 243 let set_of_closures = 244 translate_set_of_closures 245 (inconstants : Inconstant_idents.result) 246 (aliases : Alias_analysis.allocation_point Variable.Map.t) 247 (var_to_symbol_tbl : Symbol.t Variable.Tbl.t) 248 (var_to_definition_tbl: 249 Alias_analysis.constant_defining_value Variable.Tbl.t) 250 (set_of_closures : Flambda.set_of_closures) 251 in 252 Flambda.Set_of_closures set_of_closures) 253 constant_defining_values 254 255let find_original_set_of_closure 256 (aliases : Alias_analysis.allocation_point Variable.Map.t) 257 (var_to_symbol_tbl : Symbol.t Variable.Tbl.t) 258 (var_to_definition_tbl: 259 Alias_analysis.constant_defining_value Variable.Tbl.t) 260 project_closure_map 261 var = 262 let rec loop var = 263 match Variable.Map.find var aliases with 264 | Variable var -> 265 begin match Variable.Tbl.find var_to_definition_tbl var with 266 | Project_closure { set_of_closures = var } 267 | Move_within_set_of_closures { closure = var } -> 268 loop var 269 | Set_of_closures _ -> begin 270 match Variable.Tbl.find var_to_symbol_tbl var with 271 | s -> 272 s 273 | exception Not_found -> 274 Format.eprintf "var: %a@." Variable.print var; 275 assert false 276 end 277 | _ -> assert false 278 end 279 | Symbol s -> 280 match Symbol.Map.find s project_closure_map with 281 | exception Not_found -> 282 Misc.fatal_errorf "find_original_set_of_closure: cannot find \ 283 symbol %a in the project-closure map" 284 Symbol.print s 285 | s -> s 286 in 287 loop var 288 289let translate_definition_and_resolve_alias inconstants 290 (aliases : Alias_analysis.allocation_point Variable.Map.t) 291 (var_to_symbol_tbl : Symbol.t Variable.Tbl.t) 292 (var_to_definition_tbl : 293 Alias_analysis.constant_defining_value Variable.Tbl.t) 294 (symbol_definition_map : Flambda.constant_defining_value Symbol.Map.t) 295 (project_closure_map : Symbol.t Symbol.Map.t) 296 (definition : Alias_analysis.constant_defining_value) 297 ~(backend : (module Backend_intf.S)) 298 : Flambda.constant_defining_value option = 299 let resolve_float_array_involving_variables 300 ~(mutability : Asttypes.mutable_flag) ~vars = 301 (* Resolve an [Allocated_const] of the form: 302 [Array (Pfloatarray, _, _)] 303 (which references its contents via variables; it does not contain 304 manifest floats). *) 305 let find_float_var_definition var = 306 match Variable.Tbl.find var_to_definition_tbl var with 307 | Allocated_const (Normal (Float f)) -> f 308 | const_defining_value -> 309 Misc.fatal_errorf "Bad definition for float array member %a: %a" 310 Variable.print var 311 Alias_analysis.print_constant_defining_value 312 const_defining_value 313 in 314 let find_float_symbol_definition sym = 315 match Symbol.Map.find sym symbol_definition_map with 316 | Allocated_const (Float f) -> f 317 | const_defining_value -> 318 Misc.fatal_errorf "Bad definition for float array member %a: %a" 319 Symbol.print sym 320 Flambda.print_constant_defining_value 321 const_defining_value 322 in 323 let floats = 324 List.map (fun var -> 325 match Variable.Map.find var aliases with 326 | exception Not_found -> find_float_var_definition var 327 | Variable var -> find_float_var_definition var 328 | Symbol sym -> find_float_symbol_definition sym) 329 vars 330 in 331 let const : Allocated_const.t = 332 match mutability with 333 | Immutable -> Immutable_float_array floats 334 | Mutable -> Float_array floats 335 in 336 Some (Flambda.Allocated_const const) 337 in 338 match definition with 339 | Block (tag, fields) -> 340 Some (Flambda.Block (tag, 341 List.map (resolve_variable aliases var_to_symbol_tbl 342 var_to_definition_tbl) 343 fields)) 344 | Allocated_const (Normal const) -> Some (Flambda.Allocated_const const) 345 | Allocated_const (Duplicate_array (Pfloatarray, mutability, var)) -> 346 (* CR-someday mshinwell: This next section could do with cleanup. 347 What happens is: 348 - Duplicate contains a variable, which is resolved to 349 a float array thing full of variables; 350 - We send that value back through this function again so the 351 individual members of that array are resolved from variables to 352 floats. 353 - Then we can build the Flambda.name term containing the 354 Allocated_const (full of floats). 355 We should maybe factor out the code from the 356 Allocated_const (Array (...)) case below so this function doesn't have 357 to be recursive. *) 358 let (constant_defining_value : Alias_analysis.constant_defining_value) = 359 match Variable.Map.find var aliases with 360 | exception Not_found -> 361 Variable.Tbl.find var_to_definition_tbl var 362 | Variable var -> 363 Variable.Tbl.find var_to_definition_tbl var 364 | Symbol sym -> 365 match Symbol.Map.find sym symbol_definition_map with 366 | Allocated_const ((Immutable_float_array _) as const) -> 367 Alias_analysis.Allocated_const (Normal const) 368 | (Allocated_const _ | Block _ | Set_of_closures _ 369 | Project_closure _) as wrong -> 370 Misc.fatal_errorf 371 "Lift_constants.translate_definition_and_resolve_alias: \ 372 Duplicate Pfloatarray %a with symbol %a mapping to \ 373 wrong constant defining value %a" 374 Variable.print var 375 Alias_analysis.print_constant_defining_value definition 376 Flambda.print_constant_defining_value wrong 377 | exception Not_found -> 378 let module Backend = (val backend) in 379 match (Backend.import_symbol sym).descr with 380 | Value_unresolved _ -> 381 Misc.fatal_errorf 382 "Lift_constants.translate_definition_and_resolve_alias: \ 383 Duplicate Pfloatarray %a with unknown symbol: %a" 384 Variable.print var 385 Alias_analysis.print_constant_defining_value definition 386 | Value_float_array value_float_array -> 387 let contents = 388 Simple_value_approx.float_array_as_constant value_float_array 389 in 390 begin match contents with 391 | None -> 392 Misc.fatal_errorf 393 "Lift_constants.translate_definition_and_resolve_alias: \ 394 Duplicate Pfloatarray %a with not completely known float \ 395 array from symbol: %a" 396 Variable.print var 397 Alias_analysis.print_constant_defining_value definition 398 | Some l -> 399 Alias_analysis.Allocated_const (Normal (Immutable_float_array l)) 400 end 401 | wrong -> 402 (* CR-someday mshinwell: we might hit this if we ever duplicate 403 a mutable array across compilation units (e.g. "snapshotting" 404 an array). We do not currently generate such code. *) 405 Misc.fatal_errorf 406 "Lift_constants.translate_definition_and_resolve_alias: \ 407 Duplicate Pfloatarray %a with symbol %a that does not \ 408 have an export description of an immutable array" 409 Variable.print var 410 Alias_analysis.print_constant_defining_value definition 411 Simple_value_approx.print_descr wrong 412 in 413 begin match constant_defining_value with 414 | Allocated_const (Normal (Float_array _)) -> 415 (* This example from pchambart illustrates why we do not allow 416 the duplication of mutable arrays: 417 418 {| 419 let_symbol a = Allocated_const (Immutable_float_array [|0.|]) 420 initialize_symbol b = Duparray(Mutable, a) 421 effect b.(0) <- 1. 422 initialize_symbol c = Duparray(Mutable, b) 423 |} 424 425 This will be converted to: 426 {| 427 let_symbol a = Allocated_const (Immutable_float_array [|0.|]) 428 let_symbol b = Allocated_const (Float_array [|0.|]) 429 effect b.(0) <- 1. 430 let_symbol c = Allocated_const (Float_array [|0.|]) 431 |} 432 433 We can't encounter that currently, but it's scary. 434 *) 435 Misc.fatal_error "Pduparray is not allowed on mutable arrays" 436 | Allocated_const (Normal (Immutable_float_array floats)) -> 437 let const : Allocated_const.t = 438 match mutability with 439 | Immutable -> Immutable_float_array floats 440 | Mutable -> Float_array floats 441 in 442 Some (Flambda.Allocated_const const) 443 | Allocated_const (Array (Pfloatarray, _, vars)) -> 444 (* Important: [mutability] is from the [Duplicate_array] 445 construction above. *) 446 resolve_float_array_involving_variables ~mutability ~vars 447 | const -> 448 Misc.fatal_errorf 449 "Lift_constants.translate_definition_and_resolve_alias: \ 450 Duplicate Pfloatarray %a with wrong argument: %a" 451 Variable.print var 452 Alias_analysis.print_constant_defining_value const 453 end 454 | Allocated_const (Duplicate_array (_, _, _)) -> 455 Misc.fatal_errorf "Lift_constants.translate_definition_and_resolve_alias: \ 456 Duplicate_array with non-Pfloatarray kind: %a" 457 Alias_analysis.print_constant_defining_value definition 458 | Allocated_const (Array (Pfloatarray, mutability, vars)) -> 459 resolve_float_array_involving_variables ~mutability ~vars 460 | Allocated_const (Array (_, _, _)) -> 461 Misc.fatal_errorf "Lift_constants.translate_definition_and_resolve_alias: \ 462 Array with non-Pfloatarray kind: %a" 463 Alias_analysis.print_constant_defining_value definition 464 | Project_closure { set_of_closures; closure_id } -> 465 begin match Variable.Map.find set_of_closures aliases with 466 | Symbol s -> 467 Some (Flambda.Project_closure (s, closure_id)) 468 (* If a closure projection is a constant, the set of closures must 469 be assigned to a symbol. *) 470 | exception Not_found -> 471 assert false 472 | Variable v -> 473 match Variable.Tbl.find var_to_symbol_tbl v with 474 | s -> 475 Some (Flambda.Project_closure (s, closure_id)) 476 | exception Not_found -> 477 Format.eprintf "var: %a@." Variable.print v; 478 assert false 479 end 480 | Move_within_set_of_closures { closure; move_to } -> 481 let set_of_closure_symbol = 482 find_original_set_of_closure 483 aliases 484 var_to_symbol_tbl 485 var_to_definition_tbl 486 project_closure_map 487 closure 488 in 489 Some (Flambda.Project_closure (set_of_closure_symbol, move_to)) 490 | Set_of_closures set_of_closures -> 491 let set_of_closures = 492 translate_set_of_closures 493 inconstants 494 aliases 495 var_to_symbol_tbl 496 var_to_definition_tbl 497 set_of_closures 498 in 499 Some (Flambda.Set_of_closures set_of_closures) 500 | Project_var _ -> None 501 | Field (_,_) | Symbol_field _ -> None 502 | Const _ -> None 503 | Symbol _ -> None 504 | Variable _ -> None 505 506let translate_definitions_and_resolve_alias 507 inconstants 508 (aliases : Alias_analysis.allocation_point Variable.Map.t) 509 (var_to_symbol_tbl : Symbol.t Variable.Tbl.t) 510 (var_to_definition_tbl: 511 Alias_analysis.constant_defining_value Variable.Tbl.t) 512 symbol_definition_map 513 project_closure_map 514 ~backend = 515 Variable.Tbl.fold (fun var def map -> 516 match 517 translate_definition_and_resolve_alias inconstants aliases ~backend 518 var_to_symbol_tbl var_to_definition_tbl symbol_definition_map 519 project_closure_map def 520 with 521 | None -> map 522 | Some def -> 523 let symbol = Variable.Tbl.find var_to_symbol_tbl var in 524 Symbol.Map.add symbol def map) 525 var_to_definition_tbl Symbol.Map.empty 526 527(* Resorting of graph including Initialize_symbol *) 528let constant_dependencies ~backend:_ 529 (const : Flambda.constant_defining_value) = 530 match const with 531 | Allocated_const _ -> Symbol.Set.empty 532 | Block (_, fields) -> 533 let symbol_fields = 534 Misc.Stdlib.List.filter_map 535 (function 536 | (Symbol s : Flambda.constant_defining_value_block_field) -> Some s 537 | Flambda.Const _ -> None) 538 fields 539 in 540 Symbol.Set.of_list symbol_fields 541 | Set_of_closures set_of_closures -> 542 Flambda.free_symbols_named (Set_of_closures set_of_closures) 543 | Project_closure (s, _) -> 544 Symbol.Set.singleton s 545 546let program_graph ~backend imported_symbols symbol_to_constant 547 (initialize_symbol_tbl : 548 (Tag.t * Flambda.t list * Symbol.t option) Symbol.Tbl.t) 549 (effect_tbl : (Flambda.t * Symbol.t option) Symbol.Tbl.t) = 550 let expression_symbol_dependencies expr = Flambda.free_symbols expr in 551 let graph_with_only_constant_parts = 552 Symbol.Map.map (fun const -> 553 Symbol.Set.diff (constant_dependencies ~backend const) 554 imported_symbols) 555 symbol_to_constant 556 in 557 let graph_with_initialisation = 558 Symbol.Tbl.fold (fun sym (_tag, fields, previous) -> 559 let order_dep = 560 match previous with 561 | None -> Symbol.Set.empty 562 | Some previous -> Symbol.Set.singleton previous 563 in 564 let deps = List.fold_left (fun set field -> 565 Symbol.Set.union (expression_symbol_dependencies field) set) 566 order_dep fields 567 in 568 let deps = Symbol.Set.diff deps imported_symbols in 569 Symbol.Map.add sym deps) 570 initialize_symbol_tbl graph_with_only_constant_parts 571 in 572 let graph = 573 Symbol.Tbl.fold (fun sym (expr, previous) -> 574 let order_dep = 575 match previous with 576 | None -> Symbol.Set.empty 577 | Some previous -> Symbol.Set.singleton previous 578 in 579 let deps = 580 Symbol.Set.union (expression_symbol_dependencies expr) order_dep 581 in 582 let deps = Symbol.Set.diff deps imported_symbols in 583 Symbol.Map.add sym deps 584 ) 585 effect_tbl graph_with_initialisation 586 in 587 let module Symbol_SCC = Strongly_connected_components.Make (Symbol) in 588 let components = 589 Symbol_SCC.connected_components_sorted_from_roots_to_leaf 590 graph 591 in 592 components 593 594(* rebuilding the program *) 595let add_definition_of_symbol constant_definitions 596 (initialize_symbol_tbl : 597 (Tag.t * Flambda.t list * Symbol.t option) Symbol.Tbl.t) 598 (effect_tbl : (Flambda.t * Symbol.t option) Symbol.Tbl.t) 599 (program : Flambda.program_body) component : Flambda.program_body = 600 let symbol_declaration sym = 601 (* A symbol declared through an Initialize_symbol construct 602 cannot be recursive, this is not allowed in the construction. 603 This also couldn't have been introduced by this pass, so we can 604 safely assert that this is not possible here *) 605 assert(not (Symbol.Tbl.mem initialize_symbol_tbl sym)); 606 (sym, Symbol.Map.find sym constant_definitions) 607 in 608 let module Symbol_SCC = Strongly_connected_components.Make (Symbol) in 609 match component with 610 | Symbol_SCC.Has_loop l -> 611 let l = List.map symbol_declaration l in 612 Let_rec_symbol (l, program) 613 | Symbol_SCC.No_loop sym -> 614 match Symbol.Tbl.find initialize_symbol_tbl sym with 615 | (tag, fields, _previous) -> 616 Initialize_symbol (sym, tag, fields, program) 617 | exception Not_found -> 618 match Symbol.Tbl.find effect_tbl sym with 619 | (expr, _previous) -> 620 Effect (expr, program) 621 | exception Not_found -> 622 let decl = Symbol.Map.find sym constant_definitions in 623 Let_symbol (sym, decl, program) 624 625let add_definitions_of_symbols constant_definitions initialize_symbol_tbl 626 effect_tbl program components = 627 Array.fold_left 628 (add_definition_of_symbol constant_definitions initialize_symbol_tbl 629 effect_tbl) 630 program components 631 632let introduce_free_variables_in_set_of_closures 633 (var_to_block_field_tbl : 634 Flambda.constant_defining_value_block_field Variable.Tbl.t) 635 ({ Flambda.function_decls; free_vars; specialised_args; 636 direct_call_surrogates; } 637 as set_of_closures) = 638 let add_definition_and_make_substitution var (expr, subst) = 639 let searched_var = 640 match Variable.Map.find var specialised_args with 641 | exception Not_found -> var 642 | external_var -> 643 (* specialised arguments bound to constant can be rewritten *) 644 external_var.var 645 in 646 match Variable.Tbl.find var_to_block_field_tbl searched_var with 647 | def -> 648 let fresh = Variable.rename var in 649 let named : Flambda.named = match def with 650 | Symbol sym -> Symbol sym 651 | Const c -> Const c 652 in 653 (Flambda.create_let fresh named expr), Variable.Map.add var fresh subst 654 | exception Not_found -> 655 (* The variable is bound by the closure or the arguments or not 656 constant. In either case it does not need to be bound *) 657 expr, subst 658 in 659 let done_something = ref false in 660 let function_decls : Flambda.function_declarations = 661 Flambda.update_function_declarations function_decls 662 ~funs:(Variable.Map.map 663 (fun (func_decl : Flambda.function_declaration) -> 664 let variables_to_bind = 665 (* Closures from the same set must not be bound. *) 666 Variable.Set.diff func_decl.free_variables 667 (Variable.Map.keys function_decls.funs) 668 in 669 let body, subst = 670 Variable.Set.fold add_definition_and_make_substitution 671 variables_to_bind 672 (func_decl.body, Variable.Map.empty) 673 in 674 if Variable.Map.is_empty subst then begin 675 func_decl 676 end else begin 677 done_something := true; 678 let body = Flambda_utils.toplevel_substitution subst body in 679 Flambda.create_function_declaration 680 ~params:func_decl.params 681 ~body 682 ~stub:func_decl.stub 683 ~dbg:func_decl.dbg 684 ~inline:func_decl.inline 685 ~specialise:func_decl.specialise 686 ~is_a_functor:func_decl.is_a_functor 687 end) 688 function_decls.funs) 689 in 690 let free_vars = 691 (* Keep only those that are not rewritten to constants. *) 692 Variable.Map.filter (fun v _ -> 693 let keep = not (Variable.Tbl.mem var_to_block_field_tbl v) in 694 if not keep then done_something := true; 695 keep) 696 free_vars 697 in 698 let free_vars = 699 Flambda_utils.clean_projections ~which_variables:free_vars 700 in 701 let specialised_args = 702 (* Keep only those that are not rewritten to constants. *) 703 Variable.Map.filter (fun _ (spec_to : Flambda.specialised_to) -> 704 let keep = 705 not (Variable.Tbl.mem var_to_block_field_tbl spec_to.var) 706 in 707 if not keep then begin 708 done_something := true 709 end; 710 keep) 711 specialised_args 712 in 713 let specialised_args = 714 Flambda_utils.clean_projections ~which_variables:specialised_args 715 in 716 if not !done_something then 717 set_of_closures 718 else 719 Flambda.create_set_of_closures ~function_decls ~free_vars 720 ~specialised_args ~direct_call_surrogates 721 722let rewrite_project_var 723 (var_to_block_field_tbl 724 : Flambda.constant_defining_value_block_field Variable.Tbl.t) 725 (project_var : Flambda.project_var) ~original : Flambda.named = 726 let var = Var_within_closure.unwrap project_var.var in 727 match Variable.Tbl.find var_to_block_field_tbl var with 728 | exception Not_found -> original 729 | Symbol sym -> Symbol sym 730 | Const const -> Const const 731 732let introduce_free_variables_in_sets_of_closures 733 (var_to_block_field_tbl: 734 Flambda.constant_defining_value_block_field Variable.Tbl.t) 735 (translate_definition : Flambda.constant_defining_value Symbol.Map.t) = 736 Symbol.Map.map (fun (def : Flambda.constant_defining_value) -> 737 match def with 738 | Allocated_const _ 739 | Block _ 740 | Project_closure _ -> def 741 | Set_of_closures set_of_closures -> 742 Flambda.Set_of_closures 743 (introduce_free_variables_in_set_of_closures 744 var_to_block_field_tbl 745 set_of_closures)) 746 translate_definition 747 748let var_to_block_field 749 (aliases : Alias_analysis.allocation_point Variable.Map.t) 750 (var_to_symbol_tbl : Symbol.t Variable.Tbl.t) 751 (var_to_definition_tbl : 752 Alias_analysis.constant_defining_value Variable.Tbl.t) = 753 let var_to_block_field_tbl = Variable.Tbl.create 42 in 754 Variable.Tbl.iter (fun var _ -> 755 let def = 756 resolve_variable aliases var_to_symbol_tbl var_to_definition_tbl var 757 in 758 Variable.Tbl.add var_to_block_field_tbl var def) 759 var_to_definition_tbl; 760 var_to_block_field_tbl 761 762let program_symbols ~backend (program : Flambda.program) = 763 let new_fake_symbol = 764 let r = ref 0 in 765 fun () -> 766 incr r; 767 Symbol.create (Compilation_unit.get_current_exn ()) 768 (Linkage_name.create ("fake_effect_symbol_" ^ string_of_int !r)) 769 in 770 let initialize_symbol_tbl = Symbol.Tbl.create 42 in 771 let effect_tbl = Symbol.Tbl.create 42 in 772 let symbol_definition_tbl = Symbol.Tbl.create 42 in 773 let add_project_closure_definitions def_symbol 774 (const : Flambda.constant_defining_value) = 775 match const with 776 | Set_of_closures { function_decls = { funs } } -> 777 Variable.Map.iter (fun fun_var _ -> 778 let closure_id = Closure_id.wrap fun_var in 779 let closure_symbol = closure_symbol ~backend closure_id in 780 let project_closure = 781 Flambda.Project_closure (def_symbol, closure_id) 782 in 783 Symbol.Tbl.add symbol_definition_tbl closure_symbol 784 project_closure) 785 funs 786 | Project_closure _ 787 | Allocated_const _ 788 | Block _ -> () 789 in 790 let rec loop (program : Flambda.program_body) previous_effect = 791 match program with 792 | Flambda.Let_symbol (symbol, def, program) -> 793 add_project_closure_definitions symbol def; 794 Symbol.Tbl.add symbol_definition_tbl symbol def; 795 loop program previous_effect 796 | Flambda.Let_rec_symbol (defs, program) -> 797 List.iter (fun (symbol, def) -> 798 add_project_closure_definitions symbol def; 799 Symbol.Tbl.add symbol_definition_tbl symbol def) 800 defs; 801 loop program previous_effect 802 | Flambda.Initialize_symbol (symbol, tag, fields, program) -> 803 (* previous_effect is used to keep the order of initialize and effect 804 values. Their effects order must be kept ordered. 805 it is used as an extra dependency when sorting the symbols. *) 806 (* CR-someday pchambart: if the fields expressions are pure, we could 807 drop this dependency 808 mshinwell: deferred CR *) 809 Symbol.Tbl.add initialize_symbol_tbl symbol 810 (tag, fields, previous_effect); 811 loop program (Some symbol) 812 | Flambda.Effect (expr, program) -> 813 (* Used to ensure that effects are correctly ordered *) 814 let fake_effect_symbol = new_fake_symbol () in 815 Symbol.Tbl.add effect_tbl fake_effect_symbol (expr, previous_effect); 816 loop program (Some fake_effect_symbol) 817 | Flambda.End _ -> () 818 in 819 loop program.program_body None; 820 initialize_symbol_tbl, symbol_definition_tbl, effect_tbl 821 822let replace_definitions_in_initialize_symbol_and_effects 823 (inconstants : Inconstant_idents.result) 824 (aliases : Alias_analysis.allocation_point Variable.Map.t) 825 (var_to_symbol_tbl : Symbol.t Variable.Tbl.t) 826 (var_to_definition_tbl : 827 Alias_analysis.constant_defining_value Variable.Tbl.t) 828 (initialize_symbol_tbl : 829 (Tag.t * Flambda.t list * Symbol.t option) Symbol.Tbl.t) 830 (effect_tbl : (Flambda.t * Symbol.t option) Symbol.Tbl.t) = 831 let rewrite_expr expr = 832 Flambda_iterators.map_all_immutable_let_and_let_rec_bindings expr 833 ~f:(fun var (named : Flambda.named) : Flambda.named -> 834 if Inconstant_idents.variable var inconstants then 835 named 836 else 837 let resolved = 838 resolve_variable 839 aliases 840 var_to_symbol_tbl 841 var_to_definition_tbl 842 var 843 in 844 match named, resolved with 845 | Symbol s1, Symbol s2 -> 846 assert (s1 == s2); (* physical equality for speed *) 847 named; 848 | Const c1, Const c2 -> 849 assert (c1 == c2); 850 named 851 | _, Symbol s -> Symbol s 852 | _, Const c -> Const c) 853 in 854 (* This is safe because we only [replace] the current key during 855 iteration (cf. https://github.com/ocaml/ocaml/pull/337) *) 856 Symbol.Tbl.iter 857 (fun symbol (tag, fields, previous) -> 858 let fields = List.map rewrite_expr fields in 859 Symbol.Tbl.replace initialize_symbol_tbl symbol (tag, fields, previous)) 860 initialize_symbol_tbl; 861 Symbol.Tbl.iter 862 (fun symbol (expr, previous) -> 863 Symbol.Tbl.replace effect_tbl symbol (rewrite_expr expr, previous)) 864 effect_tbl 865 866(* CR-soon mshinwell: Update the name of [project_closure_map]. *) 867let project_closure_map symbol_definition_map = 868 Symbol.Map.fold (fun sym (const : Flambda.constant_defining_value) acc -> 869 match const with 870 | Project_closure (set_of_closures, _) -> 871 Symbol.Map.add sym set_of_closures acc 872 | Set_of_closures _ -> 873 Symbol.Map.add sym sym acc 874 | Allocated_const _ 875 | Block _ -> acc) 876 symbol_definition_map 877 Symbol.Map.empty 878 879let the_dead_constant_index = ref 0 880 881let lift_constants (program : Flambda.program) ~backend = 882 let the_dead_constant = 883 let index = !the_dead_constant_index in 884 incr the_dead_constant_index; 885 let name = Printf.sprintf "the_dead_constant_%d" index in 886 Symbol.create (Compilation_unit.get_current_exn ()) 887 (Linkage_name.create name) 888 in 889 let program_body : Flambda.program_body = 890 Let_symbol (the_dead_constant, Allocated_const (Nativeint 0n), 891 program.program_body) 892 in 893 let program : Flambda.program = 894 { program with program_body; } 895 in 896 let inconstants = 897 Inconstant_idents.inconstants_on_program program ~backend 898 ~compilation_unit:(Compilation_unit.get_current_exn ()) 899 in 900 let initialize_symbol_tbl, symbol_definition_tbl, effect_tbl = 901 program_symbols ~backend program 902 in 903 let var_to_symbol_tbl, var_to_definition_tbl, let_symbol_to_definition_tbl, 904 initialize_symbol_to_definition_tbl = 905 assign_symbols_and_collect_constant_definitions ~backend ~program 906 ~inconstants 907 in 908 let aliases = 909 Alias_analysis.run var_to_definition_tbl 910 initialize_symbol_to_definition_tbl 911 let_symbol_to_definition_tbl 912 ~the_dead_constant 913 in 914 replace_definitions_in_initialize_symbol_and_effects 915 (inconstants : Inconstant_idents.result) 916 (aliases : Alias_analysis.allocation_point Variable.Map.t) 917 (var_to_symbol_tbl : Symbol.t Variable.Tbl.t) 918 (var_to_definition_tbl 919 : Alias_analysis.constant_defining_value Variable.Tbl.t) 920 initialize_symbol_tbl 921 effect_tbl; 922 let symbol_definition_map = 923 translate_constant_set_of_closures 924 (inconstants : Inconstant_idents.result) 925 (aliases : Alias_analysis.allocation_point Variable.Map.t) 926 (var_to_symbol_tbl : Symbol.t Variable.Tbl.t) 927 (var_to_definition_tbl 928 : Alias_analysis.constant_defining_value Variable.Tbl.t) 929 (Symbol.Tbl.to_map symbol_definition_tbl) 930 in 931 let project_closure_map = project_closure_map symbol_definition_map in 932 let translated_definitions = 933 translate_definitions_and_resolve_alias 934 inconstants 935 (aliases : Alias_analysis.allocation_point Variable.Map.t) 936 (var_to_symbol_tbl : Symbol.t Variable.Tbl.t) 937 (var_to_definition_tbl 938 : Alias_analysis.constant_defining_value Variable.Tbl.t) 939 symbol_definition_map 940 project_closure_map 941 ~backend 942 in 943 let var_to_block_field_tbl = 944 var_to_block_field 945 (aliases : Alias_analysis.allocation_point Variable.Map.t) 946 (var_to_symbol_tbl : Symbol.t Variable.Tbl.t) 947 (var_to_definition_tbl 948 : Alias_analysis.constant_defining_value Variable.Tbl.t) 949 in 950 let translated_definitions = 951 introduce_free_variables_in_sets_of_closures var_to_block_field_tbl 952 translated_definitions 953 in 954 let constant_definitions = 955 (* Add previous Let_symbol to the newly discovered ones *) 956 Symbol.Map.union 957 (fun _sym 958 (c1:Flambda.constant_defining_value) 959 (c2:Flambda.constant_defining_value) -> 960 match c1, c2 with 961 | Project_closure (s1, closure_id1), 962 Project_closure (s2, closure_id2) when 963 Symbol.equal s1 s2 && 964 Closure_id.equal closure_id1 closure_id2 -> 965 Some c1 966 | Project_closure (s1, closure_id1), 967 Project_closure (s2, closure_id2) -> 968 Format.eprintf "not equal project closure@. s %a %a@. cid %a %a@." 969 Symbol.print s1 Symbol.print s2 970 Closure_id.print closure_id1 Closure_id.print closure_id2; 971 assert false 972 | _ -> 973 assert false 974 ) 975 symbol_definition_map 976 translated_definitions 977 in 978 (* Upon the [Initialize_symbol]s, the [Effect]s and the constant definitions, 979 do the following: 980 1. Introduce [Let]s to bind variables that are going to be replaced 981 by constants. 982 2. If a variable bound by a closure gets replaced by a symbol and 983 thus eliminated from the [free_vars] set of the closure, we need to 984 rewrite any subsequent [Project_var] expressions that project that 985 variable. *) 986 let rewrite_expr expr = 987 Flambda_iterators.map_named (function 988 | (Set_of_closures set_of_closures) as named -> 989 let new_set_of_closures = 990 introduce_free_variables_in_set_of_closures 991 var_to_block_field_tbl set_of_closures 992 in 993 if new_set_of_closures == set_of_closures then 994 named 995 else 996 Set_of_closures new_set_of_closures 997 | (Project_var project_var) as original -> 998 rewrite_project_var var_to_block_field_tbl project_var ~original 999 | (Symbol _ | Const _ | Allocated_const _ | Project_closure _ 1000 | Move_within_set_of_closures _ | Prim _ | Expr _ 1001 | Read_mutable _ | Read_symbol_field _) as named -> named) 1002 expr 1003 in 1004 let constant_definitions = 1005 Symbol.Map.map (fun (const : Flambda.constant_defining_value) -> 1006 match const with 1007 | Allocated_const _ | Block _ | Project_closure _ -> const 1008 | Set_of_closures set_of_closures -> 1009 let set_of_closures = 1010 Flambda_iterators.map_function_bodies set_of_closures 1011 ~f:rewrite_expr 1012 in 1013 Flambda.Set_of_closures 1014 (introduce_free_variables_in_set_of_closures 1015 var_to_block_field_tbl set_of_closures)) 1016 constant_definitions 1017 in 1018 let effect_tbl = 1019 Symbol.Tbl.map effect_tbl (fun (effect, dep) -> rewrite_expr effect, dep) 1020 in 1021 let initialize_symbol_tbl = 1022 Symbol.Tbl.map initialize_symbol_tbl (fun (tag, fields, dep) -> 1023 let fields = List.map rewrite_expr fields in 1024 tag, fields, dep) 1025 in 1026 let imported_symbols = Flambda_utils.imported_symbols program in 1027 let components = 1028 program_graph ~backend imported_symbols constant_definitions 1029 initialize_symbol_tbl effect_tbl 1030 in 1031 let program_body = 1032 add_definitions_of_symbols constant_definitions 1033 initialize_symbol_tbl 1034 effect_tbl 1035 (End (Flambda_utils.root_symbol program)) 1036 components 1037 in 1038 Flambda_utils.introduce_needed_import_symbols { program with program_body; } 1039