Skip to main content

cranelift_frontend/frontend/
safepoints.rs

1//! Support for safepoints and stack maps.
2
3use super::*;
4use crate::{HashMap, HashSet};
5use core::ops::{Index, IndexMut};
6
7#[derive(Clone, Copy)]
8#[repr(u8)]
9enum SlotSize {
10    Size8 = 0,
11    Size16 = 1,
12    Size32 = 2,
13    Size64 = 3,
14    Size128 = 4,
15    // If adding support for more slot sizes, update `SLOT_SIZE_LEN` below.
16}
17const SLOT_SIZE_LEN: usize = 5;
18
19impl TryFrom<ir::Type> for SlotSize {
20    type Error = &'static str;
21
22    fn try_from(ty: ir::Type) -> Result<Self, Self::Error> {
23        Self::new(ty.bytes()).ok_or("type is not supported in stack maps")
24    }
25}
26
27impl SlotSize {
28    fn new(bytes: u32) -> Option<Self> {
29        match bytes {
30            1 => Some(Self::Size8),
31            2 => Some(Self::Size16),
32            4 => Some(Self::Size32),
33            8 => Some(Self::Size64),
34            16 => Some(Self::Size128),
35            _ => None,
36        }
37    }
38
39    fn unwrap_new(bytes: u32) -> Self {
40        Self::new(bytes).unwrap_or_else(|| panic!("cannot create a `SlotSize` for {bytes} bytes"))
41    }
42}
43
44/// A map from every `SlotSize` to a `T`.
45struct SlotSizeMap<T>([T; SLOT_SIZE_LEN]);
46
47impl<T> Default for SlotSizeMap<T>
48where
49    T: Default,
50{
51    fn default() -> Self {
52        Self::new()
53    }
54}
55
56impl<T> Index<SlotSize> for SlotSizeMap<T> {
57    type Output = T;
58    fn index(&self, index: SlotSize) -> &Self::Output {
59        self.get(index)
60    }
61}
62
63impl<T> IndexMut<SlotSize> for SlotSizeMap<T> {
64    fn index_mut(&mut self, index: SlotSize) -> &mut Self::Output {
65        self.get_mut(index)
66    }
67}
68
69impl<T> SlotSizeMap<T> {
70    fn new() -> Self
71    where
72        T: Default,
73    {
74        Self([
75            T::default(),
76            T::default(),
77            T::default(),
78            T::default(),
79            T::default(),
80        ])
81    }
82
83    fn clear(&mut self)
84    where
85        T: Default,
86    {
87        *self = Self::new();
88    }
89
90    fn get(&self, size: SlotSize) -> &T {
91        let index = size as u8 as usize;
92        &self.0[index]
93    }
94
95    fn get_mut(&mut self, size: SlotSize) -> &mut T {
96        let index = size as u8 as usize;
97        &mut self.0[index]
98    }
99}
100
101/// A set of live values.
102///
103/// Make sure to copy to a vec and sort, or something, before iterating over the
104/// values to ensure deterministic output.
105type LiveSet = HashSet<ir::Value>;
106
107/// A worklist of blocks' post-order indices that we need to process.
108#[derive(Default)]
109struct Worklist {
110    /// Stack of blocks to process.
111    stack: Vec<u32>,
112
113    /// The set of blocks that need to be updated.
114    ///
115    /// This is a subset of the elements present in `self.stack`, *not* the
116    /// exact same elements. `self.stack` is allowed to have duplicates, and
117    /// once we pop the first occurrence of a duplicate, we remove it from this
118    /// set, since it no longer needs updates at that point. This potentially
119    /// uses more stack space than necessary, but prefers processing immediate
120    /// predecessors, and therefore inner loop bodies before continuing to
121    /// process outer loop bodies. This ultimately results in fewer iterations
122    /// required to reach a fixed point.
123    need_updates: HashSet<u32>,
124}
125
126impl Extend<u32> for Worklist {
127    fn extend<T>(&mut self, iter: T)
128    where
129        T: IntoIterator<Item = u32>,
130    {
131        for block_index in iter {
132            self.push(block_index);
133        }
134    }
135}
136
137impl Worklist {
138    fn clear(&mut self) {
139        let Worklist {
140            stack,
141            need_updates,
142        } = self;
143        stack.clear();
144        need_updates.clear();
145    }
146
147    fn reserve(&mut self, capacity: usize) {
148        let Worklist {
149            stack,
150            need_updates,
151        } = self;
152        stack.reserve(capacity);
153        need_updates.reserve(capacity);
154    }
155
156    fn push(&mut self, block_index: u32) {
157        // Mark this block as needing an update. If it wasn't in `self.stack`,
158        // now it is and it needs an update. If it was already in `self.stack`,
159        // then pushing this copy logically hoists it to the top of the
160        // stack. See the above note about processing inner-most loops first.
161        self.need_updates.insert(block_index);
162        self.stack.push(block_index);
163    }
164
165    fn pop(&mut self) -> Option<u32> {
166        while let Some(block_index) = self.stack.pop() {
167            // If this block was pushed multiple times, we only need to update
168            // it once, so remove it from the need-updates set. In other words
169            // it was logically hoisted up to the top of the stack, while this
170            // entry was left behind, and we already popped the hoisted
171            // copy. See the above note about processing inner-most loops first.
172            if self.need_updates.remove(&block_index) {
173                return Some(block_index);
174            }
175        }
176        None
177    }
178}
179
180/// A simple liveness analysis.
181///
182/// This analysis is used to determine which needs-stack-map values are live
183/// across safepoint instructions.
184///
185/// This is a backwards analysis, from uses (which mark values live) to defs
186/// (which remove values from the live set) and from successor blocks to
187/// predecessor blocks.
188///
189/// We compute two live sets for each block:
190///
191/// 1. The live-in set, which is the set of values that are live when control
192///    enters the block.
193///
194/// 2. The live-out set, which is the set of values that are live when control
195///    exits the block.
196///
197/// A block's live-out set is the union of its successors' live-in sets. A
198/// block's live-in set is the set of values that are still live after the
199/// block's instructions have been processed.
200///
201/// ```text
202/// live_in(block) = union(live_out(s) for s in successors(block))
203/// live_out(block) = live_in(block) - defs(block) + uses(block)
204/// ```
205///
206/// Whenever we update a block's live-in set, we must reprocess all of its
207/// predecessors, because those predecessors' live-out sets depend on this
208/// block's live-in set. Processing continues until the live sets stop changing
209/// and we've reached a fixed-point. Each time we process a block, its live sets
210/// can only grow monotonically, and therefore we know that the computation will
211/// reach its fixed-point and terminate. This fixed-point is implemented with a
212/// classic worklist algorithm.
213///
214/// The worklist is seeded such that we initially process blocks in post-order,
215/// which ensures that, when we have a loop-free control-flow graph, we only
216/// process each block once. We pop a block off the worklist for
217/// processing. Whenever a block's live-in set is updated during processing, we
218/// push its predecessors onto the worklist so that their live-in sets can be
219/// updated. Once the worklist is empty, there are no more blocks needing
220/// updates, and we've reached the fixed-point.
221///
222/// Note: For simplicity, we do not flow liveness from block parameters back to
223/// branch arguments, and instead always consider branch arguments live.
224///
225/// Furthermore, we do not differentiate between uses of a needs-stack-map value
226/// that ultimately flow into a side-effecting operation versus uses that
227/// themselves are not live. This could be tightened up in the future, but we're
228/// starting with the easiest, simplest thing. It also means that we do not need
229/// `O(all values)` space, only `O(needs-stack-map values)`. Finally, none of
230/// our mid-end optimization passes have run at this point in time yet, so there
231/// probably isn't much, if any, dead code.
232///
233/// After we've computed the live-in and -out sets for each block, we pass once
234/// more over each block, processing its instructions again. This time, we
235/// record the precise set of needs-stack-map values that are live across each
236/// safepoint instruction inside the block, which is the final output of this
237/// analysis.
238pub(crate) struct LivenessAnalysis {
239    /// Reusable depth-first search state for traversing a function's blocks.
240    dfs: Dfs,
241
242    /// The cached post-order traversal of the function's blocks.
243    post_order: Vec<ir::Block>,
244
245    /// A secondary map from each block to its index in `post_order`.
246    block_to_index: SecondaryMap<ir::Block, u32>,
247
248    /// A mapping from each block's post-order index to the post-order indices
249    /// of its direct (non-transitive) predecessors.
250    predecessors: Vec<SmallVec<[u32; 4]>>,
251
252    /// A worklist of blocks to process. Used to determine which blocks need
253    /// updates cascaded to them and when we reach a fixed-point.
254    worklist: Worklist,
255
256    /// A map from a block's post-order index to its live-in set.
257    live_ins: Vec<LiveSet>,
258
259    /// A map from a block's post-order index to its live-out set.
260    live_outs: Vec<LiveSet>,
261
262    /// The set of each needs-stack-map value that is currently live while
263    /// processing a block.
264    currently_live: LiveSet,
265
266    /// A mapping from each safepoint instruction to the set of needs-stack-map
267    /// values that are live across it.
268    safepoints: HashMap<ir::Inst, SmallVec<[ir::Value; 4]>>,
269
270    /// The set of values that are live across *any* safepoint in the function,
271    /// i.e. the union of all the values in the `safepoints` map.
272    live_across_any_safepoint: EntitySet<ir::Value>,
273}
274
275impl Default for LivenessAnalysis {
276    fn default() -> Self {
277        Self {
278            dfs: Default::default(),
279            post_order: Default::default(),
280            block_to_index: SecondaryMap::with_default(u32::MAX),
281            predecessors: Default::default(),
282            worklist: Default::default(),
283            live_ins: Default::default(),
284            live_outs: Default::default(),
285            currently_live: Default::default(),
286            safepoints: Default::default(),
287            live_across_any_safepoint: Default::default(),
288        }
289    }
290}
291
292#[derive(Clone, Copy, PartialEq, Eq)]
293enum RecordSafepoints {
294    Yes,
295    No,
296}
297
298impl LivenessAnalysis {
299    /// Clear and reset all internal state such that this analysis is ready for
300    /// reuse with a new function.
301    pub fn clear(&mut self) {
302        let LivenessAnalysis {
303            dfs,
304            post_order,
305            block_to_index,
306            predecessors,
307            worklist,
308            live_ins,
309            live_outs,
310            currently_live,
311            safepoints,
312            live_across_any_safepoint,
313        } = self;
314        dfs.clear();
315        post_order.clear();
316        block_to_index.clear();
317        predecessors.clear();
318        worklist.clear();
319        live_ins.clear();
320        live_outs.clear();
321        currently_live.clear();
322        safepoints.clear();
323        live_across_any_safepoint.clear();
324    }
325
326    /// Given that we've initialized `self.post_order`, reserve capacity for the
327    /// various data structures we use during our analysis.
328    fn reserve_capacity(&mut self, func: &Function) {
329        let LivenessAnalysis {
330            dfs: _,
331            post_order,
332            block_to_index,
333            predecessors,
334            worklist,
335            live_ins,
336            live_outs,
337            currently_live: _,
338            safepoints: _,
339            live_across_any_safepoint: _,
340        } = self;
341
342        block_to_index.resize(func.dfg.num_blocks());
343
344        let capacity = post_order.len();
345        worklist.reserve(capacity);
346        predecessors.resize(capacity, Default::default());
347        live_ins.resize(capacity, Default::default());
348        live_outs.resize(capacity, Default::default());
349    }
350
351    fn initialize_block_to_index_map(&mut self) {
352        for (block_index, block) in self.post_order.iter().enumerate() {
353            self.block_to_index[*block] = u32::try_from(block_index).unwrap();
354        }
355    }
356
357    fn initialize_predecessors_map(&mut self, func: &Function) {
358        for (block_index, block) in self.post_order.iter().enumerate() {
359            let block_index = u32::try_from(block_index).unwrap();
360            for succ in func.block_successors(*block) {
361                let succ_index = self.block_to_index[succ];
362                debug_assert_ne!(succ_index, u32::MAX);
363                let succ_index = usize::try_from(succ_index).unwrap();
364                self.predecessors[succ_index].push(block_index);
365            }
366        }
367    }
368
369    /// Process a value's definition, removing it from the currently-live set.
370    fn process_def(&mut self, func: &Function, val: ir::Value) {
371        debug_assert!(!func.dfg.value_is_alias(val));
372        if self.currently_live.remove(&val) {
373            log::trace!("liveness:   defining {val:?}, removing it from the live set");
374        }
375    }
376
377    /// Record the live set of needs-stack-map values at the given safepoint.
378    fn record_safepoint(&mut self, func: &Function, inst: Inst) {
379        log::trace!(
380            "liveness:   found safepoint: {inst:?}: {}",
381            func.dfg.display_inst(inst)
382        );
383        log::trace!("liveness:     live set = {:?}", self.currently_live);
384
385        let mut live = self.currently_live.iter().copied().collect::<SmallVec<_>>();
386        // Keep order deterministic since we add stack map entries in this
387        // order.
388        live.sort();
389        debug_assert!(live.iter().all(|v| !func.dfg.value_is_alias(*v)));
390
391        self.live_across_any_safepoint.extend(live.iter().copied());
392        self.safepoints.insert(inst, live);
393    }
394
395    /// Process a use of a needs-stack-map value, inserting it into the
396    /// currently-live set.
397    fn process_use(&mut self, func: &Function, inst: Inst, val: Value) {
398        debug_assert!(!func.dfg.value_is_alias(val));
399        if self.currently_live.insert(val) {
400            log::trace!(
401                "liveness:   found use of {val:?}, marking it live: {inst:?}: {}",
402                func.dfg.display_inst(inst)
403            );
404        }
405    }
406
407    /// Process all the instructions in a block in reverse order.
408    fn process_block(
409        &mut self,
410        func: &mut Function,
411        stack_map_values: &EntitySet<ir::Value>,
412        block_index: usize,
413        record_safepoints: RecordSafepoints,
414    ) {
415        let block = self.post_order[block_index];
416        log::trace!("liveness: traversing {block:?}");
417
418        // Reset the currently-live set to this block's live-out set.
419        self.currently_live.clear();
420        self.currently_live
421            .extend(self.live_outs[block_index].iter().copied());
422
423        // Now process this block's instructions, incrementally building its
424        // live-in set inside the currently-live set.
425        let mut option_inst = func.layout.last_inst(block);
426        while let Some(inst) = option_inst {
427            // 1. Process needs-stack-map block-call args. These are
428            // uses that logically happen *after* the instruction --
429            // this matters for `try_call` in particular which is also
430            // a callsite (safepoint).
431            for block_call in func.dfg.insts[inst]
432                .branch_destination(&func.dfg.jump_tables, &func.dfg.exception_tables)
433            {
434                for val in block_call
435                    .args(&func.dfg.value_lists)
436                    .filter_map(|arg| arg.as_value())
437                {
438                    let val = func.dfg.resolve_aliases(val);
439                    if stack_map_values.contains(val) {
440                        self.process_use(func, inst, val);
441                    }
442                }
443            }
444
445            // 2. Process needs-stack-map defs.
446            for val in func.dfg.inst_results(inst) {
447                self.process_def(func, *val);
448            }
449
450            // 3. If this instruction is a safepoint and we've been
451            // asked to record safepoints, then do so.
452            let opcode = func.dfg.insts[inst].opcode();
453            if record_safepoints == RecordSafepoints::Yes && opcode.is_safepoint() {
454                self.record_safepoint(func, inst);
455            }
456
457            // 4. Process any other needs-stack-map uses.
458            for val in func.dfg.inst_args(inst) {
459                let val = func.dfg.resolve_aliases(*val);
460                if stack_map_values.contains(val) {
461                    self.process_use(func, inst, val);
462                }
463            }
464
465            // We do not support GC refs as exception-context values
466            // (Wasmtime uses its vmctx as the
467            // exception-context). Reject this explicitly; this slot
468            // is not handled by the analysis.
469            if let Some(et) = func.dfg.insts[inst].exception_table() {
470                for item in func.dfg.exception_tables[et].items() {
471                    let ir::ExceptionTableItem::Context(ctx) = item else {
472                        continue;
473                    };
474                    let ctx = func.dfg.resolve_aliases(ctx);
475                    assert!(
476                        !stack_map_values.contains(ctx),
477                        "exception-table context {ctx:?} on {inst:?} must not be a \
478                         needs-stack-map value"
479                    );
480                }
481            }
482
483            option_inst = func.layout.prev_inst(inst);
484        }
485
486        // After we've processed this block's instructions, remove its
487        // parameters from the live set. This is part of step (1).
488        for val in func.dfg.block_params(block) {
489            self.process_def(func, *val);
490        }
491    }
492
493    /// Run the liveness analysis on the given function.
494    pub fn run(&mut self, func: &mut Function, stack_map_values: &EntitySet<ir::Value>) {
495        self.clear();
496        self.post_order.extend(self.dfs.post_order_iter(func));
497        self.reserve_capacity(func);
498        self.initialize_block_to_index_map();
499        self.initialize_predecessors_map(func);
500
501        // Initially enqueue all blocks for processing. We push them in reverse
502        // post-order (which yields them in post-order when popped) because if
503        // there are no back-edges in the control-flow graph, post-order will
504        // result in only a single pass over the blocks.
505        self.worklist
506            .extend((0..u32::try_from(self.post_order.len()).unwrap()).rev());
507
508        // Pump the worklist until we reach a fixed-point.
509        while let Some(block_index) = self.worklist.pop() {
510            let block_index = usize::try_from(block_index).unwrap();
511
512            // Because our live sets grow monotonically, we just need to see if
513            // the size changed to determine whether the whole set changed.
514            let initial_live_in_len = self.live_ins[block_index].len();
515
516            // The live-out set for a block is the union of the live-in sets of
517            // its successors.
518            for successor in func.block_successors(self.post_order[block_index]) {
519                let successor_index = self.block_to_index[successor];
520                debug_assert_ne!(successor_index, u32::MAX);
521                let successor_index = usize::try_from(successor_index).unwrap();
522                self.live_outs[block_index].extend(self.live_ins[successor_index].iter().copied());
523            }
524            debug_assert!(
525                self.live_outs[block_index]
526                    .iter()
527                    .all(|v| !func.dfg.value_is_alias(*v))
528            );
529
530            // Process the block to compute its live-in set, but do not record
531            // safepoints yet, as we haven't yet processed loop back edges (see
532            // below).
533            self.process_block(func, stack_map_values, block_index, RecordSafepoints::No);
534
535            // The live-in set for a block is the set of values that are still
536            // live after the block's instructions have been processed.
537            self.live_ins[block_index].extend(self.currently_live.iter().copied());
538            debug_assert!(
539                self.live_ins[block_index]
540                    .iter()
541                    .all(|v| !func.dfg.value_is_alias(*v))
542            );
543
544            // If the live-in set changed, then we need to revisit all this
545            // block's predecessors.
546            if self.live_ins[block_index].len() != initial_live_in_len {
547                self.worklist
548                    .extend(self.predecessors[block_index].iter().copied());
549            }
550        }
551
552        // Once we've reached a fixed-point, compute the actual live set for
553        // each safepoint instruction in each block, backwards from the block's
554        // live-out set.
555        for block_index in 0..self.post_order.len() {
556            self.process_block(func, stack_map_values, block_index, RecordSafepoints::Yes);
557
558            debug_assert_eq!(
559                self.currently_live, self.live_ins[block_index],
560                "when we recompute the live-in set for a block as part of \
561                 computing live sets at each safepoint, we should get the same \
562                 result we computed in the fixed-point"
563            );
564        }
565    }
566}
567
568/// A mapping from each needs-stack-map value to its associated stack slot.
569///
570/// Internally maintains free lists for stack slots that won't be used again, so
571/// that we can reuse them and minimize the number of stack slots we need to
572/// allocate.
573#[derive(Default)]
574struct StackSlots {
575    /// A mapping from each needs-stack-map value that is live across some
576    /// safepoint to the stack slot that it resides within. Note that if a
577    /// needs-stack-map value is never live across a safepoint, then we won't
578    /// ever add it to this map, it can remain in a virtual register for the
579    /// duration of its lifetime, and we won't replace all its uses with reloads
580    /// and all that stuff.
581    stack_slots: HashMap<ir::Value, ir::StackSlot>,
582
583    /// A map from slot size to free stack slots that are not being used
584    /// anymore. This allows us to reuse stack slots across multiple values
585    /// helps cut down on the ultimate size of our stack frames.
586    free_stack_slots: SlotSizeMap<SmallVec<[ir::StackSlot; 4]>>,
587}
588
589impl StackSlots {
590    fn clear(&mut self) {
591        let StackSlots {
592            stack_slots,
593            free_stack_slots,
594        } = self;
595        stack_slots.clear();
596        free_stack_slots.clear();
597    }
598
599    fn get(&self, val: ir::Value) -> Option<ir::StackSlot> {
600        self.stack_slots.get(&val).copied()
601    }
602
603    fn get_or_create_stack_slot(&mut self, func: &mut Function, val: ir::Value) -> ir::StackSlot {
604        debug_assert!(!func.dfg.value_is_alias(val));
605        *self.stack_slots.entry(val).or_insert_with(|| {
606            log::trace!("rewriting:     {val:?} needs a stack slot");
607            let ty = func.dfg.value_type(val);
608            let size = ty.bytes();
609            match self.free_stack_slots[SlotSize::unwrap_new(size)].pop() {
610                Some(slot) => {
611                    log::trace!("rewriting:       reusing free stack slot {slot:?} for {val:?}");
612                    slot
613                }
614                None => {
615                    debug_assert!(size.is_power_of_two());
616                    let log2_size = size.ilog2();
617                    let slot = func.create_sized_stack_slot(ir::StackSlotData::new(
618                        ir::StackSlotKind::ExplicitSlot,
619                        size,
620                        log2_size.try_into().unwrap(),
621                    ));
622                    log::trace!("rewriting:       created new stack slot {slot:?} for {val:?}");
623                    slot
624                }
625            }
626        })
627    }
628
629    fn free_stack_slot(&mut self, size: SlotSize, slot: ir::StackSlot) {
630        log::trace!("rewriting:     returning {slot:?} to the free list");
631        self.free_stack_slots[size].push(slot);
632    }
633}
634
635/// A pass to rewrite a function's instructions to spill and reload values that
636/// are live across safepoints.
637///
638/// A single `SafepointSpiller` instance may be reused to rewrite many
639/// functions, amortizing the cost of its internal allocations and avoiding
640/// repeated `malloc` and `free` calls.
641#[derive(Default)]
642pub(super) struct SafepointSpiller {
643    liveness: LivenessAnalysis,
644    stack_slots: StackSlots,
645
646    /// Optional embedder callback used to assign an alias region to the loads
647    /// and stores emitted when spilling and reloading values that are live
648    /// across safepoints.
649    pub(super) make_alias_region: Option<
650        Box<
651            dyn Fn(&mut ir::AliasRegionSet, ir::Type, ir::StackSlot, u32) -> Option<ir::AliasRegion>
652                + Send
653                + Sync,
654        >,
655    >,
656}
657
658impl SafepointSpiller {
659    /// Clear and reset all internal state such that this pass is ready to run
660    /// on a new function.
661    pub fn clear(&mut self) {
662        let SafepointSpiller {
663            liveness,
664            stack_slots,
665            make_alias_region: _,
666        } = self;
667        liveness.clear();
668        stack_slots.clear();
669    }
670
671    /// Identify needs-stack-map values that are live across safepoints, and
672    /// rewrite the function's instructions to spill and reload them as
673    /// necessary.
674    pub fn run(
675        &mut self,
676        func: &mut Function,
677        stack_map_values: &EntitySet<ir::Value>,
678        pointer_type: ir::Type,
679    ) {
680        log::trace!("values needing inclusion in stack maps: {stack_map_values:?}");
681        log::trace!(
682            "before inserting safepoint spills and reloads:\n{}",
683            func.display()
684        );
685
686        self.clear();
687        self.liveness.run(func, stack_map_values);
688        self.rewrite(func, pointer_type);
689
690        log::trace!(
691            "after inserting safepoint spills and reloads:\n{}",
692            func.display()
693        );
694    }
695
696    /// Spill this value to a stack slot if it has been declared that it must be
697    /// included in stack maps and is live across any safepoints.
698    ///
699    /// The given cursor must point just after this value's definition.
700    fn rewrite_def(&mut self, pos: &mut FuncCursor<'_>, val: ir::Value, pointer_type: ir::Type) {
701        debug_assert!(!pos.func.dfg.value_is_alias(val));
702        if let Some(slot) = self.stack_slots.get(val) {
703            let ty = pos.func.dfg.value_type(val);
704
705            let mut flags = ir::MemFlagsData::trusted();
706            flags.set_notrap();
707            if let Some(make_alias_region) = &self.make_alias_region {
708                if let Some(region) =
709                    make_alias_region(&mut pos.func.dfg.alias_regions, ty, slot, 0)
710                {
711                    flags.set_alias_region(Some(region));
712                }
713            }
714
715            let addr = pos.ins().stack_addr(pointer_type, slot, 0);
716            let inst = pos.ins().store(flags, val, addr, 0);
717            log::trace!(
718                "rewriting:   spilling {val:?} to {slot:?}: {}",
719                pos.func.dfg.display_inst(inst)
720            );
721
722            // Now that we've defined this value, there cannot be any more uses
723            // of it, and therefore this stack slot is now available for reuse.
724            let size = SlotSize::try_from(ty).unwrap();
725            self.stack_slots.free_stack_slot(size, slot);
726        }
727    }
728
729    /// Add a stack map entry for each needs-stack-map value that is live across
730    /// the given safepoint instruction.
731    ///
732    /// This will additionally assign stack slots to needs-stack-map values, if
733    /// no such assignment has already been made.
734    fn rewrite_safepoint(&mut self, func: &mut Function, inst: ir::Inst) {
735        log::trace!(
736            "rewriting:   found safepoint: {inst:?}: {}",
737            func.dfg.display_inst(inst)
738        );
739
740        let live = self
741            .liveness
742            .safepoints
743            .get(&inst)
744            .expect("should only call `rewrite_safepoint` on safepoint instructions");
745
746        for val in live {
747            debug_assert!(!func.dfg.value_is_alias(*val));
748
749            // Get or create the stack slot for this live needs-stack-map value.
750            let slot = self.stack_slots.get_or_create_stack_slot(func, *val);
751
752            log::trace!(
753                "rewriting:     adding stack map entry for {val:?} at {slot:?}: {}",
754                func.dfg.display_inst(inst)
755            );
756            let ty = func.dfg.value_type(*val);
757            func.dfg.append_user_stack_map_entry(
758                inst,
759                ir::UserStackMapEntry {
760                    ty,
761                    slot,
762                    offset: 0,
763                },
764            );
765        }
766    }
767
768    /// If `val` is a needs-stack-map value that has been spilled to a stack
769    /// slot, then rewrite `val` to be a load from its associated stack
770    /// slot.
771    ///
772    /// Returns `true` if `val` was rewritten, `false` if not.
773    ///
774    /// The given cursor must point just before the use of the value that we are
775    /// replacing.
776    fn rewrite_use(
777        &mut self,
778        pos: &mut FuncCursor<'_>,
779        val: &mut ir::Value,
780        pointer_type: ir::Type,
781    ) -> bool {
782        debug_assert!(!pos.func.dfg.value_is_alias(*val));
783        if !self.liveness.live_across_any_safepoint.contains(*val) {
784            return false;
785        }
786
787        let old_val = *val;
788        log::trace!("rewriting:     found use of {old_val:?}");
789
790        let ty = pos.func.dfg.value_type(*val);
791        let slot = self.stack_slots.get_or_create_stack_slot(pos.func, *val);
792
793        let mut flags = ir::MemFlagsData::trusted();
794        flags.set_notrap();
795        if let Some(make_alias_region) = &self.make_alias_region {
796            if let Some(region) = make_alias_region(&mut pos.func.dfg.alias_regions, ty, slot, 0) {
797                flags.set_alias_region(Some(region));
798            }
799        }
800
801        let addr = pos.ins().stack_addr(pointer_type, slot, 0);
802        *val = pos.ins().load(ty, flags, addr, 0);
803
804        log::trace!(
805            "rewriting:     reloading {old_val:?}: {}",
806            pos.func
807                .dfg
808                .display_inst(pos.func.dfg.value_def(*val).unwrap_inst())
809        );
810
811        true
812    }
813
814    /// Rewrite the outgoing edges of a safepoint instruction that is
815    /// also a branch (i.e. a `try_call` or `try_call_indirect`) such
816    /// that any spilled needs-stack-map values passed as block-call
817    /// arguments are reloaded *after* the safepoint. These uses
818    /// logically occur after the call, not before, so we need to
819    /// split edges to insert these reloads.
820    fn rewrite_safepoint_edges(
821        &mut self,
822        func: &mut Function,
823        block: ir::Block,
824        inst: ir::Inst,
825        pointer_type: ir::Type,
826    ) {
827        let num_dests = func.dfg.insts[inst]
828            .branch_destination(&func.dfg.jump_tables, &func.dfg.exception_tables)
829            .len();
830        let mut insert_after = block;
831
832        for i in 0..num_dests {
833            let block_call = func.dfg.insts[inst]
834                .branch_destination(&func.dfg.jump_tables, &func.dfg.exception_tables)[i];
835            let target = block_call.block(&func.dfg.value_lists);
836            let args: SmallVec<[_; 8]> = block_call.args(&func.dfg.value_lists).collect();
837
838            // Only split this edge if it passes some spilled value.
839            let needs_reload = args.iter().any(|arg| {
840                arg.as_value().is_some_and(|val| {
841                    let val = func.dfg.resolve_aliases(val);
842                    self.liveness.live_across_any_safepoint.contains(val)
843                })
844            });
845            if !needs_reload {
846                continue;
847            }
848
849            let landing = func.dfg.make_block();
850            func.layout.insert_block_after(landing, insert_after);
851            if func.layout.is_cold(target) {
852                func.layout.set_cold(landing);
853            }
854            insert_after = landing;
855            log::trace!(
856                "rewriting:     splitting edge {block:?} -> {target:?} with {landing:?} \
857                 to reload values after {inst:?}"
858            );
859
860            let mut jump_args: SmallVec<[_; 8]> = SmallVec::new();
861            let mut landing_args: SmallVec<[_; 8]> = SmallVec::new();
862            for (j, arg) in args.iter().enumerate() {
863                match *arg {
864                    ir::BlockArg::Value(val) => {
865                        let mut val = func.dfg.resolve_aliases(val);
866                        let mut pos = FuncCursor::new(func).at_bottom(landing);
867                        self.rewrite_use(&mut pos, &mut val, pointer_type);
868                        jump_args.push(ir::BlockArg::Value(val));
869                    }
870                    ir::BlockArg::TryCallRet(_) | ir::BlockArg::TryCallExn(_) => {
871                        let ty = func.dfg.value_type(func.dfg.block_params(target)[j]);
872                        let param = func.dfg.append_block_param(landing, ty);
873                        landing_args.push(*arg);
874                        jump_args.push(ir::BlockArg::Value(param));
875                    }
876                }
877            }
878
879            let mut pos = FuncCursor::new(func).at_bottom(landing);
880            pos.ins().jump(target, &jump_args);
881            let new_block_call =
882                ir::BlockCall::new(landing, landing_args, &mut func.dfg.value_lists);
883            let dfg = &mut func.dfg;
884            dfg.insts[inst]
885                .branch_destination_mut(&mut dfg.jump_tables, &mut dfg.exception_tables)[i] =
886                new_block_call;
887        }
888    }
889
890    /// Rewrite the function's instructions to spill and reload values that are
891    /// live across safepoints:
892    ///
893    /// 1. Definitions of needs-stack-map values that are live across some
894    ///    safepoint need to be spilled to their assigned stack slot.
895    ///
896    /// 2. Instructions that are themselves safepoints must have stack map
897    ///    entries added for the needs-stack-map values that are live across
898    ///    them.
899    ///
900    /// 3. Uses of needs-stack-map values that have been spilled to a stack slot
901    ///    need to be replaced with reloads from the slot.
902    fn rewrite(&mut self, func: &mut Function, pointer_type: ir::Type) {
903        // Shared temporary storage for operand and result lists.
904        let mut vals: SmallVec<[_; 8]> = Default::default();
905
906        // Rewrite the function's instructions in post-order. This ensures that
907        // we rewrite uses before defs, and therefore once we see a def we know
908        // its stack slot will never be used for that value again. Therefore,
909        // the slot can be reappropriated for a new needs-stack-map value with a
910        // non-overlapping live range. See `rewrite_def` and `free_stack_slots`
911        // for more details.
912        for block_index in 0..self.liveness.post_order.len() {
913            let block = self.liveness.post_order[block_index];
914            log::trace!("rewriting: processing {block:?}");
915
916            // Eagerly allocate all stack slots for values that are live-out in
917            // this block. This reserves a loop-invariant value's stack slot
918            // across the whole loop, rather than just at the first use site we
919            // see within the loop, which could otherwise lead to incorrect
920            // stack slot reuse.
921            vals.extend(
922                self.liveness.live_outs[block_index]
923                    .iter()
924                    .copied()
925                    .filter(|val| self.liveness.live_across_any_safepoint.contains(*val)),
926            );
927            vals.sort_unstable();
928            for val in vals.drain(..) {
929                self.stack_slots.get_or_create_stack_slot(func, val);
930            }
931
932            let mut option_inst = func.layout.last_inst(block);
933            while let Some(inst) = option_inst {
934                // If this instruction defines a needs-stack-map value that is
935                // live across a safepoint, then spill the value to its stack
936                // slot.
937                let mut pos = FuncCursor::new(func).after_inst(inst);
938                vals.extend_from_slice(pos.func.dfg.inst_results(inst));
939                for val in vals.drain(..) {
940                    self.rewrite_def(&mut pos, val, pointer_type);
941                }
942
943                // If this instruction is a safepoint, then we must add stack
944                // map entries for the needs-stack-map values that are live
945                // across it.
946                if self.liveness.safepoints.contains_key(&inst) {
947                    self.rewrite_safepoint(func, inst);
948
949                    // If this safepoint is also a branch (i.e. a `try_call`),
950                    // then values it passes along its edges must be reloaded
951                    // *after* the call, on the edge, rather than before it.
952                    self.rewrite_safepoint_edges(func, block, inst, pointer_type);
953                }
954
955                // Replace all uses of needs-stack-map values with loads from
956                // the value's associated stack slot.
957                let mut pos = FuncCursor::new(func).at_inst(inst);
958                vals.extend(pos.func.dfg.inst_values(inst));
959                let mut replaced_any = false;
960                for val in &mut vals {
961                    *val = pos.func.dfg.resolve_aliases(*val);
962                    replaced_any |= self.rewrite_use(&mut pos, val, pointer_type);
963                }
964                if replaced_any {
965                    pos.func.dfg.overwrite_inst_values(inst, vals.drain(..));
966                    log::trace!(
967                        "rewriting:     updated {inst:?} operands with reloaded values: {}",
968                        pos.func.dfg.display_inst(inst)
969                    );
970                } else {
971                    vals.clear();
972                }
973
974                option_inst = func.layout.prev_inst(inst);
975            }
976
977            // Spill needs-stack-map values defined by block parameters to their
978            // associated stack slots.
979            let mut pos = FuncCursor::new(func).at_position(CursorPosition::Before(block));
980            pos.next_inst(); // Advance to the first instruction in the block.
981            vals.clear();
982            vals.extend_from_slice(pos.func.dfg.block_params(block));
983            for val in vals.drain(..) {
984                self.rewrite_def(&mut pos, val, pointer_type);
985            }
986        }
987    }
988}
989
990#[cfg(test)]
991mod tests {
992    use super::*;
993    use alloc::string::ToString;
994    use cranelift_codegen::ir::{BlockCall, ExceptionTableData};
995    use cranelift_codegen::isa::{CallConv, TargetFrontendConfig};
996
997    fn systemv_frontend_config() -> TargetFrontendConfig {
998        TargetFrontendConfig {
999            default_call_conv: CallConv::SystemV,
1000            pointer_width: target_lexicon::PointerWidth::U64,
1001            page_size_align_log2: 12,
1002        }
1003    }
1004
1005    #[test]
1006    fn needs_stack_map_and_loop() {
1007        let mut sig = Signature::new(CallConv::SystemV);
1008        sig.params.push(AbiParam::new(ir::types::I32));
1009        sig.params.push(AbiParam::new(ir::types::I32));
1010
1011        let mut fn_ctx = FunctionBuilderContext::new();
1012        let mut func = Function::with_name_signature(ir::UserFuncName::testcase("sample"), sig);
1013        let mut builder = FunctionBuilder::new(&mut func, &mut fn_ctx);
1014
1015        let name = builder
1016            .func
1017            .declare_imported_user_function(ir::UserExternalName {
1018                namespace: 0,
1019                index: 0,
1020            });
1021        let mut sig = Signature::new(CallConv::SystemV);
1022        sig.params.push(AbiParam::new(ir::types::I32));
1023        let signature = builder.func.import_signature(sig);
1024        let func_ref = builder.import_function(ir::ExtFuncData {
1025            name: ir::ExternalName::user(name),
1026            signature,
1027            colocated: true,
1028            patchable: false,
1029        });
1030
1031        // Here the value `v1` is technically not live but our single-pass liveness
1032        // analysis treats every branch argument to a block as live to avoid
1033        // needing to do a fixed-point loop.
1034        //
1035        //     block0(v0, v1):
1036        //       call $foo(v0)
1037        //       jump block0(v0, v1)
1038        let block0 = builder.create_block();
1039        builder.append_block_params_for_function_params(block0);
1040        let a = builder.func.dfg.block_params(block0)[0];
1041        let b = builder.func.dfg.block_params(block0)[1];
1042        builder.declare_value_needs_stack_map(a);
1043        builder.declare_value_needs_stack_map(b);
1044        builder.switch_to_block(block0);
1045        builder.ins().call(func_ref, &[a]);
1046        builder.ins().jump(block0, &[a.into(), b.into()]);
1047        builder.seal_all_blocks();
1048        builder.finalize(systemv_frontend_config());
1049
1050        assert_eq_output!(
1051            func.display().to_string(),
1052            r#"
1053function %sample(i32, i32) system_v {
1054    ss0 = explicit_slot 4, align = 4
1055    ss1 = explicit_slot 4, align = 4
1056    sig0 = (i32) system_v
1057    fn0 = colocated u0:0 sig0
1058
1059block0(v0: i32, v1: i32):
1060    v8 = stack_addr.i64 ss0
1061    store notrap aligned v0, v8
1062    v9 = stack_addr.i64 ss1
1063    store notrap aligned v1, v9
1064    v6 = stack_addr.i64 ss0
1065    v7 = load.i32 notrap aligned v6
1066    call fn0(v7), stack_map=[i32 @ ss0+0, i32 @ ss1+0]
1067    v2 = stack_addr.i64 ss0
1068    v3 = load.i32 notrap aligned v2
1069    v4 = stack_addr.i64 ss1
1070    v5 = load.i32 notrap aligned v4
1071    jump block0(v3, v5)
1072}
1073            "#
1074        );
1075    }
1076
1077    #[test]
1078    fn needs_stack_map_simple() {
1079        let _ = env_logger::try_init();
1080
1081        let sig = Signature::new(CallConv::SystemV);
1082
1083        let mut fn_ctx = FunctionBuilderContext::new();
1084        let mut func = Function::with_name_signature(ir::UserFuncName::testcase("sample"), sig);
1085        let mut builder = FunctionBuilder::new(&mut func, &mut fn_ctx);
1086
1087        let name = builder
1088            .func
1089            .declare_imported_user_function(ir::UserExternalName {
1090                namespace: 0,
1091                index: 0,
1092            });
1093        let mut sig = Signature::new(CallConv::SystemV);
1094        sig.params.push(AbiParam::new(ir::types::I32));
1095        let signature = builder.func.import_signature(sig);
1096        let func_ref = builder.import_function(ir::ExtFuncData {
1097            name: ir::ExternalName::user(name),
1098            signature,
1099            colocated: true,
1100            patchable: false,
1101        });
1102
1103        // At each `call` we are losing one more value as no longer live, so
1104        // each stack map should be one smaller than the last. `v3` is never
1105        // live across a safepoint, so should never appear in a stack map. Note
1106        // that a value that is an argument to the call, but is not live after
1107        // the call, should not appear in the stack map. This is why `v0`
1108        // appears in the first call's stack map, but not the second call's
1109        // stack map.
1110        //
1111        //     block0:
1112        //       v0 = needs stack map
1113        //       v1 = needs stack map
1114        //       v2 = needs stack map
1115        //       v3 = needs stack map
1116        //       call $foo(v3)
1117        //       call $foo(v0)
1118        //       call $foo(v1)
1119        //       call $foo(v2)
1120        //       return
1121        let block0 = builder.create_block();
1122        builder.append_block_params_for_function_params(block0);
1123        builder.switch_to_block(block0);
1124        let v0 = builder.ins().iconst(ir::types::I32, 0);
1125        builder.declare_value_needs_stack_map(v0);
1126        let v1 = builder.ins().iconst(ir::types::I32, 1);
1127        builder.declare_value_needs_stack_map(v1);
1128        let v2 = builder.ins().iconst(ir::types::I32, 2);
1129        builder.declare_value_needs_stack_map(v2);
1130        let v3 = builder.ins().iconst(ir::types::I32, 3);
1131        builder.declare_value_needs_stack_map(v3);
1132        builder.ins().call(func_ref, &[v3]);
1133        builder.ins().call(func_ref, &[v0]);
1134        builder.ins().call(func_ref, &[v1]);
1135        builder.ins().call(func_ref, &[v2]);
1136        builder.ins().return_(&[]);
1137        builder.seal_all_blocks();
1138        builder.finalize(systemv_frontend_config());
1139
1140        assert_eq_output!(
1141            func.display().to_string(),
1142            r#"
1143function %sample() system_v {
1144    ss0 = explicit_slot 4, align = 4
1145    ss1 = explicit_slot 4, align = 4
1146    ss2 = explicit_slot 4, align = 4
1147    sig0 = (i32) system_v
1148    fn0 = colocated u0:0 sig0
1149
1150block0:
1151    v0 = iconst.i32 0
1152    v12 = stack_addr.i64 ss2
1153    store notrap aligned v0, v12  ; v0 = 0
1154    v1 = iconst.i32 1
1155    v11 = stack_addr.i64 ss1
1156    store notrap aligned v1, v11  ; v1 = 1
1157    v2 = iconst.i32 2
1158    v10 = stack_addr.i64 ss0
1159    store notrap aligned v2, v10  ; v2 = 2
1160    v3 = iconst.i32 3
1161    call fn0(v3), stack_map=[i32 @ ss2+0, i32 @ ss1+0, i32 @ ss0+0]  ; v3 = 3
1162    v8 = stack_addr.i64 ss2
1163    v9 = load.i32 notrap aligned v8
1164    call fn0(v9), stack_map=[i32 @ ss1+0, i32 @ ss0+0]
1165    v6 = stack_addr.i64 ss1
1166    v7 = load.i32 notrap aligned v6
1167    call fn0(v7), stack_map=[i32 @ ss0+0]
1168    v4 = stack_addr.i64 ss0
1169    v5 = load.i32 notrap aligned v4
1170    call fn0(v5)
1171    return
1172}
1173            "#
1174        );
1175    }
1176
1177    #[test]
1178    fn needs_stack_map_and_post_order_early_return() {
1179        let mut sig = Signature::new(CallConv::SystemV);
1180        sig.params.push(AbiParam::new(ir::types::I32));
1181
1182        let mut fn_ctx = FunctionBuilderContext::new();
1183        let mut func = Function::with_name_signature(ir::UserFuncName::testcase("sample"), sig);
1184        let mut builder = FunctionBuilder::new(&mut func, &mut fn_ctx);
1185
1186        let name = builder
1187            .func
1188            .declare_imported_user_function(ir::UserExternalName {
1189                namespace: 0,
1190                index: 0,
1191            });
1192        let signature = builder
1193            .func
1194            .import_signature(Signature::new(CallConv::SystemV));
1195        let func_ref = builder.import_function(ir::ExtFuncData {
1196            name: ir::ExternalName::user(name),
1197            signature,
1198            colocated: true,
1199            patchable: false,
1200        });
1201
1202        // Here we rely on the post-order to make sure that we never visit block
1203        // 4 and add `v1` to our live set, then visit block 2 and add `v1` to
1204        // its stack map even though `v1` is not in scope. Thanksfully, that
1205        // sequence is impossible because it would be an invalid post-order
1206        // traversal. The only valid post-order traversals are [3, 1, 2, 0] and
1207        // [2, 3, 1, 0].
1208        //
1209        //     block0(v0):
1210        //       brif v0, block1, block2
1211        //
1212        //     block1:
1213        //       <stuff>
1214        //       v1 = get some gc ref
1215        //       jump block3
1216        //
1217        //     block2:
1218        //       call $needs_safepoint_accidentally
1219        //       return
1220        //
1221        //     block3:
1222        //       stuff keeping v1 live
1223        //       return
1224        let block0 = builder.create_block();
1225        let block1 = builder.create_block();
1226        let block2 = builder.create_block();
1227        let block3 = builder.create_block();
1228        builder.append_block_params_for_function_params(block0);
1229
1230        builder.switch_to_block(block0);
1231        let v0 = builder.func.dfg.block_params(block0)[0];
1232        builder.ins().brif(v0, block1, &[], block2, &[]);
1233
1234        builder.switch_to_block(block1);
1235        let v1 = builder.ins().iconst(ir::types::I64, 0x12345678);
1236        builder.declare_value_needs_stack_map(v1);
1237        builder.ins().jump(block3, &[]);
1238
1239        builder.switch_to_block(block2);
1240        builder.ins().call(func_ref, &[]);
1241        builder.ins().return_(&[]);
1242
1243        builder.switch_to_block(block3);
1244        // NB: Our simplistic liveness analysis conservatively treats any use of
1245        // a value as keeping it live, regardless if the use has side effects or
1246        // is otherwise itself live, so an `iadd_imm` suffices to keep `v1` live
1247        // here.
1248        builder.ins().iadd_imm_s(v1, 0);
1249        builder.ins().return_(&[]);
1250
1251        builder.seal_all_blocks();
1252        builder.finalize(systemv_frontend_config());
1253
1254        assert_eq_output!(
1255            func.display().to_string(),
1256            r#"
1257function %sample(i32) system_v {
1258    sig0 = () system_v
1259    fn0 = colocated u0:0 sig0
1260
1261block0(v0: i32):
1262    brif v0, block1, block2
1263
1264block1:
1265    v1 = iconst.i64 0x1234_5678
1266    jump block3
1267
1268block2:
1269    call fn0()
1270    return
1271
1272block3:
1273    v2 = iconst.i64 0
1274    v3 = iadd.i64 v1, v2  ; v1 = 0x1234_5678, v2 = 0
1275    return
1276}
1277            "#
1278        );
1279    }
1280
1281    #[test]
1282    fn needs_stack_map_conditional_branches_and_liveness() {
1283        let mut sig = Signature::new(CallConv::SystemV);
1284        sig.params.push(AbiParam::new(ir::types::I32));
1285
1286        let mut fn_ctx = FunctionBuilderContext::new();
1287        let mut func = Function::with_name_signature(ir::UserFuncName::testcase("sample"), sig);
1288        let mut builder = FunctionBuilder::new(&mut func, &mut fn_ctx);
1289
1290        let name = builder
1291            .func
1292            .declare_imported_user_function(ir::UserExternalName {
1293                namespace: 0,
1294                index: 0,
1295            });
1296        let signature = builder
1297            .func
1298            .import_signature(Signature::new(CallConv::SystemV));
1299        let func_ref = builder.import_function(ir::ExtFuncData {
1300            name: ir::ExternalName::user(name),
1301            signature,
1302            colocated: true,
1303            patchable: false,
1304        });
1305
1306        // We should not have a stack map entry for `v1` in block 1 because it
1307        // is not live across the call.
1308        //
1309        //     block0(v0):
1310        //       v1 = needs stack map
1311        //       brif v0, block1, block2
1312        //
1313        //     block1:
1314        //       call $foo()
1315        //       return
1316        //
1317        //     block2:
1318        //       keep v1 alive
1319        //       return
1320        let block0 = builder.create_block();
1321        let block1 = builder.create_block();
1322        let block2 = builder.create_block();
1323        builder.append_block_params_for_function_params(block0);
1324
1325        builder.switch_to_block(block0);
1326        let v0 = builder.func.dfg.block_params(block0)[0];
1327        let v1 = builder.ins().iconst(ir::types::I64, 0x12345678);
1328        builder.declare_value_needs_stack_map(v1);
1329        builder.ins().brif(v0, block1, &[], block2, &[]);
1330
1331        builder.switch_to_block(block1);
1332        builder.ins().call(func_ref, &[]);
1333        builder.ins().return_(&[]);
1334
1335        builder.switch_to_block(block2);
1336        // NB: Our simplistic liveness analysis conservatively treats any use of
1337        // a value as keeping it live, regardless if the use has side effects or
1338        // is otherwise itself live, so an `iadd_imm` suffices to keep `v1` live
1339        // here.
1340        builder.ins().iadd_imm_s(v1, 0);
1341        builder.ins().return_(&[]);
1342
1343        builder.seal_all_blocks();
1344        builder.finalize(systemv_frontend_config());
1345
1346        assert_eq_output!(
1347            func.display().to_string(),
1348            r#"
1349function %sample(i32) system_v {
1350    sig0 = () system_v
1351    fn0 = colocated u0:0 sig0
1352
1353block0(v0: i32):
1354    v1 = iconst.i64 0x1234_5678
1355    brif v0, block1, block2
1356
1357block1:
1358    call fn0()
1359    return
1360
1361block2:
1362    v2 = iconst.i64 0
1363    v3 = iadd.i64 v1, v2  ; v1 = 0x1234_5678, v2 = 0
1364    return
1365}
1366            "#
1367        );
1368
1369        // Now do the same test but with block 1 and 2 swapped so that we
1370        // exercise what we are trying to exercise, regardless of which
1371        // post-order traversal we happen to take.
1372        func.clear();
1373        fn_ctx.clear();
1374
1375        let mut sig = Signature::new(CallConv::SystemV);
1376        sig.params.push(AbiParam::new(ir::types::I32));
1377
1378        func.signature = sig;
1379        let mut builder = FunctionBuilder::new(&mut func, &mut fn_ctx);
1380
1381        let name = builder
1382            .func
1383            .declare_imported_user_function(ir::UserExternalName {
1384                namespace: 0,
1385                index: 0,
1386            });
1387        let signature = builder
1388            .func
1389            .import_signature(Signature::new(CallConv::SystemV));
1390        let func_ref = builder.import_function(ir::ExtFuncData {
1391            name: ir::ExternalName::user(name),
1392            signature,
1393            colocated: true,
1394            patchable: false,
1395        });
1396
1397        let block0 = builder.create_block();
1398        let block1 = builder.create_block();
1399        let block2 = builder.create_block();
1400        builder.append_block_params_for_function_params(block0);
1401
1402        builder.switch_to_block(block0);
1403        let v0 = builder.func.dfg.block_params(block0)[0];
1404        let v1 = builder.ins().iconst(ir::types::I64, 0x12345678);
1405        builder.declare_value_needs_stack_map(v1);
1406        builder.ins().brif(v0, block1, &[], block2, &[]);
1407
1408        builder.switch_to_block(block1);
1409        builder.ins().iadd_imm_s(v1, 0);
1410        builder.ins().return_(&[]);
1411
1412        builder.switch_to_block(block2);
1413        builder.ins().call(func_ref, &[]);
1414        builder.ins().return_(&[]);
1415
1416        builder.seal_all_blocks();
1417        builder.finalize(systemv_frontend_config());
1418
1419        assert_eq_output!(
1420            func.display().to_string(),
1421            r#"
1422function u0:0(i32) system_v {
1423    sig0 = () system_v
1424    fn0 = colocated u0:0 sig0
1425
1426block0(v0: i32):
1427    v1 = iconst.i64 0x1234_5678
1428    brif v0, block1, block2
1429
1430block1:
1431    v2 = iconst.i64 0
1432    v3 = iadd.i64 v1, v2  ; v1 = 0x1234_5678, v2 = 0
1433    return
1434
1435block2:
1436    call fn0()
1437    return
1438}
1439            "#
1440        );
1441    }
1442
1443    #[test]
1444    fn needs_stack_map_and_tail_calls() {
1445        let mut sig = Signature::new(CallConv::SystemV);
1446        sig.params.push(AbiParam::new(ir::types::I32));
1447
1448        let mut fn_ctx = FunctionBuilderContext::new();
1449        let mut func = Function::with_name_signature(ir::UserFuncName::testcase("sample"), sig);
1450        let mut builder = FunctionBuilder::new(&mut func, &mut fn_ctx);
1451
1452        let name = builder
1453            .func
1454            .declare_imported_user_function(ir::UserExternalName {
1455                namespace: 0,
1456                index: 0,
1457            });
1458        let signature = builder
1459            .func
1460            .import_signature(Signature::new(CallConv::SystemV));
1461        let func_ref = builder.import_function(ir::ExtFuncData {
1462            name: ir::ExternalName::user(name),
1463            signature,
1464            colocated: true,
1465            patchable: false,
1466        });
1467
1468        // Depending on which post-order traversal we take, we might consider
1469        // `v1` live inside `block1`. But nothing is live after a tail call so
1470        // we shouldn't spill `v1` either way here.
1471        //
1472        //     block0(v0):
1473        //       v1 = needs stack map
1474        //       brif v0, block1, block2
1475        //
1476        //     block1:
1477        //       return_call $foo()
1478        //
1479        //     block2:
1480        //       keep v1 alive
1481        //       return
1482        let block0 = builder.create_block();
1483        let block1 = builder.create_block();
1484        let block2 = builder.create_block();
1485        builder.append_block_params_for_function_params(block0);
1486
1487        builder.switch_to_block(block0);
1488        let v0 = builder.func.dfg.block_params(block0)[0];
1489        let v1 = builder.ins().iconst(ir::types::I64, 0x12345678);
1490        builder.declare_value_needs_stack_map(v1);
1491        builder.ins().brif(v0, block1, &[], block2, &[]);
1492
1493        builder.switch_to_block(block1);
1494        builder.ins().return_call(func_ref, &[]);
1495
1496        builder.switch_to_block(block2);
1497        // NB: Our simplistic liveness analysis conservatively treats any use of
1498        // a value as keeping it live, regardless if the use has side effects or
1499        // is otherwise itself live, so an `iadd_imm` suffices to keep `v1` live
1500        // here.
1501        builder.ins().iadd_imm_s(v1, 0);
1502        builder.ins().return_(&[]);
1503
1504        builder.seal_all_blocks();
1505        builder.finalize(systemv_frontend_config());
1506
1507        assert_eq_output!(
1508            func.display().to_string(),
1509            r#"
1510function %sample(i32) system_v {
1511    sig0 = () system_v
1512    fn0 = colocated u0:0 sig0
1513
1514block0(v0: i32):
1515    v1 = iconst.i64 0x1234_5678
1516    brif v0, block1, block2
1517
1518block1:
1519    return_call fn0()
1520
1521block2:
1522    v2 = iconst.i64 0
1523    v3 = iadd.i64 v1, v2  ; v1 = 0x1234_5678, v2 = 0
1524    return
1525}
1526            "#
1527        );
1528
1529        // Do the same test but with block 1 and 2 swapped so that we exercise
1530        // what we are trying to exercise, regardless of which post-order
1531        // traversal we happen to take.
1532        func.clear();
1533        fn_ctx.clear();
1534
1535        let mut sig = Signature::new(CallConv::SystemV);
1536        sig.params.push(AbiParam::new(ir::types::I32));
1537        func.signature = sig;
1538
1539        let mut builder = FunctionBuilder::new(&mut func, &mut fn_ctx);
1540
1541        let name = builder
1542            .func
1543            .declare_imported_user_function(ir::UserExternalName {
1544                namespace: 0,
1545                index: 0,
1546            });
1547        let signature = builder
1548            .func
1549            .import_signature(Signature::new(CallConv::SystemV));
1550        let func_ref = builder.import_function(ir::ExtFuncData {
1551            name: ir::ExternalName::user(name),
1552            signature,
1553            colocated: true,
1554            patchable: false,
1555        });
1556
1557        let block0 = builder.create_block();
1558        let block1 = builder.create_block();
1559        let block2 = builder.create_block();
1560        builder.append_block_params_for_function_params(block0);
1561
1562        builder.switch_to_block(block0);
1563        let v0 = builder.func.dfg.block_params(block0)[0];
1564        let v1 = builder.ins().iconst(ir::types::I64, 0x12345678);
1565        builder.declare_value_needs_stack_map(v1);
1566        builder.ins().brif(v0, block1, &[], block2, &[]);
1567
1568        builder.switch_to_block(block1);
1569        builder.ins().iadd_imm_s(v1, 0);
1570        builder.ins().return_(&[]);
1571
1572        builder.switch_to_block(block2);
1573        builder.ins().return_call(func_ref, &[]);
1574
1575        builder.seal_all_blocks();
1576        builder.finalize(systemv_frontend_config());
1577
1578        assert_eq_output!(
1579            func.display().to_string(),
1580            r#"
1581function u0:0(i32) system_v {
1582    sig0 = () system_v
1583    fn0 = colocated u0:0 sig0
1584
1585block0(v0: i32):
1586    v1 = iconst.i64 0x1234_5678
1587    brif v0, block1, block2
1588
1589block1:
1590    v2 = iconst.i64 0
1591    v3 = iadd.i64 v1, v2  ; v1 = 0x1234_5678, v2 = 0
1592    return
1593
1594block2:
1595    return_call fn0()
1596}
1597            "#
1598        );
1599    }
1600
1601    #[test]
1602    fn needs_stack_map_and_cfg_diamond() {
1603        let mut sig = Signature::new(CallConv::SystemV);
1604        sig.params.push(AbiParam::new(ir::types::I32));
1605
1606        let mut fn_ctx = FunctionBuilderContext::new();
1607        let mut func = Function::with_name_signature(ir::UserFuncName::testcase("sample"), sig);
1608        let mut builder = FunctionBuilder::new(&mut func, &mut fn_ctx);
1609
1610        let name = builder
1611            .func
1612            .declare_imported_user_function(ir::UserExternalName {
1613                namespace: 0,
1614                index: 0,
1615            });
1616        let signature = builder
1617            .func
1618            .import_signature(Signature::new(CallConv::SystemV));
1619        let func_ref = builder.import_function(ir::ExtFuncData {
1620            name: ir::ExternalName::user(name),
1621            signature,
1622            colocated: true,
1623            patchable: false,
1624        });
1625
1626        // Create an if/else CFG diamond that and check that various things get
1627        // spilled as needed.
1628        //
1629        //     block0(v0):
1630        //       brif v0, block1, block2
1631        //
1632        //     block1:
1633        //       v1 = needs stack map
1634        //       v2 = needs stack map
1635        //       call $foo()
1636        //       jump block3(v1, v2)
1637        //
1638        //     block2:
1639        //       v3 = needs stack map
1640        //       v4 = needs stack map
1641        //       call $foo()
1642        //       jump block3(v3, v3)  ;; Note: v4 is not live
1643        //
1644        //     block3(v5, v6):
1645        //       call $foo()
1646        //       keep v5 alive, but not v6
1647        let block0 = builder.create_block();
1648        let block1 = builder.create_block();
1649        let block2 = builder.create_block();
1650        let block3 = builder.create_block();
1651        builder.append_block_params_for_function_params(block0);
1652
1653        builder.switch_to_block(block0);
1654        let v0 = builder.func.dfg.block_params(block0)[0];
1655        builder.ins().brif(v0, block1, &[], block2, &[]);
1656
1657        builder.switch_to_block(block1);
1658        let v1 = builder.ins().iconst(ir::types::I64, 1);
1659        builder.declare_value_needs_stack_map(v1);
1660        let v2 = builder.ins().iconst(ir::types::I64, 2);
1661        builder.declare_value_needs_stack_map(v2);
1662        builder.ins().call(func_ref, &[]);
1663        builder.ins().jump(block3, &[v1.into(), v2.into()]);
1664
1665        builder.switch_to_block(block2);
1666        let v3 = builder.ins().iconst(ir::types::I64, 3);
1667        builder.declare_value_needs_stack_map(v3);
1668        let v4 = builder.ins().iconst(ir::types::I64, 4);
1669        builder.declare_value_needs_stack_map(v4);
1670        builder.ins().call(func_ref, &[]);
1671        builder.ins().jump(block3, &[v3.into(), v3.into()]);
1672
1673        builder.switch_to_block(block3);
1674        builder.append_block_param(block3, ir::types::I64);
1675        builder.append_block_param(block3, ir::types::I64);
1676        builder.ins().call(func_ref, &[]);
1677        // NB: Our simplistic liveness analysis conservatively treats any use of
1678        // a value as keeping it live, regardless if the use has side effects or
1679        // is otherwise itself live, so an `iadd_imm` suffices to keep `v1` live
1680        // here.
1681        builder.ins().iadd_imm_s(v1, 0);
1682        builder.ins().return_(&[]);
1683
1684        builder.seal_all_blocks();
1685        builder.finalize(systemv_frontend_config());
1686
1687        assert_eq_output!(
1688            func.display().to_string(),
1689            r#"
1690function %sample(i32) system_v {
1691    ss0 = explicit_slot 8, align = 8
1692    ss1 = explicit_slot 8, align = 8
1693    sig0 = () system_v
1694    fn0 = colocated u0:0 sig0
1695
1696block0(v0: i32):
1697    brif v0, block1, block2
1698
1699block1:
1700    v1 = iconst.i64 1
1701    v16 = stack_addr.i64 ss0
1702    store notrap aligned v1, v16  ; v1 = 1
1703    v2 = iconst.i64 2
1704    v15 = stack_addr.i64 ss1
1705    store notrap aligned v2, v15  ; v2 = 2
1706    call fn0(), stack_map=[i64 @ ss0+0, i64 @ ss1+0]
1707    v11 = stack_addr.i64 ss0
1708    v12 = load.i64 notrap aligned v11
1709    v13 = stack_addr.i64 ss1
1710    v14 = load.i64 notrap aligned v13
1711    jump block3(v12, v14)
1712
1713block2:
1714    v3 = iconst.i64 3
1715    v21 = stack_addr.i64 ss0
1716    store notrap aligned v3, v21  ; v3 = 3
1717    v4 = iconst.i64 4
1718    call fn0(), stack_map=[i64 @ ss0+0, i64 @ ss0+0]
1719    v17 = stack_addr.i64 ss0
1720    v18 = load.i64 notrap aligned v17
1721    v19 = stack_addr.i64 ss0
1722    v20 = load.i64 notrap aligned v19
1723    jump block3(v18, v20)
1724
1725block3(v5: i64, v6: i64):
1726    call fn0(), stack_map=[i64 @ ss0+0]
1727    v7 = iconst.i64 0
1728    v9 = stack_addr.i64 ss0
1729    v10 = load.i64 notrap aligned v9
1730    v8 = iadd v10, v7  ; v7 = 0
1731    return
1732}
1733            "#
1734        );
1735    }
1736
1737    #[test]
1738    fn needs_stack_map_and_heterogeneous_types() {
1739        let _ = env_logger::try_init();
1740
1741        let mut sig = Signature::new(CallConv::SystemV);
1742        for ty in [
1743            ir::types::I8,
1744            ir::types::I16,
1745            ir::types::I32,
1746            ir::types::I64,
1747            ir::types::I128,
1748            ir::types::F32,
1749            ir::types::F64,
1750            ir::types::I8X16,
1751            ir::types::I16X8,
1752        ] {
1753            sig.params.push(AbiParam::new(ty));
1754            sig.returns.push(AbiParam::new(ty));
1755        }
1756
1757        let mut fn_ctx = FunctionBuilderContext::new();
1758        let mut func = Function::with_name_signature(ir::UserFuncName::testcase("sample"), sig);
1759        let mut builder = FunctionBuilder::new(&mut func, &mut fn_ctx);
1760
1761        let name = builder
1762            .func
1763            .declare_imported_user_function(ir::UserExternalName {
1764                namespace: 0,
1765                index: 0,
1766            });
1767        let signature = builder
1768            .func
1769            .import_signature(Signature::new(CallConv::SystemV));
1770        let func_ref = builder.import_function(ir::ExtFuncData {
1771            name: ir::ExternalName::user(name),
1772            signature,
1773            colocated: true,
1774            patchable: false,
1775        });
1776
1777        // Test that we support stack maps of heterogeneous types and properly
1778        // coalesce types into stack slots based on their size.
1779        //
1780        //     block0(v0, v1, v2, v3, v4, v5, v6, v7, v8):
1781        //       call $foo()
1782        //       return v0, v1, v2, v3, v4, v5, v6, v7, v8
1783        let block0 = builder.create_block();
1784        builder.append_block_params_for_function_params(block0);
1785
1786        builder.switch_to_block(block0);
1787        let params = builder.func.dfg.block_params(block0).to_vec();
1788        for val in &params {
1789            builder.declare_value_needs_stack_map(*val);
1790        }
1791        builder.ins().call(func_ref, &[]);
1792        builder.ins().return_(&params);
1793
1794        builder.seal_all_blocks();
1795        builder.finalize(systemv_frontend_config());
1796
1797        assert_eq_output!(
1798            func.display().to_string(),
1799            r#"
1800function %sample(i8, i16, i32, i64, i128, f32, f64, i8x16, i16x8) -> i8, i16, i32, i64, i128, f32, f64, i8x16, i16x8 system_v {
1801    ss0 = explicit_slot 1
1802    ss1 = explicit_slot 2, align = 2
1803    ss2 = explicit_slot 4, align = 4
1804    ss3 = explicit_slot 8, align = 8
1805    ss4 = explicit_slot 16, align = 16
1806    ss5 = explicit_slot 4, align = 4
1807    ss6 = explicit_slot 8, align = 8
1808    ss7 = explicit_slot 16, align = 16
1809    ss8 = explicit_slot 16, align = 16
1810    sig0 = () system_v
1811    fn0 = colocated u0:0 sig0
1812
1813block0(v0: i8, v1: i16, v2: i32, v3: i64, v4: i128, v5: f32, v6: f64, v7: i8x16, v8: i16x8):
1814    v27 = stack_addr.i64 ss0
1815    store notrap aligned v0, v27
1816    v28 = stack_addr.i64 ss1
1817    store notrap aligned v1, v28
1818    v29 = stack_addr.i64 ss2
1819    store notrap aligned v2, v29
1820    v30 = stack_addr.i64 ss3
1821    store notrap aligned v3, v30
1822    v31 = stack_addr.i64 ss4
1823    store notrap aligned v4, v31
1824    v32 = stack_addr.i64 ss5
1825    store notrap aligned v5, v32
1826    v33 = stack_addr.i64 ss6
1827    store notrap aligned v6, v33
1828    v34 = stack_addr.i64 ss7
1829    store notrap aligned v7, v34
1830    v35 = stack_addr.i64 ss8
1831    store notrap aligned v8, v35
1832    call fn0(), stack_map=[i8 @ ss0+0, i16 @ ss1+0, i32 @ ss2+0, i64 @ ss3+0, i128 @ ss4+0, f32 @ ss5+0, f64 @ ss6+0, i8x16 @ ss7+0, i16x8 @ ss8+0]
1833    v9 = stack_addr.i64 ss0
1834    v10 = load.i8 notrap aligned v9
1835    v11 = stack_addr.i64 ss1
1836    v12 = load.i16 notrap aligned v11
1837    v13 = stack_addr.i64 ss2
1838    v14 = load.i32 notrap aligned v13
1839    v15 = stack_addr.i64 ss3
1840    v16 = load.i64 notrap aligned v15
1841    v17 = stack_addr.i64 ss4
1842    v18 = load.i128 notrap aligned v17
1843    v19 = stack_addr.i64 ss5
1844    v20 = load.f32 notrap aligned v19
1845    v21 = stack_addr.i64 ss6
1846    v22 = load.f64 notrap aligned v21
1847    v23 = stack_addr.i64 ss7
1848    v24 = load.i8x16 notrap aligned v23
1849    v25 = stack_addr.i64 ss8
1850    v26 = load.i16x8 notrap aligned v25
1851    return v10, v12, v14, v16, v18, v20, v22, v24, v26
1852}
1853            "#
1854        );
1855    }
1856
1857    #[test]
1858    fn series_of_non_overlapping_live_ranges_needs_stack_map() {
1859        let sig = Signature::new(CallConv::SystemV);
1860
1861        let mut fn_ctx = FunctionBuilderContext::new();
1862        let mut func = Function::with_name_signature(ir::UserFuncName::testcase("sample"), sig);
1863        let mut builder = FunctionBuilder::new(&mut func, &mut fn_ctx);
1864
1865        let name = builder
1866            .func
1867            .declare_imported_user_function(ir::UserExternalName {
1868                namespace: 0,
1869                index: 0,
1870            });
1871        let signature = builder
1872            .func
1873            .import_signature(Signature::new(CallConv::SystemV));
1874        let foo_func_ref = builder.import_function(ir::ExtFuncData {
1875            name: ir::ExternalName::user(name),
1876            signature,
1877            colocated: true,
1878            patchable: false,
1879        });
1880
1881        let name = builder
1882            .func
1883            .declare_imported_user_function(ir::UserExternalName {
1884                namespace: 0,
1885                index: 1,
1886            });
1887        let mut sig = Signature::new(CallConv::SystemV);
1888        sig.params.push(AbiParam::new(ir::types::I32));
1889        let signature = builder.func.import_signature(sig);
1890        let consume_func_ref = builder.import_function(ir::ExtFuncData {
1891            name: ir::ExternalName::user(name),
1892            signature,
1893            colocated: true,
1894            patchable: false,
1895        });
1896
1897        // Create a series of needs-stack-map values that do not have
1898        // overlapping live ranges, but which do appear in stack maps for calls
1899        // to `$foo`:
1900        //
1901        //     block0:
1902        //       v0 = needs stack map
1903        //       call $foo()
1904        //       call consume(v0)
1905        //       v1 = needs stack map
1906        //       call $foo()
1907        //       call consume(v1)
1908        //       v2 = needs stack map
1909        //       call $foo()
1910        //       call consume(v2)
1911        //       v3 = needs stack map
1912        //       call $foo()
1913        //       call consume(v3)
1914        //       return
1915        let block0 = builder.create_block();
1916        builder.append_block_params_for_function_params(block0);
1917        builder.switch_to_block(block0);
1918        let v0 = builder.ins().iconst(ir::types::I32, 0);
1919        builder.declare_value_needs_stack_map(v0);
1920        builder.ins().call(foo_func_ref, &[]);
1921        builder.ins().call(consume_func_ref, &[v0]);
1922        let v1 = builder.ins().iconst(ir::types::I32, 1);
1923        builder.declare_value_needs_stack_map(v1);
1924        builder.ins().call(foo_func_ref, &[]);
1925        builder.ins().call(consume_func_ref, &[v1]);
1926        let v2 = builder.ins().iconst(ir::types::I32, 2);
1927        builder.declare_value_needs_stack_map(v2);
1928        builder.ins().call(foo_func_ref, &[]);
1929        builder.ins().call(consume_func_ref, &[v2]);
1930        let v3 = builder.ins().iconst(ir::types::I32, 3);
1931        builder.declare_value_needs_stack_map(v3);
1932        builder.ins().call(foo_func_ref, &[]);
1933        builder.ins().call(consume_func_ref, &[v3]);
1934        builder.ins().return_(&[]);
1935        builder.seal_all_blocks();
1936        builder.finalize(systemv_frontend_config());
1937
1938        assert_eq_output!(
1939            func.display().to_string(),
1940            r#"
1941function %sample() system_v {
1942    ss0 = explicit_slot 4, align = 4
1943    sig0 = () system_v
1944    sig1 = (i32) system_v
1945    fn0 = colocated u0:0 sig0
1946    fn1 = colocated u0:1 sig1
1947
1948block0:
1949    v0 = iconst.i32 0
1950    v15 = stack_addr.i64 ss0
1951    store notrap aligned v0, v15  ; v0 = 0
1952    call fn0(), stack_map=[i32 @ ss0+0]
1953    v13 = stack_addr.i64 ss0
1954    v14 = load.i32 notrap aligned v13
1955    call fn1(v14)
1956    v1 = iconst.i32 1
1957    v12 = stack_addr.i64 ss0
1958    store notrap aligned v1, v12  ; v1 = 1
1959    call fn0(), stack_map=[i32 @ ss0+0]
1960    v10 = stack_addr.i64 ss0
1961    v11 = load.i32 notrap aligned v10
1962    call fn1(v11)
1963    v2 = iconst.i32 2
1964    v9 = stack_addr.i64 ss0
1965    store notrap aligned v2, v9  ; v2 = 2
1966    call fn0(), stack_map=[i32 @ ss0+0]
1967    v7 = stack_addr.i64 ss0
1968    v8 = load.i32 notrap aligned v7
1969    call fn1(v8)
1970    v3 = iconst.i32 3
1971    v6 = stack_addr.i64 ss0
1972    store notrap aligned v3, v6  ; v3 = 3
1973    call fn0(), stack_map=[i32 @ ss0+0]
1974    v4 = stack_addr.i64 ss0
1975    v5 = load.i32 notrap aligned v4
1976    call fn1(v5)
1977    return
1978}
1979            "#
1980        );
1981    }
1982
1983    #[test]
1984    fn vars_block_params_and_needs_stack_map() {
1985        let _ = env_logger::try_init();
1986
1987        let mut sig = Signature::new(CallConv::SystemV);
1988        sig.params.push(AbiParam::new(ir::types::I32));
1989        sig.returns.push(AbiParam::new(ir::types::I32));
1990
1991        let mut fn_ctx = FunctionBuilderContext::new();
1992        let mut func = Function::with_name_signature(ir::UserFuncName::testcase("sample"), sig);
1993        let mut builder = FunctionBuilder::new(&mut func, &mut fn_ctx);
1994
1995        let name = builder
1996            .func
1997            .declare_imported_user_function(ir::UserExternalName {
1998                namespace: 0,
1999                index: 0,
2000            });
2001        let mut sig = Signature::new(CallConv::SystemV);
2002        sig.params.push(AbiParam::new(ir::types::I32));
2003        let signature = builder.func.import_signature(sig);
2004        let func_ref = builder.import_function(ir::ExtFuncData {
2005            name: ir::ExternalName::user(name),
2006            signature,
2007            colocated: true,
2008            patchable: false,
2009        });
2010
2011        // Use a variable, create a control flow diamond so that the variable
2012        // forces a block parameter on the control join point, and make sure
2013        // that we get stack maps for all the appropriate uses of the variable
2014        // in all blocks, as well as that we are reusing stack slots for each of
2015        // the values.
2016        //
2017        //                        block0:
2018        //                          x := needs stack map
2019        //                          call $foo(x)
2020        //                          br_if v0, block1, block2
2021        //
2022        //
2023        //     block1:                                     block2:
2024        //       call $foo(x)                                call $foo(x)
2025        //       call $foo(x)                                call $foo(x)
2026        //       x := new needs stack map                    x := new needs stack map
2027        //       call $foo(x)                                call $foo(x)
2028        //       jump block3                                 jump block3
2029        //
2030        //
2031        //                        block3:
2032        //                          call $foo(x)
2033        //                          return x
2034
2035        let x = builder.declare_var(ir::types::I32);
2036        builder.declare_var_needs_stack_map(x);
2037
2038        let block0 = builder.create_block();
2039        let block1 = builder.create_block();
2040        let block2 = builder.create_block();
2041        let block3 = builder.create_block();
2042
2043        builder.append_block_params_for_function_params(block0);
2044        builder.switch_to_block(block0);
2045        let v0 = builder.func.dfg.block_params(block0)[0];
2046        let val = builder.ins().iconst(ir::types::I32, 42);
2047        builder.def_var(x, val);
2048        {
2049            let x = builder.use_var(x);
2050            builder.ins().call(func_ref, &[x]);
2051        }
2052        builder.ins().brif(v0, block1, &[], block2, &[]);
2053
2054        builder.switch_to_block(block1);
2055        {
2056            let x = builder.use_var(x);
2057            builder.ins().call(func_ref, &[x]);
2058            builder.ins().call(func_ref, &[x]);
2059        }
2060        let val = builder.ins().iconst(ir::types::I32, 36);
2061        builder.def_var(x, val);
2062        {
2063            let x = builder.use_var(x);
2064            builder.ins().call(func_ref, &[x]);
2065        }
2066        builder.ins().jump(block3, &[]);
2067
2068        builder.switch_to_block(block2);
2069        {
2070            let x = builder.use_var(x);
2071            builder.ins().call(func_ref, &[x]);
2072            builder.ins().call(func_ref, &[x]);
2073        }
2074        let val = builder.ins().iconst(ir::types::I32, 36);
2075        builder.def_var(x, val);
2076        {
2077            let x = builder.use_var(x);
2078            builder.ins().call(func_ref, &[x]);
2079        }
2080        builder.ins().jump(block3, &[]);
2081
2082        builder.switch_to_block(block3);
2083        let x = builder.use_var(x);
2084        builder.ins().call(func_ref, &[x]);
2085        builder.ins().return_(&[x]);
2086
2087        builder.seal_all_blocks();
2088        builder.finalize(systemv_frontend_config());
2089
2090        assert_eq_output!(
2091            func.display().to_string(),
2092            r#"
2093function %sample(i32) -> i32 system_v {
2094    ss0 = explicit_slot 4, align = 4
2095    ss1 = explicit_slot 4, align = 4
2096    sig0 = (i32) system_v
2097    fn0 = colocated u0:0 sig0
2098
2099block0(v0: i32):
2100    v1 = iconst.i32 42
2101    v2 -> v1
2102    v4 -> v1
2103    v32 = stack_addr.i64 ss0
2104    store notrap aligned v1, v32  ; v1 = 42
2105    v30 = stack_addr.i64 ss0
2106    v31 = load.i32 notrap aligned v30
2107    call fn0(v31), stack_map=[i32 @ ss0+0]
2108    brif v0, block1, block2
2109
2110block1:
2111    v19 = stack_addr.i64 ss0
2112    v20 = load.i32 notrap aligned v19
2113    call fn0(v20), stack_map=[i32 @ ss0+0]
2114    v17 = stack_addr.i64 ss0
2115    v18 = load.i32 notrap aligned v17
2116    call fn0(v18)
2117    v3 = iconst.i32 36
2118    v16 = stack_addr.i64 ss0
2119    store notrap aligned v3, v16  ; v3 = 36
2120    v14 = stack_addr.i64 ss0
2121    v15 = load.i32 notrap aligned v14
2122    call fn0(v15), stack_map=[i32 @ ss0+0]
2123    v12 = stack_addr.i64 ss0
2124    v13 = load.i32 notrap aligned v12
2125    jump block3(v13)
2126
2127block2:
2128    v28 = stack_addr.i64 ss0
2129    v29 = load.i32 notrap aligned v28
2130    call fn0(v29), stack_map=[i32 @ ss0+0]
2131    v26 = stack_addr.i64 ss0
2132    v27 = load.i32 notrap aligned v26
2133    call fn0(v27)
2134    v5 = iconst.i32 36
2135    v25 = stack_addr.i64 ss1
2136    store notrap aligned v5, v25  ; v5 = 36
2137    v23 = stack_addr.i64 ss1
2138    v24 = load.i32 notrap aligned v23
2139    call fn0(v24), stack_map=[i32 @ ss1+0]
2140    v21 = stack_addr.i64 ss1
2141    v22 = load.i32 notrap aligned v21
2142    jump block3(v22)
2143
2144block3(v6: i32):
2145    v11 = stack_addr.i64 ss0
2146    store notrap aligned v6, v11
2147    v9 = stack_addr.i64 ss0
2148    v10 = load.i32 notrap aligned v9
2149    call fn0(v10), stack_map=[i32 @ ss0+0]
2150    v7 = stack_addr.i64 ss0
2151    v8 = load.i32 notrap aligned v7
2152    return v8
2153}
2154            "#
2155        );
2156    }
2157
2158    #[test]
2159    fn var_needs_stack_map() {
2160        let mut sig = Signature::new(CallConv::SystemV);
2161        sig.params
2162            .push(AbiParam::new(cranelift_codegen::ir::types::I32));
2163        sig.returns
2164            .push(AbiParam::new(cranelift_codegen::ir::types::I32));
2165
2166        let mut fn_ctx = FunctionBuilderContext::new();
2167        let mut func = Function::with_name_signature(ir::UserFuncName::testcase("sample"), sig);
2168        let mut builder = FunctionBuilder::new(&mut func, &mut fn_ctx);
2169
2170        let var = builder.declare_var(cranelift_codegen::ir::types::I32);
2171        builder.declare_var_needs_stack_map(var);
2172
2173        let name = builder
2174            .func
2175            .declare_imported_user_function(ir::UserExternalName {
2176                namespace: 0,
2177                index: 0,
2178            });
2179        let signature = builder
2180            .func
2181            .import_signature(Signature::new(CallConv::SystemV));
2182        let func_ref = builder.import_function(ir::ExtFuncData {
2183            name: ir::ExternalName::user(name),
2184            signature,
2185            colocated: true,
2186            patchable: false,
2187        });
2188
2189        let block0 = builder.create_block();
2190        builder.append_block_params_for_function_params(block0);
2191        builder.switch_to_block(block0);
2192
2193        let arg = builder.func.dfg.block_params(block0)[0];
2194        builder.def_var(var, arg);
2195
2196        builder.ins().call(func_ref, &[]);
2197
2198        let val = builder.use_var(var);
2199        builder.ins().return_(&[val]);
2200
2201        builder.seal_all_blocks();
2202        builder.finalize(systemv_frontend_config());
2203
2204        assert_eq_output!(
2205            func.display().to_string(),
2206            r#"
2207function %sample(i32) -> i32 system_v {
2208    ss0 = explicit_slot 4, align = 4
2209    sig0 = () system_v
2210    fn0 = colocated u0:0 sig0
2211
2212block0(v0: i32):
2213    v3 = stack_addr.i64 ss0
2214    store notrap aligned v0, v3
2215    call fn0(), stack_map=[i32 @ ss0+0]
2216    v1 = stack_addr.i64 ss0
2217    v2 = load.i32 notrap aligned v1
2218    return v2
2219}
2220            "#
2221        );
2222    }
2223
2224    #[test]
2225    fn first_inst_defines_needs_stack_map() {
2226        let mut sig = Signature::new(CallConv::SystemV);
2227        sig.params
2228            .push(AbiParam::new(cranelift_codegen::ir::types::I32));
2229        sig.returns
2230            .push(AbiParam::new(cranelift_codegen::ir::types::I32));
2231        sig.returns
2232            .push(AbiParam::new(cranelift_codegen::ir::types::I32));
2233
2234        let mut fn_ctx = FunctionBuilderContext::new();
2235        let mut func = Function::with_name_signature(ir::UserFuncName::testcase("sample"), sig);
2236        let mut builder = FunctionBuilder::new(&mut func, &mut fn_ctx);
2237
2238        let name = builder
2239            .func
2240            .declare_imported_user_function(ir::UserExternalName {
2241                namespace: 0,
2242                index: 0,
2243            });
2244        let signature = builder
2245            .func
2246            .import_signature(Signature::new(CallConv::SystemV));
2247        let func_ref = builder.import_function(ir::ExtFuncData {
2248            name: ir::ExternalName::user(name),
2249            signature,
2250            colocated: true,
2251            patchable: false,
2252        });
2253
2254        // Regression test found via fuzzing in
2255        // https://github.com/bytecodealliance/wasmtime/pull/8941 involving the
2256        // combination of cursor positions after we have block parameters that
2257        // need inclusion in stack maps and when the first instruction in a
2258        // block defines a value that needs inclusion in stack maps.
2259        //
2260        // block0(v0: i32):
2261        //   v1 = iconst.i32 42
2262        //   call $foo()
2263        //   return v0, v1
2264
2265        let block0 = builder.create_block();
2266        builder.append_block_params_for_function_params(block0);
2267        builder.switch_to_block(block0);
2268
2269        let arg = builder.func.dfg.block_params(block0)[0];
2270        builder.declare_value_needs_stack_map(arg);
2271
2272        let val = builder.ins().iconst(ir::types::I32, 42);
2273        builder.declare_value_needs_stack_map(val);
2274
2275        builder.ins().call(func_ref, &[]);
2276
2277        builder.ins().return_(&[arg, val]);
2278
2279        builder.seal_all_blocks();
2280        builder.finalize(systemv_frontend_config());
2281
2282        assert_eq_output!(
2283            func.display().to_string(),
2284            r#"
2285function %sample(i32) -> i32, i32 system_v {
2286    ss0 = explicit_slot 4, align = 4
2287    ss1 = explicit_slot 4, align = 4
2288    sig0 = () system_v
2289    fn0 = colocated u0:0 sig0
2290
2291block0(v0: i32):
2292    v7 = stack_addr.i64 ss0
2293    store notrap aligned v0, v7
2294    v1 = iconst.i32 42
2295    v6 = stack_addr.i64 ss1
2296    store notrap aligned v1, v6  ; v1 = 42
2297    call fn0(), stack_map=[i32 @ ss0+0, i32 @ ss1+0]
2298    v2 = stack_addr.i64 ss0
2299    v3 = load.i32 notrap aligned v2
2300    v4 = stack_addr.i64 ss1
2301    v5 = load.i32 notrap aligned v4
2302    return v3, v5
2303}
2304            "#
2305        );
2306    }
2307
2308    #[test]
2309    fn needs_stack_map_and_loops_and_partially_live_values() {
2310        let _ = env_logger::try_init();
2311
2312        let mut sig = Signature::new(CallConv::SystemV);
2313        sig.params.push(AbiParam::new(ir::types::I32));
2314
2315        let mut fn_ctx = FunctionBuilderContext::new();
2316        let mut func =
2317            Function::with_name_signature(ir::UserFuncName::testcase("sample"), sig.clone());
2318        let mut builder = FunctionBuilder::new(&mut func, &mut fn_ctx);
2319
2320        let name = builder
2321            .func
2322            .declare_imported_user_function(ir::UserExternalName {
2323                namespace: 0,
2324                index: 0,
2325            });
2326        let signature = builder
2327            .func
2328            .import_signature(Signature::new(CallConv::SystemV));
2329        let foo_func_ref = builder.import_function(ir::ExtFuncData {
2330            name: ir::ExternalName::user(name),
2331            signature,
2332            colocated: true,
2333            patchable: false,
2334        });
2335
2336        let name = builder
2337            .func
2338            .declare_imported_user_function(ir::UserExternalName {
2339                namespace: 1,
2340                index: 1,
2341            });
2342        let signature = builder.func.import_signature(sig);
2343        let bar_func_ref = builder.import_function(ir::ExtFuncData {
2344            name: ir::ExternalName::user(name),
2345            signature,
2346            colocated: true,
2347            patchable: false,
2348        });
2349
2350        // Test that we support stack maps in loops and that we properly handle
2351        // value that are only live for part of the loop body on each iteration,
2352        // but are live across the whole loop because they will be used again
2353        // the next iteration. Note that `v0` below, which is a GC value, is not
2354        // live *within a single iteration of the loop* after the call to `bar`,
2355        // but is actually live across the whole loop because it will be used
2356        // again in the *next iteration of the loop*:
2357        //
2358        //     block0(v0: i32):
2359        //       jump block1
2360        //
2361        //     block1:
2362        //       call $foo()
2363        //       call $bar(v0)
2364        //       call $foo()
2365        //       jump block1
2366        let block0 = builder.create_block();
2367        let block1 = builder.create_block();
2368        builder.append_block_params_for_function_params(block0);
2369
2370        builder.switch_to_block(block0);
2371        builder.ins().jump(block1, &[]);
2372
2373        builder.switch_to_block(block1);
2374        let v0 = builder.func.dfg.block_params(block0)[0];
2375        builder.declare_value_needs_stack_map(v0);
2376        builder.ins().call(foo_func_ref, &[]);
2377        builder.ins().call(bar_func_ref, &[v0]);
2378        builder.ins().call(foo_func_ref, &[]);
2379        builder.ins().jump(block1, &[]);
2380
2381        builder.seal_all_blocks();
2382        builder.finalize(systemv_frontend_config());
2383
2384        assert_eq_output!(
2385            func.display().to_string(),
2386            r#"
2387function %sample(i32) system_v {
2388    ss0 = explicit_slot 4, align = 4
2389    sig0 = () system_v
2390    sig1 = (i32) system_v
2391    fn0 = colocated u0:0 sig0
2392    fn1 = colocated u1:1 sig1
2393
2394block0(v0: i32):
2395    v3 = stack_addr.i64 ss0
2396    store notrap aligned v0, v3
2397    jump block1
2398
2399block1:
2400    call fn0(), stack_map=[i32 @ ss0+0]
2401    v1 = stack_addr.i64 ss0
2402    v2 = load.i32 notrap aligned v1
2403    call fn1(v2), stack_map=[i32 @ ss0+0]
2404    call fn0(), stack_map=[i32 @ ss0+0]
2405    jump block1
2406}
2407            "#,
2408        );
2409    }
2410
2411    #[test]
2412    fn needs_stack_map_and_irreducible_loops() {
2413        let _ = env_logger::try_init();
2414
2415        let mut sig = Signature::new(CallConv::SystemV);
2416        sig.params.push(AbiParam::new(ir::types::I32));
2417        sig.params.push(AbiParam::new(ir::types::I32));
2418
2419        let mut fn_ctx = FunctionBuilderContext::new();
2420        let mut func = Function::with_name_signature(ir::UserFuncName::testcase("sample"), sig);
2421        let mut builder = FunctionBuilder::new(&mut func, &mut fn_ctx);
2422
2423        let name = builder
2424            .func
2425            .declare_imported_user_function(ir::UserExternalName {
2426                namespace: 0,
2427                index: 0,
2428            });
2429        let signature = builder
2430            .func
2431            .import_signature(Signature::new(CallConv::SystemV));
2432        let foo_func_ref = builder.import_function(ir::ExtFuncData {
2433            name: ir::ExternalName::user(name),
2434            signature,
2435            colocated: true,
2436            patchable: false,
2437        });
2438
2439        let name = builder
2440            .func
2441            .declare_imported_user_function(ir::UserExternalName {
2442                namespace: 1,
2443                index: 1,
2444            });
2445        let mut sig = Signature::new(CallConv::SystemV);
2446        sig.params.push(AbiParam::new(ir::types::I32));
2447        let signature = builder.func.import_signature(sig);
2448        let bar_func_ref = builder.import_function(ir::ExtFuncData {
2449            name: ir::ExternalName::user(name),
2450            signature,
2451            colocated: true,
2452            patchable: false,
2453        });
2454
2455        // Test an irreducible loop with multiple entry points, both block1 and
2456        // block2, in this case:
2457        //
2458        //     block0(v0: i32, v1: i32):
2459        //       brif v0, block1, block2
2460        //
2461        //     block1:
2462        //       jump block3
2463        //
2464        //     block2:
2465        //       jump block4
2466        //
2467        //     block3:
2468        //       call $foo()
2469        //       call $bar(v1)
2470        //       call $foo()
2471        //       jump block2
2472        //
2473        //     block4:
2474        //       call $foo()
2475        //       call $bar(v1)
2476        //       call $foo()
2477        //       jump block1
2478        let block0 = builder.create_block();
2479        let block1 = builder.create_block();
2480        let block2 = builder.create_block();
2481        let block3 = builder.create_block();
2482        let block4 = builder.create_block();
2483        builder.append_block_params_for_function_params(block0);
2484
2485        builder.switch_to_block(block0);
2486        let v0 = builder.func.dfg.block_params(block0)[0];
2487        let v1 = builder.func.dfg.block_params(block0)[1];
2488        builder.declare_value_needs_stack_map(v1);
2489        builder.ins().brif(v0, block1, &[], block2, &[]);
2490
2491        builder.switch_to_block(block1);
2492        builder.ins().jump(block3, &[]);
2493
2494        builder.switch_to_block(block2);
2495        builder.ins().jump(block4, &[]);
2496
2497        builder.switch_to_block(block3);
2498        builder.ins().call(foo_func_ref, &[]);
2499        builder.ins().call(bar_func_ref, &[v1]);
2500        builder.ins().call(foo_func_ref, &[]);
2501        builder.ins().jump(block2, &[]);
2502
2503        builder.switch_to_block(block4);
2504        builder.ins().call(foo_func_ref, &[]);
2505        builder.ins().call(bar_func_ref, &[v1]);
2506        builder.ins().call(foo_func_ref, &[]);
2507        builder.ins().jump(block1, &[]);
2508
2509        builder.seal_all_blocks();
2510        builder.finalize(systemv_frontend_config());
2511
2512        assert_eq_output!(
2513            func.display().to_string(),
2514            r#"
2515function %sample(i32, i32) system_v {
2516    ss0 = explicit_slot 4, align = 4
2517    sig0 = () system_v
2518    sig1 = (i32) system_v
2519    fn0 = colocated u0:0 sig0
2520    fn1 = colocated u1:1 sig1
2521
2522block0(v0: i32, v1: i32):
2523    v6 = stack_addr.i64 ss0
2524    store notrap aligned v1, v6
2525    brif v0, block1, block2
2526
2527block1:
2528    jump block3
2529
2530block2:
2531    jump block4
2532
2533block3:
2534    call fn0(), stack_map=[i32 @ ss0+0]
2535    v4 = stack_addr.i64 ss0
2536    v5 = load.i32 notrap aligned v4
2537    call fn1(v5), stack_map=[i32 @ ss0+0]
2538    call fn0(), stack_map=[i32 @ ss0+0]
2539    jump block2
2540
2541block4:
2542    call fn0(), stack_map=[i32 @ ss0+0]
2543    v2 = stack_addr.i64 ss0
2544    v3 = load.i32 notrap aligned v2
2545    call fn1(v3), stack_map=[i32 @ ss0+0]
2546    call fn0(), stack_map=[i32 @ ss0+0]
2547    jump block1
2548}
2549            "#,
2550        );
2551    }
2552
2553    #[test]
2554    fn needs_stack_map_and_back_edge_to_back_edge() {
2555        let _ = env_logger::try_init();
2556
2557        let mut sig = Signature::new(CallConv::SystemV);
2558        sig.params.push(AbiParam::new(ir::types::I32));
2559        sig.params.push(AbiParam::new(ir::types::I32));
2560        sig.params.push(AbiParam::new(ir::types::I32));
2561        sig.params.push(AbiParam::new(ir::types::I32));
2562
2563        let mut fn_ctx = FunctionBuilderContext::new();
2564        let mut func = Function::with_name_signature(ir::UserFuncName::testcase("sample"), sig);
2565        let mut builder = FunctionBuilder::new(&mut func, &mut fn_ctx);
2566
2567        let name = builder
2568            .func
2569            .declare_imported_user_function(ir::UserExternalName {
2570                namespace: 0,
2571                index: 0,
2572            });
2573        let signature = builder
2574            .func
2575            .import_signature(Signature::new(CallConv::SystemV));
2576        let foo_func_ref = builder.import_function(ir::ExtFuncData {
2577            name: ir::ExternalName::user(name),
2578            signature,
2579            colocated: true,
2580            patchable: false,
2581        });
2582
2583        let name = builder
2584            .func
2585            .declare_imported_user_function(ir::UserExternalName {
2586                namespace: 1,
2587                index: 1,
2588            });
2589        let mut sig = Signature::new(CallConv::SystemV);
2590        sig.params.push(AbiParam::new(ir::types::I32));
2591        let signature = builder.func.import_signature(sig);
2592        let bar_func_ref = builder.import_function(ir::ExtFuncData {
2593            name: ir::ExternalName::user(name),
2594            signature,
2595            colocated: true,
2596            patchable: false,
2597        });
2598
2599        // Test that we detect the `block1 -> block2 -> block3 -> block2 ->
2600        // block1` loop in our liveness analysis and keep `v{0,1,2}` live across
2601        // the whole loop body.
2602        //
2603        //     block0(v0, v1, v2, v3):
2604        //       jump block1(v3)
2605        //
2606        //     block1(v4):
2607        //       call foo_func_ref()
2608        //       call bar_func_ref(v0)
2609        //       call foo_func_ref()
2610        //       jump block2
2611        //
2612        //     block2:
2613        //       call foo_func_ref()
2614        //       call bar_func_ref(v1)
2615        //       call foo_func_ref()
2616        //       v5 = iadd_imm v4, -1
2617        //       brif v4, block1(v5), block3
2618        //
2619        //     block3:
2620        //       call foo_func_ref()
2621        //       call bar_func_ref(v2)
2622        //       call foo_func_ref()
2623        //       jump block2
2624
2625        let block0 = builder.create_block();
2626        let block1 = builder.create_block();
2627        let block2 = builder.create_block();
2628        let block3 = builder.create_block();
2629
2630        builder.append_block_params_for_function_params(block0);
2631
2632        builder.switch_to_block(block0);
2633
2634        let v0 = builder.func.dfg.block_params(block0)[0];
2635        builder.declare_value_needs_stack_map(v0);
2636        let v1 = builder.func.dfg.block_params(block0)[1];
2637        builder.declare_value_needs_stack_map(v1);
2638        let v2 = builder.func.dfg.block_params(block0)[2];
2639        builder.declare_value_needs_stack_map(v2);
2640        let v3 = builder.func.dfg.block_params(block0)[3];
2641
2642        builder.ins().jump(block1, &[v3.into()]);
2643
2644        builder.switch_to_block(block1);
2645        let v4 = builder.append_block_param(block1, ir::types::I32);
2646        builder.ins().call(foo_func_ref, &[]);
2647        builder.ins().call(bar_func_ref, &[v0]);
2648        builder.ins().call(foo_func_ref, &[]);
2649        builder.ins().jump(block2, &[]);
2650
2651        builder.switch_to_block(block2);
2652        builder.ins().call(foo_func_ref, &[]);
2653        builder.ins().call(bar_func_ref, &[v1]);
2654        builder.ins().call(foo_func_ref, &[]);
2655        let v5 = builder.ins().iadd_imm_s(v4, -1);
2656        builder.ins().brif(v4, block1, &[v5.into()], block3, &[]);
2657
2658        builder.switch_to_block(block3);
2659        builder.ins().call(foo_func_ref, &[]);
2660        builder.ins().call(bar_func_ref, &[v2]);
2661        builder.ins().call(foo_func_ref, &[]);
2662        builder.ins().jump(block2, &[]);
2663
2664        builder.seal_all_blocks();
2665        builder.finalize(systemv_frontend_config());
2666
2667        assert_eq_output!(
2668            func.display().to_string(),
2669            r#"
2670function %sample(i32, i32, i32, i32) system_v {
2671    ss0 = explicit_slot 4, align = 4
2672    ss1 = explicit_slot 4, align = 4
2673    ss2 = explicit_slot 4, align = 4
2674    sig0 = () system_v
2675    sig1 = (i32) system_v
2676    fn0 = colocated u0:0 sig0
2677    fn1 = colocated u1:1 sig1
2678
2679block0(v0: i32, v1: i32, v2: i32, v3: i32):
2680    v13 = stack_addr.i64 ss0
2681    store notrap aligned v0, v13
2682    v14 = stack_addr.i64 ss1
2683    store notrap aligned v1, v14
2684    v15 = stack_addr.i64 ss2
2685    store notrap aligned v2, v15
2686    jump block1(v3)
2687
2688block1(v4: i32):
2689    call fn0(), stack_map=[i32 @ ss0+0, i32 @ ss1+0, i32 @ ss2+0]
2690    v11 = stack_addr.i64 ss0
2691    v12 = load.i32 notrap aligned v11
2692    call fn1(v12), stack_map=[i32 @ ss0+0, i32 @ ss1+0, i32 @ ss2+0]
2693    call fn0(), stack_map=[i32 @ ss0+0, i32 @ ss1+0, i32 @ ss2+0]
2694    jump block2
2695
2696block2:
2697    call fn0(), stack_map=[i32 @ ss0+0, i32 @ ss1+0, i32 @ ss2+0]
2698    v9 = stack_addr.i64 ss1
2699    v10 = load.i32 notrap aligned v9
2700    call fn1(v10), stack_map=[i32 @ ss0+0, i32 @ ss1+0, i32 @ ss2+0]
2701    call fn0(), stack_map=[i32 @ ss0+0, i32 @ ss1+0, i32 @ ss2+0]
2702    v5 = iconst.i32 -1
2703    v6 = iadd.i32 v4, v5  ; v5 = -1
2704    brif.i32 v4, block1(v6), block3
2705
2706block3:
2707    call fn0(), stack_map=[i32 @ ss0+0, i32 @ ss1+0, i32 @ ss2+0]
2708    v7 = stack_addr.i64 ss2
2709    v8 = load.i32 notrap aligned v7
2710    call fn1(v8), stack_map=[i32 @ ss0+0, i32 @ ss1+0, i32 @ ss2+0]
2711    call fn0(), stack_map=[i32 @ ss0+0, i32 @ ss1+0, i32 @ ss2+0]
2712    jump block2
2713}
2714            "#,
2715        );
2716    }
2717
2718    fn import_func(
2719        builder: &mut FunctionBuilder,
2720        params: impl IntoIterator<Item = ir::Type>,
2721        results: impl IntoIterator<Item = ir::Type>,
2722    ) -> ir::FuncRef {
2723        let index = u32::try_from(builder.func.dfg.ext_funcs.len()).unwrap();
2724
2725        let name = builder
2726            .func
2727            .declare_imported_user_function(ir::UserExternalName {
2728                namespace: 0,
2729                index,
2730            });
2731        let name = ir::ExternalName::user(name);
2732
2733        let mut signature = Signature::new(CallConv::SystemV);
2734        signature
2735            .params
2736            .extend(params.into_iter().map(|ty| AbiParam::new(ty)));
2737        signature
2738            .returns
2739            .extend(results.into_iter().map(|ty| AbiParam::new(ty)));
2740        let signature = builder.func.import_signature(signature);
2741
2742        builder.import_function(ir::ExtFuncData {
2743            name,
2744            signature,
2745            colocated: true,
2746            patchable: false,
2747        })
2748    }
2749
2750    #[test]
2751    fn issue_10397_stack_map_vars_and_indirect_block_params() {
2752        let _ = env_logger::try_init();
2753
2754        let mut sig = Signature::new(CallConv::SystemV);
2755        if false {
2756            sig.params.push(AbiParam::new(ir::types::I32));
2757        }
2758
2759        let mut fn_ctx = FunctionBuilderContext::new();
2760        let mut func = Function::with_name_signature(ir::UserFuncName::testcase("f"), sig);
2761        let mut builder = FunctionBuilder::new(&mut func, &mut fn_ctx);
2762
2763        let alloc_struct = import_func(&mut builder, None, Some(ir::types::I32));
2764        let alloc_array = import_func(&mut builder, None, Some(ir::types::I32));
2765        let array_init_elem = import_func(&mut builder, Some(ir::types::I32), None);
2766        let type_of = import_func(&mut builder, Some(ir::types::I32), Some(ir::types::I32));
2767        let ref_test = import_func(
2768            &mut builder,
2769            vec![ir::types::I32, ir::types::I32],
2770            Some(ir::types::I32),
2771        );
2772        let access_array = import_func(&mut builder, Some(ir::types::I32), None);
2773        let should_continue_inner_loop = import_func(&mut builder, None, Some(ir::types::I32));
2774        let access_struct = import_func(&mut builder, Some(ir::types::I32), None);
2775        let should_return = import_func(&mut builder, None, Some(ir::types::I32));
2776
2777        // This test exercises the combination of declaring stack maps and when
2778        // we need to insert block params during SSA construction. The following
2779        // CLIF uses vars in such a way that it will ultimately require that we
2780        // insert block params. However, some of the inserted block params are
2781        // only ever used in such a way that they are passed as arguments to
2782        // *other* blocks. That is, there are regions where we never use them
2783        // directly, and therefore, if we are only declaring that the variable's
2784        // values need inclusion in stack maps on direct uses, we will
2785        // completely miss these inserted block params. That leads to incomplete
2786        // stack maps, which leads to collecting objects too early, which leads
2787        // to use-after-free bugs.
2788        //
2789        // After inserting safepoint spills and stack map annotations to the
2790        // following pseudo-CLIF, we should have stack map entries for the
2791        // `live: {...}` annotations (or an over-approximation of that
2792        // annotation).
2793        //
2794        //     block_entry:
2795        //       v0 = call alloc_struct()                          ;; live: {}
2796        //       define var_struct = v0
2797        //       jump block_outer_loop_head
2798        //
2799        //     block_outer_loop_head:
2800        //       v1 = call alloc_array()                           ;; live: {struct}
2801        //       define var_array = v1
2802        //       v2 = iconst.i32 0
2803        //       jump block_array_init_loop_head(v2)
2804        //
2805        //     block_array_init_loop_head(v3):
2806        //       v4 = iconst.i32 1
2807        //       v5 = icmp ult v3, v4
2808        //       br_if v5, block_array_init_loop_body, block_array_init_loop_done
2809        //
2810        //     block_array_init_loop_body:
2811        //       call array_init_elem(v1, v4)                      ;; live: {struct, array}
2812        //       v6 = iconst.i32 1
2813        //       v7 = iadd v4, v6
2814        //       jump block_array_init_loop_head(v7)
2815        //
2816        //     block_array_init_loop_done:
2817        //       jump block_inner_loop_head
2818        //
2819        //     block_inner_loop_head:
2820        //       v8 = use var_array
2821        //       v9 = iconst.i32 0
2822        //       v10 = icmp eq v8, v9
2823        //       br_if v10, block_ref_test_done(v9), block_ref_test_non_null
2824        //
2825        //     block_ref_test_non_null:
2826        //       v11 = call type_of(v8)                            ;; live: {struct, array}
2827        //       v12 = iconst.i32 0xbeefbeef
2828        //       v13 = icmp eq v11, v12
2829        //       v14 = iconst.i31 1
2830        //       br_if v13, block_ref_test_done(v14), block_ref_test_slow
2831        //
2832        //     block_ref_test_slow:
2833        //       v15 = call ref_test(v8, v12)                      ;; live: {struct, array}
2834        //       jump block_ref_test_done(v15)
2835        //
2836        //     block_ref_test_done(v16):
2837        //       trapz v16, user1
2838        //       define var_array = v8
2839        //       call access_array(v8)                             ;; live: {struct}
2840        //       v17 = call should_continue_inner_loop()           ;; live: {struct}
2841        //       br_if v17, block_inner_loop_head, block_after_inner_loop
2842        //
2843        //     block_after_inner_loop:
2844        //       v18 = use var_struct
2845        //       call access_struct(v18)                           ;; live: {struct}
2846        //       v19 = call should_return()                        ;; live: {struct}
2847        //       br_if v19, block_return, block_outer_loop_head
2848        //
2849        //     block_return:
2850        //       return
2851
2852        let var_struct = builder.declare_var(cranelift_codegen::ir::types::I32);
2853        builder.declare_var_needs_stack_map(var_struct);
2854
2855        let var_array = builder.declare_var(cranelift_codegen::ir::types::I32);
2856        builder.declare_var_needs_stack_map(var_array);
2857
2858        let block_entry = builder.create_block();
2859        let block_outer_loop_head = builder.create_block();
2860        let block_array_init_loop_head = builder.create_block();
2861        let block_array_init_loop_body = builder.create_block();
2862        let block_array_init_loop_done = builder.create_block();
2863        let block_inner_loop_head = builder.create_block();
2864        let block_ref_test_non_null = builder.create_block();
2865        let block_ref_test_slow = builder.create_block();
2866        let block_ref_test_done = builder.create_block();
2867        let block_after_inner_loop = builder.create_block();
2868        let block_return = builder.create_block();
2869
2870        builder.append_block_params_for_function_params(block_entry);
2871        builder.switch_to_block(block_entry);
2872        builder.seal_block(block_entry);
2873        let call_inst = builder.ins().call(alloc_struct, &[]);
2874        let v0 = builder.func.dfg.first_result(call_inst);
2875        builder.def_var(var_struct, v0);
2876        builder.ins().jump(block_outer_loop_head, &[]);
2877
2878        builder.switch_to_block(block_outer_loop_head);
2879        let call_inst = builder.ins().call(alloc_array, &[]);
2880        let v1 = builder.func.dfg.first_result(call_inst);
2881        builder.def_var(var_array, v1);
2882        let v2 = builder.ins().iconst(ir::types::I32, 0);
2883        builder.ins().jump(block_array_init_loop_head, &[v2.into()]);
2884
2885        builder.switch_to_block(block_array_init_loop_head);
2886        let v3 = builder.append_block_param(block_array_init_loop_head, ir::types::I32);
2887        let v4 = builder.ins().iconst(ir::types::I32, 1);
2888        let v5 = builder
2889            .ins()
2890            .icmp(ir::condcodes::IntCC::UnsignedLessThan, v3, v4);
2891        builder.ins().brif(
2892            v5,
2893            block_array_init_loop_body,
2894            &[],
2895            block_array_init_loop_done,
2896            &[],
2897        );
2898
2899        builder.switch_to_block(block_array_init_loop_body);
2900        builder.seal_block(block_array_init_loop_body);
2901        builder.ins().call(array_init_elem, &[v1, v4]);
2902        let v6 = builder.ins().iconst(ir::types::I32, 1);
2903        let v7 = builder.ins().iadd(v4, v6);
2904        builder.ins().jump(block_array_init_loop_head, &[v7.into()]);
2905        builder.seal_block(block_array_init_loop_head);
2906
2907        builder.switch_to_block(block_array_init_loop_done);
2908        builder.seal_block(block_array_init_loop_done);
2909        builder.ins().jump(block_inner_loop_head, &[]);
2910
2911        builder.switch_to_block(block_inner_loop_head);
2912        let v8 = builder.use_var(var_array);
2913        let v9 = builder.ins().iconst(ir::types::I32, 0);
2914        let v10 = builder.ins().icmp(ir::condcodes::IntCC::Equal, v8, v9);
2915        builder.ins().brif(
2916            v10,
2917            block_ref_test_done,
2918            &[v9.into()],
2919            block_ref_test_non_null,
2920            &[],
2921        );
2922
2923        builder.switch_to_block(block_ref_test_non_null);
2924        builder.seal_block(block_ref_test_non_null);
2925        let call_inst = builder.ins().call(type_of, &[v8]);
2926        let v11 = builder.func.dfg.first_result(call_inst);
2927        let v12 = builder.ins().iconst(ir::types::I32, 0xbeefbeef);
2928        let v13 = builder.ins().icmp(ir::condcodes::IntCC::Equal, v11, v12);
2929        let v14 = builder.ins().iconst(ir::types::I32, 1);
2930        builder.ins().brif(
2931            v13,
2932            block_ref_test_done,
2933            &[v14.into()],
2934            block_ref_test_slow,
2935            &[],
2936        );
2937
2938        builder.switch_to_block(block_ref_test_slow);
2939        builder.seal_block(block_ref_test_slow);
2940        let call_inst = builder.ins().call(ref_test, &[v8, v12]);
2941        let v15 = builder.func.dfg.first_result(call_inst);
2942        builder.ins().jump(block_ref_test_done, &[v15.into()]);
2943
2944        builder.switch_to_block(block_ref_test_done);
2945        let v16 = builder.append_block_param(block_ref_test_done, ir::types::I32);
2946        builder.seal_block(block_ref_test_done);
2947        builder.ins().trapz(v16, ir::TrapCode::user(1).unwrap());
2948        builder.def_var(var_array, v8);
2949        builder.ins().call(access_array, &[v8]);
2950        let call_inst = builder.ins().call(should_continue_inner_loop, &[]);
2951        let v17 = builder.func.dfg.first_result(call_inst);
2952        builder
2953            .ins()
2954            .brif(v17, block_inner_loop_head, &[], block_after_inner_loop, &[]);
2955        builder.seal_block(block_inner_loop_head);
2956
2957        builder.switch_to_block(block_after_inner_loop);
2958        builder.seal_block(block_after_inner_loop);
2959        let v18 = builder.use_var(var_struct);
2960        builder.ins().call(access_struct, &[v18]);
2961        let call_inst = builder.ins().call(should_return, &[]);
2962        let v19 = builder.func.dfg.first_result(call_inst);
2963        builder
2964            .ins()
2965            .brif(v19, block_return, &[], block_outer_loop_head, &[]);
2966        builder.seal_block(block_outer_loop_head);
2967
2968        builder.switch_to_block(block_return);
2969        builder.seal_block(block_return);
2970        builder.ins().return_(&[]);
2971
2972        builder.finalize(systemv_frontend_config());
2973        assert_eq_output!(
2974            func.display().to_string(),
2975            r#"
2976function %f() system_v {
2977    ss0 = explicit_slot 4, align = 4
2978    ss1 = explicit_slot 4, align = 4
2979    ss2 = explicit_slot 4, align = 4
2980    sig0 = () -> i32 system_v
2981    sig1 = () -> i32 system_v
2982    sig2 = (i32) system_v
2983    sig3 = (i32) -> i32 system_v
2984    sig4 = (i32, i32) -> i32 system_v
2985    sig5 = (i32) system_v
2986    sig6 = () -> i32 system_v
2987    sig7 = (i32) system_v
2988    sig8 = () -> i32 system_v
2989    fn0 = colocated u0:0 sig0
2990    fn1 = colocated u0:1 sig1
2991    fn2 = colocated u0:2 sig2
2992    fn3 = colocated u0:3 sig3
2993    fn4 = colocated u0:4 sig4
2994    fn5 = colocated u0:5 sig5
2995    fn6 = colocated u0:6 sig6
2996    fn7 = colocated u0:7 sig7
2997    fn8 = colocated u0:8 sig8
2998
2999block0:
3000    v0 = call fn0()
3001    jump block1(v0)
3002
3003block1(v22: i32):
3004    v21 -> v22
3005    v44 = stack_addr.i64 ss1
3006    store notrap aligned v22, v44
3007    v1 = call fn1(), stack_map=[i32 @ ss1+0]
3008    v8 -> v1
3009    v18 -> v1
3010    v43 = stack_addr.i64 ss0
3011    store notrap aligned v1, v43
3012    v2 = iconst.i32 0
3013    jump block2(v2)  ; v2 = 0
3014
3015block2(v3: i32):
3016    v4 = iconst.i32 1
3017    v5 = icmp ult v3, v4  ; v4 = 1
3018    brif v5, block3, block4
3019
3020block3:
3021    v24 = stack_addr.i64 ss0
3022    v25 = load.i32 notrap aligned v24
3023    call fn2(v25, v4), stack_map=[i32 @ ss0+0, i32 @ ss1+0]  ; v4 = 1
3024    v6 = iconst.i32 1
3025    v7 = iadd.i32 v4, v6  ; v4 = 1, v6 = 1
3026    jump block2(v7)
3027
3028block4:
3029    v41 = stack_addr.i64 ss1
3030    v42 = load.i32 notrap aligned v41
3031    jump block5(v42)
3032
3033block5(v20: i32):
3034    v19 -> v20
3035    v40 = stack_addr.i64 ss2
3036    store notrap aligned v20, v40
3037    v9 = iconst.i32 0
3038    v38 = stack_addr.i64 ss0
3039    v39 = load.i32 notrap aligned v38
3040    v10 = icmp eq v39, v9  ; v9 = 0
3041    brif v10, block8(v9), block6  ; v9 = 0
3042
3043block6:
3044    v36 = stack_addr.i64 ss0
3045    v37 = load.i32 notrap aligned v36
3046    v11 = call fn3(v37), stack_map=[i32 @ ss0+0, i32 @ ss2+0]
3047    v12 = iconst.i32 -1091584273
3048    v13 = icmp eq v11, v12  ; v12 = -1091584273
3049    v14 = iconst.i32 1
3050    brif v13, block8(v14), block7  ; v14 = 1
3051
3052block7:
3053    v34 = stack_addr.i64 ss0
3054    v35 = load.i32 notrap aligned v34
3055    v15 = call fn4(v35, v12), stack_map=[i32 @ ss0+0, i32 @ ss2+0]  ; v12 = -1091584273
3056    jump block8(v15)
3057
3058block8(v16: i32):
3059    trapz v16, user1
3060    v32 = stack_addr.i64 ss0
3061    v33 = load.i32 notrap aligned v32
3062    call fn5(v33), stack_map=[i32 @ ss0+0, i32 @ ss2+0]
3063    v17 = call fn6(), stack_map=[i32 @ ss0+0, i32 @ ss2+0]
3064    v30 = stack_addr.i64 ss2
3065    v31 = load.i32 notrap aligned v30
3066    brif v17, block5(v31), block9
3067
3068block9:
3069    v28 = stack_addr.i64 ss2
3070    v29 = load.i32 notrap aligned v28
3071    call fn7(v29), stack_map=[i32 @ ss2+0]
3072    v23 = call fn8(), stack_map=[i32 @ ss2+0]
3073    v26 = stack_addr.i64 ss2
3074    v27 = load.i32 notrap aligned v26
3075    brif v23, block10, block1(v27)
3076
3077block10:
3078    return
3079}
3080            "#,
3081        );
3082    }
3083
3084    #[test]
3085    fn needs_stack_map_across_try_call() {
3086        let _ = env_logger::try_init();
3087
3088        // Test that values needing stack maps get stack_map entries on
3089        // try_call instructions, not just regular call instructions.
3090        //
3091        //     block0:
3092        //       v0 = call fn0()       ;; returns a gc ref
3093        //       v1 = call fn0()       ;; returns another gc ref
3094        //       try_call fn0(), sig0, block1(), [default: block2()]
3095        //                             ;; v0 and v1 should be in the stack map here
3096        //     block1:
3097        //       call fn1(v0)          ;; uses v0, v1 is dead
3098        //       return
3099        //     block2(v2: i64):
3100        //       return
3101
3102        let sig = Signature::new(CallConv::SystemV);
3103
3104        let mut fn_ctx = FunctionBuilderContext::new();
3105        let mut func = Function::with_name_signature(ir::UserFuncName::testcase("sample"), sig);
3106        let mut builder = FunctionBuilder::new(&mut func, &mut fn_ctx);
3107
3108        // fn0: () -> i32 (returns a gc ref)
3109        let name0 = builder
3110            .func
3111            .declare_imported_user_function(ir::UserExternalName {
3112                namespace: 0,
3113                index: 0,
3114            });
3115        let mut sig0 = Signature::new(CallConv::SystemV);
3116        sig0.returns.push(AbiParam::new(ir::types::I32));
3117        let signature0 = builder.func.import_signature(sig0);
3118        let func_ref0 = builder.import_function(ir::ExtFuncData {
3119            name: ir::ExternalName::user(name0),
3120            signature: signature0,
3121            colocated: true,
3122            patchable: false,
3123        });
3124
3125        // fn1: (i32) -> () (consumes a gc ref)
3126        let name1 = builder
3127            .func
3128            .declare_imported_user_function(ir::UserExternalName {
3129                namespace: 0,
3130                index: 1,
3131            });
3132        let mut sig1 = Signature::new(CallConv::SystemV);
3133        sig1.params.push(AbiParam::new(ir::types::I32));
3134        let signature1 = builder.func.import_signature(sig1);
3135        let func_ref1 = builder.import_function(ir::ExtFuncData {
3136            name: ir::ExternalName::user(name1),
3137            signature: signature1,
3138            colocated: true,
3139            patchable: false,
3140        });
3141
3142        let block0 = builder.create_block();
3143        let block1 = builder.create_block();
3144        let block2 = builder.create_block();
3145
3146        builder.switch_to_block(block0);
3147
3148        // v0 = call fn0()  -- a gc ref that's live across the try_call
3149        let call0 = builder.ins().call(func_ref0, &[]);
3150        let v0 = builder.func.dfg.inst_results(call0)[0];
3151        builder.declare_value_needs_stack_map(v0);
3152
3153        // v1 = call fn0()  -- another gc ref, also live across the try_call
3154        let call1 = builder.ins().call(func_ref0, &[]);
3155        let v1 = builder.func.dfg.inst_results(call1)[0];
3156        builder.declare_value_needs_stack_map(v1);
3157
3158        // try_call fn0() -> block1, exception -> block2
3159        let normal_return = BlockCall::new(block1, [], &mut builder.func.dfg.value_lists);
3160        let exception_table = builder
3161            .func
3162            .dfg
3163            .exception_tables
3164            .push(ExceptionTableData::new(signature0, normal_return, []));
3165        builder.ins().try_call(func_ref0, &[], exception_table);
3166
3167        // block1: use v0 (so it's live across the try_call)
3168        builder.switch_to_block(block1);
3169        builder.ins().call(func_ref1, &[v0]);
3170        builder.ins().return_(&[]);
3171
3172        // block2: exception handler
3173        builder.switch_to_block(block2);
3174        builder.ins().return_(&[]);
3175
3176        builder.seal_all_blocks();
3177        builder.finalize(systemv_frontend_config());
3178
3179        // The try_call should have a stack_map with v0 (and v1 should NOT
3180        // be in it since v1 is not used after the try_call).
3181        // The second call (which produces v1) should have v0 in its stack_map
3182        // since v0 is live across that call too.
3183        let output = func.display().to_string();
3184        assert!(
3185            output.contains("try_call fn0(), sig0, block1, [], stack_map=[i32 @ ss0+0]"),
3186            "try_call should have stack_map entry for v0 (spilled to ss0), got:\n{output}"
3187        );
3188    }
3189
3190    #[test]
3191    fn needs_stack_map_try_call_edge_args() {
3192        let _ = env_logger::try_init();
3193
3194        // Test that values needing stack maps which are passed as block-call
3195        // arguments on a `try_call`'s edges are live across the `try_call`,
3196        // and are reloaded *after* it, on each edge (see #14452).
3197        //
3198        //     block0:
3199        //       v0 = call fn0()   ;; returns a gc ref
3200        //       try_call fn0(), sig0, block1(ret0, v0), [ default: block2(exn0, v0) ]
3201        //                         ;; v0 should be in the stack map here
3202        //     block1(v1: i32, v2: i32):
3203        //       return
3204        //     block2(v3: i64, v4: i32):
3205        //       return
3206
3207        let sig = Signature::new(CallConv::SystemV);
3208
3209        let mut fn_ctx = FunctionBuilderContext::new();
3210        let mut func = Function::with_name_signature(ir::UserFuncName::testcase("sample"), sig);
3211        let mut builder = FunctionBuilder::new(&mut func, &mut fn_ctx);
3212
3213        // fn0: () -> i32 (returns a gc ref)
3214        let name0 = builder
3215            .func
3216            .declare_imported_user_function(ir::UserExternalName {
3217                namespace: 0,
3218                index: 0,
3219            });
3220        let mut sig0 = Signature::new(CallConv::SystemV);
3221        sig0.returns.push(AbiParam::new(ir::types::I32));
3222        let signature0 = builder.func.import_signature(sig0);
3223        let func_ref0 = builder.import_function(ir::ExtFuncData {
3224            name: ir::ExternalName::user(name0),
3225            signature: signature0,
3226            colocated: true,
3227            patchable: false,
3228        });
3229
3230        let block0 = builder.create_block();
3231        let block1 = builder.create_block();
3232        let block2 = builder.create_block();
3233        builder.append_block_param(block1, ir::types::I32);
3234        builder.append_block_param(block1, ir::types::I32);
3235        builder.append_block_param(block2, ir::types::I64);
3236        builder.append_block_param(block2, ir::types::I32);
3237
3238        builder.switch_to_block(block0);
3239        let call0 = builder.ins().call(func_ref0, &[]);
3240        let v0 = builder.func.dfg.inst_results(call0)[0];
3241        builder.declare_value_needs_stack_map(v0);
3242
3243        let normal_return = BlockCall::new(
3244            block1,
3245            [ir::BlockArg::TryCallRet(0), ir::BlockArg::Value(v0)],
3246            &mut builder.func.dfg.value_lists,
3247        );
3248        let handler = BlockCall::new(
3249            block2,
3250            [ir::BlockArg::TryCallExn(0), ir::BlockArg::Value(v0)],
3251            &mut builder.func.dfg.value_lists,
3252        );
3253        let exception_table = builder
3254            .func
3255            .dfg
3256            .exception_tables
3257            .push(ExceptionTableData::new(
3258                signature0,
3259                normal_return,
3260                [ir::ExceptionTableItem::Default(handler)],
3261            ));
3262        builder.ins().try_call(func_ref0, &[], exception_table);
3263
3264        builder.switch_to_block(block1);
3265        builder.ins().return_(&[]);
3266
3267        builder.switch_to_block(block2);
3268        builder.ins().return_(&[]);
3269
3270        builder.seal_all_blocks();
3271        builder.finalize(systemv_frontend_config());
3272
3273        assert_eq_output!(
3274            func.display().to_string(),
3275            r#"
3276function %sample() system_v {
3277    ss0 = explicit_slot 4, align = 4
3278    sig0 = () -> i32 system_v
3279    fn0 = colocated u0:0 sig0
3280
3281block0:
3282    v4 = call fn0()
3283    v11 = stack_addr.i64 ss0
3284    store notrap aligned v4, v11
3285    try_call fn0(), sig0, block4(ret0), [ default: block3(exn0) ], stack_map=[i32 @ ss0+0]
3286
3287block3(v5: i64):
3288    v6 = stack_addr.i64 ss0
3289    v7 = load.i32 notrap aligned v6
3290    jump block2(v5, v7)
3291
3292block4(v8: i32):
3293    v9 = stack_addr.i64 ss0
3294    v10 = load.i32 notrap aligned v9
3295    jump block1(v8, v10)
3296
3297block1(v0: i32, v1: i32):
3298    return
3299
3300block2(v2: i64, v3: i32):
3301    return
3302}
3303            "#
3304        );
3305    }
3306
3307    #[test]
3308    fn rewrite_uses_of_alias_values() {
3309        let _ = env_logger::try_init();
3310
3311        let mut sig = Signature::new(CallConv::SystemV);
3312        sig.params.push(AbiParam::new(ir::types::I32));
3313        let mut fn_ctx = FunctionBuilderContext::new();
3314        let mut func = Function::with_name_signature(ir::UserFuncName::testcase("sample"), sig);
3315        let mut builder = FunctionBuilder::new(&mut func, &mut fn_ctx);
3316
3317        // fn0: () -> i32 (returns a gc ref)
3318        let name0 = builder
3319            .func
3320            .declare_imported_user_function(ir::UserExternalName {
3321                namespace: 0,
3322                index: 0,
3323            });
3324        let mut sig = Signature::new(CallConv::SystemV);
3325        sig.returns.push(AbiParam::new(ir::types::I32));
3326        let signature = builder.func.import_signature(sig);
3327        let fn0 = builder.import_function(ir::ExtFuncData {
3328            name: ir::ExternalName::user(name0),
3329            signature,
3330            colocated: true,
3331            patchable: false,
3332        });
3333
3334        // fn1: (i32) -> () (consumes a gc ref)
3335        let name1 = builder
3336            .func
3337            .declare_imported_user_function(ir::UserExternalName {
3338                namespace: 0,
3339                index: 1,
3340            });
3341        let mut sig = Signature::new(CallConv::SystemV);
3342        sig.params.push(AbiParam::new(ir::types::I32));
3343        let signature = builder.func.import_signature(sig);
3344        let fn1 = builder.import_function(ir::ExtFuncData {
3345            name: ir::ExternalName::user(name1),
3346            signature,
3347            colocated: true,
3348            patchable: false,
3349        });
3350
3351        let var = builder.declare_var(ir::types::I32);
3352        builder.declare_var_needs_stack_map(var);
3353
3354        let block0 = builder.create_block();
3355        let block1 = builder.create_block();
3356        let block2 = builder.create_block();
3357
3358        builder.append_block_params_for_function_params(block0);
3359        builder.switch_to_block(block0);
3360        let v0 = builder.func.dfg.block_params(block0)[0];
3361        let inst = builder.ins().call(fn0, &[]);
3362        let v1 = builder.func.dfg.first_result(inst);
3363        builder.def_var(var, v1);
3364        builder.ins().brif(v0, block1, &[], block2, &[]);
3365
3366        builder.switch_to_block(block1);
3367        builder.def_var(var, v1);
3368        builder.ins().jump(block2, &[]);
3369
3370        builder.switch_to_block(block2);
3371        let v2 = builder.use_var(var);
3372        builder.ins().call(fn1, &[v2]);
3373        let v3 = builder.use_var(var);
3374        builder.ins().call(fn1, &[v3]);
3375        builder.ins().return_(&[]);
3376
3377        builder.seal_all_blocks();
3378
3379        builder.finalize(systemv_frontend_config());
3380
3381        // The SSA construction / `Variable` infrastructure makes this value
3382        // into an alias. But it is also an alias of a value that needs
3383        // inclusion in stack maps, so we should reload it from the stack before
3384        // it is passed into `fn1` in the disassembly below, rather than pass
3385        // the original or aliased value directly.
3386        assert!(func.dfg.value_is_alias(v3));
3387        assert_eq_output!(
3388            func.display().to_string(),
3389            r#"
3390function %sample(i32) system_v {
3391    ss0 = explicit_slot 4, align = 4
3392    sig0 = () -> i32 system_v
3393    sig1 = (i32) system_v
3394    fn0 = colocated u0:0 sig0
3395    fn1 = colocated u0:1 sig1
3396
3397block0(v0: i32):
3398    v1 = call fn0()
3399    v2 -> v1
3400    v7 = stack_addr.i64 ss0
3401    store notrap aligned v1, v7
3402    brif v0, block1, block2
3403
3404block1:
3405    jump block2
3406
3407block2:
3408    v5 = stack_addr.i64 ss0
3409    v6 = load.i32 notrap aligned v5
3410    call fn1(v6), stack_map=[i32 @ ss0+0]
3411    v3 = stack_addr.i64 ss0
3412    v4 = load.i32 notrap aligned v3
3413    call fn1(v4)
3414    return
3415}
3416            "#,
3417        );
3418    }
3419
3420    /// Regression test for a missed stack-map slot when a stack-map
3421    /// variable is re-defined in the same block while an earlier SSA
3422    /// value bound to it is still live across a safepoint.
3423    ///
3424    /// Previously, `SSABuilder::variables` stored only the *latest*
3425    /// `Value` per `(Variable, Block)`, and `values_for_var` iterated
3426    /// that map, so the earlier `Value` was dropped from the set the
3427    /// safepoint pass propagates stack-map flags onto. The expected
3428    /// output below requires `v1` to be spilled and listed in the
3429    /// `call` instruction's stack map alongside `v2`.
3430    #[test]
3431    fn var_redefined_in_same_block_keeps_earlier_value_in_stack_map() {
3432        let mut sig = Signature::new(CallConv::SystemV);
3433        sig.returns
3434            .push(AbiParam::new(cranelift_codegen::ir::types::I32));
3435        sig.returns
3436            .push(AbiParam::new(cranelift_codegen::ir::types::I32));
3437
3438        let mut fn_ctx = FunctionBuilderContext::new();
3439        let mut func = Function::with_name_signature(ir::UserFuncName::testcase("sample"), sig);
3440        let mut builder = FunctionBuilder::new(&mut func, &mut fn_ctx);
3441
3442        let var = builder.declare_var(cranelift_codegen::ir::types::I32);
3443        builder.declare_var_needs_stack_map(var);
3444
3445        let name = builder
3446            .func
3447            .declare_imported_user_function(ir::UserExternalName {
3448                namespace: 0,
3449                index: 0,
3450            });
3451        let signature = builder
3452            .func
3453            .import_signature(Signature::new(CallConv::SystemV));
3454        let func_ref = builder.import_function(ir::ExtFuncData {
3455            name: ir::ExternalName::user(name),
3456            signature,
3457            colocated: true,
3458            patchable: false,
3459        });
3460
3461        let block0 = builder.create_block();
3462        builder.switch_to_block(block0);
3463
3464        // First definition of `var`. Read it out (giving us an SSA
3465        // value `v1` that we'll keep live across the safepoint below).
3466        let v1 = builder.ins().iconst(ir::types::I32, 11);
3467        builder.def_var(var, v1);
3468        let v1_use = builder.use_var(var);
3469
3470        // Re-define `var` to a different SSA value within the same
3471        // block. Without the fix, this overwrites the only place
3472        // `v1` was recorded, so `values_for_var(var)` no longer
3473        // returns it and the safepoint pass omits it from the
3474        // upcoming stack map.
3475        let v2 = builder.ins().iconst(ir::types::I32, 22);
3476        builder.def_var(var, v2);
3477        let v2_use = builder.use_var(var);
3478
3479        // Safepoint at which both `v1_use` (= v1) and `v2_use` (= v2)
3480        // must be live and listed in the stack map.
3481        builder.ins().call(func_ref, &[]);
3482
3483        builder.ins().return_(&[v1_use, v2_use]);
3484
3485        builder.seal_all_blocks();
3486        builder.finalize(systemv_frontend_config());
3487
3488        assert_eq_output!(
3489            func.display().to_string(),
3490            r#"
3491function %sample() -> i32, i32 system_v {
3492    ss0 = explicit_slot 4, align = 4
3493    ss1 = explicit_slot 4, align = 4
3494    sig0 = () system_v
3495    fn0 = colocated u0:0 sig0
3496
3497block0:
3498    v0 = iconst.i32 11
3499    v7 = stack_addr.i64 ss0
3500    store notrap aligned v0, v7  ; v0 = 11
3501    v1 = iconst.i32 22
3502    v6 = stack_addr.i64 ss1
3503    store notrap aligned v1, v6  ; v1 = 22
3504    call fn0(), stack_map=[i32 @ ss0+0, i32 @ ss1+0]
3505    v2 = stack_addr.i64 ss0
3506    v3 = load.i32 notrap aligned v2
3507    v4 = stack_addr.i64 ss1
3508    v5 = load.i32 notrap aligned v4
3509    return v3, v5
3510}
3511            "#
3512        );
3513    }
3514
3515    #[test]
3516    fn loop_invariant_value_needs_stack_map() {
3517        let sig = Signature::new(CallConv::SystemV);
3518        let mut fn_ctx = FunctionBuilderContext::new();
3519        let mut func = Function::with_name_signature(ir::UserFuncName::testcase("sample"), sig);
3520        let mut builder = FunctionBuilder::new(&mut func, &mut fn_ctx);
3521
3522        // alloc: () -> i32
3523        let alloc_name = builder
3524            .func
3525            .declare_imported_user_function(ir::UserExternalName {
3526                namespace: 0,
3527                index: 1,
3528            });
3529        let mut sig = Signature::new(CallConv::SystemV);
3530        sig.returns.push(AbiParam::new(ir::types::I32));
3531        let signature = builder.func.import_signature(sig);
3532        let alloc = builder.import_function(ir::ExtFuncData {
3533            name: ir::ExternalName::user(alloc_name),
3534            signature,
3535            colocated: true,
3536            patchable: false,
3537        });
3538
3539        // observe: (i32) -> ()
3540        let observe_name = builder
3541            .func
3542            .declare_imported_user_function(ir::UserExternalName {
3543                namespace: 0,
3544                index: 1,
3545            });
3546        let mut sig = Signature::new(CallConv::SystemV);
3547        sig.params.push(AbiParam::new(ir::types::I32));
3548        let signature = builder.func.import_signature(sig);
3549        let observe = builder.import_function(ir::ExtFuncData {
3550            name: ir::ExternalName::user(observe_name),
3551            signature,
3552            colocated: true,
3553            patchable: false,
3554        });
3555
3556        let block0 = builder.create_block();
3557        let block1 = builder.create_block();
3558        let block2 = builder.create_block();
3559
3560        // block0: Allocate a GC reference, then enter the loop.
3561        builder.switch_to_block(block0);
3562        let call_inst = builder.ins().call(alloc, &[]);
3563        let v0 = builder.func.dfg.first_result(call_inst);
3564        builder.declare_value_needs_stack_map(v0);
3565        builder.ins().jump(block1, &[]);
3566
3567        // block1: A loop that uses the loop-invariant GC ref and then allocates
3568        // another GC ref. The two GC refs' live ranges are overlapping (since
3569        // `v0` is live across the whole loop) even if they "don't" overlap
3570        // within a single loop iteration.
3571        builder.switch_to_block(block1);
3572        builder.ins().call(observe, &[v0]);
3573        let call_inst = builder.ins().call(alloc, &[]);
3574        let v1 = builder.func.dfg.first_result(call_inst);
3575        builder.declare_value_needs_stack_map(v1);
3576        // This call should have stack map entries for both `v0` and `v1`.
3577        builder.ins().call(observe, &[v1]);
3578        builder.ins().brif(v1, block2, &[], block1, &[]);
3579
3580        // block2: Keep `v1` alive across a safepoint and return.
3581        builder.switch_to_block(block2);
3582        builder.ins().call(observe, &[v1]);
3583        builder.ins().call(observe, &[v1]);
3584        builder.ins().return_(&[]);
3585
3586        builder.seal_all_blocks();
3587        builder.finalize(systemv_frontend_config());
3588
3589        // `v0` and `v1` should be spilled and reloaded from different stack
3590        // slots.
3591        assert_eq_output!(
3592            func.display().to_string(),
3593            r#"
3594function %sample() system_v {
3595    ss0 = explicit_slot 4, align = 4
3596    ss1 = explicit_slot 4, align = 4
3597    sig0 = () -> i32 system_v
3598    sig1 = (i32) system_v
3599    fn0 = colocated u0:1 sig0
3600    fn1 = colocated u0:1 sig1
3601
3602block0:
3603    v0 = call fn0()
3604    v13 = stack_addr.i64 ss1
3605    store notrap aligned v0, v13
3606    jump block1
3607
3608block1:
3609    v11 = stack_addr.i64 ss1
3610    v12 = load.i32 notrap aligned v11
3611    call fn1(v12), stack_map=[i32 @ ss1+0]
3612    v1 = call fn0(), stack_map=[i32 @ ss1+0]
3613    v10 = stack_addr.i64 ss0
3614    store notrap aligned v1, v10
3615    v8 = stack_addr.i64 ss0
3616    v9 = load.i32 notrap aligned v8
3617    call fn1(v9), stack_map=[i32 @ ss1+0, i32 @ ss0+0]
3618    v6 = stack_addr.i64 ss0
3619    v7 = load.i32 notrap aligned v6
3620    brif v7, block2, block1
3621
3622block2:
3623    v4 = stack_addr.i64 ss0
3624    v5 = load.i32 notrap aligned v4
3625    call fn1(v5), stack_map=[i32 @ ss0+0]
3626    v2 = stack_addr.i64 ss0
3627    v3 = load.i32 notrap aligned v2
3628    call fn1(v3)
3629    return
3630}
3631            "#
3632        );
3633    }
3634
3635    #[test]
3636    fn stack_map_alias_region() {
3637        let _ = env_logger::try_init();
3638
3639        let sig = Signature::new(CallConv::SystemV);
3640
3641        let mut fn_ctx = FunctionBuilderContext::new();
3642        let mut func = Function::with_name_signature(ir::UserFuncName::testcase("sample"), sig);
3643        let mut builder = FunctionBuilder::new(&mut func, &mut fn_ctx);
3644
3645        // Place every safepoint spill and reload into a single deduplicated
3646        // alias region.
3647        builder.make_stack_map_alias_region(Box::new(|regions, _ty, _slot, _offset| {
3648            Some(regions.insert(ir::AliasRegionData {
3649                user_id: 0,
3650                description: "stack map".into(),
3651            }))
3652        }));
3653
3654        let name = builder
3655            .func
3656            .declare_imported_user_function(ir::UserExternalName {
3657                namespace: 0,
3658                index: 0,
3659            });
3660        let mut sig = Signature::new(CallConv::SystemV);
3661        sig.params.push(AbiParam::new(ir::types::I32));
3662        let signature = builder.func.import_signature(sig);
3663        let func_ref = builder.import_function(ir::ExtFuncData {
3664            name: ir::ExternalName::user(name),
3665            signature,
3666            colocated: true,
3667            patchable: false,
3668        });
3669
3670        // `v0` is live across the first `call` safepoint (it is used again by
3671        // the second call), so it is spilled at its definition and reloaded at
3672        // each use. The spill `store` and reload `load`s should all be tagged
3673        // with the configured alias region.
3674        //
3675        //     block0:
3676        //       v0 = iconst.i32 42  ; needs stack map
3677        //       call $foo(v0)
3678        //       call $foo(v0)
3679        //       return
3680        let block0 = builder.create_block();
3681        builder.append_block_params_for_function_params(block0);
3682        builder.switch_to_block(block0);
3683        let v0 = builder.ins().iconst(ir::types::I32, 42);
3684        builder.declare_value_needs_stack_map(v0);
3685        builder.ins().call(func_ref, &[v0]);
3686        builder.ins().call(func_ref, &[v0]);
3687        builder.ins().return_(&[]);
3688        builder.seal_all_blocks();
3689        builder.finalize(systemv_frontend_config());
3690
3691        assert_eq_output!(
3692            func.display().to_string(),
3693            r#"
3694function %sample() system_v {
3695    ss0 = explicit_slot 4, align = 4
3696    region0 = 0 "stack map"
3697    sig0 = (i32) system_v
3698    fn0 = colocated u0:0 sig0
3699
3700block0:
3701    v0 = iconst.i32 42
3702    v5 = stack_addr.i64 ss0
3703    store notrap aligned region0 v0, v5  ; v0 = 42
3704    v3 = stack_addr.i64 ss0
3705    v4 = load.i32 notrap aligned region0 v3
3706    call fn0(v4), stack_map=[i32 @ ss0+0]
3707    v1 = stack_addr.i64 ss0
3708    v2 = load.i32 notrap aligned region0 v1
3709    call fn0(v2)
3710    return
3711}
3712            "#
3713        );
3714    }
3715}