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}