Skip to main content

charon_lib/transform/control_flow/
ullbc_to_llbc.rs

1//! ULLBC to LLBC
2//!
3//! We reconstruct the control-flow in the Unstructured LLBC.
4//!
5//! The reconstruction algorithm is not written to be efficient (its complexity
6//! is probably very bad), but it was not written to be: this is still an early
7//! stage and we want the algorithm to generate the best reconstruction as
8//! possible. We still need to test the algorithm on more interesting examples,
9//! and will consider making it more efficient once it is a bit mature and well
10//! tested.
11//! Also note that we more importantly focus on making the algorithm sound: the
12//! reconstructed program must always be equivalent to the original MIR program,
13//! and the fact that the reconstruction preserves this property must be obvious.
14use itertools::Itertools;
15use petgraph::algo::dijkstra;
16use petgraph::algo::dominators::{Dominators, simple_fast};
17use petgraph::graphmap::DiGraphMap;
18use petgraph::visit::{Dfs, DfsPostOrder, EdgeFiltered, EdgeRef, Walker};
19use rustc_hash::{FxHashMap as HashMap, FxHashSet as HashSet};
20use smallvec::SmallVec;
21use std::cmp::Reverse;
22use std::mem;
23
24use crate::ids::IndexVec;
25use crate::llbc_ast::{self as tgt, StatementId};
26use crate::transform::TransformCtx;
27use crate::transform::ctx::TransformPass;
28use crate::ullbc_ast::{self as src, BlockId};
29use crate::utils::{DfsWithPath, ensure_sufficient_stack};
30use crate::{ast::*, register_error};
31
32/// The amount of "flow" reaching a block.
33#[derive(Debug, Clone, Copy, Default)]
34struct Flow(f64);
35
36impl Ord for Flow {
37    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
38        self.0.total_cmp(&other.0)
39    }
40}
41impl PartialOrd for Flow {
42    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
43        Some(self.cmp(other))
44    }
45}
46impl PartialEq for Flow {
47    fn eq(&self, other: &Self) -> bool {
48        self.cmp(other).is_eq()
49    }
50}
51impl Eq for Flow {}
52impl std::ops::AddAssign<Flow> for Flow {
53    fn add_assign(&mut self, other: Flow) {
54        self.0 += other.0;
55    }
56}
57impl Flow {
58    const ZERO: Flow = Flow(0.);
59    const ONE: Flow = Flow(1.);
60    fn divided_by(self, n: usize) -> Flow {
61        Flow(self.0 / (n as f64))
62    }
63}
64
65/// Control-Flow Graph
66type Cfg = DiGraphMap<src::BlockId, (), rustc_hash::FxBuildHasher>;
67
68/// Information precomputed about a function's CFG.
69#[derive(Debug)]
70struct CfgInfo {
71    /// The CFG
72    pub cfg: Cfg,
73    /// The CFG where all the backward edges have been removed. Aka "forward CFG".
74    pub fwd_cfg: Cfg,
75    /// We consider the destination of the backward edges to be loop entries and
76    /// store them here.
77    pub loop_entries: HashSet<src::BlockId>,
78    /// The blocks whose terminators are a switch are stored here.
79    pub switch_blocks: HashSet<src::BlockId>,
80    /// Tree of which nodes dominates which other nodes.
81    #[expect(unused)]
82    pub dominator_tree: Dominators<BlockId>,
83    /// Computed data about each block.
84    pub block_data: IndexVec<BlockId, Box<BlockData>>,
85}
86
87#[derive(Debug)]
88struct BlockData {
89    pub id: BlockId,
90    pub span: Span,
91    /// The (unique) entrypoints of each loop. Unique because we error on irreducible cfgs.
92    pub is_loop_header: bool,
93    /// Whether this block is a switch.
94    pub is_switch: bool,
95    /// Whether this block is only reachable by taking an unwind edge.
96    pub is_unwind: bool,
97    /// Whether this block has multiple incoming control-flow edges in the forward graph.
98    pub is_merge_target: bool,
99    /// Order in a reverse postorder numbering. `None` if the block is unreachable.
100    pub reverse_postorder: Option<u32>,
101    /// Nodes that this block immediately dominates. Sorted by reverse_postorder_id, with largest
102    /// id first.
103    pub immediately_dominates: SmallVec<[BlockId; 2]>,
104    /// The nodes from `immediately_dominates` that are also merge targets. Sorted in the same
105    /// order.
106    pub immediately_dominated_merge_targets: SmallVec<[BlockId; 2]>,
107    /// List of loops inside of which this node is (loops are identified by their header). A node
108    /// is considered inside a loop if it is reachable from the loop header and if it can reach the
109    /// loop header using only the backwards edges into it (i.e. we don't count a path that enters
110    /// the loop header through a forward edge).
111    ///
112    /// Note that we might have to take a backward edge to reach the loop header, e.g.:
113    ///   'a: loop {
114    ///       // ...
115    ///       'b: loop {
116    ///           // ...
117    ///           if true {
118    ///               continue 'a;
119    ///           } else {
120    ///               if true {
121    ///                   break 'a;
122    ///               }
123    ///               // This node has to take two backward edges in order to reach the start of `'a`.
124    ///           }
125    ///       }
126    ///   }
127    ///
128    /// The restriction on backwards edges is for the following case:
129    ///   loop {
130    ///     loop {
131    ///       ..
132    ///     }
133    ///     // Not in inner loop
134    ///   }
135    ///
136    /// This is sorted by path order from the graph root.
137    pub within_loops: SmallVec<[BlockId; 2]>,
138    /// Node from where we can only reach error nodes (panic, etc.)
139    // TODO: track more nicely the set of targets reachable from a node: panic, return, exit loop,
140    // continue loop (this is a partial order).
141    pub only_reach_error: bool,
142    /// List of reachable nodes, with the length of shortest path to them. Includes the current
143    /// node.
144    pub shortest_paths: hashbrown::HashMap<BlockId, usize>,
145    /// Let's say we put a quantity of water equal to 1 on the block, and the water flows downards.
146    /// Whenever there is a branching, the quantity of water gets equally divided between the
147    /// branches. When the control flows join, we put the water back together. The set below
148    /// computes the amount of water received by each descendant of the node.
149    ///
150    /// TODO: there must be a known algorithm which computes this, right?...
151    /// This is exactly this problems:
152    /// <https://stackoverflow.com/questions/78221666/algorithm-for-total-flow-through-weighted-directed-acyclic-graph>
153    /// TODO: the way I compute this is not efficient.
154    pub flow: IndexVec<BlockId, Flow>,
155    /// Reconstructed information about loops and switches.
156    pub exit_info: ExitInfo,
157}
158
159#[derive(Debug, Default, Clone)]
160struct ExitInfo {
161    /// The loop exit
162    loop_exit: Option<src::BlockId>,
163    /// The switch exit.
164    switch_exit: Option<src::BlockId>,
165}
166
167/// Error indicating that the control-flow graph is not reducible. The contained block id is a
168/// block involved in an irreducible subgraph.
169struct Irreducible(BlockId);
170
171impl CfgInfo {
172    /// Build the CFGs (the "regular" CFG and the CFG without backward edges) and precompute a
173    /// bunch of graph information about the CFG.
174    fn build(ctx: &TransformCtx, body: &src::BodyContents) -> Result<Self, Irreducible> {
175        // The steps in this function follow a precise order, as each step typically requires the
176        // previous one.
177        let start_block = BlockId::ZERO;
178
179        let empty_flow = body.map_ref(|_| Flow::ZERO);
180        let mut block_data: IndexVec<BlockId, _> = body.map_ref_indexed(|id, contents| {
181            Box::new(BlockData {
182                id,
183                span: contents.terminator.span,
184                is_loop_header: false,
185                is_switch: false,
186                is_unwind: false,
187                is_merge_target: false,
188                reverse_postorder: None,
189                immediately_dominates: Default::default(),
190                immediately_dominated_merge_targets: Default::default(),
191                within_loops: Default::default(),
192                only_reach_error: false,
193                shortest_paths: Default::default(),
194                flow: empty_flow.clone(),
195                exit_info: Default::default(),
196            })
197        });
198
199        // Build the node graph.
200        let mut cfg = Cfg::new();
201        let mut cfg_without_unwind = Cfg::new();
202        for (block_id, block) in body.iter_enumerated() {
203            cfg.add_node(block_id);
204            cfg_without_unwind.add_node(block_id);
205            for tgt in block.targets() {
206                cfg.add_edge(block_id, tgt, ());
207            }
208            for tgt in block.targets_ignoring_unwind() {
209                cfg_without_unwind.add_edge(block_id, tgt, ());
210            }
211        }
212
213        // Compute the dominator tree.
214        let dominator_tree = simple_fast(&cfg, start_block);
215
216        // Compute reverse postorder numbering.
217        for (i, block_id) in DfsPostOrder::new(&cfg, start_block).iter(&cfg).enumerate() {
218            let rev_post_id = body.len() - i;
219            block_data[block_id].reverse_postorder = Some(rev_post_id.try_into().unwrap());
220
221            // Store the dominator tree in `block_data`.
222            if let Some(dominator) = dominator_tree.immediate_dominator(block_id) {
223                block_data[dominator].immediately_dominates.push(block_id);
224            }
225        }
226
227        // Compute the forward graph (without backward edges).
228        let mut fwd_cfg = Cfg::new();
229        let mut loop_entries = HashSet::default();
230        let mut switch_blocks = HashSet::default();
231        for block_id in Dfs::new(&cfg, start_block).iter(&cfg) {
232            fwd_cfg.add_node(block_id);
233
234            if body[block_id].terminator.kind.is_switch() {
235                switch_blocks.insert(block_id);
236                block_data[block_id].is_switch = true;
237            }
238
239            // Iterate over edges into this node (so that we can determine whether this node is a
240            // loop header).
241            let mut incoming_fwd_edges = 0;
242            for from in cfg.neighbors_directed(block_id, petgraph::Direction::Incoming) {
243                // Check if the edge is a backward edge.
244                if block_data[from].reverse_postorder >= block_data[block_id].reverse_postorder {
245                    // This is a backward edge
246                    block_data[block_id].is_loop_header = true;
247                    loop_entries.insert(block_id);
248                    // A cfg is reducible iff the target of every back edge dominates the
249                    // edge's source.
250                    if !dominator_tree.dominators(from).unwrap().contains(&block_id) {
251                        return Err(Irreducible(from));
252                    }
253                } else {
254                    incoming_fwd_edges += 1;
255                    fwd_cfg.add_edge(from, block_id, ());
256                }
257            }
258
259            // Detect merge targets.
260            if incoming_fwd_edges >= 2 {
261                block_data[block_id].is_merge_target = true;
262            }
263        }
264
265        let reachable_without_unwind: HashSet<BlockId> = Dfs::new(&cfg_without_unwind, start_block)
266            .iter(&cfg_without_unwind)
267            .collect();
268
269        // Finish filling in information.
270        for block_id in DfsPostOrder::new(&fwd_cfg, start_block).iter(&fwd_cfg) {
271            let block = &body[block_id];
272            let targets = cfg.neighbors(block_id).collect_vec();
273            let fwd_targets = fwd_cfg.neighbors(block_id).collect_vec();
274
275            // Compute the nodes that are part of unwind paths.
276            block_data[block_id].is_unwind = !reachable_without_unwind.contains(&block_id);
277
278            // Compute the nodes that can only reach error nodes.
279            // The node can only reach error nodes if:
280            // - it is an error node;
281            // - or it has neighbors and they all lead to errors.
282            // Note that if there is a backward edge, `only_reach_error` cannot contain this
283            // node yet. In other words, this does not consider infinite loops as reaching an
284            // error node.
285            if block.terminator.is_error()
286                || (!targets.is_empty()
287                    && targets.iter().all(|&tgt| block_data[tgt].only_reach_error))
288            {
289                block_data[block_id].only_reach_error = true;
290            }
291
292            // Compute the flows between each pair of nodes.
293            let mut flow: IndexVec<src::BlockId, Flow> = mem::take(&mut block_data[block_id].flow);
294            // The flow to self is 1.
295            flow[block_id] = Flow::ONE;
296            // If a block has both regular and unwind targets, don't let the unwind path dilute
297            // the normal-control-flow heuristic used to pick switch exits. Unwind-only subgraphs
298            // still flow through their own targets.
299            let mut flow_targets = fwd_targets
300                .iter()
301                .copied()
302                .filter(|&child| !block_data[child].is_unwind)
303                .collect_vec();
304            if flow_targets.is_empty() {
305                flow_targets = fwd_targets;
306            }
307            // Divide the flow from each child to a given target block by the number of children.
308            // This is a sparse matrix multiplication and could be implemented using a linalg
309            // library.
310            let num_children = flow_targets.len();
311            for child in flow_targets {
312                for grandchild in block_data[child].reachable_including_self() {
313                    // Flow from `child` to `grandchild`
314                    flow[grandchild] += block_data[child].flow[grandchild].divided_by(num_children);
315                }
316            }
317            block_data[block_id].flow = flow;
318
319            // Compute shortest paths to all reachable nodes in the forward graph.
320            block_data[block_id].shortest_paths = dijkstra(&fwd_cfg, block_id, None, |_| 1usize);
321
322            // Fill in the rest of the domination data.
323            let mut dominatees = mem::take(&mut block_data[block_id].immediately_dominates);
324            dominatees.sort_by_key(|&child| block_data[child].reverse_postorder);
325            dominatees.reverse();
326            block_data[block_id].immediately_dominates = dominatees;
327            block_data[block_id].immediately_dominated_merge_targets = block_data[block_id]
328                .immediately_dominates
329                .iter()
330                .copied()
331                .filter(|&child| block_data[child].is_merge_target)
332                .collect();
333        }
334
335        // Fill in the within_loop information. See the docs of `within_loops` to understand what
336        // we're computing.
337        let mut path_dfs = DfsWithPath::new(&cfg, start_block);
338        while let Some(block_id) = path_dfs.next(&cfg) {
339            // Store all the loops on the path to this
340            // node.
341            let mut within_loops: SmallVec<_> = path_dfs
342                .path
343                .iter()
344                .copied()
345                .filter(|&loop_id| block_data[loop_id].is_loop_header)
346                .collect();
347            // The loops that we can reach by taking a single backward edge.
348            let loops_directly_within = within_loops
349                .iter()
350                .copied()
351                .filter(|&loop_header| {
352                    cfg.neighbors_directed(loop_header, petgraph::Direction::Incoming)
353                        .any(|bid| block_data[block_id].shortest_paths.contains_key(&bid))
354                })
355                .collect_vec();
356            // The loops that we can reach by taking any number of backward edges.
357            let loops_indirectly_within: HashSet<_> = loops_directly_within
358                .iter()
359                .copied()
360                .flat_map(|loop_header| &block_data[loop_header].within_loops)
361                .chain(&loops_directly_within)
362                .copied()
363                .collect();
364            within_loops.retain(|id| loops_indirectly_within.contains(id));
365            block_data[block_id].within_loops = within_loops;
366        }
367
368        let mut cfg = CfgInfo {
369            cfg,
370            fwd_cfg,
371            loop_entries,
372            switch_blocks,
373            dominator_tree,
374            block_data,
375        };
376
377        // Pick an exit block for each loop, if we find one.
378        ExitInfo::compute_loop_exits(ctx, &mut cfg);
379
380        // Pick an exit block for each switch, if we find one.
381        ExitInfo::compute_switch_exits(&mut cfg);
382
383        Ok(cfg)
384    }
385
386    fn block_data(&self, block_id: BlockId) -> &BlockData {
387        &self.block_data[block_id]
388    }
389    // fn can_reach(&self, src: BlockId, tgt: BlockId) -> bool {
390    //     self.block_data[src].shortest_paths.contains_key(&tgt)
391    // }
392    fn topo_rank(&self, block_id: BlockId) -> u32 {
393        self.block_data[block_id].reverse_postorder.unwrap()
394    }
395    #[expect(unused)]
396    fn is_backward_edge(&self, src: BlockId, tgt: BlockId) -> bool {
397        self.block_data[src].reverse_postorder >= self.block_data[tgt].reverse_postorder
398            && self.cfg.contains_edge(src, tgt)
399    }
400
401    /// Check if the node is within the given loop.
402    fn is_within_loop(&self, loop_header: src::BlockId, block_id: src::BlockId) -> bool {
403        self.block_data[block_id]
404            .within_loops
405            .contains(&loop_header)
406    }
407
408    /// Check if all paths from `src` to nodes in `target_set` go through `through_node`. If `src`
409    /// is already in `target_set`, we ignore that empty path.
410    fn all_paths_go_through(
411        &self,
412        src: src::BlockId,
413        through_node: src::BlockId,
414        target_set: &HashSet<src::BlockId>,
415    ) -> bool {
416        let graph = EdgeFiltered::from_fn(&self.fwd_cfg, |edge| edge.source() != through_node);
417        !Dfs::new(&graph, src)
418            .iter(&graph)
419            .skip(1) // skip src
420            .any(|bid| target_set.contains(&bid))
421    }
422}
423
424impl BlockData {
425    fn shortest_paths_including_self(&self) -> impl Iterator<Item = (BlockId, usize)> {
426        self.shortest_paths.iter().map(|(bid, d)| (*bid, *d))
427    }
428    fn shortest_paths_excluding_self(&self) -> impl Iterator<Item = (BlockId, usize)> {
429        self.shortest_paths_including_self()
430            .filter(move |&(bid, _)| bid != self.id)
431    }
432    fn reachable_including_self(&self) -> impl Iterator<Item = BlockId> {
433        self.shortest_paths_including_self().map(|(bid, _)| bid)
434    }
435    fn reachable_excluding_self(&self) -> impl Iterator<Item = BlockId> {
436        self.shortest_paths_excluding_self().map(|(bid, _)| bid)
437    }
438    #[expect(unused)]
439    fn can_reach_excluding_self(&self, other: BlockId) -> bool {
440        self.shortest_paths.contains_key(&other) && self.id != other
441    }
442}
443
444/// See [`ExitInfo::compute_loop_exit_ranks`].
445#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
446struct LoopExitRank {
447    /// Prefer regular blocks over blocks reachable only through unwind edges.
448    is_unwind_block: Reverse<bool>,
449    /// Number of paths we found going to this exit.
450    path_count: usize,
451    /// Distance from the loop header.
452    distance_from_header: Reverse<usize>,
453}
454
455impl ExitInfo {
456    /// Compute the first node on each path that exits the loop.
457    fn compute_loop_exit_starting_points(
458        cfg: &CfgInfo,
459        loop_header: src::BlockId,
460    ) -> Vec<src::BlockId> {
461        let mut loop_exits = Vec::new();
462        // Do a dfs from the loop header while keeping track of the path from the loop header to
463        // the current node.
464        let mut dfs = Dfs::new(&cfg.fwd_cfg, loop_header);
465        while let Some(block_id) = dfs.next(&cfg.fwd_cfg) {
466            // If we've exited all the loops after and including the target one, this node is an
467            // exit node for the target loop.
468            if !cfg.is_within_loop(loop_header, block_id) {
469                loop_exits.push(block_id);
470                // Don't explore any more paths from this node.
471                dfs.discovered.extend(cfg.fwd_cfg.neighbors(block_id));
472            }
473        }
474        loop_exits
475    }
476
477    /// Compute the loop exit candidates along with a rank.
478    ///
479    /// In the simple case, there is one exit node through which all exit paths go. We want to be
480    /// sure to catch that case, and when that's not possible we want to still find a node through
481    /// which a lot of exit paths go.
482    ///
483    /// To do that, we first count for each exit node how many exit paths go through it, and pick
484    /// the node with most occurrences. If there are many such nodes, we pick the one with shortest
485    /// distance from the loop header. Finally if there are still many such nodes, we keep the
486    /// first node found (the order in which we explore the graph is deterministic, and we use an
487    /// insertion-order hash map).
488    ///
489    /// Note that exit candidates will typically be referenced more than once for one loop. This
490    /// comes from the fact that whenever we reach a node outside the current loop, we register
491    /// this node as well as all its children as exit candidates.
492    /// Consider the following example:
493    /// ```text
494    /// while i < max {
495    ///     if cond {
496    ///         break;
497    ///     }
498    ///     s += i;
499    ///     i += 1
500    /// }
501    /// // All the below nodes are exit candidates (each of them is referenced twice)
502    /// s += 1;
503    /// return s;
504    /// ```
505    fn compute_loop_exit_ranks(
506        cfg: &CfgInfo,
507        loop_header: src::BlockId,
508    ) -> SeqHashMap<src::BlockId, LoopExitRank> {
509        let mut loop_exits: SeqHashMap<BlockId, LoopExitRank> = SeqHashMap::new();
510        for block_id in Self::compute_loop_exit_starting_points(cfg, loop_header) {
511            for bid in cfg.block_data(block_id).reachable_including_self() {
512                loop_exits
513                    .entry(bid)
514                    .or_insert_with(|| LoopExitRank {
515                        is_unwind_block: Reverse(cfg.block_data[bid].is_unwind),
516                        path_count: 0,
517                        distance_from_header: Reverse(
518                            cfg.block_data[loop_header].shortest_paths[&bid],
519                        ),
520                    })
521                    .path_count += 1;
522            }
523        }
524        loop_exits
525    }
526
527    /// A loop exit is any block reachable from the loop header that isn't inside the loop.
528    /// This function choses an exit for every loop. See `compute_loop_exit_ranks` for how we
529    /// select them.
530    ///
531    /// For example:
532    /// ```text
533    /// while ... {
534    ///    ...
535    ///    if ... {
536    ///        // We can't reach the loop entry from here: this is an exit
537    ///        // candidate
538    ///        return;
539    ///    }
540    /// }
541    /// // This is another exit candidate - and this is the one we want to use
542    /// // as the "real" exit...
543    /// ...
544    /// ```
545    ///
546    /// Once we listed all the exit candidates, we find the "best" one for every loop. The best
547    /// exit is the following one:
548    /// - it is the one which is used the most times (note that there can be
549    ///   several candidates which are referenced strictly more than once: see the
550    ///   comment below)
551    /// - if several exits have the same number of occurrences, we choose the one
552    ///   for which we goto the "earliest" (earliest meaning that the goto is close to
553    ///   the loop entry node in the AST). The reason is that all the loops should
554    ///   have an outer if ... then ... else ... which executes the loop body or goes
555    ///   to the exit (note that this is not necessarily the first
556    ///   if ... then ... else ... we find: loop conditions can be arbitrary
557    ///   expressions, containing branchings).
558    ///
559    /// # Several candidates for a loop exit:
560    /// =====================================
561    /// There used to be a sanity check to ensure there are no two different
562    /// candidates with exactly the same number of occurrences and distance from
563    /// the entry of the loop, if the number of occurrences is > 1.
564    ///
565    /// We removed it because it does happen, for instance here (the match
566    /// introduces an `unreachable` node, and it has the same number of
567    /// occurrences and the same distance to the loop entry as the `panic`
568    /// node):
569    ///
570    /// ```text
571    /// pub fn list_nth_mut_loop_pair<'a, T>(
572    ///     mut ls: &'a mut List<T>,
573    ///     mut i: u32,
574    /// ) -> &'a mut T {
575    ///     loop {
576    ///         match ls {
577    ///             List::Nil => {
578    ///                 panic!() // <-- best candidate
579    ///             }
580    ///             List::Cons(x, tl) => {
581    ///                 if i == 0 {
582    ///                     return x;
583    ///                 } else {
584    ///                     ls = tl;
585    ///                     i -= 1;
586    ///                 }
587    ///             }
588    ///             _ => {
589    ///               // Note that Rustc always introduces an unreachable branch after
590    ///               // desugaring matches.
591    ///               unreachable!(), // <-- best candidate
592    ///             }
593    ///         }
594    ///     }
595    /// }
596    /// ```
597    ///
598    /// When this happens we choose an exit candidate whose edges don't necessarily
599    /// lead to an error (above there are none, so we don't choose any exits). Note
600    /// that this last condition is important to prevent loops from being unnecessarily
601    /// nested:
602    ///
603    /// ```text
604    /// pub fn nested_loops_enum(step_out: usize, step_in: usize) -> usize {
605    ///     let mut s = 0;
606    ///
607    ///     for _ in 0..128 { // We don't want this loop to be nested with the loops below
608    ///         s += 1;
609    ///     }
610    ///
611    ///     for _ in 0..(step_out) {
612    ///         for _ in 0..(step_in) {
613    ///             s += 1;
614    ///         }
615    ///     }
616    ///
617    ///     s
618    /// }
619    /// ```
620    fn compute_loop_exits(_ctx: &TransformCtx, cfg: &mut CfgInfo) {
621        for &loop_id in &cfg.loop_entries {
622            // Compute the candidates.
623            let loop_exits: SeqHashMap<BlockId, LoopExitRank> =
624                Self::compute_loop_exit_ranks(cfg, loop_id);
625            // We choose the exit with:
626            // - the most occurrences
627            // - the least total distance (if there are several possibilities)
628            // - doesn't necessarily lead to an error (panic, unreachable)
629            let best_exits: Vec<(BlockId, LoopExitRank)> =
630                loop_exits.into_iter().max_set_by_key(|&(_, rank)| rank);
631            // If there is exactly one best candidate, use it. Otherwise we need to split further.
632            let chosen_exit = match best_exits.into_iter().map(|(bid, _)| bid).exactly_one() {
633                Ok(best_exit) => Some(best_exit),
634                Err(best_exits) => {
635                    // Remove the candidates which only lead to errors (panic or unreachable).
636                    // If there is exactly one candidate we select it, otherwise we do not select any
637                    // exit.
638                    // We don't want to select any exit if we are in the below situation
639                    // (all paths lead to errors). We added a sanity check below to
640                    // catch the situations where there are several exits which don't
641                    // lead to errors.
642                    //
643                    // Example:
644                    // ========
645                    // ```
646                    // loop {
647                    //     match ls {
648                    //         List::Nil => {
649                    //             panic!() // <-- best candidate
650                    //         }
651                    //         List::Cons(x, tl) => {
652                    //             if i == 0 {
653                    //                 return x;
654                    //             } else {
655                    //                 ls = tl;
656                    //                 i -= 1;
657                    //             }
658                    //         }
659                    //         _ => {
660                    //           unreachable!(); // <-- best candidate (Rustc introduces an `unreachable` case)
661                    //         }
662                    //     }
663                    // }
664                    // ```
665                    // The `exactly_one` can fail, see `tests/ui/control-flow/ambiguous-loop-exit.rs`
666                    best_exits
667                        .filter(|&bid| !cfg.block_data[bid].only_reach_error)
668                        .exactly_one()
669                        .ok()
670                }
671            };
672            cfg.block_data[loop_id].exit_info.loop_exit = chosen_exit;
673        }
674    }
675
676    /// Let's consider the following piece of code:
677    /// ```text
678    /// if cond1 { ... } else { ... };
679    /// if cond2 { ... } else { ... };
680    /// ```
681    /// Once converted to MIR, the control-flow is destructured, which means we
682    /// have gotos everywhere. When reconstructing the control-flow, we have
683    /// to be careful about the point where we should join the two branches of
684    /// the first if.
685    /// For instance, if we don't notice they should be joined at some point (i.e,
686    /// whatever the branch we take, there is a moment when we go to the exact
687    /// same place, just before the second if), we might generate code like
688    /// this, with some duplicata:
689    /// ```text
690    /// if cond1 { ...; if cond2 { ... } else { ...} }
691    /// else { ...; if cond2 { ... } else { ...} }
692    /// ```
693    ///
694    /// Such a reconstructed program is valid, but it is definitely non-optimal:
695    /// it is very different from the original program (making it less clean and
696    /// clear), more bloated, and might involve duplicating the proof effort.
697    ///
698    /// For this reason, we need to find the "exit" of the first switch, which is
699    /// the point where the two branches join. Note that this can be a bit tricky,
700    /// because there may be more than two branches (if we do `switch(x) { ... }`),
701    /// and some of them might not join (if they contain a `break`, `panic`,
702    /// `return`, etc.).
703    ///
704    /// In order to compute the switch exits, we simply recursively compute a
705    /// topologically ordered set of "filtered successors" as follows (note
706    /// that we work in the CFG *without* back edges):
707    /// - for a block which doesn't branch (only one successor), the filtered
708    ///   successors is the set of reachable nodes.
709    /// - for a block which branches, we compute the nodes reachable from all
710    ///   the children, and find the "best" intersection between those.
711    ///   Note that we find the "best" intersection (a pair of branches which
712    ///   maximize the intersection of filtered successors) because some branches
713    ///   might never join the control-flow of the other branches, if they contain
714    ///   a `break`, `return`, `panic`, etc., like here:
715    ///   ```text
716    ///   if b { x = 3; } { return; }
717    ///   y += x;
718    ///   ...
719    ///   ```
720    /// Note that with nested switches, the branches of the inner switches might
721    /// goto the exits of the outer switches: for this reason, we give precedence
722    /// to the outer switches.
723    fn compute_switch_exits(cfg: &mut CfgInfo) {
724        // We need to give precedence to the outer switches: we thus iterate
725        // over the switch blocks in topological order.
726        let mut exits_set = HashSet::default();
727        for bid in cfg
728            .switch_blocks
729            .iter()
730            .copied()
731            .sorted_unstable_by_key(|&bid| (cfg.topo_rank(bid), bid))
732        {
733            let block_data = &cfg.block_data[bid];
734            // Find the best successor: this is the node with the highest flow, and the lowest
735            // topological rank. If several nodes have the same flow, we want to take the highest
736            // one in the hierarchy: hence the use of the topological rank.
737            //
738            // Ex.:
739            // ```text
740            // A  -- we start here
741            // |
742            // |---------------------------------------
743            // |            |            |            |
744            // B:(0.25,-1)  C:(0.25,-2)  D:(0.25,-3)  E:(0.25,-4)
745            // |            |            |
746            // |--------------------------
747            // |
748            // F:(0.75,-5)
749            // |
750            // |
751            // G:(0.75,-6)
752            // ```
753            // The "best" node (with the highest (flow, rank) in the graph above is F.
754            // If the switch is inside a loop, we also only consider exists that are inside that
755            // same loop. There must be one, otherwise the switch entry would not be inside the
756            // loop.
757            let current_loop = block_data.within_loops.last().copied();
758            let best_exit: Option<BlockId> = block_data
759                .reachable_excluding_self()
760                .filter(|&b| {
761                    current_loop.is_none_or(|current_loop| cfg.is_within_loop(current_loop, b))
762                })
763                .max_by_key(|&id| {
764                    let is_unwind_block = Reverse(cfg.block_data[id].is_unwind);
765                    let flow = &block_data.flow[id];
766                    let rank = Reverse(cfg.topo_rank(id));
767                    ((is_unwind_block, flow, rank), id)
768                });
769            // We have an exit candidate: we first check that it was not already taken by an
770            // external switch.
771            //
772            // We then check that we can't reach the exit of an external switch from one of the
773            // branches, without going through the exit candidate. We do this by simply checking
774            // that we can't reach any of the exits of outer switches.
775            //
776            // The reason is that it can lead to code like the following:
777            // ```
778            // if ... { // if #1
779            //   if ... { // if #2
780            //     ...
781            //     // here, we have a `goto b1`, where b1 is the exit
782            //     // of if #2: we thus stop translating the blocks.
783            //   }
784            //   else {
785            //     ...
786            //     // here, we have a `goto b2`, where b2 is the exit
787            //     // of if #1: we thus stop translating the blocks.
788            //   }
789            //   // We insert code for the block b1 here (which is the exit of
790            //   // the exit of if #2). However, this block should only
791            //   // be executed in the branch "then" of the if #2, not in
792            //   // the branch "else".
793            //   ...
794            // }
795            // else {
796            //   ...
797            // }
798            // ```
799            if let Some(exit_id) = best_exit
800                && !exits_set.contains(&exit_id)
801                && cfg.all_paths_go_through(bid, exit_id, &exits_set)
802            {
803                exits_set.insert(exit_id);
804                cfg.block_data[bid].exit_info.switch_exit = Some(exit_id);
805            }
806        }
807    }
808}
809
810/// Iter over the last non-switch statements that may be executed on any branch of this block.
811/// Skips over `Nop`s.
812fn iter_tail_statements(block: &mut tgt::Block, f: &mut impl FnMut(&mut tgt::Statement)) {
813    let Some(st) = block
814        .statements
815        .iter_mut()
816        .rev()
817        .find(|st| !st.kind.is_nop())
818    else {
819        return;
820    };
821    if let tgt::StatementKind::Switch { branches, .. } = &mut st.kind {
822        for block in branches {
823            iter_tail_statements(block, f);
824        }
825    } else {
826        f(st)
827    };
828}
829
830type Depth = usize;
831
832#[derive(Debug, Clone, Copy)]
833enum SpecialJumpKind {
834    /// This block can be reached by a `continue` to the given depth.
835    LoopContinue(Depth),
836    /// This block can be reached by a `break` to the given depth. This comes from a loop.
837    LoopBreak(Depth),
838    /// This block can be reached by a `break` to the given depth. This is a `loop` context
839    /// introduced only for forward jumps.
840    ForwardBreak(Depth),
841    /// This block can be reached by doing nothing as this is the next block that will be
842    /// translated. Only applies if this block is at the top of the stack.
843    NextBlock,
844}
845
846#[derive(Clone, Copy)]
847struct SpecialJump {
848    /// The relevant block.
849    target_block: BlockId,
850    /// How to translate a jump to the target block.
851    kind: SpecialJumpKind,
852}
853
854impl SpecialJump {
855    fn new(target_block: BlockId, kind: SpecialJumpKind) -> Self {
856        Self { target_block, kind }
857    }
858}
859
860impl std::fmt::Debug for SpecialJump {
861    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
862        write!(f, "SpecialJump({}, {:?})", self.target_block, self.kind)
863    }
864}
865
866/// How to handle blocks reachable from multiple branches.
867enum ReconstructMode {
868    /// Duplicate blocks reachable from multiple branches.
869    Duplicate,
870    /// Insert loops for the purpose of breaking forward to them to implement DAG-like control-flow
871    /// without duplicating blocks.
872    /// Based on the algorithm from "Beyond Relooper" (https://dl.acm.org/doi/10.1145/3547621).
873    ForwardBreak,
874}
875
876struct ReconstructCtx<'a> {
877    cfg: CfgInfo,
878    body: &'a src::ExprBody,
879    /// The depth of `loop` contexts we may `break`/`continue` to.
880    break_context_depth: Depth,
881    /// Stack of block ids that should be translated to special jumps (`break`/`continue`/do
882    /// nothing) in the current context.
883    /// The block where control-flow continues naturally after this block is kept at the top of the
884    /// stack.
885    special_jump_stack: Vec<SpecialJump>,
886    mode: ReconstructMode,
887}
888
889impl<'a> ReconstructCtx<'a> {
890    fn build(ctx: &TransformCtx, src_body: &'a src::ExprBody) -> Result<Self, Irreducible> {
891        // Compute all sorts of graph-related information about the control-flow graph, including
892        // reachability, the dominator tree, loop entries, and loop/switch exits.
893        let cfg = CfgInfo::build(ctx, &src_body.body)?;
894
895        // Translate the body by reconstructing the loops and the
896        // conditional branchings.
897        let allow_duplication = true;
898        Ok(ReconstructCtx {
899            cfg,
900            body: src_body,
901            break_context_depth: 0,
902            special_jump_stack: Vec::new(),
903            mode: if allow_duplication {
904                ReconstructMode::Duplicate
905            } else {
906                ReconstructMode::ForwardBreak
907            },
908        })
909    }
910
911    fn translate_statement(&self, st: &src::Statement) -> tgt::Statement {
912        let span = st.span;
913        let st = match st.kind.clone() {
914            src::StatementKind::Assign(place, rvalue) => tgt::StatementKind::Assign(place, rvalue),
915            src::StatementKind::SetDiscriminant(place, variant_id) => {
916                tgt::StatementKind::SetDiscriminant(place, variant_id)
917            }
918            src::StatementKind::StorageLive(var_id) => tgt::StatementKind::StorageLive(var_id),
919            src::StatementKind::StorageDead(var_id) => tgt::StatementKind::StorageDead(var_id),
920            src::StatementKind::Assert { assert, on_failure } => tgt::StatementKind::Assert {
921                assert,
922                on_failure: on_failure.clone(),
923                on_unwind: tgt::Statement::new(span, tgt::StatementKind::UnwindResume).into_block(),
924            },
925            src::StatementKind::PlaceMention(place) => tgt::StatementKind::PlaceMention(place),
926            src::StatementKind::Borrowck(statement) => tgt::StatementKind::Borrowck(statement),
927            src::StatementKind::Nop => tgt::StatementKind::Nop,
928        };
929        tgt::Statement::new(span, st)
930    }
931
932    /// Translate a jump to the given block. The span is used to create the jump statement, if any.
933    #[tracing::instrument(skip(self), ret, fields(stack = ?self.special_jump_stack))]
934    fn translate_jump(&mut self, span: Span, target_block: src::BlockId) -> tgt::Block {
935        match self
936            .special_jump_stack
937            .iter()
938            .rev()
939            .enumerate()
940            .find(|(_, j)| j.target_block == target_block)
941        {
942            Some((i, jump_target)) => {
943                let mk_block = |kind| tgt::Statement::new(span, kind).into_block();
944                match jump_target.kind {
945                    // The top of the stack is where control-flow goes naturally, no need to add a
946                    // `break`/`continue`. We only do that in `ForwardBreak` mode to avoid breaking aeneas.
947                    SpecialJumpKind::LoopContinue(_)
948                    | SpecialJumpKind::ForwardBreak(_)
949                    | SpecialJumpKind::NextBlock
950                        if i == 0 && matches!(self.mode, ReconstructMode::ForwardBreak) =>
951                    {
952                        mk_block(tgt::StatementKind::Nop)
953                    }
954                    SpecialJumpKind::LoopContinue(depth) => mk_block(tgt::StatementKind::Continue(
955                        self.break_context_depth - depth,
956                    )),
957                    SpecialJumpKind::ForwardBreak(depth) | SpecialJumpKind::LoopBreak(depth) => {
958                        mk_block(tgt::StatementKind::Break(self.break_context_depth - depth))
959                    }
960                    SpecialJumpKind::NextBlock if i == 0 => mk_block(tgt::StatementKind::Nop),
961                    // Translate the block without a jump.
962                    SpecialJumpKind::NextBlock => self.translate_block(target_block),
963                }
964            }
965            // Translate the block without a jump.
966            None => self.translate_block(target_block),
967        }
968    }
969
970    fn translate_terminator(&mut self, terminator: &src::Terminator) -> tgt::Block {
971        let src_span = terminator.span;
972
973        match &terminator.kind {
974            src::TerminatorKind::Abort(kind) => {
975                tgt::Statement::new(src_span, tgt::StatementKind::Abort(kind.clone())).into_block()
976            }
977            src::TerminatorKind::Return => {
978                tgt::Statement::new(src_span, tgt::StatementKind::Return).into_block()
979            }
980            src::TerminatorKind::UnwindResume => {
981                tgt::Statement::new(src_span, tgt::StatementKind::UnwindResume).into_block()
982            }
983            src::TerminatorKind::Call {
984                call,
985                target,
986                on_unwind,
987            } => {
988                let on_unwind = self.translate_block(*on_unwind);
989                let st = tgt::Statement::new(
990                    src_span,
991                    tgt::StatementKind::Call {
992                        call: call.clone(),
993                        on_unwind,
994                    },
995                );
996                let mut block = self.translate_jump(terminator.span, *target);
997                block.statements.insert(0, st);
998                block
999            }
1000            src::TerminatorKind::Drop {
1001                kind,
1002                place,
1003                fn_ptr,
1004                target,
1005                on_unwind,
1006            } => {
1007                let on_unwind = self.translate_block(*on_unwind);
1008                let st = tgt::Statement::new(
1009                    src_span,
1010                    tgt::StatementKind::Drop {
1011                        place: place.clone(),
1012                        fn_ptr: fn_ptr.clone(),
1013                        kind: *kind,
1014                        on_unwind,
1015                    },
1016                );
1017                let mut block = self.translate_jump(terminator.span, *target);
1018                block.statements.insert(0, st);
1019                block
1020            }
1021            src::TerminatorKind::Assert {
1022                assert,
1023                target,
1024                on_unwind,
1025            } => {
1026                let on_unwind = self.translate_block(*on_unwind);
1027                let st = tgt::StatementKind::Assert {
1028                    assert: assert.clone(),
1029                    on_failure: AbortKind::Panic(None),
1030                    on_unwind,
1031                };
1032                let target = self.translate_jump(terminator.span, *target);
1033                tgt::Statement::new(src_span, st).into_block().merge(target)
1034            }
1035            src::TerminatorKind::InlineAsm {
1036                asm,
1037                targets,
1038                on_unwind,
1039            } => {
1040                let targets = targets
1041                    .iter()
1042                    .map(|target| self.translate_jump(terminator.span, *target))
1043                    .collect();
1044                let on_unwind = self.translate_block(*on_unwind);
1045                let st = tgt::StatementKind::InlineAsm {
1046                    asm: asm.clone(),
1047                    targets,
1048                    on_unwind,
1049                };
1050                tgt::Statement::new(src_span, st).into_block()
1051            }
1052            src::TerminatorKind::Goto { target } => self.translate_jump(terminator.span, *target),
1053            src::TerminatorKind::Switch { data, branches } => {
1054                let mut branches =
1055                    branches.map_ref(|target| self.translate_jump(terminator.span, *target));
1056                // If we considered the match to be exhaustive, the fallback block is still in
1057                // the list of branches and can be removed.
1058                if data.fallback.is_none()
1059                    && let Some(last_branch_id) = branches.indices().next_back()
1060                    && !data
1061                        .branches
1062                        .iter()
1063                        .map(|(_, branch_id)| *branch_id)
1064                        .contains(&last_branch_id)
1065                {
1066                    branches.pop();
1067                }
1068
1069                // Return
1070                let span = combine_span_iter(branches.iter().map(|branch| &branch.span));
1071                let span = combine_span(&src_span, &span);
1072                let st = tgt::StatementKind::Switch {
1073                    data: data.clone(),
1074                    branches,
1075                };
1076                tgt::Statement::new(span, st).into_block()
1077            }
1078        }
1079    }
1080
1081    /// Translate just the block statements and terminator.
1082    fn translate_block_itself(&mut self, block_id: BlockId) -> tgt::Block {
1083        let block = &self.body.body[block_id];
1084        // Translate the statements inside the block
1085        let statements = block
1086            .statements
1087            .iter()
1088            .map(|st| self.translate_statement(st))
1089            .collect_vec();
1090        // Translate the terminator.
1091        let terminator = self.translate_terminator(&block.terminator);
1092        // Prepend the statements to the terminator.
1093        if let Some(st) = tgt::Block::from_seq(statements) {
1094            st.merge(terminator)
1095        } else {
1096            terminator
1097        }
1098    }
1099
1100    /// Translate a block including surrounding control-flow like looping.
1101    #[tracing::instrument(skip(self), fields(stack = ?self.special_jump_stack))]
1102    fn translate_block(&mut self, block_id: src::BlockId) -> tgt::Block {
1103        ensure_sufficient_stack(|| self.translate_block_inner(block_id))
1104    }
1105    fn translate_block_inner(&mut self, block_id: src::BlockId) -> tgt::Block {
1106        // Some of the blocks we might jump to inside this tree can't be translated as normal
1107        // blocks: the loop backward edges must become `continue`s and the merge nodes may need
1108        // some care if we're jumping to them from distant locations.
1109        // For this purpose, we push to the `special_jump_stack` the block ids that must be
1110        // translated specially. In `translate_jump` we check the stack. At the end of this
1111        // function we restore the stack to its previous state.
1112        let old_context_depth = self.special_jump_stack.len();
1113        let block_data = &self.cfg.block_data[block_id];
1114        let span = block_data.span;
1115
1116        // Catch jumps to the loop header or loop exit.
1117        if block_data.is_loop_header {
1118            self.break_context_depth += 1;
1119            if let Some(exit_id) = block_data.exit_info.loop_exit {
1120                self.special_jump_stack.push(SpecialJump::new(
1121                    exit_id,
1122                    SpecialJumpKind::LoopBreak(self.break_context_depth),
1123                ));
1124            }
1125            // Put the next block at the top of the stack.
1126            self.special_jump_stack.push(SpecialJump::new(
1127                block_id,
1128                SpecialJumpKind::LoopContinue(self.break_context_depth),
1129            ));
1130        }
1131
1132        // Catch jumps to a merge node.
1133        let merge_children = &block_data.immediately_dominated_merge_targets;
1134        if let ReconstructMode::ForwardBreak = self.mode {
1135            // We support forward-jumps using `break`
1136            // The child with highest postorder numbering is nested outermost in this scheme.
1137            for &child in merge_children {
1138                self.break_context_depth += 1;
1139                self.special_jump_stack.push(SpecialJump::new(
1140                    child,
1141                    SpecialJumpKind::ForwardBreak(self.break_context_depth),
1142                ));
1143            }
1144        }
1145
1146        if let Some(bid) = block_data.exit_info.switch_exit
1147            && !block_data.is_loop_header
1148            && !(matches!(self.mode, ReconstructMode::ForwardBreak)
1149                && merge_children.contains(&bid))
1150        {
1151            // Move some code that would be inside one or several switch branches to be after the
1152            // switch intead.
1153            self.special_jump_stack
1154                .push(SpecialJump::new(bid, SpecialJumpKind::NextBlock));
1155        }
1156
1157        // Translate this block. Any jumps to a loop header or a merge node will be replaced with
1158        // `continue`/`break`.
1159        let mut block = self.translate_block_itself(block_id);
1160
1161        // Reset the state to what it was previously, and translate what remains.
1162        let new_statement = move |kind| tgt::Statement::new(block.span, kind);
1163        while self.special_jump_stack.len() > old_context_depth {
1164            let special_jump = self.special_jump_stack.pop().unwrap();
1165            match &special_jump.kind {
1166                SpecialJumpKind::LoopContinue(_) => {
1167                    self.break_context_depth -= 1;
1168                    if let ReconstructMode::ForwardBreak = self.mode {
1169                        // We add `continue` at the end for users that don't know that the default
1170                        // behavior at the end of a loop block is `continue`. Not needed for
1171                        // `Duplicate` mode because we use explicit `continue`s there. TODO: clean
1172                        // that up.
1173                        block
1174                            .statements
1175                            .push(new_statement(tgt::StatementKind::Continue(0)));
1176                    }
1177                    block = new_statement(tgt::StatementKind::Loop(block)).into_block();
1178                }
1179                SpecialJumpKind::ForwardBreak(_) => {
1180                    self.break_context_depth -= 1;
1181                    // Remove unneeded `break`s in branches leading up to that final one.
1182                    iter_tail_statements(&mut block, &mut |st| {
1183                        if matches!(st.kind, tgt::StatementKind::Break(0)) {
1184                            st.kind = tgt::StatementKind::Nop;
1185                        }
1186                    });
1187                    // We add a `loop { ...; break }` so that we can use `break` to jump forward.
1188                    block
1189                        .statements
1190                        .push(new_statement(tgt::StatementKind::Break(0)));
1191                    block = new_statement(tgt::StatementKind::Loop(block)).into_block();
1192                    // We must translate the merge nodes after the block used for forward jumps to
1193                    // them.
1194                    let next_block = self.translate_jump(span, special_jump.target_block);
1195                    block = block.merge(next_block);
1196                }
1197                SpecialJumpKind::NextBlock | SpecialJumpKind::LoopBreak(..) => {
1198                    let next_block = self.translate_jump(span, special_jump.target_block);
1199                    block = block.merge(next_block);
1200                }
1201            }
1202        }
1203        block
1204    }
1205}
1206
1207fn remove_useless_jump_blocks(body: &mut tgt::ExprBody) {
1208    use tgt::StatementKind;
1209    #[derive(Default)]
1210    struct Count {
1211        continue_count: u32,
1212        break_count: u32,
1213    }
1214    #[derive(Default, Visitor)]
1215    struct CountJumpsVisitor {
1216        counts: HashMap<StatementId, Count>,
1217        loop_stack: Vec<StatementId>,
1218    }
1219    #[derive(Visitor)]
1220    struct RemoveUselessJumpsVisitor {
1221        counts: HashMap<StatementId, Count>,
1222        /// For every loop we encounter, whether we're keeping it or removing it.
1223        loop_stack: Vec<bool>,
1224    }
1225
1226    impl VisitBodyMut for CountJumpsVisitor {
1227        fn visit_llbc_statement(&mut self, st: &mut tgt::Statement) -> ControlFlow<Self::Break> {
1228            if let StatementKind::Loop(_) = &st.kind {
1229                self.loop_stack.push(st.id);
1230            }
1231            match &st.kind {
1232                StatementKind::Break(depth) => {
1233                    let loop_id = self.loop_stack[self.loop_stack.len() - 1 - depth];
1234                    self.counts.entry(loop_id).or_default().break_count += 1;
1235                }
1236                StatementKind::Continue(depth) => {
1237                    let loop_id = self.loop_stack[self.loop_stack.len() - 1 - depth];
1238                    self.counts.entry(loop_id).or_default().continue_count += 1;
1239                }
1240                _ => {}
1241            }
1242            self.visit_inner(st)?;
1243            if let StatementKind::Loop(_) = &st.kind {
1244                self.loop_stack.pop();
1245            }
1246            ControlFlow::Continue(())
1247        }
1248    }
1249
1250    impl VisitBodyMut for RemoveUselessJumpsVisitor {
1251        fn visit_llbc_block(&mut self, block: &mut tgt::Block) -> ControlFlow<Self::Break> {
1252            for mut st in mem::take(&mut block.statements) {
1253                if let tgt::StatementKind::Loop(block) = &mut st.kind {
1254                    let counts = &self.counts[&st.id];
1255                    let remove = counts.continue_count == 0
1256                        && counts.break_count == 1
1257                        && matches!(
1258                            block.statements.last().unwrap().kind,
1259                            StatementKind::Break(0)
1260                        );
1261                    self.loop_stack.push(!remove);
1262                }
1263                self.visit(&mut st)?;
1264                if st.kind.is_loop() && !self.loop_stack.pop().unwrap() {
1265                    // Remove the loop.
1266                    let StatementKind::Loop(mut inner_block) = st.kind else {
1267                        unreachable!()
1268                    };
1269                    inner_block.statements.last_mut().unwrap().kind = StatementKind::Nop;
1270                    block.statements.extend(inner_block.statements);
1271                } else {
1272                    block.statements.push(st);
1273                }
1274            }
1275            ControlFlow::Continue(())
1276        }
1277        fn enter_llbc_statement(&mut self, st: &mut tgt::Statement) {
1278            match &st.kind {
1279                StatementKind::Break(depth) => {
1280                    let new_depth = self.loop_stack[self.loop_stack.len() - depth..]
1281                        .iter()
1282                        .filter(|&&keep| keep)
1283                        .count();
1284                    st.kind = StatementKind::Break(new_depth);
1285                }
1286                StatementKind::Continue(depth) => {
1287                    let new_depth = self.loop_stack[self.loop_stack.len() - depth..]
1288                        .iter()
1289                        .filter(|&&keep| keep)
1290                        .count();
1291                    st.kind = StatementKind::Continue(new_depth);
1292                }
1293                _ => {}
1294            }
1295        }
1296    }
1297
1298    let mut v = CountJumpsVisitor::default();
1299    body.body.drive_body_mut(&mut v);
1300    let mut v = RemoveUselessJumpsVisitor {
1301        counts: v.counts,
1302        loop_stack: Default::default(),
1303    };
1304    body.body.drive_body_mut(&mut v);
1305}
1306
1307fn translate_body(ctx: &mut TransformCtx, body: &mut Body) {
1308    use Body::{Structured, Unstructured};
1309    let Unstructured(src_body) = body else {
1310        panic!("Called `ullbc_to_llbc` on an already restructured body")
1311    };
1312    trace!("About to translate to ullbc: {:?}", src_body.span);
1313    // Report the time spent per body size, since this pass is superlinear in the number of blocks.
1314    let _guard = crate::timing::scope_lazy("ullbc_to_llbc-body", || {
1315        format!(
1316            "{:05} blocks or less",
1317            src_body.body.len().next_power_of_two()
1318        )
1319    });
1320
1321    // Calculate info about the graph and heuristically determine loop and switch exit blocks.
1322    let start_block = BlockId::ZERO;
1323    let mut ctx = match ReconstructCtx::build(ctx, src_body) {
1324        Ok(ctx) => ctx,
1325        Err(Irreducible(bid)) => {
1326            let span = src_body.body[bid].terminator.span;
1327            register_error!(
1328                ctx,
1329                span,
1330                "the control-flow graph of this function is not reducible"
1331            );
1332            panic!("can't reconstruct irreducible control-flow")
1333        }
1334    };
1335    // Translate the blocks using the computed data.
1336    let tgt_body = ctx.translate_block(start_block);
1337
1338    let mut tgt_body = tgt::ExprBody {
1339        span: src_body.span,
1340        locals: src_body.locals.clone(),
1341        bound_body_regions: src_body.bound_body_regions,
1342        body: tgt_body,
1343        comments: src_body.comments.clone(),
1344    };
1345    remove_useless_jump_blocks(&mut tgt_body);
1346
1347    *body = Structured(tgt_body);
1348}
1349
1350pub struct Transform;
1351impl TransformPass for Transform {
1352    fn transform_ctx(&self, ctx: &mut TransformCtx) {
1353        // Translate the bodies one at a time.
1354        ctx.for_each_body(|ctx, body| {
1355            translate_body(ctx, body);
1356        });
1357
1358        if ctx.options.print_built_llbc {
1359            eprintln!("# LLBC resulting from control-flow reconstruction:\n\n{ctx}\n",);
1360        } else {
1361            trace!("# LLBC resulting from control-flow reconstruction:\n\n{ctx}\n",);
1362        }
1363    }
1364}