1use alloc::collections::VecDeque;
2use core::ops;
3
4use crate::{IndexMap, IndexSet, error::InternalError};
5use malachite_bigint::BigInt;
6use num_traits::{ToPrimitive, Zero};
7
8use rustpython_compiler_core::{
9 OneIndexed, SourceLocation,
10 bytecode::{
11 AnyInstruction, Arg, CO_FAST_CELL, CO_FAST_FREE, CO_FAST_HIDDEN, CO_FAST_LOCAL, CodeFlags,
12 CodeObject, CodeUnit, CodeUnits, ConstantData, ExceptionTableEntry, InstrDisplayContext,
13 Instruction, InstructionMetadata, Label, OpArg, PseudoInstruction, PyCodeLocationInfoKind,
14 encode_exception_table, oparg,
15 },
16 varint::{write_signed_varint, write_varint},
17};
18
19#[derive(Clone, Copy, Debug, PartialEq, Eq)]
21struct LineTableLocation {
22 line: i32,
23 end_line: i32,
24 col: i32,
25 end_col: i32,
26}
27
28const MAX_INT_SIZE_BITS: u64 = 128;
29const MIN_CONST_SEQUENCE_SIZE: usize = 3;
30
31#[derive(Clone, Debug)]
34pub struct CodeUnitMetadata {
35 pub name: String, pub qualname: Option<String>, pub consts: IndexSet<ConstantData>, pub names: IndexSet<String>, pub varnames: IndexSet<String>, pub cellvars: IndexSet<String>, pub freevars: IndexSet<String>, pub fast_hidden: IndexMap<String, bool>, pub argcount: u32, pub posonlyargcount: u32, pub kwonlyargcount: u32, pub firstlineno: OneIndexed, }
48#[derive(Copy, Clone, PartialEq, Eq, Debug)]
51pub struct BlockIdx(u32);
52
53impl BlockIdx {
54 pub const NULL: Self = Self::new(u32::MAX);
55
56 #[must_use]
58 pub const fn new(value: u32) -> Self {
59 Self(value)
60 }
61
62 #[must_use]
64 pub const fn idx(self) -> usize {
65 self.0 as usize
66 }
67}
68
69impl From<BlockIdx> for u32 {
70 fn from(block_idx: BlockIdx) -> Self {
71 block_idx.0
72 }
73}
74
75impl ops::Index<BlockIdx> for [Block] {
76 type Output = Block;
77
78 fn index(&self, idx: BlockIdx) -> &Block {
79 &self[idx.idx()]
80 }
81}
82
83impl ops::IndexMut<BlockIdx> for [Block] {
84 fn index_mut(&mut self, idx: BlockIdx) -> &mut Block {
85 &mut self[idx.idx()]
86 }
87}
88
89impl ops::Index<BlockIdx> for Vec<Block> {
90 type Output = Block;
91
92 fn index(&self, idx: BlockIdx) -> &Block {
93 &self[idx.idx()]
94 }
95}
96
97impl ops::IndexMut<BlockIdx> for Vec<Block> {
98 fn index_mut(&mut self, idx: BlockIdx) -> &mut Block {
99 &mut self[idx.idx()]
100 }
101}
102
103#[derive(Clone, Copy, Debug)]
104pub struct InstructionInfo {
105 pub instr: AnyInstruction,
106 pub arg: OpArg,
107 pub target: BlockIdx,
108 pub location: SourceLocation,
109 pub end_location: SourceLocation,
110 pub except_handler: Option<ExceptHandlerInfo>,
111 pub lineno_override: Option<i32>,
113 pub cache_entries: u32,
115}
116
117#[derive(Clone, Copy, Debug, PartialEq, Eq)]
119pub struct ExceptHandlerInfo {
120 pub handler_block: BlockIdx,
122 pub stack_depth: u32,
124 pub preserve_lasti: bool,
126}
127
128#[derive(Debug, Clone)]
132pub struct Block {
133 pub instructions: Vec<InstructionInfo>,
134 pub next: BlockIdx,
135 pub except_handler: bool,
138 pub preserve_lasti: bool,
140 pub start_depth: Option<u32>,
142 pub cold: bool,
144}
145
146impl Default for Block {
147 fn default() -> Self {
148 Self {
149 instructions: Vec::new(),
150 next: BlockIdx::NULL,
151 except_handler: false,
152 preserve_lasti: false,
153 start_depth: None,
154 cold: false,
155 }
156 }
157}
158
159pub struct CodeInfo {
160 pub flags: CodeFlags,
161 pub source_path: String,
162 pub private: Option<String>, pub blocks: Vec<Block>,
165 pub current_block: BlockIdx,
166
167 pub metadata: CodeUnitMetadata,
168
169 pub static_attributes: Option<IndexSet<String>>,
171
172 pub in_inlined_comp: bool,
174
175 pub fblock: Vec<crate::compile::FBlockInfo>,
177
178 pub symbol_table_index: usize,
180
181 pub in_conditional_block: u32,
184
185 pub next_conditional_annotation_index: u32,
188}
189
190impl CodeInfo {
191 pub fn finalize_code(
192 mut self,
193 opts: &crate::compile::CompileOpts,
194 ) -> crate::InternalResult<CodeObject> {
195 self.fold_binop_constants();
197 self.remove_nops();
198 self.fold_unary_negative();
199 self.fold_binop_constants(); self.remove_nops(); self.fold_tuple_constants();
202 self.fold_list_constants();
203 self.fold_set_constants();
204 self.remove_nops(); self.fold_const_iterable_for_iter();
206 self.convert_to_load_small_int();
207 self.remove_unused_consts();
208 self.remove_nops();
209
210 self.dce();
212 self.peephole_optimize();
214
215 split_blocks_at_jumps(&mut self.blocks);
218 mark_except_handlers(&mut self.blocks);
219 label_exception_targets(&mut self.blocks);
220 jump_threading(&mut self.blocks);
222 self.eliminate_unreachable_blocks();
223 self.remove_nops();
224 push_cold_blocks_to_end(&mut self.blocks);
226
227 normalize_jumps(&mut self.blocks);
229 inline_small_or_no_lineno_blocks(&mut self.blocks);
230 self.dce(); self.eliminate_unreachable_blocks();
232 resolve_line_numbers(&mut self.blocks);
233 duplicate_end_returns(&mut self.blocks);
234 self.dce(); self.eliminate_unreachable_blocks(); self.optimize_load_fast_borrow();
238 self.optimize_load_global_push_null();
239
240 let max_stackdepth = self.max_stackdepth()?;
241
242 let Self {
243 flags,
244 source_path,
245 private: _, mut blocks,
248 current_block: _,
249 metadata,
250 static_attributes: _,
251 in_inlined_comp: _,
252 fblock: _,
253 symbol_table_index: _,
254 in_conditional_block: _,
255 next_conditional_annotation_index: _,
256 } = self;
257
258 let CodeUnitMetadata {
259 name: obj_name,
260 qualname,
261 consts: constants,
262 names: name_cache,
263 varnames: varname_cache,
264 cellvars: cellvar_cache,
265 freevars: freevar_cache,
266 fast_hidden,
267 argcount: arg_count,
268 posonlyargcount: posonlyarg_count,
269 kwonlyargcount: kwonlyarg_count,
270 firstlineno: first_line_number,
271 } = metadata;
272
273 let mut instructions = Vec::new();
274 let mut locations = Vec::new();
275 let mut linetable_locations: Vec<LineTableLocation> = Vec::new();
276
277 let cellfixedoffsets =
279 build_cellfixedoffsets(&varname_cache, &cellvar_cache, &freevar_cache);
280 convert_pseudo_ops(&mut blocks, &cellfixedoffsets);
282 fixup_deref_opargs(&mut blocks, &cellfixedoffsets);
283 let mut block_order = Vec::new();
286 let mut current = BlockIdx(0);
287 while current != BlockIdx::NULL {
288 block_order.push(current);
289 current = blocks[current.idx()].next;
290 }
291 for block_idx in block_order {
292 let bi = block_idx.idx();
293 let mut src_instructions = core::mem::take(&mut blocks[bi].instructions);
294 let mut kept = Vec::with_capacity(src_instructions.len());
295 let mut prev_lineno = -1i32;
296
297 for src in 0..src_instructions.len() {
298 let instr = src_instructions[src];
299 let lineno = instr
300 .lineno_override
301 .unwrap_or_else(|| instr.location.line.get() as i32);
302 let mut remove = false;
303
304 if matches!(instr.instr.real(), Some(Instruction::Nop)) {
305 if lineno < 0 || prev_lineno == lineno {
307 remove = true;
308 }
309 else if src < src_instructions.len() - 1 {
311 let next_lineno =
312 src_instructions[src + 1]
313 .lineno_override
314 .unwrap_or_else(|| {
315 src_instructions[src + 1].location.line.get() as i32
316 });
317 if next_lineno == lineno {
318 remove = true;
319 } else if next_lineno < 0 {
320 src_instructions[src + 1].lineno_override = Some(lineno);
321 remove = true;
322 }
323 }
324 else {
327 let mut next = blocks[bi].next;
328 while next != BlockIdx::NULL && blocks[next.idx()].instructions.is_empty() {
329 next = blocks[next.idx()].next;
330 }
331 if next != BlockIdx::NULL {
332 let mut next_lineno = None;
333 for next_instr in &blocks[next.idx()].instructions {
334 let line = next_instr
335 .lineno_override
336 .unwrap_or_else(|| next_instr.location.line.get() as i32);
337 if matches!(next_instr.instr.real(), Some(Instruction::Nop))
338 && line < 0
339 {
340 continue;
341 }
342 next_lineno = Some(line);
343 break;
344 }
345 if next_lineno.is_some_and(|line| line == lineno) {
346 remove = true;
347 }
348 }
349 }
350 }
351
352 if !remove {
353 kept.push(instr);
354 prev_lineno = lineno;
355 }
356 }
357
358 blocks[bi].instructions = kept;
359 }
360
361 for block in blocks.iter_mut() {
364 if let Some(pos) = block
365 .instructions
366 .iter()
367 .position(|ins| ins.instr.is_scope_exit() || ins.instr.is_unconditional_jump())
368 {
369 block.instructions.truncate(pos + 1);
370 }
371 }
372
373 for block in blocks.iter_mut() {
375 for instr in &mut block.instructions {
376 if let AnyInstruction::Real(op) = instr.instr {
377 instr.cache_entries = op.cache_entries() as u32;
378 }
379 }
380 }
381
382 let mut block_to_offset = vec![Label::from_u32(0); blocks.len()];
383 let mut block_to_index = vec![0u32; blocks.len()];
386 const END_SEND_OFFSET: u32 = 5;
388 loop {
389 let mut num_instructions = 0;
390 for (idx, block) in iter_blocks(&blocks) {
391 block_to_offset[idx.idx()] = Label::from_u32(num_instructions as u32);
392 block_to_index[idx.idx()] = num_instructions as u32;
396 for instr in &block.instructions {
397 num_instructions += instr.arg.instr_size() + instr.cache_entries as usize;
398 }
399 }
400
401 instructions.reserve_exact(num_instructions);
402 locations.reserve_exact(num_instructions);
403
404 let mut recompile = false;
405 let mut next_block = BlockIdx(0);
406 while next_block != BlockIdx::NULL {
407 let block = &mut blocks[next_block];
408 let mut current_offset = block_to_offset[next_block.idx()].as_u32();
410 for info in &mut block.instructions {
411 let target = info.target;
412 let mut op = info.instr.expect_real();
413 let old_arg_size = info.arg.instr_size();
414 let old_cache_entries = info.cache_entries;
415 let offset_after = current_offset + old_arg_size as u32 + old_cache_entries;
418
419 if target != BlockIdx::NULL {
420 let target_offset = block_to_offset[target.idx()].as_u32();
421 op = match op {
425 Instruction::JumpForward { .. } if target_offset <= current_offset => {
426 Instruction::JumpBackward {
427 delta: Arg::marker(),
428 }
429 }
430 Instruction::JumpBackward { .. } if target_offset > current_offset => {
431 Instruction::JumpForward {
432 delta: Arg::marker(),
433 }
434 }
435 Instruction::JumpBackwardNoInterrupt { .. }
436 if target_offset > current_offset =>
437 {
438 Instruction::JumpForward {
439 delta: Arg::marker(),
440 }
441 }
442 _ => op,
443 };
444 info.instr = op.into();
445 let updated_cache = op.cache_entries() as u32;
446 recompile |= updated_cache != old_cache_entries;
447 info.cache_entries = updated_cache;
448 let new_arg = if matches!(op, Instruction::EndAsyncFor) {
449 let arg = offset_after
450 .checked_sub(target_offset + END_SEND_OFFSET)
451 .expect("END_ASYNC_FOR target must be before instruction");
452 OpArg::new(arg)
453 } else if matches!(
454 op,
455 Instruction::JumpBackward { .. }
456 | Instruction::JumpBackwardNoInterrupt { .. }
457 ) {
458 let arg = offset_after
459 .checked_sub(target_offset)
460 .expect("backward jump target must be before instruction");
461 OpArg::new(arg)
462 } else {
463 let arg = target_offset
464 .checked_sub(offset_after)
465 .expect("forward jump target must be after instruction");
466 OpArg::new(arg)
467 };
468 recompile |= new_arg.instr_size() != old_arg_size;
469 info.arg = new_arg;
470 }
471
472 let cache_count = info.cache_entries as usize;
473 let (extras, lo_arg) = info.arg.split();
474 let loc_pair = (info.location, info.end_location);
475 locations.extend(core::iter::repeat_n(
476 loc_pair,
477 info.arg.instr_size() + cache_count,
478 ));
479 let lt_loc = LineTableLocation {
481 line: info
482 .lineno_override
483 .unwrap_or_else(|| info.location.line.get() as i32),
484 end_line: info.end_location.line.get() as i32,
485 col: info.location.character_offset.to_zero_indexed() as i32,
486 end_col: info.end_location.character_offset.to_zero_indexed() as i32,
487 };
488 linetable_locations.extend(core::iter::repeat_n(lt_loc, info.arg.instr_size()));
489 if cache_count > 0 {
491 linetable_locations.extend(core::iter::repeat_n(lt_loc, cache_count));
492 }
493 instructions.extend(
494 extras
495 .map(|byte| CodeUnit::new(Instruction::ExtendedArg, byte))
496 .chain([CodeUnit { op, arg: lo_arg }]),
497 );
498 if cache_count > 0 {
500 instructions.extend(core::iter::repeat_n(
501 CodeUnit::new(Instruction::Cache, 0.into()),
502 cache_count,
503 ));
504 }
505 current_offset = offset_after;
506 }
507 next_block = block.next;
508 }
509
510 if !recompile {
511 break;
512 }
513
514 instructions.clear();
515 locations.clear();
516 linetable_locations.clear();
517 }
518
519 let linetable = generate_linetable(
521 &linetable_locations,
522 first_line_number.get() as i32,
523 opts.debug_ranges,
524 );
525
526 let exceptiontable = generate_exception_table(&blocks, &block_to_index);
528
529 let nlocals = varname_cache.len();
531 let ncells = cellvar_cache.len();
532 let nfrees = freevar_cache.len();
533 let numdropped = cellvar_cache
534 .iter()
535 .filter(|cv| varname_cache.contains(cv.as_str()))
536 .count();
537 let nlocalsplus = nlocals + ncells - numdropped + nfrees;
538 let mut localspluskinds = vec![0u8; nlocalsplus];
539 for kind in localspluskinds.iter_mut().take(nlocals) {
541 *kind = CO_FAST_LOCAL;
542 }
543 for (i, cellvar) in cellvar_cache.iter().enumerate() {
545 let idx = cellfixedoffsets[i] as usize;
546 if varname_cache.contains(cellvar.as_str()) {
547 localspluskinds[idx] |= CO_FAST_CELL; } else {
549 localspluskinds[idx] = CO_FAST_CELL;
550 }
551 }
552 for i in 0..nfrees {
554 let idx = cellfixedoffsets[ncells + i] as usize;
555 localspluskinds[idx] = CO_FAST_FREE;
556 }
557 for (name, &hidden) in &fast_hidden {
559 if hidden && let Some(idx) = varname_cache.get_index_of(name.as_str()) {
560 localspluskinds[idx] |= CO_FAST_HIDDEN;
561 }
562 }
563
564 Ok(CodeObject {
565 flags,
566 posonlyarg_count,
567 arg_count,
568 kwonlyarg_count,
569 source_path,
570 first_line_number: Some(first_line_number),
571 obj_name: obj_name.clone(),
572 qualname: qualname.unwrap_or(obj_name),
573
574 max_stackdepth,
575 instructions: CodeUnits::from(instructions),
576 locations: locations.into_boxed_slice(),
577 constants: constants.into_iter().collect(),
578 names: name_cache.into_iter().collect(),
579 varnames: varname_cache.into_iter().collect(),
580 cellvars: cellvar_cache.into_iter().collect(),
581 freevars: freevar_cache.into_iter().collect(),
582 localspluskinds: localspluskinds.into_boxed_slice(),
583 linetable,
584 exceptiontable,
585 })
586 }
587
588 fn dce(&mut self) {
589 for block in &mut self.blocks {
591 let mut last_instr = None;
592 for (i, ins) in block.instructions.iter().enumerate() {
593 if ins.instr.is_scope_exit() || ins.instr.is_unconditional_jump() {
594 last_instr = Some(i);
595 break;
596 }
597 }
598 if let Some(i) = last_instr {
599 block.instructions.truncate(i + 1);
600 }
601 }
602 }
603
604 fn eliminate_unreachable_blocks(&mut self) {
607 let mut reachable = vec![false; self.blocks.len()];
608 reachable[0] = true;
609
610 let mut changed = true;
612 while changed {
613 changed = false;
614 for i in 0..self.blocks.len() {
615 if !reachable[i] {
616 continue;
617 }
618 for ins in &self.blocks[i].instructions {
620 if ins.target != BlockIdx::NULL && !reachable[ins.target.idx()] {
621 reachable[ins.target.idx()] = true;
622 changed = true;
623 }
624 if let Some(eh) = &ins.except_handler
625 && !reachable[eh.handler_block.idx()]
626 {
627 reachable[eh.handler_block.idx()] = true;
628 changed = true;
629 }
630 }
631 let next = self.blocks[i].next;
633 if next != BlockIdx::NULL
634 && !reachable[next.idx()]
635 && !self.blocks[i].instructions.last().is_some_and(|ins| {
636 ins.instr.is_scope_exit() || ins.instr.is_unconditional_jump()
637 })
638 {
639 reachable[next.idx()] = true;
640 changed = true;
641 }
642 }
643 }
644
645 for (i, block) in self.blocks.iter_mut().enumerate() {
646 if !reachable[i] {
647 block.instructions.clear();
648 }
649 }
650 }
651
652 fn fold_unary_negative(&mut self) {
654 for block in &mut self.blocks {
655 let mut i = 0;
656 while i + 1 < block.instructions.len() {
657 let next = &block.instructions[i + 1];
658 let Some(Instruction::UnaryNegative) = next.instr.real() else {
659 i += 1;
660 continue;
661 };
662 let curr = &block.instructions[i];
663 let value = match curr.instr.real() {
664 Some(Instruction::LoadConst { .. }) => {
665 let idx = u32::from(curr.arg) as usize;
666 match self.metadata.consts.get_index(idx) {
667 Some(ConstantData::Integer { value }) => {
668 Some(ConstantData::Integer { value: -value })
669 }
670 Some(ConstantData::Float { value }) => {
671 Some(ConstantData::Float { value: -value })
672 }
673 _ => None,
674 }
675 }
676 Some(Instruction::LoadSmallInt { .. }) => {
677 let v = u32::from(curr.arg) as i32;
678 Some(ConstantData::Integer {
679 value: BigInt::from(-v),
680 })
681 }
682 _ => None,
683 };
684 if let Some(neg_const) = value {
685 let (const_idx, _) = self.metadata.consts.insert_full(neg_const);
686 let load_location = block.instructions[i].location;
688 block.instructions[i].instr = Instruction::LoadConst {
689 consti: Arg::marker(),
690 }
691 .into();
692 block.instructions[i].arg = OpArg::new(const_idx as u32);
693 block.instructions[i + 1].instr = Instruction::Nop.into();
696 block.instructions[i + 1].location = load_location;
697 block.instructions[i + 1].end_location = block.instructions[i].end_location;
698 i += 2;
700 } else {
701 i += 1;
702 }
703 }
704 }
705 }
706
707 fn fold_binop_constants(&mut self) {
711 use oparg::BinaryOperator as BinOp;
712
713 for block in &mut self.blocks {
714 let mut i = 0;
715 while i + 2 < block.instructions.len() {
716 let Some(Instruction::BinaryOp { .. }) = block.instructions[i + 2].instr.real()
718 else {
719 i += 1;
720 continue;
721 };
722
723 let op_raw = u32::from(block.instructions[i + 2].arg);
724 let Ok(op) = BinOp::try_from(op_raw) else {
725 i += 1;
726 continue;
727 };
728
729 let left = Self::get_const_value_from(&self.metadata, &block.instructions[i]);
730 let right = Self::get_const_value_from(&self.metadata, &block.instructions[i + 1]);
731
732 let (Some(left_val), Some(right_val)) = (left, right) else {
733 i += 1;
734 continue;
735 };
736
737 let result = Self::eval_binop(&left_val, &right_val, op);
738
739 if let Some(result_const) = result {
740 if Self::const_too_big(&result_const) {
742 i += 1;
743 continue;
744 }
745 let (const_idx, _) = self.metadata.consts.insert_full(result_const);
746 block.instructions[i].instr = Instruction::LoadConst {
748 consti: Arg::marker(),
749 }
750 .into();
751 block.instructions[i].arg = OpArg::new(const_idx as u32);
752 let loc = block.instructions[i].location;
754 let end_loc = block.instructions[i].end_location;
755 block.instructions[i + 1].instr = Instruction::Nop.into();
756 block.instructions[i + 1].location = loc;
757 block.instructions[i + 1].end_location = end_loc;
758 block.instructions[i + 2].instr = Instruction::Nop.into();
759 block.instructions[i + 2].location = loc;
760 block.instructions[i + 2].end_location = end_loc;
761 i = i.saturating_sub(1); } else {
765 i += 1;
766 }
767 }
768 }
769 }
770
771 fn get_const_value_from(
772 metadata: &CodeUnitMetadata,
773 info: &InstructionInfo,
774 ) -> Option<ConstantData> {
775 match info.instr.real() {
776 Some(Instruction::LoadConst { .. }) => {
777 let idx = u32::from(info.arg) as usize;
778 metadata.consts.get_index(idx).cloned()
779 }
780 Some(Instruction::LoadSmallInt { .. }) => {
781 let v = u32::from(info.arg) as i32;
782 Some(ConstantData::Integer {
783 value: BigInt::from(v),
784 })
785 }
786 _ => None,
787 }
788 }
789
790 fn eval_binop(
791 left: &ConstantData,
792 right: &ConstantData,
793 op: oparg::BinaryOperator,
794 ) -> Option<ConstantData> {
795 use oparg::BinaryOperator as BinOp;
796 match (left, right) {
797 (ConstantData::Integer { value: l }, ConstantData::Integer { value: r }) => {
798 let result = match op {
799 BinOp::Add => l + r,
800 BinOp::Subtract => l - r,
801 BinOp::Multiply => {
802 if !l.is_zero() && !r.is_zero() && l.bits() + r.bits() > MAX_INT_SIZE_BITS {
803 return None;
804 }
805 l * r
806 }
807 BinOp::FloorDivide => {
808 if r.is_zero() {
809 return None;
810 }
811 let (q, rem) = (l.clone() / r.clone(), l.clone() % r.clone());
813 if !rem.is_zero() && (rem < BigInt::from(0)) != (*r < BigInt::from(0)) {
814 q - 1
815 } else {
816 q
817 }
818 }
819 BinOp::Remainder => {
820 if r.is_zero() {
821 return None;
822 }
823 let rem = l.clone() % r.clone();
825 if !rem.is_zero() && (rem < BigInt::from(0)) != (*r < BigInt::from(0)) {
826 rem + r
827 } else {
828 rem
829 }
830 }
831 BinOp::Power => {
832 let exp: u64 = r.try_into().ok()?;
833 let exp_usize = usize::try_from(exp).ok()?;
834 if !l.is_zero() && exp > 0 && l.bits() > MAX_INT_SIZE_BITS / exp {
835 return None;
836 }
837 num_traits::pow::pow(l.clone(), exp_usize)
838 }
839 BinOp::Lshift => {
840 let shift: u64 = r.try_into().ok()?;
841 let shift_usize = usize::try_from(shift).ok()?;
842 if shift > MAX_INT_SIZE_BITS
843 || (!l.is_zero() && l.bits() > MAX_INT_SIZE_BITS - shift)
844 {
845 return None;
846 }
847 l << shift_usize
848 }
849 BinOp::Rshift => {
850 let shift: u32 = r.try_into().ok()?;
851 l >> (shift as usize)
852 }
853 BinOp::And => l & r,
854 BinOp::Or => l | r,
855 BinOp::Xor => l ^ r,
856 _ => return None,
857 };
858 Some(ConstantData::Integer { value: result })
859 }
860 (ConstantData::Float { value: l }, ConstantData::Float { value: r }) => {
861 let result = match op {
862 BinOp::Add => l + r,
863 BinOp::Subtract => l - r,
864 BinOp::Multiply => l * r,
865 BinOp::TrueDivide => {
866 if *r == 0.0 {
867 return None;
868 }
869 l / r
870 }
871 BinOp::FloorDivide => {
872 return None;
874 }
875 BinOp::Remainder => {
876 return None;
878 }
879 BinOp::Power => l.powf(*r),
880 _ => return None,
881 };
882 if !result.is_finite() {
883 return None;
884 }
885 Some(ConstantData::Float { value: result })
886 }
887 (ConstantData::Integer { value: l }, ConstantData::Float { value: r }) => {
889 let l_f = l.to_f64()?;
890 Self::eval_binop(
891 &ConstantData::Float { value: l_f },
892 &ConstantData::Float { value: *r },
893 op,
894 )
895 }
896 (ConstantData::Float { value: l }, ConstantData::Integer { value: r }) => {
897 let r_f = r.to_f64()?;
898 Self::eval_binop(
899 &ConstantData::Float { value: *l },
900 &ConstantData::Float { value: r_f },
901 op,
902 )
903 }
904 (ConstantData::Str { value: l }, ConstantData::Str { value: r })
906 if matches!(op, BinOp::Add) =>
907 {
908 let mut result = l.to_string();
909 result.push_str(&r.to_string());
910 Some(ConstantData::Str {
911 value: result.into(),
912 })
913 }
914 (ConstantData::Str { value: s }, ConstantData::Integer { value: n })
915 if matches!(op, BinOp::Multiply) =>
916 {
917 let n: usize = n.try_into().ok()?;
918 if n > 4096 {
919 return None;
920 }
921 let result = s.to_string().repeat(n);
922 Some(ConstantData::Str {
923 value: result.into(),
924 })
925 }
926 _ => None,
927 }
928 }
929
930 fn const_too_big(c: &ConstantData) -> bool {
931 match c {
932 ConstantData::Integer { value } => value.bits() > 4096 * 8,
933 ConstantData::Str { value } => value.len() > 4096,
934 _ => false,
935 }
936 }
937
938 fn fold_tuple_constants(&mut self) {
941 for block in &mut self.blocks {
942 let mut i = 0;
943 while i < block.instructions.len() {
944 let instr = &block.instructions[i];
945 let Some(Instruction::BuildTuple { .. }) = instr.instr.real() else {
947 i += 1;
948 continue;
949 };
950
951 let tuple_size = u32::from(instr.arg) as usize;
952 if tuple_size == 0 {
953 let (const_idx, _) = self.metadata.consts.insert_full(ConstantData::Tuple {
955 elements: Vec::new(),
956 });
957 block.instructions[i].instr = Instruction::LoadConst {
958 consti: Arg::marker(),
959 }
960 .into();
961 block.instructions[i].arg = OpArg::new(const_idx as u32);
962 i += 1;
963 continue;
964 }
965 if i < tuple_size {
966 i += 1;
967 continue;
968 }
969
970 let start_idx = i - tuple_size;
972 let mut elements = Vec::with_capacity(tuple_size);
973 let mut all_const = true;
974
975 for j in start_idx..i {
976 let load_instr = &block.instructions[j];
977 match load_instr.instr.real() {
978 Some(Instruction::LoadConst { .. }) => {
979 let const_idx = u32::from(load_instr.arg) as usize;
980 if let Some(constant) =
981 self.metadata.consts.get_index(const_idx).cloned()
982 {
983 elements.push(constant);
984 } else {
985 all_const = false;
986 break;
987 }
988 }
989 Some(Instruction::LoadSmallInt { .. }) => {
990 let value = u32::from(load_instr.arg) as i32;
992 elements.push(ConstantData::Integer {
993 value: BigInt::from(value),
994 });
995 }
996 _ => {
997 all_const = false;
998 break;
999 }
1000 }
1001 }
1002
1003 if !all_const {
1004 i += 1;
1005 continue;
1006 }
1007
1008 let tuple_const = ConstantData::Tuple { elements };
1014 let (const_idx, _) = self.metadata.consts.insert_full(tuple_const);
1015
1016 let folded_loc = block.instructions[i].location;
1019 for j in start_idx..i {
1020 block.instructions[j].instr = Instruction::Nop.into();
1021 block.instructions[j].location = folded_loc;
1022 }
1023
1024 block.instructions[i].instr = Instruction::LoadConst {
1026 consti: Arg::marker(),
1027 }
1028 .into();
1029 block.instructions[i].arg = OpArg::new(const_idx as u32);
1030
1031 i += 1;
1032 }
1033 }
1034 }
1035
1036 fn fold_list_constants(&mut self) {
1039 for block in &mut self.blocks {
1040 let mut i = 0;
1041 while i < block.instructions.len() {
1042 let instr = &block.instructions[i];
1043 let Some(Instruction::BuildList { .. }) = instr.instr.real() else {
1044 i += 1;
1045 continue;
1046 };
1047
1048 let list_size = u32::from(instr.arg) as usize;
1049 if list_size == 0 || i < list_size {
1050 i += 1;
1051 continue;
1052 }
1053
1054 let start_idx = i - list_size;
1055 let mut elements = Vec::with_capacity(list_size);
1056 let mut all_const = true;
1057
1058 for j in start_idx..i {
1059 let load_instr = &block.instructions[j];
1060 match load_instr.instr.real() {
1061 Some(Instruction::LoadConst { .. }) => {
1062 let const_idx = u32::from(load_instr.arg) as usize;
1063 if let Some(constant) =
1064 self.metadata.consts.get_index(const_idx).cloned()
1065 {
1066 elements.push(constant);
1067 } else {
1068 all_const = false;
1069 break;
1070 }
1071 }
1072 Some(Instruction::LoadSmallInt { .. }) => {
1073 let value = u32::from(load_instr.arg) as i32;
1074 elements.push(ConstantData::Integer {
1075 value: BigInt::from(value),
1076 });
1077 }
1078 _ => {
1079 all_const = false;
1080 break;
1081 }
1082 }
1083 }
1084
1085 if !all_const || list_size < MIN_CONST_SEQUENCE_SIZE {
1086 i += 1;
1087 continue;
1088 }
1089
1090 let tuple_const = ConstantData::Tuple { elements };
1091 let (const_idx, _) = self.metadata.consts.insert_full(tuple_const);
1092
1093 let folded_loc = block.instructions[i].location;
1094 let end_loc = block.instructions[i].end_location;
1095 let eh = block.instructions[i].except_handler;
1096
1097 block.instructions[start_idx].instr = Instruction::BuildList {
1099 count: Arg::marker(),
1100 }
1101 .into();
1102 block.instructions[start_idx].arg = OpArg::new(0);
1103 block.instructions[start_idx].location = folded_loc;
1104 block.instructions[start_idx].end_location = end_loc;
1105 block.instructions[start_idx].except_handler = eh;
1106
1107 block.instructions[start_idx + 1].instr = Instruction::LoadConst {
1109 consti: Arg::marker(),
1110 }
1111 .into();
1112 block.instructions[start_idx + 1].arg = OpArg::new(const_idx as u32);
1113 block.instructions[start_idx + 1].location = folded_loc;
1114 block.instructions[start_idx + 1].end_location = end_loc;
1115 block.instructions[start_idx + 1].except_handler = eh;
1116
1117 for j in (start_idx + 2)..i {
1119 block.instructions[j].instr = Instruction::Nop.into();
1120 block.instructions[j].location = folded_loc;
1121 }
1122
1123 block.instructions[i].instr = Instruction::ListExtend { i: Arg::marker() }.into();
1125 block.instructions[i].arg = OpArg::new(1);
1126
1127 i += 1;
1128 }
1129 }
1130 }
1131
1132 fn fold_const_iterable_for_iter(&mut self) {
1136 for block in &mut self.blocks {
1137 let mut i = 0;
1138 while i + 1 < block.instructions.len() {
1139 let is_build = matches!(
1140 block.instructions[i].instr.real(),
1141 Some(Instruction::BuildList { .. })
1142 ) && u32::from(block.instructions[i].arg) == 0;
1143
1144 let is_const = matches!(
1145 block
1146 .instructions
1147 .get(i + 1)
1148 .and_then(|instr| instr.instr.real()),
1149 Some(Instruction::LoadConst { .. })
1150 );
1151
1152 let is_extend = matches!(
1153 block
1154 .instructions
1155 .get(i + 2)
1156 .and_then(|instr| instr.instr.real()),
1157 Some(Instruction::ListExtend { .. })
1158 ) && block
1159 .instructions
1160 .get(i + 2)
1161 .is_some_and(|instr| u32::from(instr.arg) == 1);
1162
1163 let is_iter = matches!(
1164 block
1165 .instructions
1166 .get(i + 3)
1167 .and_then(|instr| instr.instr.real()),
1168 Some(Instruction::GetIter)
1169 );
1170
1171 if is_build && is_const && is_extend && is_iter {
1172 let loc = block.instructions[i].location;
1174 block.instructions[i].instr = Instruction::Nop.into();
1175 block.instructions[i].location = loc;
1176 block.instructions[i + 2].instr = Instruction::Nop.into();
1177 block.instructions[i + 2].location = loc;
1178 i += 4;
1179 } else if matches!(
1180 block.instructions[i].instr.real(),
1181 Some(Instruction::BuildList { .. })
1182 ) && matches!(
1183 block.instructions[i + 1].instr.real(),
1184 Some(Instruction::GetIter)
1185 ) {
1186 let seq_size = u32::from(block.instructions[i].arg) as usize;
1187
1188 if seq_size != 0 && i >= seq_size {
1189 let start_idx = i - seq_size;
1190 let mut elements = Vec::with_capacity(seq_size);
1191 let mut all_const = true;
1192
1193 for j in start_idx..i {
1194 match Self::get_const_value_from(&self.metadata, &block.instructions[j])
1195 {
1196 Some(constant) => elements.push(constant),
1197 None => {
1198 all_const = false;
1199 break;
1200 }
1201 }
1202 }
1203
1204 if all_const {
1205 let const_data = ConstantData::Tuple { elements };
1206 let (const_idx, _) = self.metadata.consts.insert_full(const_data);
1207 let folded_loc = block.instructions[i].location;
1208
1209 for j in start_idx..i {
1210 block.instructions[j].instr = Instruction::Nop.into();
1211 block.instructions[j].location = folded_loc;
1212 }
1213
1214 block.instructions[i].instr = Instruction::LoadConst {
1215 consti: Arg::marker(),
1216 }
1217 .into();
1218 block.instructions[i].arg = OpArg::new(const_idx as u32);
1219 i += 2;
1220 continue;
1221 }
1222 }
1223
1224 block.instructions[i].instr = Instruction::BuildTuple {
1225 count: Arg::marker(),
1226 }
1227 .into();
1228 i += 2;
1229 } else {
1230 i += 1;
1231 }
1232 }
1233 }
1234 }
1235
1236 fn fold_set_constants(&mut self) {
1239 for block in &mut self.blocks {
1240 let mut i = 0;
1241 while i < block.instructions.len() {
1242 let instr = &block.instructions[i];
1243 let Some(Instruction::BuildSet { .. }) = instr.instr.real() else {
1244 i += 1;
1245 continue;
1246 };
1247
1248 let set_size = u32::from(instr.arg) as usize;
1249 if set_size < 3 || i < set_size {
1250 i += 1;
1251 continue;
1252 }
1253
1254 let start_idx = i - set_size;
1255 let mut elements = Vec::with_capacity(set_size);
1256 let mut all_const = true;
1257
1258 for j in start_idx..i {
1259 let load_instr = &block.instructions[j];
1260 match load_instr.instr.real() {
1261 Some(Instruction::LoadConst { .. }) => {
1262 let const_idx = u32::from(load_instr.arg) as usize;
1263 if let Some(constant) =
1264 self.metadata.consts.get_index(const_idx).cloned()
1265 {
1266 elements.push(constant);
1267 } else {
1268 all_const = false;
1269 break;
1270 }
1271 }
1272 Some(Instruction::LoadSmallInt { .. }) => {
1273 let value = u32::from(load_instr.arg) as i32;
1274 elements.push(ConstantData::Integer {
1275 value: BigInt::from(value),
1276 });
1277 }
1278 _ => {
1279 all_const = false;
1280 break;
1281 }
1282 }
1283 }
1284
1285 if !all_const {
1286 i += 1;
1287 continue;
1288 }
1289
1290 let const_data = ConstantData::Tuple { elements };
1292 let (const_idx, _) = self.metadata.consts.insert_full(const_data);
1293
1294 let folded_loc = block.instructions[i].location;
1295 let end_loc = block.instructions[i].end_location;
1296 let eh = block.instructions[i].except_handler;
1297
1298 block.instructions[start_idx].instr = Instruction::BuildSet {
1299 count: Arg::marker(),
1300 }
1301 .into();
1302 block.instructions[start_idx].arg = OpArg::new(0);
1303 block.instructions[start_idx].location = folded_loc;
1304 block.instructions[start_idx].end_location = end_loc;
1305 block.instructions[start_idx].except_handler = eh;
1306
1307 block.instructions[start_idx + 1].instr = Instruction::LoadConst {
1308 consti: Arg::marker(),
1309 }
1310 .into();
1311 block.instructions[start_idx + 1].arg = OpArg::new(const_idx as u32);
1312 block.instructions[start_idx + 1].location = folded_loc;
1313 block.instructions[start_idx + 1].end_location = end_loc;
1314 block.instructions[start_idx + 1].except_handler = eh;
1315
1316 for j in (start_idx + 2)..i {
1317 block.instructions[j].instr = Instruction::Nop.into();
1318 block.instructions[j].location = folded_loc;
1319 }
1320
1321 block.instructions[i].instr = Instruction::SetUpdate { i: Arg::marker() }.into();
1322 block.instructions[i].arg = OpArg::new(1);
1323
1324 i += 1;
1325 }
1326 }
1327 }
1328
1329 fn peephole_optimize(&mut self) {
1331 for block in &mut self.blocks {
1332 let mut i = 0;
1333 while i + 1 < block.instructions.len() {
1334 let combined = {
1335 let curr = &block.instructions[i];
1336 let next = &block.instructions[i + 1];
1337
1338 let (Some(curr_instr), Some(next_instr)) =
1340 (curr.instr.real(), next.instr.real())
1341 else {
1342 i += 1;
1343 continue;
1344 };
1345
1346 match (curr_instr, next_instr) {
1347 (Instruction::LoadFast { .. }, Instruction::LoadFast { .. }) => {
1349 let idx1 = u32::from(curr.arg);
1350 let idx2 = u32::from(next.arg);
1351 if idx1 < 16 && idx2 < 16 {
1352 let packed = (idx1 << 4) | idx2;
1353 Some((
1354 Instruction::LoadFastLoadFast {
1355 var_nums: Arg::marker(),
1356 },
1357 OpArg::new(packed),
1358 ))
1359 } else {
1360 None
1361 }
1362 }
1363 (Instruction::StoreFast { .. }, Instruction::StoreFast { .. }) => {
1365 let idx1 = u32::from(curr.arg);
1366 let idx2 = u32::from(next.arg);
1367 if idx1 < 16 && idx2 < 16 {
1368 let packed = (idx1 << 4) | idx2;
1369 Some((
1370 Instruction::StoreFastStoreFast {
1371 var_nums: Arg::marker(),
1372 },
1373 OpArg::new(packed),
1374 ))
1375 } else {
1376 None
1377 }
1378 }
1379 (Instruction::LoadConst { consti }, Instruction::ToBool) => {
1383 let consti = consti.get(curr.arg);
1384 let constant = &self.metadata.consts[consti.as_usize()];
1385 if let ConstantData::Boolean { .. } = constant {
1386 Some((curr_instr, OpArg::from(consti.as_u32())))
1387 } else {
1388 None
1389 }
1390 }
1391 (Instruction::LoadConst { consti }, Instruction::UnaryNot) => {
1392 let constant = &self.metadata.consts[consti.get(curr.arg).as_usize()];
1393 match constant {
1394 ConstantData::Boolean { value } => {
1395 let (const_idx, _) = self
1396 .metadata
1397 .consts
1398 .insert_full(ConstantData::Boolean { value: !value });
1399 Some((
1400 (Instruction::LoadConst {
1401 consti: Arg::marker(),
1402 }),
1403 OpArg::new(const_idx as u32),
1404 ))
1405 }
1406 _ => None,
1407 }
1408 }
1409 _ => None,
1410 }
1411 };
1412
1413 if let Some((new_instr, new_arg)) = combined {
1414 block.instructions[i].instr = new_instr.into();
1416 block.instructions[i].arg = new_arg;
1417 block.instructions.remove(i + 1);
1419 } else {
1421 i += 1;
1422 }
1423 }
1424 }
1425 }
1426
1427 fn optimize_load_global_push_null(&mut self) {
1429 for block in &mut self.blocks {
1430 let mut i = 0;
1431 while i + 1 < block.instructions.len() {
1432 let curr = &block.instructions[i];
1433 let next = &block.instructions[i + 1];
1434
1435 let (Some(Instruction::LoadGlobal { .. }), Some(Instruction::PushNull)) =
1436 (curr.instr.real(), next.instr.real())
1437 else {
1438 i += 1;
1439 continue;
1440 };
1441
1442 let oparg = u32::from(block.instructions[i].arg);
1443 if (oparg & 1) != 0 {
1444 i += 1;
1445 continue;
1446 }
1447
1448 block.instructions[i].arg = OpArg::new(oparg | 1);
1449 block.instructions.remove(i + 1);
1450 }
1451 }
1452 }
1453
1454 fn convert_to_load_small_int(&mut self) {
1457 for block in &mut self.blocks {
1458 for instr in &mut block.instructions {
1459 let Some(Instruction::LoadConst { .. }) = instr.instr.real() else {
1461 continue;
1462 };
1463
1464 let const_idx = u32::from(instr.arg) as usize;
1466 let Some(constant) = self.metadata.consts.get_index(const_idx) else {
1467 continue;
1468 };
1469
1470 let ConstantData::Integer { value } = constant else {
1472 continue;
1473 };
1474
1475 if let Some(small) = value.to_i32().filter(|v| (0..=255).contains(v)) {
1477 instr.instr = Instruction::LoadSmallInt { i: Arg::marker() }.into();
1479 instr.arg = OpArg::new(small as u32);
1481 }
1482 }
1483 }
1484 }
1485
1486 fn remove_unused_consts(&mut self) {
1489 let nconsts = self.metadata.consts.len();
1490 if nconsts == 0 {
1491 return;
1492 }
1493
1494 let mut used = vec![false; nconsts];
1497 used[0] = true;
1498
1499 for block in &self.blocks {
1500 for instr in &block.instructions {
1501 if let Some(Instruction::LoadConst { .. }) = instr.instr.real() {
1502 let idx = u32::from(instr.arg) as usize;
1503 if idx < nconsts {
1504 used[idx] = true;
1505 }
1506 }
1507 }
1508 }
1509
1510 let n_used: usize = used.iter().filter(|&&u| u).count();
1512 if n_used == nconsts {
1513 return; }
1515
1516 let mut old_to_new = vec![0usize; nconsts];
1518 let mut new_idx = 0usize;
1519 for (old_idx, &is_used) in used.iter().enumerate() {
1520 if is_used {
1521 old_to_new[old_idx] = new_idx;
1522 new_idx += 1;
1523 }
1524 }
1525
1526 let old_consts: Vec<_> = self.metadata.consts.iter().cloned().collect();
1528 self.metadata.consts.clear();
1529 for (old_idx, constant) in old_consts.into_iter().enumerate() {
1530 if used[old_idx] {
1531 self.metadata.consts.insert(constant);
1532 }
1533 }
1534
1535 for block in &mut self.blocks {
1537 for instr in &mut block.instructions {
1538 if let Some(Instruction::LoadConst { .. }) = instr.instr.real() {
1539 let old_idx = u32::from(instr.arg) as usize;
1540 if old_idx < nconsts {
1541 instr.arg = OpArg::new(old_to_new[old_idx] as u32);
1542 }
1543 }
1544 }
1545 }
1546 }
1547
1548 fn remove_nops(&mut self) {
1551 for block in &mut self.blocks {
1552 let mut prev_line = None;
1553 block.instructions.retain(|ins| {
1554 if matches!(ins.instr.real(), Some(Instruction::Nop)) {
1555 let line = ins.location.line;
1556 if prev_line == Some(line) {
1557 return false;
1558 }
1559 }
1560 prev_line = Some(ins.location.line);
1561 true
1562 });
1563 }
1564 }
1565
1566 #[allow(dead_code)]
1571 fn combine_store_fast_load_fast(&mut self) {
1572 for block in &mut self.blocks {
1573 let mut i = 0;
1574 while i + 1 < block.instructions.len() {
1575 let curr = &block.instructions[i];
1576 let next = &block.instructions[i + 1];
1577 let (Some(Instruction::StoreFast { .. }), Some(Instruction::LoadFast { .. })) =
1578 (curr.instr.real(), next.instr.real())
1579 else {
1580 i += 1;
1581 continue;
1582 };
1583 let line1 = curr.location.line;
1585 let line2 = next.location.line;
1586 if line1 != line2 {
1587 i += 1;
1588 continue;
1589 }
1590 let idx1 = u32::from(curr.arg);
1591 let idx2 = u32::from(next.arg);
1592 if idx1 < 16 && idx2 < 16 {
1593 let packed = (idx1 << 4) | idx2;
1594 block.instructions[i].instr = Instruction::StoreFastLoadFast {
1595 var_nums: Arg::marker(),
1596 }
1597 .into();
1598 block.instructions[i].arg = OpArg::new(packed);
1599 block.instructions[i + 1].instr = Instruction::Nop.into();
1601 block.instructions[i + 1].arg = OpArg::new(0);
1602 i += 2; } else {
1604 i += 1;
1605 }
1606 }
1607 }
1608 }
1609
1610 fn optimize_load_fast_borrow(&mut self) {
1611 const NOT_LOCAL: usize = usize::MAX;
1613
1614 for block in &mut self.blocks {
1615 if block.instructions.is_empty() {
1616 continue;
1617 }
1618
1619 let mut unconsumed = vec![false; block.instructions.len()];
1622
1623 let mut stack: Vec<usize> = Vec::new();
1631 let mut underflow = false;
1632
1633 for (i, info) in block.instructions.iter().enumerate() {
1634 let Some(instr) = info.instr.real() else {
1635 continue;
1636 };
1637
1638 let stack_effect_info = instr.stack_effect_info(info.arg.into());
1639 let (pushes, pops) = (stack_effect_info.pushed(), stack_effect_info.popped());
1640
1641 for _ in 0..pops {
1643 if stack.pop().is_none() {
1644 underflow = true;
1647 break;
1648 }
1649 }
1650 if underflow {
1651 break;
1652 }
1653
1654 let source = match instr {
1656 Instruction::LoadFast { .. } | Instruction::LoadFastLoadFast { .. } => i,
1657 _ => NOT_LOCAL,
1658 };
1659 for _ in 0..pushes {
1660 stack.push(source);
1661 }
1662 }
1663
1664 if underflow {
1665 continue;
1666 }
1667
1668 for &src in &stack {
1670 if src != NOT_LOCAL {
1671 unconsumed[src] = true;
1672 }
1673 }
1674
1675 for (i, info) in block.instructions.iter_mut().enumerate() {
1677 if unconsumed[i] {
1678 continue;
1679 }
1680 let Some(instr) = info.instr.real() else {
1681 continue;
1682 };
1683 match instr {
1684 Instruction::LoadFast { .. } => {
1685 info.instr = Instruction::LoadFastBorrow {
1686 var_num: Arg::marker(),
1687 }
1688 .into();
1689 }
1690 Instruction::LoadFastLoadFast { .. } => {
1691 info.instr = Instruction::LoadFastBorrowLoadFastBorrow {
1692 var_nums: Arg::marker(),
1693 }
1694 .into();
1695 }
1696 _ => {}
1697 }
1698 }
1699 }
1700 }
1701
1702 fn max_stackdepth(&mut self) -> crate::InternalResult<u32> {
1703 let mut maxdepth = 0u32;
1704 let mut stack = Vec::with_capacity(self.blocks.len());
1705 let mut start_depths = vec![u32::MAX; self.blocks.len()];
1706 stackdepth_push(&mut stack, &mut start_depths, BlockIdx(0), 0);
1707 const DEBUG: bool = false;
1708 'process_blocks: while let Some(block_idx) = stack.pop() {
1709 let idx = block_idx.idx();
1710 let mut depth = start_depths[idx];
1711 if DEBUG {
1712 eprintln!("===BLOCK {}===", block_idx.0);
1713 }
1714 let block = &self.blocks[block_idx];
1715 for ins in &block.instructions {
1716 let instr = &ins.instr;
1717 let effect = instr.stack_effect(ins.arg.into());
1718 if DEBUG {
1719 let display_arg = if ins.target == BlockIdx::NULL {
1720 ins.arg
1721 } else {
1722 OpArg::new(ins.target.0)
1723 };
1724 let instr_display = instr.display(display_arg, self);
1725 eprint!("{instr_display}: {depth} {effect:+} => ");
1726 }
1727 let new_depth = depth.checked_add_signed(effect).ok_or({
1728 if effect < 0 {
1729 InternalError::StackUnderflow
1730 } else {
1731 InternalError::StackOverflow
1732 }
1733 })?;
1734 if DEBUG {
1735 eprintln!("{new_depth}");
1736 }
1737 if new_depth > maxdepth {
1738 maxdepth = new_depth
1739 }
1740 if ins.target != BlockIdx::NULL {
1742 if instr.is_block_push() {
1743 let handler_effect: u32 = match instr.pseudo() {
1749 Some(PseudoInstruction::SetupCleanup { .. }) => 2,
1750 _ => 1, };
1752 let handler_depth = depth + handler_effect;
1753 if handler_depth > maxdepth {
1754 maxdepth = handler_depth;
1755 }
1756 stackdepth_push(&mut stack, &mut start_depths, ins.target, handler_depth);
1757 } else {
1758 let jump_effect = match instr.real() {
1761 Some(Instruction::Send { .. }) => 0i32,
1762 _ => effect,
1763 };
1764 let target_depth = depth.checked_add_signed(jump_effect).ok_or({
1765 if jump_effect < 0 {
1766 InternalError::StackUnderflow
1767 } else {
1768 InternalError::StackOverflow
1769 }
1770 })?;
1771 if target_depth > maxdepth {
1772 maxdepth = target_depth
1773 }
1774 stackdepth_push(&mut stack, &mut start_depths, ins.target, target_depth);
1775 }
1776 }
1777 depth = new_depth;
1778 if instr.is_scope_exit() || instr.is_unconditional_jump() {
1779 continue 'process_blocks;
1780 }
1781 }
1782 if block.next != BlockIdx::NULL {
1784 stackdepth_push(&mut stack, &mut start_depths, block.next, depth);
1785 }
1786 }
1787 if DEBUG {
1788 eprintln!("DONE: {maxdepth}");
1789 }
1790
1791 for (block, &start_depth) in self.blocks.iter_mut().zip(&start_depths) {
1792 block.start_depth = (start_depth != u32::MAX).then_some(start_depth);
1793 }
1794
1795 for block in self.blocks.iter_mut() {
1798 for ins in &mut block.instructions {
1799 if let Some(ref mut handler) = ins.except_handler {
1800 let h_start = start_depths[handler.handler_block.idx()];
1801 if h_start != u32::MAX {
1802 let adjustment = 1 + handler.preserve_lasti as u32;
1803 debug_assert!(
1804 h_start >= adjustment,
1805 "handler start depth {h_start} too shallow for adjustment {adjustment}"
1806 );
1807 handler.stack_depth = h_start.saturating_sub(adjustment);
1808 }
1809 }
1810 }
1811 }
1812
1813 Ok(maxdepth)
1814 }
1815}
1816
1817impl InstrDisplayContext for CodeInfo {
1818 type Constant = ConstantData;
1819
1820 fn get_constant(&self, consti: oparg::ConstIdx) -> &ConstantData {
1821 &self.metadata.consts[consti.as_usize()]
1822 }
1823
1824 fn get_name(&self, i: usize) -> &str {
1825 self.metadata.names[i].as_ref()
1826 }
1827
1828 fn get_varname(&self, var_num: oparg::VarNum) -> &str {
1829 self.metadata.varnames[var_num.as_usize()].as_ref()
1830 }
1831
1832 fn get_localsplus_name(&self, var_num: oparg::VarNum) -> &str {
1833 let idx = var_num.as_usize();
1834 let nlocals = self.metadata.varnames.len();
1835 if idx < nlocals {
1836 self.metadata.varnames[idx].as_ref()
1837 } else {
1838 let cell_idx = idx - nlocals;
1839 self.metadata
1840 .cellvars
1841 .get_index(cell_idx)
1842 .unwrap_or_else(|| &self.metadata.freevars[cell_idx - self.metadata.cellvars.len()])
1843 .as_ref()
1844 }
1845 }
1846}
1847
1848fn stackdepth_push(
1849 stack: &mut Vec<BlockIdx>,
1850 start_depths: &mut [u32],
1851 target: BlockIdx,
1852 depth: u32,
1853) {
1854 let idx = target.idx();
1855 let block_depth = &mut start_depths[idx];
1856 if depth > *block_depth || *block_depth == u32::MAX {
1857 *block_depth = depth;
1858 stack.push(target);
1859 }
1860}
1861
1862fn iter_blocks(blocks: &[Block]) -> impl Iterator<Item = (BlockIdx, &Block)> + '_ {
1863 let mut next = BlockIdx(0);
1864 core::iter::from_fn(move || {
1865 if next == BlockIdx::NULL {
1866 return None;
1867 }
1868 let (idx, b) = (next, &blocks[next]);
1869 next = b.next;
1870 Some((idx, b))
1871 })
1872}
1873
1874fn generate_linetable(
1876 locations: &[LineTableLocation],
1877 first_line: i32,
1878 debug_ranges: bool,
1879) -> Box<[u8]> {
1880 if locations.is_empty() {
1881 return Box::new([]);
1882 }
1883
1884 let mut linetable = Vec::new();
1885 let mut prev_line = first_line;
1888 let mut i = 0;
1889
1890 while i < locations.len() {
1891 let loc = &locations[i];
1892
1893 let mut length = 1;
1895 while i + length < locations.len() && locations[i + length] == locations[i] {
1896 length += 1;
1897 }
1898
1899 while length > 0 {
1901 let entry_length = length.min(8);
1902
1903 let line = loc.line;
1905
1906 if line == -1 {
1908 linetable.push(
1909 0x80 | ((PyCodeLocationInfoKind::None as u8) << 3) | ((entry_length - 1) as u8),
1910 );
1911 length -= entry_length;
1913 i += entry_length;
1914 continue;
1915 }
1916
1917 let end_line = loc.end_line;
1918 let line_delta = line - prev_line;
1919 let end_line_delta = end_line - line;
1920
1921 if !debug_ranges {
1923 linetable.push(
1925 0x80 | ((PyCodeLocationInfoKind::NoColumns as u8) << 3)
1926 | ((entry_length - 1) as u8),
1927 );
1928 write_signed_varint(&mut linetable, line_delta);
1929
1930 prev_line = line;
1931 length -= entry_length;
1932 i += entry_length;
1933 continue;
1934 }
1935
1936 let col = loc.col;
1938 let end_col = loc.end_col;
1939
1940 if line_delta == 0 && end_line_delta == 0 {
1942 if col < 80 && end_col - col < 16 && end_col >= col {
1943 let code = (col / 8).min(9) as u8; linetable.push(0x80 | (code << 3) | ((entry_length - 1) as u8));
1946 let col_byte = (((col % 8) as u8) << 4) | ((end_col - col) as u8 & 0xf);
1947 linetable.push(col_byte);
1948 } else if col < 128 && end_col < 128 {
1949 linetable.push(
1951 0x80 | ((PyCodeLocationInfoKind::OneLine0 as u8) << 3)
1952 | ((entry_length - 1) as u8),
1953 );
1954 linetable.push(col as u8);
1955 linetable.push(end_col as u8);
1956 } else {
1957 linetable.push(
1959 0x80 | ((PyCodeLocationInfoKind::Long as u8) << 3)
1960 | ((entry_length - 1) as u8),
1961 );
1962 write_signed_varint(&mut linetable, 0); write_varint(&mut linetable, 0); write_varint(&mut linetable, (col as u32) + 1);
1965 write_varint(&mut linetable, (end_col as u32) + 1);
1966 }
1967 } else if line_delta > 0 && line_delta < 3 && end_line_delta == 0 {
1968 if col < 128 && end_col < 128 {
1970 let code = (PyCodeLocationInfoKind::OneLine0 as u8) + (line_delta as u8);
1971 linetable.push(0x80 | (code << 3) | ((entry_length - 1) as u8));
1972 linetable.push(col as u8);
1973 linetable.push(end_col as u8);
1974 } else {
1975 linetable.push(
1977 0x80 | ((PyCodeLocationInfoKind::Long as u8) << 3)
1978 | ((entry_length - 1) as u8),
1979 );
1980 write_signed_varint(&mut linetable, line_delta);
1981 write_varint(&mut linetable, 0); write_varint(&mut linetable, (col as u32) + 1);
1983 write_varint(&mut linetable, (end_col as u32) + 1);
1984 }
1985 } else {
1986 linetable.push(
1989 0x80 | ((PyCodeLocationInfoKind::Long as u8) << 3) | ((entry_length - 1) as u8),
1990 );
1991 write_signed_varint(&mut linetable, line_delta);
1992 write_varint(&mut linetable, end_line_delta as u32);
1993 write_varint(&mut linetable, (col as u32) + 1);
1994 write_varint(&mut linetable, (end_col as u32) + 1);
1995 }
1996
1997 prev_line = line;
1998 length -= entry_length;
1999 i += entry_length;
2000 }
2001 }
2002
2003 linetable.into_boxed_slice()
2004}
2005
2006fn generate_exception_table(blocks: &[Block], block_to_index: &[u32]) -> Box<[u8]> {
2008 let mut entries: Vec<ExceptionTableEntry> = Vec::new();
2009 let mut current_entry: Option<(ExceptHandlerInfo, u32)> = None; let mut instr_index = 0u32;
2011
2012 for (_, block) in iter_blocks(blocks) {
2016 for instr in &block.instructions {
2017 let instr_size = instr.arg.instr_size() as u32 + instr.cache_entries;
2019
2020 match (¤t_entry, instr.except_handler) {
2021 (None, None) => {}
2023
2024 (None, Some(handler)) => {
2026 current_entry = Some((handler, instr_index));
2027 }
2028
2029 (Some((curr_handler, _)), Some(handler))
2031 if curr_handler.handler_block == handler.handler_block
2032 && curr_handler.stack_depth == handler.stack_depth
2033 && curr_handler.preserve_lasti == handler.preserve_lasti => {}
2034
2035 (Some((curr_handler, start)), Some(handler)) => {
2037 let target_index = block_to_index[curr_handler.handler_block.idx()];
2038 entries.push(ExceptionTableEntry::new(
2039 *start,
2040 instr_index,
2041 target_index,
2042 curr_handler.stack_depth as u16,
2043 curr_handler.preserve_lasti,
2044 ));
2045 current_entry = Some((handler, instr_index));
2046 }
2047
2048 (Some((curr_handler, start)), None) => {
2050 let target_index = block_to_index[curr_handler.handler_block.idx()];
2051 entries.push(ExceptionTableEntry::new(
2052 *start,
2053 instr_index,
2054 target_index,
2055 curr_handler.stack_depth as u16,
2056 curr_handler.preserve_lasti,
2057 ));
2058 current_entry = None;
2059 }
2060 }
2061
2062 instr_index += instr_size; }
2064 }
2065
2066 if let Some((curr_handler, start)) = current_entry {
2068 let target_index = block_to_index[curr_handler.handler_block.idx()];
2069 entries.push(ExceptionTableEntry::new(
2070 start,
2071 instr_index,
2072 target_index,
2073 curr_handler.stack_depth as u16,
2074 curr_handler.preserve_lasti,
2075 ));
2076 }
2077
2078 encode_exception_table(&entries)
2079}
2080
2081pub(crate) fn mark_except_handlers(blocks: &mut [Block]) {
2084 for block in blocks.iter_mut() {
2086 block.except_handler = false;
2087 block.preserve_lasti = false;
2088 }
2089 let targets: Vec<usize> = blocks
2091 .iter()
2092 .flat_map(|b| b.instructions.iter())
2093 .filter(|i| i.instr.is_block_push() && i.target != BlockIdx::NULL)
2094 .map(|i| i.target.idx())
2095 .collect();
2096 for idx in targets {
2097 blocks[idx].except_handler = true;
2098 }
2099}
2100
2101fn mark_cold(blocks: &mut [Block]) {
2103 let n = blocks.len();
2104 let mut warm = vec![false; n];
2105 let mut queue = VecDeque::new();
2106
2107 warm[0] = true;
2108 queue.push_back(BlockIdx(0));
2109
2110 while let Some(block_idx) = queue.pop_front() {
2111 let block = &blocks[block_idx.idx()];
2112
2113 let has_fallthrough = block
2114 .instructions
2115 .last()
2116 .map(|ins| !ins.instr.is_scope_exit() && !ins.instr.is_unconditional_jump())
2117 .unwrap_or(true);
2118 if has_fallthrough && block.next != BlockIdx::NULL {
2119 let next_idx = block.next.idx();
2120 if !blocks[next_idx].except_handler && !warm[next_idx] {
2121 warm[next_idx] = true;
2122 queue.push_back(block.next);
2123 }
2124 }
2125
2126 for instr in &block.instructions {
2127 if instr.target != BlockIdx::NULL {
2128 let target_idx = instr.target.idx();
2129 if !blocks[target_idx].except_handler && !warm[target_idx] {
2130 warm[target_idx] = true;
2131 queue.push_back(instr.target);
2132 }
2133 }
2134 }
2135 }
2136
2137 for (i, block) in blocks.iter_mut().enumerate() {
2138 block.cold = !warm[i];
2139 }
2140}
2141
2142fn push_cold_blocks_to_end(blocks: &mut Vec<Block>) {
2144 if blocks.len() <= 1 {
2145 return;
2146 }
2147
2148 mark_cold(blocks);
2149
2150 let fixups: Vec<(BlockIdx, BlockIdx)> = iter_blocks(blocks)
2152 .filter(|(_, block)| {
2153 block.cold
2154 && block.next != BlockIdx::NULL
2155 && !blocks[block.next.idx()].cold
2156 && block
2157 .instructions
2158 .last()
2159 .map(|ins| !ins.instr.is_scope_exit() && !ins.instr.is_unconditional_jump())
2160 .unwrap_or(true)
2161 })
2162 .map(|(idx, block)| (idx, block.next))
2163 .collect();
2164
2165 for (cold_idx, warm_next) in fixups {
2166 let jump_block_idx = BlockIdx(blocks.len() as u32);
2167 let loc = blocks[cold_idx.idx()]
2168 .instructions
2169 .last()
2170 .map(|i| i.location)
2171 .unwrap_or_default();
2172 let end_loc = blocks[cold_idx.idx()]
2173 .instructions
2174 .last()
2175 .map(|i| i.end_location)
2176 .unwrap_or_default();
2177 let mut jump_block = Block {
2178 cold: true,
2179 ..Block::default()
2180 };
2181 jump_block.instructions.push(InstructionInfo {
2182 instr: PseudoInstruction::JumpNoInterrupt {
2183 delta: Arg::marker(),
2184 }
2185 .into(),
2186 arg: OpArg::new(0),
2187 target: warm_next,
2188 location: loc,
2189 end_location: end_loc,
2190 except_handler: None,
2191 lineno_override: None,
2192 cache_entries: 0,
2193 });
2194 jump_block.next = blocks[cold_idx.idx()].next;
2195 blocks[cold_idx.idx()].next = jump_block_idx;
2196 blocks.push(jump_block);
2197 }
2198
2199 let mut cold_head: BlockIdx = BlockIdx::NULL;
2201 let mut cold_tail: BlockIdx = BlockIdx::NULL;
2202 let mut current = BlockIdx(0);
2203 assert!(!blocks[0].cold);
2204
2205 while current != BlockIdx::NULL {
2206 let next = blocks[current.idx()].next;
2207 if next == BlockIdx::NULL {
2208 break;
2209 }
2210
2211 if blocks[next.idx()].cold {
2212 let cold_start = next;
2213 let mut cold_end = next;
2214 while blocks[cold_end.idx()].next != BlockIdx::NULL
2215 && blocks[blocks[cold_end.idx()].next.idx()].cold
2216 {
2217 cold_end = blocks[cold_end.idx()].next;
2218 }
2219
2220 let after_cold = blocks[cold_end.idx()].next;
2221 blocks[current.idx()].next = after_cold;
2222 blocks[cold_end.idx()].next = BlockIdx::NULL;
2223
2224 if cold_head == BlockIdx::NULL {
2225 cold_head = cold_start;
2226 } else {
2227 blocks[cold_tail.idx()].next = cold_start;
2228 }
2229 cold_tail = cold_end;
2230 } else {
2231 current = next;
2232 }
2233 }
2234
2235 if cold_head != BlockIdx::NULL {
2236 let mut last = current;
2237 while blocks[last.idx()].next != BlockIdx::NULL {
2238 last = blocks[last.idx()].next;
2239 }
2240 blocks[last.idx()].next = cold_head;
2241 }
2242}
2243
2244fn split_blocks_at_jumps(blocks: &mut Vec<Block>) {
2248 let mut bi = 0;
2249 while bi < blocks.len() {
2250 let split_at = {
2252 let block = &blocks[bi];
2253 let mut found = None;
2254 for (i, ins) in block.instructions.iter().enumerate() {
2255 if is_conditional_jump(&ins.instr)
2256 || ins.instr.is_unconditional_jump()
2257 || ins.instr.is_scope_exit()
2258 {
2259 if i + 1 < block.instructions.len() {
2260 found = Some(i + 1);
2261 }
2262 break;
2263 }
2264 }
2265 found
2266 };
2267 if let Some(pos) = split_at {
2268 let new_block_idx = BlockIdx(blocks.len() as u32);
2269 let tail: Vec<InstructionInfo> = blocks[bi].instructions.drain(pos..).collect();
2270 let old_next = blocks[bi].next;
2271 let cold = blocks[bi].cold;
2272 blocks[bi].next = new_block_idx;
2273 blocks.push(Block {
2274 instructions: tail,
2275 next: old_next,
2276 cold,
2277 ..Block::default()
2278 });
2279 } else {
2281 bi += 1;
2282 }
2283 }
2284}
2285
2286fn jump_threading(blocks: &mut [Block]) {
2290 let mut changed = true;
2291 while changed {
2292 changed = false;
2293 for bi in 0..blocks.len() {
2294 let last_idx = match blocks[bi].instructions.len().checked_sub(1) {
2295 Some(i) => i,
2296 None => continue,
2297 };
2298 let ins = &blocks[bi].instructions[last_idx];
2299 let target = ins.target;
2300 if target == BlockIdx::NULL {
2301 continue;
2302 }
2303 if !ins.instr.is_unconditional_jump() && !is_conditional_jump(&ins.instr) {
2304 continue;
2305 }
2306 let target_block = &blocks[target.idx()];
2308 if let Some(target_ins) = target_block.instructions.first()
2309 && target_ins.instr.is_unconditional_jump()
2310 && target_ins.target != BlockIdx::NULL
2311 && target_ins.target != target
2312 {
2313 let final_target = target_ins.target;
2314 blocks[bi].instructions[last_idx].target = final_target;
2315 changed = true;
2316 }
2317 }
2318 }
2319}
2320
2321fn is_conditional_jump(instr: &AnyInstruction) -> bool {
2322 matches!(
2323 instr.real(),
2324 Some(
2325 Instruction::PopJumpIfFalse { .. }
2326 | Instruction::PopJumpIfTrue { .. }
2327 | Instruction::PopJumpIfNone { .. }
2328 | Instruction::PopJumpIfNotNone { .. }
2329 )
2330 )
2331}
2332
2333fn reversed_conditional(instr: &AnyInstruction) -> Option<AnyInstruction> {
2335 Some(match instr.real()? {
2336 Instruction::PopJumpIfFalse { .. } => Instruction::PopJumpIfTrue {
2337 delta: Arg::marker(),
2338 }
2339 .into(),
2340 Instruction::PopJumpIfTrue { .. } => Instruction::PopJumpIfFalse {
2341 delta: Arg::marker(),
2342 }
2343 .into(),
2344 Instruction::PopJumpIfNone { .. } => Instruction::PopJumpIfNotNone {
2345 delta: Arg::marker(),
2346 }
2347 .into(),
2348 Instruction::PopJumpIfNotNone { .. } => Instruction::PopJumpIfNone {
2349 delta: Arg::marker(),
2350 }
2351 .into(),
2352 _ => return None,
2353 })
2354}
2355
2356fn normalize_jumps(blocks: &mut Vec<Block>) {
2358 let mut visit_order = Vec::new();
2359 let mut visited = vec![false; blocks.len()];
2360 let mut current = BlockIdx(0);
2361 while current != BlockIdx::NULL {
2362 visit_order.push(current);
2363 visited[current.idx()] = true;
2364 current = blocks[current.idx()].next;
2365 }
2366
2367 visited.fill(false);
2368
2369 for &block_idx in &visit_order {
2370 let idx = block_idx.idx();
2371 visited[idx] = true;
2372
2373 let next = blocks[idx].next;
2375 if next != BlockIdx::NULL {
2376 let last = blocks[idx].instructions.last();
2377 let is_jump_to_next = last.is_some_and(|ins| {
2378 ins.instr.is_unconditional_jump()
2379 && ins.target != BlockIdx::NULL
2380 && ins.target == next
2381 });
2382 if is_jump_to_next && let Some(last_instr) = blocks[idx].instructions.last_mut() {
2383 last_instr.instr = Instruction::Nop.into();
2384 last_instr.target = BlockIdx::NULL;
2385 }
2386 }
2387
2388 let last = blocks[idx].instructions.last();
2390 if let Some(last_ins) = last
2391 && is_conditional_jump(&last_ins.instr)
2392 && last_ins.target != BlockIdx::NULL
2393 {
2394 let target = last_ins.target;
2395 let is_forward = !visited[target.idx()];
2396
2397 if is_forward {
2398 let not_taken = InstructionInfo {
2400 instr: Instruction::NotTaken.into(),
2401 arg: OpArg::new(0),
2402 target: BlockIdx::NULL,
2403 location: last_ins.location,
2404 end_location: last_ins.end_location,
2405 except_handler: last_ins.except_handler,
2406 lineno_override: None,
2407 cache_entries: 0,
2408 };
2409 blocks[idx].instructions.push(not_taken);
2410 } else {
2411 let loc = last_ins.location;
2415 let end_loc = last_ins.end_location;
2416 let exc_handler = last_ins.except_handler;
2417
2418 if let Some(reversed) = reversed_conditional(&last_ins.instr) {
2419 let old_next = blocks[idx].next;
2420 let is_cold = blocks[idx].cold;
2421
2422 let new_block_idx = BlockIdx(blocks.len() as u32);
2424 let mut new_block = Block {
2425 cold: is_cold,
2426 ..Block::default()
2427 };
2428 new_block.instructions.push(InstructionInfo {
2429 instr: Instruction::NotTaken.into(),
2430 arg: OpArg::new(0),
2431 target: BlockIdx::NULL,
2432 location: loc,
2433 end_location: end_loc,
2434 except_handler: exc_handler,
2435 lineno_override: None,
2436 cache_entries: 0,
2437 });
2438 new_block.instructions.push(InstructionInfo {
2439 instr: PseudoInstruction::Jump {
2440 delta: Arg::marker(),
2441 }
2442 .into(),
2443 arg: OpArg::new(0),
2444 target,
2445 location: loc,
2446 end_location: end_loc,
2447 except_handler: exc_handler,
2448 lineno_override: None,
2449 cache_entries: 0,
2450 });
2451 new_block.next = old_next;
2452
2453 let last_mut = blocks[idx].instructions.last_mut().unwrap();
2455 last_mut.instr = reversed;
2456 last_mut.target = old_next;
2457
2458 blocks[idx].next = new_block_idx;
2460 blocks.push(new_block);
2461
2462 visited.push(true);
2464 }
2465 }
2466 }
2467 }
2468
2469 let mut visit_order = Vec::new();
2471 let mut current = BlockIdx(0);
2472 while current != BlockIdx::NULL {
2473 visit_order.push(current);
2474 current = blocks[current.idx()].next;
2475 }
2476
2477 for &block_idx in &visit_order {
2480 let idx = block_idx.idx();
2481 let mut replacements: Vec<(usize, Vec<InstructionInfo>)> = Vec::new();
2482 for (i, ins) in blocks[idx].instructions.iter().enumerate() {
2483 if !ins.instr.is_unconditional_jump() || ins.target == BlockIdx::NULL {
2484 continue;
2485 }
2486 let mut target_idx = ins.target.idx();
2488 while blocks[target_idx].instructions.is_empty()
2489 && blocks[target_idx].next != BlockIdx::NULL
2490 {
2491 target_idx = blocks[target_idx].next.idx();
2492 }
2493 let target_block = &blocks[target_idx];
2494 if target_block.instructions.len() == 2 {
2496 let t0 = &target_block.instructions[0];
2497 let t1 = &target_block.instructions[1];
2498 if matches!(t0.instr, AnyInstruction::Real(_))
2499 && !t0.instr.is_scope_exit()
2500 && !t0.instr.is_unconditional_jump()
2501 && matches!(t1.instr.real(), Some(Instruction::ReturnValue))
2502 {
2503 let mut load = *t0;
2504 let mut ret = *t1;
2505 load.location = ins.location;
2507 load.end_location = ins.end_location;
2508 load.except_handler = ins.except_handler;
2509 ret.location = ins.location;
2510 ret.end_location = ins.end_location;
2511 ret.except_handler = ins.except_handler;
2512 replacements.push((i, vec![load, ret]));
2513 }
2514 }
2515 }
2516 for (i, new_insts) in replacements.into_iter().rev() {
2518 blocks[idx].instructions.splice(i..i + 1, new_insts);
2519 }
2520 }
2521
2522 let mut block_order = vec![0u32; blocks.len()];
2524 for (pos, &block_idx) in visit_order.iter().enumerate() {
2525 block_order[block_idx.idx()] = pos as u32;
2526 }
2527
2528 for &block_idx in &visit_order {
2529 let source_pos = block_order[block_idx.idx()];
2530 for info in &mut blocks[block_idx.idx()].instructions {
2531 let target = info.target;
2532 if target == BlockIdx::NULL {
2533 continue;
2534 }
2535 let target_pos = block_order[target.idx()];
2536 info.instr = match info.instr {
2537 AnyInstruction::Pseudo(PseudoInstruction::Jump { .. }) => {
2538 if target_pos > source_pos {
2539 Instruction::JumpForward {
2540 delta: Arg::marker(),
2541 }
2542 .into()
2543 } else {
2544 Instruction::JumpBackward {
2545 delta: Arg::marker(),
2546 }
2547 .into()
2548 }
2549 }
2550 AnyInstruction::Pseudo(PseudoInstruction::JumpNoInterrupt { .. }) => {
2551 if target_pos > source_pos {
2552 Instruction::JumpForward {
2553 delta: Arg::marker(),
2554 }
2555 .into()
2556 } else {
2557 Instruction::JumpBackwardNoInterrupt {
2558 delta: Arg::marker(),
2559 }
2560 .into()
2561 }
2562 }
2563 other => other,
2564 };
2565 }
2566 }
2567}
2568
2569fn inline_small_or_no_lineno_blocks(blocks: &mut [Block]) {
2571 const MAX_COPY_SIZE: usize = 4;
2572
2573 let block_exits_scope = |block: &Block| {
2574 block
2575 .instructions
2576 .last()
2577 .is_some_and(|ins| ins.instr.is_scope_exit())
2578 };
2579 let block_has_no_lineno = |block: &Block| {
2580 block
2581 .instructions
2582 .iter()
2583 .all(|ins| !instruction_has_lineno(ins))
2584 };
2585
2586 loop {
2587 let mut changes = false;
2588 let mut current = BlockIdx(0);
2589 while current != BlockIdx::NULL {
2590 let next = blocks[current.idx()].next;
2591 let Some(last) = blocks[current.idx()].instructions.last().copied() else {
2592 current = next;
2593 continue;
2594 };
2595 if !last.instr.is_unconditional_jump() || last.target == BlockIdx::NULL {
2596 current = next;
2597 continue;
2598 }
2599
2600 let target = last.target;
2601 let small_exit_block = block_exits_scope(&blocks[target.idx()])
2602 && blocks[target.idx()].instructions.len() <= MAX_COPY_SIZE;
2603 let no_lineno_no_fallthrough = block_has_no_lineno(&blocks[target.idx()])
2604 && !block_has_fallthrough(&blocks[target.idx()]);
2605
2606 if small_exit_block || no_lineno_no_fallthrough {
2607 if let Some(last_instr) = blocks[current.idx()].instructions.last_mut() {
2608 last_instr.instr = Instruction::Nop.into();
2609 last_instr.arg = OpArg::new(0);
2610 last_instr.target = BlockIdx::NULL;
2611 }
2612 let appended = blocks[target.idx()].instructions.clone();
2613 blocks[current.idx()].instructions.extend(appended);
2614 changes = true;
2615 }
2616
2617 current = next;
2618 }
2619
2620 if !changes {
2621 break;
2622 }
2623 }
2624}
2625
2626fn next_nonempty_block(blocks: &[Block], mut idx: BlockIdx) -> BlockIdx {
2628 while idx != BlockIdx::NULL
2629 && blocks[idx.idx()].instructions.is_empty()
2630 && blocks[idx.idx()].next != BlockIdx::NULL
2631 {
2632 idx = blocks[idx.idx()].next;
2633 }
2634 idx
2635}
2636
2637fn instruction_lineno(instr: &InstructionInfo) -> i32 {
2638 instr
2639 .lineno_override
2640 .unwrap_or_else(|| instr.location.line.get() as i32)
2641}
2642
2643fn instruction_has_lineno(instr: &InstructionInfo) -> bool {
2644 instruction_lineno(instr) > 0
2645}
2646
2647fn block_has_fallthrough(block: &Block) -> bool {
2648 block
2649 .instructions
2650 .last()
2651 .is_none_or(|ins| !ins.instr.is_scope_exit() && !ins.instr.is_unconditional_jump())
2652}
2653
2654fn is_jump_instruction(instr: &InstructionInfo) -> bool {
2655 instr.instr.is_unconditional_jump() || is_conditional_jump(&instr.instr)
2656}
2657
2658fn is_exit_without_lineno(block: &Block) -> bool {
2659 let Some(first) = block.instructions.first() else {
2660 return false;
2661 };
2662 let Some(last) = block.instructions.last() else {
2663 return false;
2664 };
2665 !instruction_has_lineno(first) && last.instr.is_scope_exit()
2666}
2667
2668fn maybe_propagate_location(
2669 instr: &mut InstructionInfo,
2670 location: SourceLocation,
2671 end_location: SourceLocation,
2672) {
2673 if !instruction_has_lineno(instr) {
2674 instr.location = location;
2675 instr.end_location = end_location;
2676 instr.lineno_override = None;
2677 }
2678}
2679
2680fn propagate_locations_in_block(
2681 block: &mut Block,
2682 location: SourceLocation,
2683 end_location: SourceLocation,
2684) {
2685 let mut prev_location = location;
2686 let mut prev_end_location = end_location;
2687 for instr in &mut block.instructions {
2688 maybe_propagate_location(instr, prev_location, prev_end_location);
2689 prev_location = instr.location;
2690 prev_end_location = instr.end_location;
2691 }
2692}
2693
2694fn compute_predecessors(blocks: &[Block]) -> Vec<u32> {
2695 let mut predecessors = vec![0u32; blocks.len()];
2696 if blocks.is_empty() {
2697 return predecessors;
2698 }
2699
2700 predecessors[0] = 1;
2701 let mut current = BlockIdx(0);
2702 while current != BlockIdx::NULL {
2703 let block = &blocks[current.idx()];
2704 if block_has_fallthrough(block) {
2705 let next = next_nonempty_block(blocks, block.next);
2706 if next != BlockIdx::NULL {
2707 predecessors[next.idx()] += 1;
2708 }
2709 }
2710 for ins in &block.instructions {
2711 if ins.target != BlockIdx::NULL {
2712 let target = next_nonempty_block(blocks, ins.target);
2713 if target != BlockIdx::NULL {
2714 predecessors[target.idx()] += 1;
2715 }
2716 }
2717 }
2718 current = block.next;
2719 }
2720 predecessors
2721}
2722
2723fn duplicate_exits_without_lineno(blocks: &mut Vec<Block>, predecessors: &mut Vec<u32>) {
2724 let mut current = BlockIdx(0);
2725 while current != BlockIdx::NULL {
2726 let block = &blocks[current.idx()];
2727 let last = match block.instructions.last() {
2728 Some(ins) if ins.target != BlockIdx::NULL && is_jump_instruction(ins) => ins,
2729 _ => {
2730 current = blocks[current.idx()].next;
2731 continue;
2732 }
2733 };
2734
2735 let target = next_nonempty_block(blocks, last.target);
2736 if target == BlockIdx::NULL || !is_exit_without_lineno(&blocks[target.idx()]) {
2737 current = blocks[current.idx()].next;
2738 continue;
2739 }
2740 if predecessors[target.idx()] <= 1 {
2741 current = blocks[current.idx()].next;
2742 continue;
2743 }
2744
2745 let new_idx = BlockIdx(blocks.len() as u32);
2747 let mut new_block = blocks[target.idx()].clone();
2748 let jump_loc = last.location;
2749 let jump_end_loc = last.end_location;
2750 propagate_locations_in_block(&mut new_block, jump_loc, jump_end_loc);
2751 let old_next = blocks[current.idx()].next;
2752 new_block.next = old_next;
2753 blocks.push(new_block);
2754 blocks[current.idx()].next = new_idx;
2755
2756 let last_mut = blocks[current.idx()].instructions.last_mut().unwrap();
2758 last_mut.target = new_idx;
2759 predecessors[target.idx()] -= 1;
2760 predecessors.push(1);
2761
2762 current = old_next;
2764 }
2765
2766 current = BlockIdx(0);
2767 while current != BlockIdx::NULL {
2768 let block = &blocks[current.idx()];
2769 if let Some(last) = block.instructions.last()
2770 && block_has_fallthrough(block)
2771 {
2772 let target = next_nonempty_block(blocks, block.next);
2773 if target != BlockIdx::NULL
2774 && predecessors[target.idx()] == 1
2775 && is_exit_without_lineno(&blocks[target.idx()])
2776 {
2777 let last_location = last.location;
2778 let last_end_location = last.end_location;
2779 propagate_locations_in_block(
2780 &mut blocks[target.idx()],
2781 last_location,
2782 last_end_location,
2783 );
2784 }
2785 }
2786 current = blocks[current.idx()].next;
2787 }
2788}
2789
2790fn propagate_line_numbers(blocks: &mut [Block], predecessors: &[u32]) {
2791 let mut current = BlockIdx(0);
2792 while current != BlockIdx::NULL {
2793 let last = blocks[current.idx()].instructions.last().copied();
2794 if let Some(last) = last {
2795 let (next_block, has_fallthrough) = {
2796 let block = &blocks[current.idx()];
2797 (block.next, block_has_fallthrough(block))
2798 };
2799
2800 {
2801 let block = &mut blocks[current.idx()];
2802 let mut prev_location = None;
2803 for instr in &mut block.instructions {
2804 if let Some((location, end_location)) = prev_location {
2805 maybe_propagate_location(instr, location, end_location);
2806 }
2807 prev_location = Some((instr.location, instr.end_location));
2808 }
2809 }
2810
2811 if has_fallthrough {
2812 let target = next_nonempty_block(blocks, next_block);
2813 if target != BlockIdx::NULL && predecessors[target.idx()] == 1 {
2814 propagate_locations_in_block(
2815 &mut blocks[target.idx()],
2816 last.location,
2817 last.end_location,
2818 );
2819 }
2820 }
2821
2822 if is_jump_instruction(&last) {
2823 let target = next_nonempty_block(blocks, last.target);
2824 if target != BlockIdx::NULL && predecessors[target.idx()] == 1 {
2825 propagate_locations_in_block(
2826 &mut blocks[target.idx()],
2827 last.location,
2828 last.end_location,
2829 );
2830 }
2831 }
2832 }
2833 current = blocks[current.idx()].next;
2834 }
2835}
2836
2837fn resolve_line_numbers(blocks: &mut Vec<Block>) {
2838 let mut predecessors = compute_predecessors(blocks);
2839 duplicate_exits_without_lineno(blocks, &mut predecessors);
2840 propagate_line_numbers(blocks, &predecessors);
2841}
2842
2843fn duplicate_end_returns(blocks: &mut [Block]) {
2846 let mut last_block = BlockIdx::NULL;
2848 let mut current = BlockIdx(0);
2849 while current != BlockIdx::NULL {
2850 if !blocks[current.idx()].instructions.is_empty() {
2851 last_block = current;
2852 }
2853 current = blocks[current.idx()].next;
2854 }
2855 if last_block == BlockIdx::NULL {
2856 return;
2857 }
2858
2859 let last_insts = &blocks[last_block.idx()].instructions;
2860 let is_return_block = last_insts.len() == 2
2862 && matches!(
2863 last_insts[0].instr,
2864 AnyInstruction::Real(Instruction::LoadConst { .. })
2865 )
2866 && matches!(
2867 last_insts[1].instr,
2868 AnyInstruction::Real(Instruction::ReturnValue)
2869 );
2870 if !is_return_block {
2871 return;
2872 }
2873
2874 let return_insts: Vec<InstructionInfo> = last_insts[last_insts.len() - 2..].to_vec();
2876
2877 let mut blocks_to_fix = Vec::new();
2879 current = BlockIdx(0);
2880 while current != BlockIdx::NULL {
2881 let block = &blocks[current.idx()];
2882 let next = next_nonempty_block(blocks, block.next);
2883 if current != last_block && next == last_block && !block.cold && !block.except_handler {
2884 let last_ins = block.instructions.last();
2885 let has_fallthrough = last_ins
2886 .map(|ins| !ins.instr.is_scope_exit() && !ins.instr.is_unconditional_jump())
2887 .unwrap_or(true);
2888 let already_has_return = block.instructions.len() >= 2 && {
2890 let n = block.instructions.len();
2891 matches!(
2892 block.instructions[n - 2].instr,
2893 AnyInstruction::Real(Instruction::LoadConst { .. })
2894 ) && matches!(
2895 block.instructions[n - 1].instr,
2896 AnyInstruction::Real(Instruction::ReturnValue)
2897 )
2898 };
2899 if has_fallthrough && !already_has_return {
2900 blocks_to_fix.push(current);
2901 }
2902 }
2903 current = blocks[current.idx()].next;
2904 }
2905
2906 for block_idx in blocks_to_fix {
2908 blocks[block_idx.idx()]
2909 .instructions
2910 .extend_from_slice(&return_insts);
2911 }
2912}
2913
2914pub(crate) fn label_exception_targets(blocks: &mut [Block]) {
2918 #[derive(Clone)]
2919 struct ExceptEntry {
2920 handler_block: BlockIdx,
2921 preserve_lasti: bool,
2922 }
2923
2924 let num_blocks = blocks.len();
2925 if num_blocks == 0 {
2926 return;
2927 }
2928
2929 let mut visited = vec![false; num_blocks];
2930 let mut block_stacks: Vec<Option<Vec<ExceptEntry>>> = vec![None; num_blocks];
2931
2932 visited[0] = true;
2934 block_stacks[0] = Some(Vec::new());
2935
2936 let mut todo = vec![BlockIdx(0)];
2937
2938 while let Some(block_idx) = todo.pop() {
2939 let bi = block_idx.idx();
2940 let mut stack = block_stacks[bi].take().unwrap_or_default();
2941 let mut last_yield_except_depth: i32 = -1;
2942
2943 let instr_count = blocks[bi].instructions.len();
2944 for i in 0..instr_count {
2945 let target = blocks[bi].instructions[i].target;
2947 let arg = blocks[bi].instructions[i].arg;
2948 let is_push = blocks[bi].instructions[i].instr.is_block_push();
2949 let is_pop = blocks[bi].instructions[i].instr.is_pop_block();
2950
2951 if is_push {
2952 let preserve_lasti = matches!(
2954 blocks[bi].instructions[i].instr.pseudo(),
2955 Some(
2956 PseudoInstruction::SetupWith { .. }
2957 | PseudoInstruction::SetupCleanup { .. }
2958 )
2959 );
2960
2961 if preserve_lasti && target != BlockIdx::NULL {
2963 blocks[target.idx()].preserve_lasti = true;
2964 }
2965
2966 if target != BlockIdx::NULL && !visited[target.idx()] {
2968 visited[target.idx()] = true;
2969 block_stacks[target.idx()] = Some(stack.clone());
2970 todo.push(target);
2971 }
2972
2973 stack.push(ExceptEntry {
2975 handler_block: target,
2976 preserve_lasti,
2977 });
2978 } else if is_pop {
2979 debug_assert!(
2980 !stack.is_empty(),
2981 "POP_BLOCK with empty except stack at block {bi} instruction {i}"
2982 );
2983 stack.pop();
2984 blocks[bi].instructions[i].instr = Instruction::Nop.into();
2986 } else {
2987 let handler_info = stack.last().map(|e| ExceptHandlerInfo {
2990 handler_block: e.handler_block,
2991 stack_depth: 0,
2992 preserve_lasti: e.preserve_lasti,
2993 });
2994 blocks[bi].instructions[i].except_handler = handler_info;
2995
2996 if let Some(Instruction::YieldValue { .. }) =
3003 blocks[bi].instructions[i].instr.real()
3004 {
3005 last_yield_except_depth = stack.len() as i32;
3006 }
3007
3008 if let Some(Instruction::Resume { context }) =
3010 blocks[bi].instructions[i].instr.real()
3011 {
3012 let location = context.get(arg).location();
3013 match location {
3014 oparg::ResumeLocation::AtFuncStart => {}
3015 _ => {
3016 if last_yield_except_depth == 1 {
3017 blocks[bi].instructions[i].arg =
3018 OpArg::new(oparg::ResumeContext::new(location, true).as_u32());
3019 }
3020 last_yield_except_depth = -1;
3021 }
3022 }
3023 }
3024
3025 if target != BlockIdx::NULL && !visited[target.idx()] {
3027 visited[target.idx()] = true;
3028 block_stacks[target.idx()] = Some(stack.clone());
3029 todo.push(target);
3030 }
3031 }
3032 }
3033
3034 let next = blocks[bi].next;
3036 if next != BlockIdx::NULL && !visited[next.idx()] {
3037 let has_fallthrough = blocks[bi]
3038 .instructions
3039 .last()
3040 .map(|ins| !ins.instr.is_scope_exit() && !ins.instr.is_unconditional_jump())
3041 .unwrap_or(true); if has_fallthrough {
3043 visited[next.idx()] = true;
3044 block_stacks[next.idx()] = Some(stack);
3045 todo.push(next);
3046 }
3047 }
3048 }
3049}
3050
3051pub(crate) fn convert_pseudo_ops(blocks: &mut [Block], cellfixedoffsets: &[u32]) {
3054 for block in blocks.iter_mut() {
3055 for info in &mut block.instructions {
3056 let Some(pseudo) = info.instr.pseudo() else {
3057 continue;
3058 };
3059 match pseudo {
3060 PseudoInstruction::SetupCleanup { .. }
3062 | PseudoInstruction::SetupFinally { .. }
3063 | PseudoInstruction::SetupWith { .. } => {
3064 info.instr = Instruction::Nop.into();
3065 }
3066 PseudoInstruction::PopBlock => {
3069 info.instr = Instruction::Nop.into();
3070 }
3071 PseudoInstruction::LoadClosure { i } => {
3073 let cell_relative = i.get(info.arg) as usize;
3074 let new_idx = cellfixedoffsets[cell_relative];
3075 info.arg = OpArg::new(new_idx);
3076 info.instr = Instruction::LoadFast {
3077 var_num: Arg::marker(),
3078 }
3079 .into();
3080 }
3081 PseudoInstruction::Jump { .. } | PseudoInstruction::JumpNoInterrupt { .. } => {}
3083 PseudoInstruction::AnnotationsPlaceholder
3085 | PseudoInstruction::JumpIfFalse { .. }
3086 | PseudoInstruction::JumpIfTrue { .. }
3087 | PseudoInstruction::StoreFastMaybeNull { .. } => {
3088 unreachable!("Unexpected pseudo instruction in convert_pseudo_ops: {pseudo:?}")
3089 }
3090 }
3091 }
3092 }
3093}
3094
3095pub(crate) fn build_cellfixedoffsets(
3099 varnames: &IndexSet<String>,
3100 cellvars: &IndexSet<String>,
3101 freevars: &IndexSet<String>,
3102) -> Vec<u32> {
3103 let nlocals = varnames.len();
3104 let ncells = cellvars.len();
3105 let nfrees = freevars.len();
3106 let mut fixed = Vec::with_capacity(ncells + nfrees);
3107 let mut numdropped = 0usize;
3108 for (i, cellvar) in cellvars.iter().enumerate() {
3109 if let Some(local_idx) = varnames.get_index_of(cellvar) {
3110 fixed.push(local_idx as u32);
3111 numdropped += 1;
3112 } else {
3113 fixed.push((nlocals + i - numdropped) as u32);
3114 }
3115 }
3116 for i in 0..nfrees {
3117 fixed.push((nlocals + ncells - numdropped + i) as u32);
3118 }
3119 fixed
3120}
3121
3122pub(crate) fn fixup_deref_opargs(blocks: &mut [Block], cellfixedoffsets: &[u32]) {
3125 for block in blocks.iter_mut() {
3126 for info in &mut block.instructions {
3127 let Some(instr) = info.instr.real() else {
3128 continue;
3129 };
3130 let needs_fixup = matches!(
3131 instr,
3132 Instruction::LoadDeref { .. }
3133 | Instruction::StoreDeref { .. }
3134 | Instruction::DeleteDeref { .. }
3135 | Instruction::LoadFromDictOrDeref { .. }
3136 | Instruction::MakeCell { .. }
3137 );
3138 if needs_fixup {
3139 let cell_relative = u32::from(info.arg) as usize;
3140 info.arg = OpArg::new(cellfixedoffsets[cell_relative]);
3141 }
3142 }
3143 }
3144}