Skip to main content

cranelift_codegen/
inline.rs

1//! Function inlining infrastructure.
2//!
3//! This module provides "inlining as a library" to Cranelift users; it does
4//! _not_ provide a complete, off-the-shelf inlining solution. Cranelift's
5//! compilation context is per-function and does not encompass the full call
6//! graph. It does not know which functions are hot and which are cold, which
7//! have been marked the equivalent of `#[inline(never)]`, etc... Only the
8//! Cranelift user can understand these aspects of the full compilation
9//! pipeline, and these things can be very different between (say) Wasmtime and
10//! `cg_clif`. Therefore, this module does not attempt to define heuristics for
11//! when inlining a particular call is likely beneficial. This module only
12//! provides hooks for the Cranelift user to define whether a given call should
13//! be inlined or not, and the mechanics to inline a callee into a particular
14//! call site when directed to do so by the Cranelift user.
15//!
16//! The top-level inlining entry point during Cranelift compilation is
17//! [`Context::inline`][crate::Context::inline]. It takes an [`Inline`] trait
18//! implementation, which is authored by the Cranelift user and directs
19//! Cranelift whether to inline a particular call, and, when inlining, gives
20//! Cranelift the body of the callee that is to be inlined.
21
22use crate::cursor::{Cursor as _, FuncCursor};
23use crate::ir::{self, DebugTag, ExceptionTableData, ExceptionTableItem, InstBuilder as _};
24use crate::result::CodegenResult;
25use crate::trace;
26use crate::traversals::Dfs;
27use alloc::borrow::Cow;
28use alloc::vec::Vec;
29use cranelift_entity::{SecondaryMap, packed_option::PackedOption};
30use smallvec::SmallVec;
31
32type SmallValueVec = SmallVec<[ir::Value; 8]>;
33type SmallBlockArgVec = SmallVec<[ir::BlockArg; 8]>;
34type SmallBlockCallVec = SmallVec<[ir::BlockCall; 8]>;
35
36/// A command directing Cranelift whether or not to inline a particular call.
37pub enum InlineCommand<'a> {
38    /// Keep the call as-is, out-of-line, and do not inline the callee.
39    KeepCall,
40
41    /// Inline the call, using this function as the body of the callee.
42    ///
43    /// It is the `Inline` implementor's responsibility to ensure that this
44    /// function is the correct callee. Providing the wrong function may result
45    /// in panics during compilation or incorrect runtime behavior.
46    Inline {
47        /// The callee function's body.
48        callee: Cow<'a, ir::Function>,
49        /// Whether to visit any function calls within the callee body after
50        /// inlining and consider them for further inlining.
51        visit_callee: bool,
52    },
53}
54
55/// A trait for directing Cranelift whether to inline a particular call or not.
56///
57/// Used in combination with the [`Context::inline`][crate::Context::inline]
58/// method.
59pub trait Inline {
60    /// A hook invoked for each direct call instruction in a function, whose
61    /// result determines whether Cranelift should inline a given call.
62    ///
63    /// The Cranelift user is responsible for defining their own heuristics and
64    /// deciding whether inlining the call is beneficial.
65    ///
66    /// When returning a function and directing Cranelift to inline its body
67    /// into the call site, the `Inline` implementer must ensure the following:
68    ///
69    /// * The returned function's signature exactly matches the `callee`
70    ///   `FuncRef`'s signature.
71    ///
72    /// * The returned function must be legalized.
73    ///
74    /// * The returned function must be valid (i.e. it must pass the CLIF
75    ///   verifier).
76    ///
77    /// * The returned function is a correct and valid implementation of the
78    ///   `callee` according to your language's semantics.
79    ///
80    /// Failure to uphold these invariants may result in panics during
81    /// compilation or incorrect runtime behavior in the generated code.
82    fn inline(
83        &mut self,
84        caller: &ir::Function,
85        call_inst: ir::Inst,
86        call_opcode: ir::Opcode,
87        callee: ir::FuncRef,
88        call_args: &[ir::Value],
89    ) -> InlineCommand<'_>;
90}
91
92impl<'a, T> Inline for &'a mut T
93where
94    T: Inline,
95{
96    fn inline(
97        &mut self,
98        caller: &ir::Function,
99        inst: ir::Inst,
100        opcode: ir::Opcode,
101        callee: ir::FuncRef,
102        args: &[ir::Value],
103    ) -> InlineCommand<'_> {
104        (*self).inline(caller, inst, opcode, callee, args)
105    }
106}
107
108/// Walk the given function, invoke the `Inline` implementation for each call
109/// instruction, and inline the callee when directed to do so.
110///
111/// Returns whether any call was inlined.
112pub(crate) fn do_inlining(
113    func: &mut ir::Function,
114    mut inliner: impl Inline,
115) -> CodegenResult<bool> {
116    trace!("function {} before inlining: {}", func.name, func);
117
118    let mut inlined_any = false;
119    let mut allocs = InliningAllocs::default();
120
121    let mut cursor = FuncCursor::new(func);
122    'block_loop: while let Some(block) = cursor.next_block() {
123        // Always keep track of our previous cursor position. Assuming that the
124        // current position is a function call that we will inline, then the
125        // previous position is just before the inlined callee function. After
126        // inlining a call, the Cranelift user can decide whether to consider
127        // any function calls in the inlined callee for further inlining or
128        // not. When they do, then we back up to this previous cursor position
129        // so that our traversal will then continue over the inlined body.
130        let mut prev_pos;
131
132        while let Some(inst) = {
133            prev_pos = cursor.position();
134            cursor.next_inst()
135        } {
136            // Make sure that `block` is always `inst`'s block, even with all of
137            // our cursor-position-updating and block-splitting-during-inlining
138            // shenanigans below.
139            debug_assert_eq!(Some(block), cursor.func.layout.inst_block(inst));
140
141            match cursor.func.dfg.insts[inst] {
142                ir::InstructionData::Call { func_ref, .. }
143                    if cursor.func.dfg.ext_funcs[func_ref].patchable =>
144                {
145                    // Can't inline patchable calls; they need to
146                    // remain patchable and inlining the whole body is
147                    // decidedly *not* patchable!
148                }
149
150                ir::InstructionData::Call {
151                    opcode: opcode @ ir::Opcode::Call | opcode @ ir::Opcode::ReturnCall,
152                    args: _,
153                    func_ref,
154                } => {
155                    trace!(
156                        "considering call site for inlining: {inst}: {}",
157                        cursor.func.dfg.display_inst(inst),
158                    );
159                    let args = cursor.func.dfg.inst_args(inst);
160                    match inliner.inline(&cursor.func, inst, opcode, func_ref, args) {
161                        InlineCommand::KeepCall => {
162                            trace!("  --> keeping call");
163                        }
164                        InlineCommand::Inline {
165                            callee,
166                            visit_callee,
167                        } => {
168                            let last_inlined_block = inline_one(
169                                &mut allocs,
170                                cursor.func,
171                                func_ref,
172                                block,
173                                inst,
174                                opcode,
175                                &callee,
176                                None,
177                            )?;
178                            inlined_any = true;
179                            if visit_callee {
180                                cursor.set_position(prev_pos);
181                            } else {
182                                // Arrange it so that the `next_block()` loop
183                                // will continue to the next block that is not
184                                // associated with the just-inlined callee.
185                                cursor.goto_bottom(last_inlined_block);
186                                continue 'block_loop;
187                            }
188                        }
189                    }
190                }
191                ir::InstructionData::TryCall {
192                    opcode: opcode @ ir::Opcode::TryCall,
193                    args: _,
194                    func_ref,
195                    exception,
196                } => {
197                    trace!(
198                        "considering call site for inlining: {inst}: {}",
199                        cursor.func.dfg.display_inst(inst),
200                    );
201                    let args = cursor.func.dfg.inst_args(inst);
202                    match inliner.inline(&cursor.func, inst, opcode, func_ref, args) {
203                        InlineCommand::KeepCall => {
204                            trace!("  --> keeping call");
205                        }
206                        InlineCommand::Inline {
207                            callee,
208                            visit_callee,
209                        } => {
210                            let last_inlined_block = inline_one(
211                                &mut allocs,
212                                cursor.func,
213                                func_ref,
214                                block,
215                                inst,
216                                opcode,
217                                &callee,
218                                Some(exception),
219                            )?;
220                            inlined_any = true;
221                            if visit_callee {
222                                cursor.set_position(prev_pos);
223                            } else {
224                                // Arrange it so that the `next_block()` loop
225                                // will continue to the next block that is not
226                                // associated with the just-inlined callee.
227                                cursor.goto_bottom(last_inlined_block);
228                                continue 'block_loop;
229                            }
230                        }
231                    }
232                }
233                ir::InstructionData::CallIndirect { .. }
234                | ir::InstructionData::TryCallIndirect { .. } => {
235                    // Can't inline indirect calls; need to have some earlier
236                    // pass rewrite them into direct calls first, when possible.
237                }
238                _ => {
239                    debug_assert!(
240                        !cursor.func.dfg.insts[inst].opcode().is_call(),
241                        "should have matched all call instructions, but found: {inst}: {}",
242                        cursor.func.dfg.display_inst(inst),
243                    );
244                }
245            }
246        }
247    }
248
249    if inlined_any {
250        trace!("function {} after inlining: {}", func.name, func);
251    } else {
252        trace!("function {} did not have any callees inlined", func.name);
253    }
254
255    Ok(inlined_any)
256}
257
258#[derive(Default)]
259struct InliningAllocs {
260    /// Map from callee value to inlined caller value.
261    values: SecondaryMap<ir::Value, PackedOption<ir::Value>>,
262
263    /// Map from callee constant to inlined caller constant.
264    ///
265    /// Not in `EntityMap` because these are hash-consed inside the
266    /// `ir::Function`.
267    constants: SecondaryMap<ir::Constant, PackedOption<ir::Constant>>,
268
269    /// Map from callee to inlined caller external name refs.
270    ///
271    /// Not in `EntityMap` because these are hash-consed inside the
272    /// `ir::Function`.
273    user_external_name_refs:
274        SecondaryMap<ir::UserExternalNameRef, PackedOption<ir::UserExternalNameRef>>,
275
276    /// The set of _caller_ inlined call instructions that need exception table
277    /// fixups at the end of inlining.
278    ///
279    /// This includes all kinds of non-returning calls, not just the literal
280    /// `call` instruction: `call_indirect`, `try_call`, `try_call_indirect`,
281    /// etc... However, it does not include `return_call` and
282    /// `return_call_indirect` instructions because the caller cannot catch
283    /// exceptions that those calls throw because the caller is no longer on the
284    /// stack as soon as they are executed.
285    ///
286    /// Note: this is a simple `Vec`, and not an `EntitySet`, because it is very
287    /// sparse: most of the caller's instructions are not inlined call
288    /// instructions. Additionally, we require deterministic iteration order and
289    /// do not require set-membership testing, so a hash set is not a good
290    /// choice either.
291    calls_needing_exception_table_fixup: Vec<ir::Inst>,
292}
293
294impl InliningAllocs {
295    fn reset(&mut self, callee: &ir::Function) {
296        let InliningAllocs {
297            values,
298            constants,
299            user_external_name_refs,
300            calls_needing_exception_table_fixup,
301        } = self;
302
303        values.clear();
304        values.resize(callee.dfg.len_values());
305
306        constants.clear();
307        constants.resize(callee.dfg.constants.len());
308
309        user_external_name_refs.clear();
310        user_external_name_refs.resize(callee.params.user_named_funcs().len());
311
312        // Note: We do not reserve capacity for
313        // `calls_needing_exception_table_fixup` because it is a sparse set and
314        // we don't know how large it needs to be ahead of time.
315        calls_needing_exception_table_fixup.clear();
316    }
317
318    fn set_inlined_value(
319        &mut self,
320        callee: &ir::Function,
321        callee_val: ir::Value,
322        inlined_val: ir::Value,
323    ) {
324        trace!("  --> callee {callee_val:?} = inlined {inlined_val:?}");
325        debug_assert!(self.values[callee_val].is_none());
326        let resolved_callee_val = callee.dfg.resolve_aliases(callee_val);
327        debug_assert!(self.values[resolved_callee_val].is_none());
328        self.values[resolved_callee_val] = Some(inlined_val).into();
329    }
330
331    fn get_inlined_value(&self, callee: &ir::Function, callee_val: ir::Value) -> Option<ir::Value> {
332        let resolved_callee_val = callee.dfg.resolve_aliases(callee_val);
333        self.values[resolved_callee_val].expand()
334    }
335}
336
337/// Inline one particular function call.
338///
339/// Returns the last inlined block in the layout.
340fn inline_one(
341    allocs: &mut InliningAllocs,
342    func: &mut ir::Function,
343    callee_func_ref: ir::FuncRef,
344    call_block: ir::Block,
345    call_inst: ir::Inst,
346    call_opcode: ir::Opcode,
347    callee: &ir::Function,
348    call_exception_table: Option<ir::ExceptionTable>,
349) -> CodegenResult<ir::Block> {
350    trace!(
351        "Inlining call {call_inst:?}: {}\n\
352         with callee = {callee:?}",
353        func.dfg.display_inst(call_inst)
354    );
355
356    // Type check callee signature.
357    let expected_callee_sig = func.dfg.ext_funcs[callee_func_ref].signature;
358    let expected_callee_sig = &func.dfg.signatures[expected_callee_sig];
359    assert_eq!(expected_callee_sig, &callee.signature);
360
361    allocs.reset(callee);
362
363    // First, append various callee entity arenas to the end of the caller's
364    // entity arenas.
365    let entity_map = create_entities(allocs, func, callee)?;
366
367    // Inlined prologue: split the call instruction's block at the point of the
368    // call and replace the call with a jump.
369    let return_block = split_off_return_block(func, call_inst, call_opcode, callee);
370    let call_stack_map = replace_call_with_jump(allocs, func, call_inst, callee, &entity_map);
371
372    // Prepare for translating the actual instructions by inserting the inlined
373    // blocks into the caller's layout in the same order that they appear in the
374    // callee.
375    let mut last_inlined_block = inline_block_layout(func, call_block, callee, &entity_map);
376
377    // Get a copy of debug tags on the call instruction; these are
378    // prepended to debug tags on inlined instructions. Remove them
379    // from the call itself as it will be rewritten to a jump (which
380    // cannot have tags).
381    let call_debug_tags = func.debug_tags.get(call_inst).to_vec();
382    func.debug_tags.set(call_inst, []);
383
384    // Translate each instruction from the callee into the caller,
385    // appending them to their associated block in the caller.
386    //
387    // Note that we iterate over the callee with a pre-order traversal so that
388    // we see value defs before uses.
389    for callee_block in Dfs::new().pre_order_iter(callee) {
390        let inlined_block = entity_map.inlined_block(callee_block);
391        trace!(
392            "Processing instructions in callee block {callee_block:?} (inlined block {inlined_block:?}"
393        );
394
395        let mut next_callee_inst = callee.layout.first_inst(callee_block);
396        while let Some(callee_inst) = next_callee_inst {
397            trace!(
398                "Processing callee instruction {callee_inst:?}: {}",
399                callee.dfg.display_inst(callee_inst)
400            );
401
402            // Remap the callee instruction's entities and insert it into the
403            // caller's DFG.
404            let mut inst_remapper = InliningInstRemapper {
405                allocs: &allocs,
406                func,
407                callee,
408                entity_map: &entity_map,
409                error: None,
410            };
411            let inlined_inst_data = callee.dfg.insts[callee_inst].map(&mut inst_remapper);
412            if let Some(err) = inst_remapper.error.take() {
413                return Err(err);
414            }
415            let inlined_inst = func.dfg.make_inst(inlined_inst_data);
416            func.layout.append_inst(inlined_inst, inlined_block);
417
418            // Copy over debug tags, translating referenced entities
419            // as appropriate.
420            let debug_tags = callee.debug_tags.get(callee_inst);
421            // If there are tags on the inlined instruction, we always
422            // add tags, and we prepend any tags from the call
423            // instruction; but we don't add tags if only the callsite
424            // had them (this would otherwise mean that every single
425            // instruction in an inlined function body would get
426            // tags).
427            if !debug_tags.is_empty() {
428                let tags = call_debug_tags
429                    .iter()
430                    .cloned()
431                    .chain(debug_tags.iter().map(|tag| match *tag {
432                        DebugTag::User(value) => DebugTag::User(value),
433                        DebugTag::StackSlot(slot) => {
434                            DebugTag::StackSlot(entity_map.inlined_stack_slot(slot))
435                        }
436                    }))
437                    .collect::<SmallVec<[_; 4]>>();
438                func.debug_tags.set(inlined_inst, tags);
439            }
440
441            let opcode = callee.dfg.insts[callee_inst].opcode();
442            if opcode.is_return() {
443                // Instructions that return do not define any values, so we
444                // don't need to worry about that, but we do need to fix them up
445                // so that they return by jumping to our control-flow join
446                // block, rather than returning from the caller.
447                if let Some(return_block) = return_block {
448                    fixup_inst_that_returns(
449                        allocs,
450                        func,
451                        callee,
452                        &entity_map,
453                        call_opcode,
454                        inlined_inst,
455                        callee_inst,
456                        return_block,
457                        call_stack_map.as_ref().map(|es| &**es),
458                    );
459                } else {
460                    // If we are inlining a callee that was invoked via
461                    // `return_call`, we leave inlined return instructions
462                    // as-is: there is no logical caller frame on the stack to
463                    // continue to.
464                    debug_assert_eq!(call_opcode, ir::Opcode::ReturnCall);
465                }
466            } else {
467                // Make the instruction's result values.
468                let ctrl_typevar = callee.dfg.ctrl_typevar(callee_inst);
469                func.dfg.make_inst_results(inlined_inst, ctrl_typevar);
470
471                // Update the value map for this instruction's defs.
472                let callee_results = callee.dfg.inst_results(callee_inst);
473                let inlined_results = func.dfg.inst_results(inlined_inst);
474                debug_assert_eq!(callee_results.len(), inlined_results.len());
475                for (callee_val, inlined_val) in callee_results.iter().zip(inlined_results) {
476                    allocs.set_inlined_value(callee, *callee_val, *inlined_val);
477                }
478
479                if opcode.is_call() {
480                    append_stack_map_entries(
481                        func,
482                        callee,
483                        &entity_map,
484                        call_stack_map.as_deref(),
485                        inlined_inst,
486                        callee_inst,
487                    );
488
489                    // When we are inlining a `try_call` call site, we need to merge
490                    // the call site's exception table into the inlined calls'
491                    // exception tables. This can involve rewriting regular `call`s
492                    // into `try_call`s, which requires mutating the CFG because
493                    // `try_call` is a block terminator. However, we can't mutate
494                    // the CFG in the middle of this traversal because we rely on
495                    // the existence of a one-to-one mapping between the callee
496                    // layout and the inlined layout. Instead, we record the set of
497                    // inlined call instructions that will need fixing up, and
498                    // perform that possibly-CFG-mutating exception table merging in
499                    // a follow up pass, when we no longer rely on that one-to-one
500                    // layout mapping.
501                    debug_assert_eq!(
502                        call_opcode == ir::Opcode::TryCall,
503                        call_exception_table.is_some()
504                    );
505                    if call_opcode == ir::Opcode::TryCall {
506                        allocs
507                            .calls_needing_exception_table_fixup
508                            .push(inlined_inst);
509                    }
510                }
511            }
512
513            trace!(
514                "  --> inserted inlined instruction {inlined_inst:?}: {}",
515                func.dfg.display_inst(inlined_inst)
516            );
517
518            next_callee_inst = callee.layout.next_inst(callee_inst);
519        }
520    }
521
522    // We copied *all* callee blocks into the caller's layout, but only copied
523    // the callee instructions in *reachable* callee blocks into the caller's
524    // associated blocks. Therefore, any *unreachable* blocks are empty in the
525    // caller, which is invalid CLIF because all blocks must end in a
526    // terminator, so do a quick pass over the inlined blocks and remove any
527    // empty blocks from the caller's layout.
528    for block in entity_map.iter_inlined_blocks(func) {
529        if func.layout.is_block_inserted(block) && func.layout.first_inst(block).is_none() {
530            log::trace!("removing unreachable inlined block from layout: {block}");
531
532            // If the block being removed is our last-inlined block, then back
533            // it up to the previous block in the layout, which will be the new
534            // last-inlined block after this one's removal.
535            if block == last_inlined_block {
536                last_inlined_block = func.layout.prev_block(last_inlined_block).expect(
537                    "there will always at least be the block that contained the call we are \
538                     inlining",
539                );
540            }
541
542            func.layout.remove_block(block);
543        }
544    }
545
546    // Final step: fixup the exception tables of any inlined calls when we are
547    // inlining a `try_call` site.
548    //
549    // Subtly, this requires rewriting non-catching `call[_indirect]`
550    // instructions into `try_call[_indirect]` instructions so that exceptions
551    // that unwound through the original callee frame and were caught by the
552    // caller's `try_call` do not unwind past this inlined frame. And turning a
553    // `call` into a `try_call` mutates the CFG, breaking our one-to-one mapping
554    // between callee blocks and inlined blocks, so we delay these fixups to
555    // this final step, when we no longer rely on that mapping.
556    debug_assert!(
557        allocs.calls_needing_exception_table_fixup.is_empty() || call_exception_table.is_some()
558    );
559    debug_assert_eq!(
560        call_opcode == ir::Opcode::TryCall,
561        call_exception_table.is_some()
562    );
563    if let Some(call_exception_table) = call_exception_table {
564        fixup_inlined_call_exception_tables(allocs, func, call_exception_table);
565    }
566
567    debug_assert!(
568        func.layout.is_block_inserted(last_inlined_block),
569        "last_inlined_block={last_inlined_block} should be inserted in the layout"
570    );
571    Ok(last_inlined_block)
572}
573
574/// Append stack map entries from the caller and callee to the given inlined
575/// instruction.
576fn append_stack_map_entries(
577    func: &mut ir::Function,
578    callee: &ir::Function,
579    entity_map: &EntityMap,
580    call_stack_map: Option<&[ir::UserStackMapEntry]>,
581    inlined_inst: ir::Inst,
582    callee_inst: ir::Inst,
583) {
584    // Add the caller's stack map to this call. These entries
585    // already refer to caller entities and do not need further
586    // translation.
587    func.dfg.append_user_stack_map_entries(
588        inlined_inst,
589        call_stack_map
590            .iter()
591            .flat_map(|entries| entries.iter().cloned()),
592    );
593
594    // Append the callee's stack map to this call. These entries
595    // refer to callee entities and therefore do require
596    // translation into the caller's index space.
597    func.dfg.append_user_stack_map_entries(
598        inlined_inst,
599        callee
600            .dfg
601            .user_stack_map_entries(callee_inst)
602            .iter()
603            .flat_map(|entries| entries.iter())
604            .map(|entry| ir::UserStackMapEntry {
605                ty: entry.ty,
606                slot: entity_map.inlined_stack_slot(entry.slot),
607                offset: entry.offset,
608            }),
609    );
610}
611
612/// Create or update the exception tables for any inlined call instructions:
613/// when inlining at a `try_call` site, we must forward our exceptional edges
614/// into each inlined call instruction.
615fn fixup_inlined_call_exception_tables(
616    allocs: &mut InliningAllocs,
617    func: &mut ir::Function,
618    call_exception_table: ir::ExceptionTable,
619) {
620    // Split a block at a `call[_indirect]` instruction, detach the
621    // instruction's results, and alias them to the new block's parameters.
622    let split_block_for_new_try_call = |func: &mut ir::Function, inst: ir::Inst| -> ir::Block {
623        debug_assert!(func.dfg.insts[inst].opcode().is_call());
624        debug_assert!(!func.dfg.insts[inst].opcode().is_terminator());
625
626        // Split the block.
627        let next_inst = func
628            .layout
629            .next_inst(inst)
630            .expect("inst is not a terminator, should have a successor");
631        let new_block = func.dfg.blocks.add();
632        func.layout.split_block(new_block, next_inst);
633
634        // `try_call[_indirect]` instructions do not define values themselves;
635        // the normal-return block has parameters for the results. So remove
636        // this instruction's results, create an associated block parameter for
637        // each of them, and alias them to the new block parameter.
638        let old_results = SmallValueVec::from_iter(func.dfg.inst_results(inst).iter().copied());
639        func.dfg.detach_inst_results(inst);
640        for old_result in old_results {
641            let ty = func.dfg.value_type(old_result);
642            let new_block_param = func.dfg.append_block_param(new_block, ty);
643            func.dfg.change_to_alias(old_result, new_block_param);
644        }
645
646        new_block
647    };
648
649    // Clone the caller's exception table, updating it for use in the current
650    // `call[_indirect]` instruction as it becomes a `try_call[_indirect]`.
651    let clone_exception_table_for_this_call = |func: &mut ir::Function,
652                                               signature: ir::SigRef,
653                                               new_block: ir::Block|
654     -> ir::ExceptionTable {
655        let mut exception = func.stencil.dfg.exception_tables[call_exception_table]
656            .deep_clone(&mut func.stencil.dfg.value_lists);
657
658        *exception.signature_mut() = signature;
659
660        let returns_len = func.dfg.signatures[signature].returns.len();
661        let returns_len = u32::try_from(returns_len).unwrap();
662
663        *exception.normal_return_mut() = ir::BlockCall::new(
664            new_block,
665            (0..returns_len).map(|i| ir::BlockArg::TryCallRet(i)),
666            &mut func.dfg.value_lists,
667        );
668
669        func.dfg.exception_tables.push(exception)
670    };
671
672    for inst in allocs.calls_needing_exception_table_fixup.drain(..) {
673        debug_assert!(func.dfg.insts[inst].opcode().is_call());
674        debug_assert!(!func.dfg.insts[inst].opcode().is_return());
675        match func.dfg.insts[inst] {
676            //     current_block:
677            //         preds...
678            //         rets... = call f(args...)
679            //         succs...
680            //
681            // becomes
682            //
683            //     current_block:
684            //         preds...
685            //         try_call f(args...), new_block(rets...), [call_exception_table...]
686            //     new_block(rets...):
687            //         succs...
688            ir::InstructionData::Call {
689                opcode: ir::Opcode::Call,
690                args,
691                func_ref,
692            } => {
693                let new_block = split_block_for_new_try_call(func, inst);
694                let signature = func.dfg.ext_funcs[func_ref].signature;
695                let exception = clone_exception_table_for_this_call(func, signature, new_block);
696                func.dfg.insts[inst] = ir::InstructionData::TryCall {
697                    opcode: ir::Opcode::TryCall,
698                    args,
699                    func_ref,
700                    exception,
701                };
702            }
703
704            //     current_block:
705            //         preds...
706            //         rets... = call_indirect sig, val(args...)
707            //         succs...
708            //
709            // becomes
710            //
711            //     current_block:
712            //         preds...
713            //         try_call_indirect sig, val(args...), new_block(rets...), [call_exception_table...]
714            //     new_block(rets...):
715            //         succs...
716            ir::InstructionData::CallIndirect {
717                opcode: ir::Opcode::CallIndirect,
718                args,
719                sig_ref,
720            } => {
721                let new_block = split_block_for_new_try_call(func, inst);
722                let exception = clone_exception_table_for_this_call(func, sig_ref, new_block);
723                func.dfg.insts[inst] = ir::InstructionData::TryCallIndirect {
724                    opcode: ir::Opcode::TryCallIndirect,
725                    args,
726                    exception,
727                };
728            }
729
730            // For `try_call[_indirect]` instructions, we just need to merge the
731            // exception tables.
732            ir::InstructionData::TryCall {
733                opcode: ir::Opcode::TryCall,
734                exception,
735                ..
736            }
737            | ir::InstructionData::TryCallIndirect {
738                opcode: ir::Opcode::TryCallIndirect,
739                exception,
740                ..
741            } => {
742                // Construct a new exception table that consists of
743                // the inlined instruction's exception table match
744                // sequence, with the inlining site's exception table
745                // appended. This will ensure that the first-match
746                // semantics emulates the original behavior of
747                // matching in the inner frame first.
748                let sig = func.dfg.exception_tables[exception].signature();
749                let normal_return = *func.dfg.exception_tables[exception].normal_return();
750                let exception_data = ExceptionTableData::new(
751                    sig,
752                    normal_return,
753                    func.dfg.exception_tables[exception]
754                        .items()
755                        .chain(func.dfg.exception_tables[call_exception_table].items()),
756                )
757                .deep_clone(&mut func.dfg.value_lists);
758
759                func.dfg.exception_tables[exception] = exception_data;
760            }
761
762            otherwise => unreachable!("unknown non-return call instruction: {otherwise:?}"),
763        }
764    }
765}
766
767/// After having created an inlined version of a callee instruction that returns
768/// in the caller, we need to fix it up so that it doesn't actually return
769/// (since we are already in the caller's frame) and instead just jumps to the
770/// control-flow join point.
771fn fixup_inst_that_returns(
772    allocs: &mut InliningAllocs,
773    func: &mut ir::Function,
774    callee: &ir::Function,
775    entity_map: &EntityMap,
776    call_opcode: ir::Opcode,
777    inlined_inst: ir::Inst,
778    callee_inst: ir::Inst,
779    return_block: ir::Block,
780    call_stack_map: Option<&[ir::UserStackMapEntry]>,
781) {
782    debug_assert!(func.dfg.insts[inlined_inst].opcode().is_return());
783    match func.dfg.insts[inlined_inst] {
784        //     return rets...
785        //
786        // becomes
787        //
788        //     jump return_block(rets...)
789        ir::InstructionData::MultiAry {
790            opcode: ir::Opcode::Return,
791            args,
792        } => {
793            let rets = SmallBlockArgVec::from_iter(
794                args.as_slice(&func.dfg.value_lists)
795                    .iter()
796                    .copied()
797                    .map(|v| v.into()),
798            );
799            func.replace(inlined_inst).jump(return_block, &rets);
800        }
801
802        //     return_call f(args...)
803        //
804        // becomes
805        //
806        //     rets... = call f(args...)
807        //     jump return_block(rets...)
808        ir::InstructionData::Call {
809            opcode: ir::Opcode::ReturnCall,
810            args,
811            func_ref,
812        } => {
813            func.dfg.insts[inlined_inst] = ir::InstructionData::Call {
814                opcode: ir::Opcode::Call,
815                args,
816                func_ref,
817            };
818            func.dfg.make_inst_results(inlined_inst, ir::types::INVALID);
819
820            append_stack_map_entries(
821                func,
822                callee,
823                &entity_map,
824                call_stack_map,
825                inlined_inst,
826                callee_inst,
827            );
828
829            let rets = SmallBlockArgVec::from_iter(
830                func.dfg
831                    .inst_results(inlined_inst)
832                    .iter()
833                    .copied()
834                    .map(|v| v.into()),
835            );
836            let mut cursor = FuncCursor::new(func);
837            cursor.goto_after_inst(inlined_inst);
838            cursor.ins().jump(return_block, &rets);
839
840            if call_opcode == ir::Opcode::TryCall {
841                allocs
842                    .calls_needing_exception_table_fixup
843                    .push(inlined_inst);
844            }
845        }
846
847        //     return_call_indirect val(args...)
848        //
849        // becomes
850        //
851        //     rets... = call_indirect val(args...)
852        //     jump return_block(rets...)
853        ir::InstructionData::CallIndirect {
854            opcode: ir::Opcode::ReturnCallIndirect,
855            args,
856            sig_ref,
857        } => {
858            func.dfg.insts[inlined_inst] = ir::InstructionData::CallIndirect {
859                opcode: ir::Opcode::CallIndirect,
860                args,
861                sig_ref,
862            };
863            func.dfg.make_inst_results(inlined_inst, ir::types::INVALID);
864
865            append_stack_map_entries(
866                func,
867                callee,
868                &entity_map,
869                call_stack_map,
870                inlined_inst,
871                callee_inst,
872            );
873
874            let rets = SmallBlockArgVec::from_iter(
875                func.dfg
876                    .inst_results(inlined_inst)
877                    .iter()
878                    .copied()
879                    .map(|v| v.into()),
880            );
881            let mut cursor = FuncCursor::new(func);
882            cursor.goto_after_inst(inlined_inst);
883            cursor.ins().jump(return_block, &rets);
884
885            if call_opcode == ir::Opcode::TryCall {
886                allocs
887                    .calls_needing_exception_table_fixup
888                    .push(inlined_inst);
889            }
890        }
891
892        inst_data => unreachable!(
893            "should have handled all `is_return() == true` instructions above; \
894             got {inst_data:?}"
895        ),
896    }
897}
898
899/// An `InstructionMapper` implementation that remaps a callee instruction's
900/// entity references to their new indices in the caller function.
901struct InliningInstRemapper<'a> {
902    allocs: &'a InliningAllocs,
903    func: &'a mut ir::Function,
904    callee: &'a ir::Function,
905    entity_map: &'a EntityMap,
906    error: Option<crate::result::CodegenError>,
907}
908
909impl<'a> ir::instructions::InstructionMapper for InliningInstRemapper<'a> {
910    fn map_value(&mut self, value: ir::Value) -> ir::Value {
911        self.allocs.get_inlined_value(self.callee, value).expect(
912            "defs come before uses; we should have already inlined all values \
913             used by an instruction",
914        )
915    }
916
917    fn map_value_list(&mut self, value_list: ir::ValueList) -> ir::ValueList {
918        let mut inlined_list = ir::ValueList::new();
919        for callee_val in value_list.as_slice(&self.callee.dfg.value_lists) {
920            let inlined_val = self.map_value(*callee_val);
921            inlined_list.push(inlined_val, &mut self.func.dfg.value_lists);
922        }
923        inlined_list
924    }
925
926    fn map_global_value(&mut self, global_value: ir::GlobalValue) -> ir::GlobalValue {
927        self.entity_map.inlined_global_value(global_value)
928    }
929
930    fn map_jump_table(&mut self, jump_table: ir::JumpTable) -> ir::JumpTable {
931        let inlined_default =
932            self.map_block_call(self.callee.dfg.jump_tables[jump_table].default_block());
933        let inlined_table = self.callee.dfg.jump_tables[jump_table]
934            .as_slice()
935            .iter()
936            .map(|callee_block_call| self.map_block_call(*callee_block_call))
937            .collect::<SmallBlockCallVec>();
938        self.func
939            .dfg
940            .jump_tables
941            .push(ir::JumpTableData::new(inlined_default, &inlined_table))
942    }
943
944    fn map_exception_table(&mut self, exception_table: ir::ExceptionTable) -> ir::ExceptionTable {
945        let exception_table = &self.callee.dfg.exception_tables[exception_table];
946        let inlined_sig_ref = self.map_sig_ref(exception_table.signature());
947        let inlined_normal_return = self.map_block_call(*exception_table.normal_return());
948        let inlined_table = exception_table
949            .items()
950            .map(|item| match item {
951                ExceptionTableItem::Tag(tag, block_call) => {
952                    ExceptionTableItem::Tag(tag, self.map_block_call(block_call))
953                }
954                ExceptionTableItem::Default(block_call) => {
955                    ExceptionTableItem::Default(self.map_block_call(block_call))
956                }
957                ExceptionTableItem::Context(value) => {
958                    ExceptionTableItem::Context(self.map_value(value))
959                }
960            })
961            .collect::<SmallVec<[_; 8]>>();
962        self.func
963            .dfg
964            .exception_tables
965            .push(ir::ExceptionTableData::new(
966                inlined_sig_ref,
967                inlined_normal_return,
968                inlined_table,
969            ))
970    }
971
972    fn map_block_call(&mut self, block_call: ir::BlockCall) -> ir::BlockCall {
973        let callee_block = block_call.block(&self.callee.dfg.value_lists);
974        let inlined_block = self.entity_map.inlined_block(callee_block);
975        let args = block_call
976            .args(&self.callee.dfg.value_lists)
977            .map(|arg| match arg {
978                ir::BlockArg::Value(value) => self.map_value(value).into(),
979                ir::BlockArg::TryCallRet(_) | ir::BlockArg::TryCallExn(_) => arg,
980            })
981            .collect::<SmallBlockArgVec>();
982        ir::BlockCall::new(inlined_block, args, &mut self.func.dfg.value_lists)
983    }
984
985    fn map_block(&mut self, block: ir::Block) -> ir::Block {
986        self.entity_map.inlined_block(block)
987    }
988
989    fn map_func_ref(&mut self, func_ref: ir::FuncRef) -> ir::FuncRef {
990        self.entity_map.inlined_func_ref(func_ref)
991    }
992
993    fn map_sig_ref(&mut self, sig_ref: ir::SigRef) -> ir::SigRef {
994        self.entity_map.inlined_sig_ref(sig_ref)
995    }
996
997    fn map_stack_slot(&mut self, stack_slot: ir::StackSlot) -> ir::StackSlot {
998        self.entity_map.inlined_stack_slot(stack_slot)
999    }
1000
1001    fn map_dynamic_stack_slot(
1002        &mut self,
1003        dynamic_stack_slot: ir::DynamicStackSlot,
1004    ) -> ir::DynamicStackSlot {
1005        self.entity_map
1006            .inlined_dynamic_stack_slot(dynamic_stack_slot)
1007    }
1008
1009    fn map_constant(&mut self, constant: ir::Constant) -> ir::Constant {
1010        self.allocs
1011            .constants
1012            .get(constant)
1013            .and_then(|o| o.expand())
1014            .expect("should have inlined all callee constants")
1015    }
1016
1017    fn map_immediate(&mut self, immediate: ir::Immediate) -> ir::Immediate {
1018        self.entity_map.inlined_immediate(immediate)
1019    }
1020    fn map_mem_flags(&mut self, flags: ir::MemFlags) -> ir::MemFlags {
1021        let mut flags_data = self.callee.dfg.mem_flags[flags];
1022        // Remap the alias region entity from callee to caller.
1023        if let Some(callee_region) = flags_data.alias_region() {
1024            let region_data = self.callee.dfg.alias_regions[callee_region].clone();
1025            let caller_region = self.func.dfg.alias_regions.insert(region_data);
1026            flags_data.set_alias_region(Some(caller_region));
1027        }
1028        match self.func.dfg.mem_flags.insert(flags_data) {
1029            Ok(flags) => flags,
1030            Err(_) => {
1031                self.error = Some(crate::result::CodegenError::ImplLimitExceeded);
1032                self.func
1033                    .dfg
1034                    .mem_flags
1035                    .insert(ir::MemFlagsData::trusted())
1036                    .unwrap()
1037            }
1038        }
1039    }
1040}
1041
1042/// Inline the callee's layout into the caller's layout.
1043///
1044/// Returns the last inlined block in the layout.
1045fn inline_block_layout(
1046    func: &mut ir::Function,
1047    call_block: ir::Block,
1048    callee: &ir::Function,
1049    entity_map: &EntityMap,
1050) -> ir::Block {
1051    debug_assert!(func.layout.is_block_inserted(call_block));
1052
1053    // Iterate over callee blocks in layout order, inserting their associated
1054    // inlined block into the caller's layout.
1055    let mut prev_inlined_block = call_block;
1056    let mut next_callee_block = callee.layout.entry_block();
1057    while let Some(callee_block) = next_callee_block {
1058        debug_assert!(func.layout.is_block_inserted(prev_inlined_block));
1059
1060        let inlined_block = entity_map.inlined_block(callee_block);
1061        func.layout
1062            .insert_block_after(inlined_block, prev_inlined_block);
1063
1064        prev_inlined_block = inlined_block;
1065        next_callee_block = callee.layout.next_block(callee_block);
1066    }
1067
1068    debug_assert!(func.layout.is_block_inserted(prev_inlined_block));
1069    prev_inlined_block
1070}
1071
1072/// Split the call instruction's block just after the call instruction to create
1073/// the point where control-flow joins after the inlined callee "returns".
1074///
1075/// Note that tail calls do not return to the caller and therefore do not have a
1076/// control-flow join point.
1077fn split_off_return_block(
1078    func: &mut ir::Function,
1079    call_inst: ir::Inst,
1080    opcode: ir::Opcode,
1081    callee: &ir::Function,
1082) -> Option<ir::Block> {
1083    // When the `call_inst` is not a block terminator, we need to split the
1084    // block.
1085    let return_block = func.layout.next_inst(call_inst).map(|next_inst| {
1086        let return_block = func.dfg.blocks.add();
1087        func.layout.split_block(return_block, next_inst);
1088
1089        // Add block parameters for each return value and alias the call
1090        // instruction's results to them.
1091        let old_results =
1092            SmallValueVec::from_iter(func.dfg.inst_results(call_inst).iter().copied());
1093        debug_assert_eq!(old_results.len(), callee.signature.returns.len());
1094        func.dfg.detach_inst_results(call_inst);
1095        for (abi, old_val) in callee.signature.returns.iter().zip(old_results) {
1096            debug_assert_eq!(abi.value_type, func.dfg.value_type(old_val));
1097            let ret_param = func.dfg.append_block_param(return_block, abi.value_type);
1098            func.dfg.change_to_alias(old_val, ret_param);
1099        }
1100
1101        return_block
1102    });
1103
1104    // When the `call_inst` is a block terminator, then it is either a
1105    // `return_call` or a `try_call`:
1106    //
1107    // * For `return_call`s, we don't have a control-flow join point, because
1108    //   the caller permanently transfers control to the callee.
1109    //
1110    // * For `try_call`s, we probably already have a block for the control-flow
1111    //   join point, but it isn't guaranteed: the `try_call` might ignore the
1112    //   call's returns and not forward them to the normal-return block or it
1113    //   might also pass additional arguments. We can only reuse the existing
1114    //   normal-return block when the `try_call` forwards exactly our callee's
1115    //   returns to that block (and therefore that block's parameter types also
1116    //   exactly match the callee's return types). Otherwise, we must create a new
1117    //   return block that forwards to the existing normal-return
1118    //   block. (Elsewhere, at the end of inlining, we will also update any inlined
1119    //   calls to forward any raised exceptions to the caller's exception table,
1120    //   as necessary.)
1121    //
1122    //   Finally, note that reusing the normal-return's target block is just an
1123    //   optimization to emit a simpler CFG when we can, and is not
1124    //   fundamentally required for correctness. We could always insert a
1125    //   temporary block as our control-flow join point that then forwards to
1126    //   the normal-return's target block. However, at the time of writing,
1127    //   Cranelift doesn't currently do any jump-threading or branch
1128    //   simplification in the mid-end, and removing unnecessary blocks in this
1129    //   way can help some subsequent mid-end optimizations. If, in the future,
1130    //   we gain support for jump-threading optimizations in the mid-end, we can
1131    //   come back and simplify the below code a bit to always generate the
1132    //   temporary block, and then rely on the subsequent optimizations to clean
1133    //   everything up.
1134    debug_assert_eq!(
1135        return_block.is_none(),
1136        opcode == ir::Opcode::ReturnCall || opcode == ir::Opcode::TryCall,
1137    );
1138    return_block.or_else(|| match func.dfg.insts[call_inst] {
1139        ir::InstructionData::TryCall {
1140            opcode: ir::Opcode::TryCall,
1141            args: _,
1142            func_ref: _,
1143            exception,
1144        } => {
1145            let normal_return = func.dfg.exception_tables[exception].normal_return();
1146            let normal_return_block = normal_return.block(&func.dfg.value_lists);
1147
1148            // Check to see if we can reuse the existing normal-return block.
1149            {
1150                let normal_return_args = normal_return.args(&func.dfg.value_lists);
1151                if normal_return_args.len() == callee.signature.returns.len()
1152                    && normal_return_args.enumerate().all(|(i, arg)| {
1153                        let i = u32::try_from(i).unwrap();
1154                        arg == ir::BlockArg::TryCallRet(i)
1155                    })
1156                {
1157                    return Some(normal_return_block);
1158                }
1159            }
1160
1161            // Okay, we cannot reuse the normal-return block. Create a new block
1162            // that has the expected block parameter types and have it jump to
1163            // the normal-return block.
1164            let return_block = func.dfg.blocks.add();
1165            func.layout.insert_block(return_block, normal_return_block);
1166
1167            let return_block_params = callee
1168                .signature
1169                .returns
1170                .iter()
1171                .map(|abi| func.dfg.append_block_param(return_block, abi.value_type))
1172                .collect::<SmallValueVec>();
1173
1174            let normal_return_args = func.dfg.exception_tables[exception]
1175                .normal_return()
1176                .args(&func.dfg.value_lists)
1177                .collect::<SmallBlockArgVec>();
1178            let jump_args = normal_return_args
1179                .into_iter()
1180                .map(|arg| match arg {
1181                    ir::BlockArg::Value(value) => ir::BlockArg::Value(value),
1182                    ir::BlockArg::TryCallRet(i) => {
1183                        let i = usize::try_from(i).unwrap();
1184                        ir::BlockArg::Value(return_block_params[i])
1185                    }
1186                    ir::BlockArg::TryCallExn(_) => {
1187                        unreachable!("normal-return edges cannot use exceptional results")
1188                    }
1189                })
1190                .collect::<SmallBlockArgVec>();
1191
1192            let mut cursor = FuncCursor::new(func);
1193            cursor.goto_first_insertion_point(return_block);
1194            cursor.ins().jump(normal_return_block, &jump_args);
1195
1196            Some(return_block)
1197        }
1198        _ => None,
1199    })
1200}
1201
1202/// Replace the caller's call instruction with a jump to the caller's inlined
1203/// copy of the callee's entry block.
1204///
1205/// Also associates the callee's parameters with the caller's arguments in our
1206/// value map.
1207///
1208/// Returns the caller's stack map entries, if any.
1209fn replace_call_with_jump(
1210    allocs: &mut InliningAllocs,
1211    func: &mut ir::Function,
1212    call_inst: ir::Inst,
1213    callee: &ir::Function,
1214    entity_map: &EntityMap,
1215) -> Option<ir::UserStackMapEntryVec> {
1216    trace!("Replacing `call` with `jump`");
1217    trace!(
1218        "  --> call instruction: {call_inst:?}: {}",
1219        func.dfg.display_inst(call_inst)
1220    );
1221
1222    let callee_entry_block = callee
1223        .layout
1224        .entry_block()
1225        .expect("callee function should have an entry block");
1226    let callee_param_values = callee.dfg.block_params(callee_entry_block);
1227    let caller_arg_values = SmallValueVec::from_iter(func.dfg.inst_args(call_inst).iter().copied());
1228    debug_assert_eq!(callee_param_values.len(), caller_arg_values.len());
1229    debug_assert_eq!(callee_param_values.len(), callee.signature.params.len());
1230    for (abi, (callee_param_value, caller_arg_value)) in callee
1231        .signature
1232        .params
1233        .iter()
1234        .zip(callee_param_values.into_iter().zip(caller_arg_values))
1235    {
1236        debug_assert_eq!(abi.value_type, callee.dfg.value_type(*callee_param_value));
1237        debug_assert_eq!(abi.value_type, func.dfg.value_type(caller_arg_value));
1238        allocs.set_inlined_value(callee, *callee_param_value, caller_arg_value);
1239    }
1240
1241    // Replace the caller's call instruction with a jump to the caller's inlined
1242    // copy of the callee's entry block.
1243    //
1244    // Note that the call block dominates the inlined entry block (and also all
1245    // other inlined blocks) so we can reference the arguments directly, and do
1246    // not need to add block parameters to the inlined entry block.
1247    let inlined_entry_block = entity_map.inlined_block(callee_entry_block);
1248    func.replace(call_inst).jump(inlined_entry_block, &[]);
1249    trace!(
1250        "  --> replaced with jump instruction: {call_inst:?}: {}",
1251        func.dfg.display_inst(call_inst)
1252    );
1253
1254    let stack_map_entries = func.dfg.take_user_stack_map_entries(call_inst);
1255    stack_map_entries
1256}
1257
1258/// Keeps track of mapping callee entities to their associated inlined caller
1259/// entities.
1260#[derive(Default)]
1261struct EntityMap {
1262    // Rather than doing an implicit, demand-based, DCE'ing translation of
1263    // entities, which would require maps from each callee entity to its
1264    // associated caller entity, we copy all entities into the caller, remember
1265    // each entity's initial offset, and then mapping from the callee to the
1266    // inlined caller entity is just adding that initial offset to the callee's
1267    // index. This should be both faster and simpler than the alternative. Most
1268    // of these sets are relatively small, and they rarely have too much dead
1269    // code in practice, so this is a good trade off.
1270    //
1271    // Note that there are a few kinds of entities that are excluded from the
1272    // `EntityMap`, and for which we do actually take the demand-based approach:
1273    // values and value lists being the notable ones.
1274    block_offset: Option<u32>,
1275    global_value_offset: Option<u32>,
1276    sig_ref_offset: Option<u32>,
1277    func_ref_offset: Option<u32>,
1278    stack_slot_offset: Option<u32>,
1279    dynamic_type_offset: Option<u32>,
1280    dynamic_stack_slot_offset: Option<u32>,
1281    immediate_offset: Option<u32>,
1282}
1283
1284impl EntityMap {
1285    fn inlined_block(&self, callee_block: ir::Block) -> ir::Block {
1286        let offset = self
1287            .block_offset
1288            .expect("must create inlined `ir::Block`s before calling `EntityMap::inlined_block`");
1289        ir::Block::from_u32(offset + callee_block.as_u32())
1290    }
1291
1292    fn iter_inlined_blocks(&self, func: &ir::Function) -> impl Iterator<Item = ir::Block> + use<> {
1293        let start = self.block_offset.expect(
1294            "must create inlined `ir::Block`s before calling `EntityMap::iter_inlined_blocks`",
1295        );
1296
1297        let end = func.dfg.blocks.len();
1298        let end = u32::try_from(end).unwrap();
1299
1300        (start..end).map(|i| ir::Block::from_u32(i))
1301    }
1302
1303    fn inlined_global_value(&self, callee_global_value: ir::GlobalValue) -> ir::GlobalValue {
1304        let offset = self
1305            .global_value_offset
1306            .expect("must create inlined `ir::GlobalValue`s before calling `EntityMap::inlined_global_value`");
1307        ir::GlobalValue::from_u32(offset + callee_global_value.as_u32())
1308    }
1309
1310    fn inlined_sig_ref(&self, callee_sig_ref: ir::SigRef) -> ir::SigRef {
1311        let offset = self.sig_ref_offset.expect(
1312            "must create inlined `ir::SigRef`s before calling `EntityMap::inlined_sig_ref`",
1313        );
1314        ir::SigRef::from_u32(offset + callee_sig_ref.as_u32())
1315    }
1316
1317    fn inlined_func_ref(&self, callee_func_ref: ir::FuncRef) -> ir::FuncRef {
1318        let offset = self.func_ref_offset.expect(
1319            "must create inlined `ir::FuncRef`s before calling `EntityMap::inlined_func_ref`",
1320        );
1321        ir::FuncRef::from_u32(offset + callee_func_ref.as_u32())
1322    }
1323
1324    fn inlined_stack_slot(&self, callee_stack_slot: ir::StackSlot) -> ir::StackSlot {
1325        let offset = self.stack_slot_offset.expect(
1326            "must create inlined `ir::StackSlot`s before calling `EntityMap::inlined_stack_slot`",
1327        );
1328        ir::StackSlot::from_u32(offset + callee_stack_slot.as_u32())
1329    }
1330
1331    fn inlined_dynamic_type(&self, callee_dynamic_type: ir::DynamicType) -> ir::DynamicType {
1332        let offset = self.dynamic_type_offset.expect(
1333            "must create inlined `ir::DynamicType`s before calling `EntityMap::inlined_dynamic_type`",
1334        );
1335        ir::DynamicType::from_u32(offset + callee_dynamic_type.as_u32())
1336    }
1337
1338    fn inlined_dynamic_stack_slot(
1339        &self,
1340        callee_dynamic_stack_slot: ir::DynamicStackSlot,
1341    ) -> ir::DynamicStackSlot {
1342        let offset = self.dynamic_stack_slot_offset.expect(
1343            "must create inlined `ir::DynamicStackSlot`s before calling `EntityMap::inlined_dynamic_stack_slot`",
1344        );
1345        ir::DynamicStackSlot::from_u32(offset + callee_dynamic_stack_slot.as_u32())
1346    }
1347
1348    fn inlined_immediate(&self, callee_immediate: ir::Immediate) -> ir::Immediate {
1349        let offset = self.immediate_offset.expect(
1350            "must create inlined `ir::Immediate`s before calling `EntityMap::inlined_immediate`",
1351        );
1352        ir::Immediate::from_u32(offset + callee_immediate.as_u32())
1353    }
1354}
1355
1356/// Translate all of the callee's various entities into the caller, producing an
1357/// `EntityMap` that can be used to translate callee entity references into
1358/// inlined caller entity references.
1359fn create_entities(
1360    allocs: &mut InliningAllocs,
1361    func: &mut ir::Function,
1362    callee: &ir::Function,
1363) -> CodegenResult<EntityMap> {
1364    let mut entity_map = EntityMap::default();
1365
1366    entity_map.block_offset = Some(create_blocks(allocs, func, callee));
1367    entity_map.global_value_offset = Some(create_global_values(func, callee)?);
1368    entity_map.sig_ref_offset = Some(create_sig_refs(func, callee));
1369    create_user_external_name_refs(allocs, func, callee);
1370    entity_map.func_ref_offset = Some(create_func_refs(allocs, func, callee, &entity_map));
1371    entity_map.stack_slot_offset = Some(create_stack_slots(func, callee));
1372    entity_map.dynamic_type_offset = Some(create_dynamic_types(func, callee, &entity_map));
1373    entity_map.dynamic_stack_slot_offset =
1374        Some(create_dynamic_stack_slots(func, callee, &entity_map));
1375    entity_map.immediate_offset = Some(create_immediates(func, callee));
1376
1377    // `ir::ConstantData` is deduplicated, so we cannot use our offset scheme
1378    // for `ir::Constant`s. Nonetheless, we still insert them into the caller
1379    // now, at the same time as the rest of our entities.
1380    create_constants(allocs, func, callee);
1381
1382    Ok(entity_map)
1383}
1384
1385/// Create inlined blocks in the caller for every block in the callee.
1386fn create_blocks(
1387    allocs: &mut InliningAllocs,
1388    func: &mut ir::Function,
1389    callee: &ir::Function,
1390) -> u32 {
1391    let offset = func.dfg.blocks.len();
1392    let offset = u32::try_from(offset).unwrap();
1393
1394    func.dfg.blocks.reserve(callee.dfg.blocks.len());
1395    for callee_block in callee.dfg.blocks.iter() {
1396        let caller_block = func.dfg.blocks.add();
1397        trace!("Callee {callee_block:?} = inlined {caller_block:?}");
1398
1399        if callee.layout.is_cold(callee_block) {
1400            func.layout.set_cold(caller_block);
1401        }
1402
1403        // Note: the entry block does not need parameters because the only
1404        // predecessor is the call block and we associate the callee's
1405        // parameters with the caller's arguments directly.
1406        if callee.layout.entry_block() != Some(callee_block) {
1407            for callee_param in callee.dfg.blocks[callee_block].params(&callee.dfg.value_lists) {
1408                let ty = callee.dfg.value_type(*callee_param);
1409                let caller_param = func.dfg.append_block_param(caller_block, ty);
1410
1411                allocs.set_inlined_value(callee, *callee_param, caller_param);
1412            }
1413        }
1414    }
1415
1416    offset
1417}
1418
1419/// Copy and translate global values from the callee into the caller.
1420fn create_global_values(func: &mut ir::Function, callee: &ir::Function) -> CodegenResult<u32> {
1421    let gv_offset = func.global_values.len();
1422    let gv_offset = u32::try_from(gv_offset).unwrap();
1423
1424    func.global_values.reserve(callee.global_values.len());
1425    for gv in callee.global_values.values() {
1426        // Re-insert callee mem flags into the caller's DFG before constructing
1427        // the global value data, to avoid borrow conflicts.
1428        let remapped_flags = match gv {
1429            ir::GlobalValueData::Load { flags, .. } => {
1430                let mut flags_data = callee.dfg.mem_flags[*flags];
1431                // Remap alias region entity from callee to caller.
1432                if let Some(callee_region) = flags_data.alias_region() {
1433                    let region_data = callee.dfg.alias_regions[callee_region].clone();
1434                    let caller_region = func.dfg.alias_regions.insert(region_data);
1435                    flags_data.set_alias_region(Some(caller_region));
1436                }
1437                Some(
1438                    func.dfg
1439                        .mem_flags
1440                        .insert(flags_data)
1441                        .map_err(|_| crate::result::CodegenError::ImplLimitExceeded)?,
1442                )
1443            }
1444            _ => None,
1445        };
1446        func.global_values.push(match gv {
1447            // These kinds of global values reference other global values, so we
1448            // need to fixup that reference.
1449            ir::GlobalValueData::Load {
1450                base,
1451                offset,
1452                global_type,
1453                flags: _,
1454            } => ir::GlobalValueData::Load {
1455                base: ir::GlobalValue::from_u32(base.as_u32() + gv_offset),
1456                offset: *offset,
1457                global_type: *global_type,
1458                flags: remapped_flags.unwrap(),
1459            },
1460            ir::GlobalValueData::IAddImm {
1461                base,
1462                offset,
1463                global_type,
1464            } => ir::GlobalValueData::IAddImm {
1465                base: ir::GlobalValue::from_u32(base.as_u32() + gv_offset),
1466                offset: *offset,
1467                global_type: *global_type,
1468            },
1469
1470            // These kinds of global values do not reference other global
1471            // values, so we can just clone them.
1472            ir::GlobalValueData::VMContext
1473            | ir::GlobalValueData::Symbol { .. }
1474            | ir::GlobalValueData::DynScaleTargetConst { .. } => gv.clone(),
1475        });
1476    }
1477
1478    Ok(gv_offset)
1479}
1480
1481/// Copy `ir::SigRef`s from the callee into the caller.
1482fn create_sig_refs(func: &mut ir::Function, callee: &ir::Function) -> u32 {
1483    let offset = func.dfg.signatures.len();
1484    let offset = u32::try_from(offset).unwrap();
1485
1486    func.dfg.signatures.reserve(callee.dfg.signatures.len());
1487    for sig in callee.dfg.signatures.values() {
1488        func.dfg.signatures.push(sig.clone());
1489    }
1490
1491    offset
1492}
1493
1494fn create_user_external_name_refs(
1495    allocs: &mut InliningAllocs,
1496    func: &mut ir::Function,
1497    callee: &ir::Function,
1498) {
1499    for (callee_named_func_ref, name) in callee.params.user_named_funcs().iter() {
1500        let caller_named_func_ref = func.declare_imported_user_function(name.clone());
1501        allocs.user_external_name_refs[callee_named_func_ref] = Some(caller_named_func_ref).into();
1502    }
1503}
1504
1505/// Translate `ir::FuncRef`s from the callee into the caller.
1506fn create_func_refs(
1507    allocs: &InliningAllocs,
1508    func: &mut ir::Function,
1509    callee: &ir::Function,
1510    entity_map: &EntityMap,
1511) -> u32 {
1512    let offset = func.dfg.ext_funcs.len();
1513    let offset = u32::try_from(offset).unwrap();
1514
1515    func.dfg.ext_funcs.reserve(callee.dfg.ext_funcs.len());
1516    for ir::ExtFuncData {
1517        name,
1518        signature,
1519        colocated,
1520        patchable,
1521    } in callee.dfg.ext_funcs.values()
1522    {
1523        func.dfg.ext_funcs.push(ir::ExtFuncData {
1524            name: match name {
1525                ir::ExternalName::User(name_ref) => {
1526                    ir::ExternalName::User(allocs.user_external_name_refs[*name_ref].expect(
1527                        "should have translated all `ir::UserExternalNameRef`s before translating \
1528                         `ir::FuncRef`s",
1529                    ))
1530                }
1531                ir::ExternalName::TestCase(_)
1532                | ir::ExternalName::LibCall(_)
1533                | ir::ExternalName::KnownSymbol(_) => name.clone(),
1534            },
1535            signature: entity_map.inlined_sig_ref(*signature),
1536            colocated: *colocated,
1537            patchable: *patchable,
1538        });
1539    }
1540
1541    offset
1542}
1543
1544/// Copy stack slots from the callee into the caller.
1545fn create_stack_slots(func: &mut ir::Function, callee: &ir::Function) -> u32 {
1546    let offset = func.sized_stack_slots.len();
1547    let offset = u32::try_from(offset).unwrap();
1548
1549    func.sized_stack_slots
1550        .reserve(callee.sized_stack_slots.len());
1551    for slot in callee.sized_stack_slots.values() {
1552        func.sized_stack_slots.push(slot.clone());
1553    }
1554
1555    offset
1556}
1557
1558/// Copy dynamic types from the callee into the caller.
1559fn create_dynamic_types(
1560    func: &mut ir::Function,
1561    callee: &ir::Function,
1562    entity_map: &EntityMap,
1563) -> u32 {
1564    let offset = func.dynamic_stack_slots.len();
1565    let offset = u32::try_from(offset).unwrap();
1566
1567    func.dfg
1568        .dynamic_types
1569        .reserve(callee.dfg.dynamic_types.len());
1570    for ir::DynamicTypeData {
1571        base_vector_ty,
1572        dynamic_scale,
1573    } in callee.dfg.dynamic_types.values()
1574    {
1575        func.dfg.dynamic_types.push(ir::DynamicTypeData {
1576            base_vector_ty: *base_vector_ty,
1577            dynamic_scale: entity_map.inlined_global_value(*dynamic_scale),
1578        });
1579    }
1580
1581    offset
1582}
1583
1584/// Copy dynamic stack slots from the callee into the caller.
1585fn create_dynamic_stack_slots(
1586    func: &mut ir::Function,
1587    callee: &ir::Function,
1588    entity_map: &EntityMap,
1589) -> u32 {
1590    let offset = func.dynamic_stack_slots.len();
1591    let offset = u32::try_from(offset).unwrap();
1592
1593    func.dynamic_stack_slots
1594        .reserve(callee.dynamic_stack_slots.len());
1595    for ir::DynamicStackSlotData { kind, dyn_ty } in callee.dynamic_stack_slots.values() {
1596        func.dynamic_stack_slots.push(ir::DynamicStackSlotData {
1597            kind: *kind,
1598            dyn_ty: entity_map.inlined_dynamic_type(*dyn_ty),
1599        });
1600    }
1601
1602    offset
1603}
1604
1605/// Copy immediates from the callee into the caller.
1606fn create_immediates(func: &mut ir::Function, callee: &ir::Function) -> u32 {
1607    let offset = func.dfg.immediates.len();
1608    let offset = u32::try_from(offset).unwrap();
1609
1610    func.dfg.immediates.reserve(callee.dfg.immediates.len());
1611    for imm in callee.dfg.immediates.values() {
1612        func.dfg.immediates.push(imm.clone());
1613    }
1614
1615    offset
1616}
1617
1618/// Copy constants from the callee into the caller.
1619fn create_constants(allocs: &mut InliningAllocs, func: &mut ir::Function, callee: &ir::Function) {
1620    for (callee_constant, data) in callee.dfg.constants.iter() {
1621        let inlined_constant = func.dfg.constants.insert(data.clone());
1622        allocs.constants[*callee_constant] = Some(inlined_constant).into();
1623    }
1624}