Skip to main content

rustpython_compiler_core/
bytecode.rs

1//! Implement python as a virtual machine with bytecode. This module
2//! implements bytecode structure.
3
4use crate::{
5    marshal::MarshalError,
6    varint::{read_varint, read_varint_with_start, write_varint_be, write_varint_with_start},
7    {OneIndexed, SourceLocation},
8};
9use alloc::{borrow::ToOwned, boxed::Box, collections::BTreeSet, fmt, string::String, vec::Vec};
10use bitflags::bitflags;
11use core::{
12    cell::UnsafeCell,
13    hash, mem,
14    ops::{Deref, DerefMut, Index, IndexMut},
15    sync::atomic::{AtomicU8, AtomicU16, AtomicUsize, Ordering},
16};
17use itertools::Itertools;
18use malachite_bigint::BigInt;
19use num_complex::Complex64;
20use num_traits::Zero;
21use rustpython_wtf8::{Wtf8, Wtf8Buf};
22
23pub use crate::bytecode::{
24    instruction::{
25        AnyInstruction, AnyOpcode, Arg, Instruction, Opcode, PseudoInstruction, PseudoOpcode,
26        StackEffect,
27    },
28    oparg::{
29        BinaryOperator, BuildSliceArgCount, CommonConstant, ComparisonOperator, ConvertValueOparg,
30        IntrinsicFunction1, IntrinsicFunction2, Invert, Label, LoadAttr, LoadSuperAttr,
31        MakeFunctionFlag, MakeFunctionFlags, NameIdx, OpArg, OpArgByte, OpArgState, OpArgType,
32        RaiseKind, SpecialMethod, UnpackExArgs,
33    },
34};
35
36mod instruction;
37mod opcode_metadata;
38
39pub mod oparg;
40
41/// Exception table entry for zero-cost exception handling
42/// Format: (start, size, target, depth<<1|lasti)
43#[derive(Clone, Copy, Debug, PartialEq, Eq)]
44pub struct ExceptionTableEntry {
45    /// Start instruction offset (inclusive)
46    pub start: u32,
47    /// End instruction offset (exclusive)
48    pub end: u32,
49    /// Handler target offset
50    pub target: u32,
51    /// Stack depth at handler entry
52    pub depth: u16,
53    /// Whether to push lasti before exception
54    pub push_lasti: bool,
55}
56
57impl ExceptionTableEntry {
58    #[must_use]
59    pub const fn new(start: u32, end: u32, target: u32, depth: u16, push_lasti: bool) -> Self {
60        Self {
61            start,
62            end,
63            target,
64            depth,
65            push_lasti,
66        }
67    }
68}
69
70/// Encode exception table entries.
71/// Uses 6-bit varint encoding with start marker (MSB) and continuation bit.
72#[must_use]
73pub fn encode_exception_table(entries: &[ExceptionTableEntry]) -> alloc::boxed::Box<[u8]> {
74    let mut data = Vec::new();
75    for entry in entries {
76        let size = entry.end.saturating_sub(entry.start);
77        let depth_lasti = ((entry.depth as u32) << 1) | (entry.push_lasti as u32);
78
79        write_varint_with_start(&mut data, entry.start);
80        write_varint_be(&mut data, size);
81        write_varint_be(&mut data, entry.target);
82        write_varint_be(&mut data, depth_lasti);
83    }
84    data.into_boxed_slice()
85}
86
87/// Find exception handler for given instruction offset.
88#[must_use]
89pub fn find_exception_handler(table: &[u8], offset: u32) -> Option<ExceptionTableEntry> {
90    let mut pos = 0;
91    while pos < table.len() {
92        let start = read_varint_with_start(table, &mut pos)?;
93        let size = read_varint(table, &mut pos)?;
94        let target = read_varint(table, &mut pos)?;
95        let depth_lasti = read_varint(table, &mut pos)?;
96
97        let end = start + size;
98        let depth = (depth_lasti >> 1) as u16;
99        let push_lasti = (depth_lasti & 1) != 0;
100
101        if offset >= start && offset < end {
102            return Some(ExceptionTableEntry {
103                start,
104                end,
105                target,
106                depth,
107                push_lasti,
108            });
109        }
110    }
111    None
112}
113
114/// Decode all exception table entries.
115#[must_use]
116pub fn decode_exception_table(table: &[u8]) -> Vec<ExceptionTableEntry> {
117    let mut entries = Vec::new();
118    let mut pos = 0;
119    while pos < table.len() {
120        let Some(start) = read_varint_with_start(table, &mut pos) else {
121            break;
122        };
123        let Some(size) = read_varint(table, &mut pos) else {
124            break;
125        };
126        let Some(target) = read_varint(table, &mut pos) else {
127            break;
128        };
129        let Some(depth_lasti) = read_varint(table, &mut pos) else {
130            break;
131        };
132        let Some(end) = start.checked_add(size) else {
133            break;
134        };
135        entries.push(ExceptionTableEntry {
136            start,
137            end,
138            target,
139            depth: (depth_lasti >> 1) as u16,
140            push_lasti: (depth_lasti & 1) != 0,
141        });
142    }
143    entries
144}
145
146/// Parse linetable to build a boolean mask indicating which code units
147/// have NO_LOCATION (line == -1). Returns a Vec<bool> of length `num_units`.
148#[must_use]
149pub fn build_no_location_mask(linetable: &[u8], num_units: usize) -> Vec<bool> {
150    let mut mask = Vec::new();
151    mask.resize(num_units, false);
152    let mut pos = 0;
153    let mut unit_idx = 0;
154
155    while pos < linetable.len() && unit_idx < num_units {
156        let header = linetable[pos];
157        pos += 1;
158        let code = (header >> 3) & 0xf;
159        let length = ((header & 7) + 1) as usize;
160
161        let is_no_location = code == PyCodeLocationInfoKind::None as u8;
162
163        // Skip payload bytes based on location kind
164        match code {
165            0..=9 => pos += 1,   // Short forms: 1 byte payload
166            10..=12 => pos += 2, // OneLine forms: 2 bytes payload
167            13 => {
168                // NoColumns: signed varint (line delta)
169                while pos < linetable.len() {
170                    let b = linetable[pos];
171                    pos += 1;
172                    if b & 0x40 == 0 {
173                        break;
174                    }
175                }
176            }
177            14 => {
178                // Long form: signed varint (line delta) + 3 unsigned varints
179                // line_delta
180                while pos < linetable.len() {
181                    let b = linetable[pos];
182                    pos += 1;
183                    if b & 0x40 == 0 {
184                        break;
185                    }
186                }
187                // end_line_delta, col+1, end_col+1
188                for _ in 0..3 {
189                    while pos < linetable.len() {
190                        let b = linetable[pos];
191                        pos += 1;
192                        if b & 0x40 == 0 {
193                            break;
194                        }
195                    }
196                }
197            }
198            15 => {} // None: no payload
199            _ => {}
200        }
201
202        for _ in 0..length {
203            if unit_idx < num_units {
204                mask[unit_idx] = is_no_location;
205                unit_idx += 1;
206            }
207        }
208    }
209
210    mask
211}
212
213/// CPython 3.11+ linetable location info codes
214#[derive(Copy, Clone, Debug, PartialEq, Eq)]
215#[repr(u8)]
216pub enum PyCodeLocationInfoKind {
217    // Short forms are 0 to 9
218    Short0 = 0,
219    Short1 = 1,
220    Short2 = 2,
221    Short3 = 3,
222    Short4 = 4,
223    Short5 = 5,
224    Short6 = 6,
225    Short7 = 7,
226    Short8 = 8,
227    Short9 = 9,
228    // One line forms are 10 to 12
229    OneLine0 = 10,
230    OneLine1 = 11,
231    OneLine2 = 12,
232    NoColumns = 13,
233    Long = 14,
234    None = 15,
235}
236
237impl PyCodeLocationInfoKind {
238    #[must_use]
239    pub fn from_code(code: u8) -> Option<Self> {
240        match code {
241            0 => Some(Self::Short0),
242            1 => Some(Self::Short1),
243            2 => Some(Self::Short2),
244            3 => Some(Self::Short3),
245            4 => Some(Self::Short4),
246            5 => Some(Self::Short5),
247            6 => Some(Self::Short6),
248            7 => Some(Self::Short7),
249            8 => Some(Self::Short8),
250            9 => Some(Self::Short9),
251            10 => Some(Self::OneLine0),
252            11 => Some(Self::OneLine1),
253            12 => Some(Self::OneLine2),
254            13 => Some(Self::NoColumns),
255            14 => Some(Self::Long),
256            15 => Some(Self::None),
257            _ => Option::None,
258        }
259    }
260
261    #[must_use]
262    pub fn is_short(&self) -> bool {
263        (*self as u8) <= 9
264    }
265
266    #[must_use]
267    pub fn short_column_group(&self) -> Option<u8> {
268        if self.is_short() {
269            Some(*self as u8)
270        } else {
271            Option::None
272        }
273    }
274
275    #[must_use]
276    pub fn one_line_delta(&self) -> Option<i32> {
277        match self {
278            Self::OneLine0 => Some(0),
279            Self::OneLine1 => Some(1),
280            Self::OneLine2 => Some(2),
281            _ => Option::None,
282        }
283    }
284}
285
286pub trait Constant: Sized + Clone {
287    type Name: AsRef<str>;
288
289    /// Transforms the given Constant to a BorrowedConstant
290    fn borrow_constant(&self) -> BorrowedConstant<'_, Self>;
291}
292
293impl Constant for ConstantData {
294    type Name = String;
295
296    fn borrow_constant(&self) -> BorrowedConstant<'_, Self> {
297        match self {
298            Self::Integer { value } => BorrowedConstant::Integer { value },
299            Self::Float { value } => BorrowedConstant::Float { value: *value },
300            Self::Complex { value } => BorrowedConstant::Complex { value: *value },
301            Self::Boolean { value } => BorrowedConstant::Boolean { value: *value },
302            Self::Str { value } => BorrowedConstant::Str { value },
303            Self::Bytes { value } => BorrowedConstant::Bytes { value },
304            Self::Code { code } => BorrowedConstant::Code { code },
305            Self::Tuple { elements } => BorrowedConstant::Tuple { elements },
306            Self::Slice { elements } => BorrowedConstant::Slice { elements },
307            Self::Frozenset { elements } => BorrowedConstant::Frozenset { elements },
308            Self::None => BorrowedConstant::None,
309            Self::Ellipsis => BorrowedConstant::Ellipsis,
310        }
311    }
312}
313
314/// A Constant Bag
315pub trait ConstantBag: Sized + Copy {
316    type Constant: Constant;
317
318    fn make_constant<C: Constant>(&self, constant: BorrowedConstant<'_, C>) -> Self::Constant;
319
320    fn make_int(&self, value: BigInt) -> Self::Constant;
321
322    fn make_tuple(&self, elements: impl Iterator<Item = Self::Constant>) -> Self::Constant;
323
324    fn make_code(&self, code: CodeObject<Self::Constant>) -> Self::Constant;
325
326    fn make_name(&self, name: &str) -> <Self::Constant as Constant>::Name;
327}
328
329pub trait AsBag {
330    type Bag: ConstantBag;
331
332    #[allow(clippy::wrong_self_convention)]
333    fn as_bag(self) -> Self::Bag;
334}
335
336impl<Bag: ConstantBag> AsBag for Bag {
337    type Bag = Self;
338
339    fn as_bag(self) -> Self {
340        self
341    }
342}
343
344#[derive(Clone, Copy)]
345pub struct BasicBag;
346
347impl ConstantBag for BasicBag {
348    type Constant = ConstantData;
349
350    fn make_constant<C: Constant>(&self, constant: BorrowedConstant<'_, C>) -> Self::Constant {
351        constant.to_owned()
352    }
353
354    fn make_int(&self, value: BigInt) -> Self::Constant {
355        ConstantData::Integer { value }
356    }
357
358    fn make_tuple(&self, elements: impl Iterator<Item = Self::Constant>) -> Self::Constant {
359        ConstantData::Tuple {
360            elements: elements.collect(),
361        }
362    }
363
364    fn make_code(&self, code: CodeObject<Self::Constant>) -> Self::Constant {
365        ConstantData::Code {
366            code: Box::new(code),
367        }
368    }
369
370    fn make_name(&self, name: &str) -> <Self::Constant as Constant>::Name {
371        name.to_owned()
372    }
373}
374
375#[derive(Clone)]
376pub struct Constants<C: Constant>(Box<[C]>);
377
378impl<C: Constant> Deref for Constants<C> {
379    type Target = [C];
380
381    fn deref(&self) -> &Self::Target {
382        &self.0
383    }
384}
385
386impl<C: Constant> DerefMut for Constants<C> {
387    fn deref_mut(&mut self) -> &mut Self::Target {
388        &mut self.0
389    }
390}
391
392impl<C: Constant> Index<oparg::ConstIdx> for Constants<C> {
393    type Output = C;
394
395    fn index(&self, consti: oparg::ConstIdx) -> &Self::Output {
396        &self.0[consti.as_usize()]
397    }
398}
399
400impl<C: Constant> IndexMut<oparg::ConstIdx> for Constants<C> {
401    fn index_mut(&mut self, consti: oparg::ConstIdx) -> &mut Self::Output {
402        &mut self.0[consti.as_usize()]
403    }
404}
405
406impl<C: Constant> FromIterator<C> for Constants<C> {
407    fn from_iter<T: IntoIterator<Item = C>>(iter: T) -> Self {
408        Self(iter.into_iter().collect())
409    }
410}
411
412// TODO: Newtype "CodeObject.varnames". Make sure only `oparg:VarNum` can be used as index
413impl<T> Index<oparg::VarNum> for [T] {
414    type Output = T;
415
416    fn index(&self, var_num: oparg::VarNum) -> &Self::Output {
417        &self[var_num.as_usize()]
418    }
419}
420
421// TODO: Newtype "CodeObject.varnames". Make sure only `oparg:VarNum` can be used as index
422impl<T> IndexMut<oparg::VarNum> for [T] {
423    fn index_mut(&mut self, var_num: oparg::VarNum) -> &mut Self::Output {
424        &mut self[var_num.as_usize()]
425    }
426}
427
428bitflagset::bitflag! {
429    /// Per-slot kind flags for localsplus (`co_localspluskinds`).
430    /// Values are bit *positions* (`Local` is bit 5 → mask `0x20`).
431    #[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
432    #[repr(u8)]
433    pub enum CoFastFlag {
434        ArgPos = 1,
435        ArgKw = 2,
436        ArgVar = 3,
437        Hidden = 4,
438        Local = 5,
439        Cell = 6,
440        Free = 7,
441    }
442}
443
444bitflagset::bitflagset! {
445    #[derive(Copy, Clone, PartialEq, Eq)]
446    pub struct CoFastFlags(u8): CoFastFlag
447}
448
449impl CoFastFlags {
450    pub const ARG: Self =
451        Self::from_slice(&[CoFastFlag::ArgPos, CoFastFlag::ArgKw, CoFastFlag::ArgVar]);
452}
453
454pub const CO_FAST_ARG_POS: u8 = CoFastFlags::from_element(CoFastFlag::ArgPos).bits();
455pub const CO_FAST_ARG_KW: u8 = CoFastFlags::from_element(CoFastFlag::ArgKw).bits();
456pub const CO_FAST_ARG_VAR: u8 = CoFastFlags::from_element(CoFastFlag::ArgVar).bits();
457pub const CO_FAST_ARG: u8 = CoFastFlags::ARG.bits();
458pub const CO_FAST_HIDDEN: u8 = CoFastFlags::from_element(CoFastFlag::Hidden).bits();
459pub const CO_FAST_LOCAL: u8 = CoFastFlags::from_element(CoFastFlag::Local).bits();
460pub const CO_FAST_CELL: u8 = CoFastFlags::from_element(CoFastFlag::Cell).bits();
461pub const CO_FAST_FREE: u8 = CoFastFlags::from_element(CoFastFlag::Free).bits();
462
463/// Primary container of a single code object. Each python function has
464/// a code object. Also a module has a code object.
465#[derive(Clone)]
466pub struct CodeObject<C: Constant = ConstantData> {
467    pub instructions: CodeUnits,
468    pub locations: Box<[(SourceLocation, SourceLocation)]>,
469    pub flags: CodeFlags,
470    /// Number of positional-only arguments
471    pub posonlyarg_count: u32,
472    pub arg_count: u32,
473    pub kwonlyarg_count: u32,
474    pub source_path: C::Name,
475    pub first_line_number: Option<OneIndexed>,
476    pub max_stackdepth: u32,
477    /// Name of the object that created this code object
478    pub obj_name: C::Name,
479    /// Qualified name of the object (like CPython's co_qualname)
480    pub qualname: C::Name,
481    pub constants: Constants<C>,
482    pub names: Box<[C::Name]>,
483    pub varnames: Box<[C::Name]>,
484    pub cellvars: Box<[C::Name]>,
485    pub freevars: Box<[C::Name]>,
486    /// Per-slot kind flags: CO_FAST_ARG_*, CO_FAST_LOCAL, CO_FAST_CELL,
487    /// CO_FAST_FREE, CO_FAST_HIDDEN.
488    /// Length = nlocalsplus (nlocals + ncells + nfrees).
489    pub localspluskinds: Box<[u8]>,
490    /// Line number table (CPython 3.11+ format)
491    pub linetable: Box<[u8]>,
492    /// Exception handling table
493    pub exceptiontable: Box<[u8]>,
494}
495
496bitflags! {
497    #[derive(Copy, Clone, Debug, PartialEq, Eq)]
498    pub struct CodeFlags: u32 {
499        const OPTIMIZED = 0x0001;
500        const NEWLOCALS = 0x0002;
501        const VARARGS = 0x0004;
502        const VARKEYWORDS = 0x0008;
503        const NESTED = 0x0010;
504        const GENERATOR = 0x0020;
505        const COROUTINE = 0x0080;
506        const ITERABLE_COROUTINE = 0x0100;
507        const ASYNC_GENERATOR = 0x0200;
508        const FUTURE_DIVISION = 0x20000;
509        const FUTURE_ABSOLUTE_IMPORT = 0x40000;
510        const FUTURE_WITH_STATEMENT = 0x80000;
511        const FUTURE_PRINT_FUNCTION = 0x100000;
512        const FUTURE_UNICODE_LITERALS = 0x200000;
513        const FUTURE_BARRY_AS_BDFL = 0x400000;
514        const FUTURE_GENERATOR_STOP = 0x800000;
515        const FUTURE_ANNOTATIONS = 0x1000000;
516        /// If a code object represents a function and has a docstring,
517        /// this bit is set and the first item in co_consts is the docstring.
518        const HAS_DOCSTRING = 0x4000000;
519        const METHOD = 0x8000000;
520    }
521}
522
523const _: () = {
524    assert!(core::mem::size_of::<CodeFlags>() == core::mem::size_of::<i32>());
525    assert!(core::mem::align_of::<CodeFlags>() == core::mem::align_of::<i32>());
526};
527
528impl CodeFlags {
529    /// The `__future__` flags that `compile()` accepts and that a compiled code
530    /// object inherits from its caller. Mirrors `PyCF_MASK`.
531    pub const FUTURE_MASK: Self = Self::FUTURE_DIVISION
532        .union(Self::FUTURE_ABSOLUTE_IMPORT)
533        .union(Self::FUTURE_WITH_STATEMENT)
534        .union(Self::FUTURE_PRINT_FUNCTION)
535        .union(Self::FUTURE_UNICODE_LITERALS)
536        .union(Self::FUTURE_BARRY_AS_BDFL)
537        .union(Self::FUTURE_GENERATOR_STOP)
538        .union(Self::FUTURE_ANNOTATIONS);
539}
540
541/// `align(2)` so that the whole unit can be read with one aligned
542/// `AtomicU16` access (see `CodeUnits::read_unit`) without relying on the
543/// allocator to happen to align the instruction array.
544#[repr(C, align(2))]
545#[derive(Copy, Clone, Debug)]
546pub struct CodeUnit {
547    pub op: Instruction,
548    pub arg: OpArgByte,
549}
550
551const _: () = assert!(mem::size_of::<CodeUnit>() == 2);
552const _: () = assert!(mem::align_of::<CodeUnit>() == 2);
553
554/// Adaptive specialization: number of executions before attempting specialization.
555///
556/// Matches CPython's `_Py_BackoffCounter` encoding.
557pub const ADAPTIVE_WARMUP_VALUE: u16 = adaptive_counter_bits(1, 1);
558/// Adaptive specialization: cooldown counter after a successful specialization.
559///
560/// Value/backoff = (52, 0), matching CPython's ADAPTIVE_COOLDOWN bits.
561pub const ADAPTIVE_COOLDOWN_VALUE: u16 = adaptive_counter_bits(52, 0);
562/// Initial JUMP_BACKWARD counter bits (value/backoff = 4095/12).
563pub const JUMP_BACKWARD_INITIAL_VALUE: u16 = adaptive_counter_bits(4095, 12);
564
565const BACKOFF_BITS: u16 = 4;
566const MAX_BACKOFF: u16 = 12;
567const UNREACHABLE_BACKOFF: u16 = 15;
568
569/// Encode an adaptive counter as `(value << 4) | backoff`.
570#[must_use]
571pub const fn adaptive_counter_bits(value: u16, backoff: u16) -> u16 {
572    (value << BACKOFF_BITS) | backoff
573}
574
575/// True when the adaptive counter should trigger specialization.
576#[inline]
577#[must_use]
578pub const fn adaptive_counter_triggers(counter: u16) -> bool {
579    counter < UNREACHABLE_BACKOFF
580}
581
582/// Decrement adaptive counter by one countdown step.
583#[inline]
584#[must_use]
585pub const fn advance_adaptive_counter(counter: u16) -> u16 {
586    counter.wrapping_sub(1 << BACKOFF_BITS)
587}
588
589/// Reset adaptive counter with exponential backoff.
590#[inline]
591#[must_use]
592pub const fn adaptive_counter_backoff(counter: u16) -> u16 {
593    let backoff = counter & ((1 << BACKOFF_BITS) - 1);
594    if backoff < MAX_BACKOFF {
595        adaptive_counter_bits((1 << (backoff + 1)) - 1, backoff + 1)
596    } else {
597        adaptive_counter_bits((1 << MAX_BACKOFF) - 1, MAX_BACKOFF)
598    }
599}
600
601impl CodeUnit {
602    #[must_use]
603    pub const fn new(op: Instruction, arg: OpArgByte) -> Self {
604        Self { op, arg }
605    }
606}
607
608impl TryFrom<&[u8]> for CodeUnit {
609    type Error = MarshalError;
610
611    fn try_from(value: &[u8]) -> Result<Self, Self::Error> {
612        match value.len() {
613            2 => Ok(Self::new(value[0].try_into()?, value[1].into())),
614            _ => Err(Self::Error::InvalidBytecode),
615        }
616    }
617}
618
619impl TryFrom<[u8; 2]> for CodeUnit {
620    type Error = MarshalError;
621
622    fn try_from(value: [u8; 2]) -> Result<Self, Self::Error> {
623        Ok(Self::new(value[0].try_into()?, value[1].into()))
624    }
625}
626
627pub struct CodeUnits {
628    units: UnsafeCell<Box<[CodeUnit]>>,
629    adaptive_counters: Box<[AtomicU16]>,
630    /// Pointer-sized cache entries for descriptor pointers.
631    /// Single atomic load/store prevents torn reads when multiple threads
632    /// specialize the same instruction concurrently.
633    pointer_cache: Box<[AtomicUsize]>,
634}
635
636// SAFETY: All cache operations use atomic read/write instructions.
637// - replace_op / compare_exchange_op: AtomicU8 store/CAS (Release)
638// - cache read/write: AtomicU16 load/store (Relaxed)
639// - adaptive counter: AtomicU16 load/store (Relaxed)
640// Ordering is established by:
641// - replace_op (Release) ↔ dispatch loop read_op (Acquire) for cache data visibility
642// - tp_version_tag (Acquire) for descriptor pointer validity
643unsafe impl Sync for CodeUnits {}
644
645impl Clone for CodeUnits {
646    fn clone(&self) -> Self {
647        // SAFETY: No concurrent mutation during clone — cloning is only done
648        // during code object construction or marshaling, not while instrumented.
649        let units = unsafe { &*self.units.get() }.clone();
650        let adaptive_counters = self
651            .adaptive_counters
652            .iter()
653            .map(|c| AtomicU16::new(c.load(Ordering::Relaxed)))
654            .collect();
655        let pointer_cache = self
656            .pointer_cache
657            .iter()
658            .map(|c| AtomicUsize::new(c.load(Ordering::Relaxed)))
659            .collect();
660        Self {
661            units: UnsafeCell::new(units),
662            adaptive_counters,
663            pointer_cache,
664        }
665    }
666}
667
668impl fmt::Debug for CodeUnits {
669    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
670        // SAFETY: Debug formatting doesn't race with replace_op
671        let inner = unsafe { &*self.units.get() };
672        f.debug_tuple("CodeUnits").field(inner).finish()
673    }
674}
675
676impl TryFrom<&[u8]> for CodeUnits {
677    type Error = MarshalError;
678
679    fn try_from(value: &[u8]) -> Result<Self, Self::Error> {
680        let (chunks, []) = value.as_chunks::<2>() else {
681            return Err(Self::Error::InvalidBytecode);
682        };
683
684        let units = chunks
685            .iter()
686            .copied()
687            .map(CodeUnit::try_from)
688            .collect::<Result<Vec<_>, _>>()?;
689        Ok(units.into())
690    }
691}
692
693impl<const N: usize> From<[CodeUnit; N]> for CodeUnits {
694    fn from(value: [CodeUnit; N]) -> Self {
695        Self::from(Vec::from(value))
696    }
697}
698
699impl From<Vec<CodeUnit>> for CodeUnits {
700    fn from(value: Vec<CodeUnit>) -> Self {
701        let units = value.into_boxed_slice();
702        let len = units.len();
703        let adaptive_counters = (0..len)
704            .map(|_| AtomicU16::new(0))
705            .collect::<Vec<_>>()
706            .into_boxed_slice();
707        let pointer_cache = (0..len)
708            .map(|_| AtomicUsize::new(0))
709            .collect::<Vec<_>>()
710            .into_boxed_slice();
711        Self {
712            units: UnsafeCell::new(units),
713            adaptive_counters,
714            pointer_cache,
715        }
716    }
717}
718
719impl FromIterator<CodeUnit> for CodeUnits {
720    fn from_iter<T: IntoIterator<Item = CodeUnit>>(iter: T) -> Self {
721        Self::from(iter.into_iter().collect::<Vec<_>>())
722    }
723}
724
725impl Deref for CodeUnits {
726    type Target = [CodeUnit];
727
728    fn deref(&self) -> &Self::Target {
729        // SAFETY: Shared references to the slice are valid even while replace_op
730        // may update individual opcode bytes — readers tolerate stale opcodes
731        // (they will re-read on the next iteration).
732        unsafe { &*self.units.get() }
733    }
734}
735
736impl CodeUnits {
737    /// Disable adaptive specialization by setting all counters to unreachable.
738    /// Used for CPython-compiled bytecode where specialization may not be safe.
739    pub fn disable_specialization(&self) {
740        for counter in &self.adaptive_counters {
741            counter.store(UNREACHABLE_BACKOFF, Ordering::Relaxed);
742        }
743    }
744
745    /// Replace the opcode at `index` in-place without changing the arg byte.
746    /// Uses atomic Release store to ensure prior cache writes are visible
747    /// to threads that subsequently read the new opcode with Acquire.
748    ///
749    /// # Safety
750    /// - `index` must be in bounds.
751    /// - `new_op` must have the same arg semantics as the original opcode.
752    pub unsafe fn replace_op(&self, index: usize, new_op: Instruction) {
753        let units = unsafe { &*self.units.get() };
754        let ptr = units.as_ptr().wrapping_add(index) as *const AtomicU8;
755        unsafe { &*ptr }.store(new_op.into(), Ordering::Release);
756    }
757
758    /// Atomically replace opcode only if it still matches `expected`.
759    /// Returns true on success. Uses Release ordering on success.
760    ///
761    /// # Safety
762    /// - `index` must be in bounds.
763    pub unsafe fn compare_exchange_op(
764        &self,
765        index: usize,
766        expected: Instruction,
767        new_op: Instruction,
768    ) -> bool {
769        let units = unsafe { &*self.units.get() };
770        let ptr = units.as_ptr().wrapping_add(index) as *const AtomicU8;
771        unsafe { &*ptr }
772            .compare_exchange(
773                expected.into(),
774                new_op.into(),
775                Ordering::Release,
776                Ordering::Relaxed,
777            )
778            .is_ok()
779    }
780
781    /// Atomically read the opcode at `index` with Acquire ordering.
782    /// Pairs with `replace_op` (Release) to ensure cache data visibility.
783    pub fn read_op(&self, index: usize) -> Instruction {
784        let units = unsafe { &*self.units.get() };
785        let ptr = units.as_ptr().wrapping_add(index) as *const AtomicU8;
786        let byte = unsafe { &*ptr }.load(Ordering::Acquire);
787        // SAFETY: Only valid Instruction values are stored via replace_op/compare_exchange_op.
788        unsafe { mem::transmute::<u8, Instruction>(byte) }
789    }
790
791    /// Base pointer of the instruction array.
792    ///
793    /// The array is allocated once when the code object is built and is never
794    /// reallocated -- specialization only rewrites units in place -- so the
795    /// pointer stays valid for as long as this `CodeUnits` does and the eval
796    /// loop may hoist it out of the dispatch loop.
797    #[inline(always)]
798    #[must_use]
799    pub fn units_ptr(&self) -> *const CodeUnit {
800        unsafe { &*self.units.get() }.as_ptr()
801    }
802
803    /// `read_unit` from a base pointer previously obtained with `units_ptr`.
804    ///
805    /// # Safety
806    /// `base` must come from `units_ptr` on a live `CodeUnits`, and `index`
807    /// must be within that array (the dispatch loop only ever reads indices
808    /// the code object itself produced).
809    #[inline(always)]
810    #[must_use]
811    pub unsafe fn read_unit_from(base: *const CodeUnit, index: usize) -> CodeUnit {
812        let ptr = base.wrapping_add(index) as *const AtomicU16;
813        let [op, arg] = unsafe { &*ptr }.load(Ordering::Acquire).to_ne_bytes();
814        // SAFETY: only valid Instruction values are ever stored into the
815        // instruction array (see `read_op`).
816        CodeUnit {
817            op: unsafe { mem::transmute::<u8, Instruction>(op) },
818            arg: OpArgByte::from(arg),
819        }
820    }
821
822    /// Atomically read the opcode and its arg byte at `index` as a single
823    /// Acquire-ordered 16-bit load.
824    ///
825    /// Equivalent to `read_op` followed by `read_arg`, but the one access
826    /// keeps the eval loop from re-loading the instruction array pointer
827    /// across the acquire barrier, and it is what the dispatch loop uses.
828    /// Acquire pairs with `replace_op` (Release) exactly as `read_op` does.
829    ///
830    /// `to_ne_bytes` yields the bytes in memory order, so `[0]` is the `op`
831    /// field and `[1]` the `arg` field of the `repr(C)` unit on either
832    /// endianness.
833    #[inline(always)]
834    #[must_use]
835    pub fn read_unit(&self, index: usize) -> CodeUnit {
836        // SAFETY: `index` is in bounds for every caller of this method.
837        unsafe { Self::read_unit_from(self.units_ptr(), index) }
838    }
839
840    /// Atomically read the arg byte at `index` with Relaxed ordering.
841    pub fn read_arg(&self, index: usize) -> OpArgByte {
842        let units = unsafe { &*self.units.get() };
843        let ptr = units.as_ptr().wrapping_add(index) as *const u8;
844        let arg_ptr = unsafe { ptr.add(1) } as *const AtomicU8;
845        OpArgByte::from(unsafe { &*arg_ptr }.load(Ordering::Relaxed))
846    }
847
848    /// Write a u16 value into a CACHE code unit at `index`.
849    /// Each CodeUnit is 2 bytes (#[repr(C)]: op u8 + arg u8), so one u16 fits exactly.
850    /// Uses Relaxed atomic store; ordering is provided by replace_op (Release).
851    ///
852    /// # Safety
853    /// - `index` must be in bounds and point to a CACHE entry.
854    pub unsafe fn write_cache_u16(&self, index: usize, value: u16) {
855        let units = unsafe { &*self.units.get() };
856        let ptr = units.as_ptr().wrapping_add(index) as *const AtomicU16;
857        unsafe { &*ptr }.store(value, Ordering::Relaxed);
858    }
859
860    /// Read a u16 value from a CACHE code unit at `index`.
861    /// Uses Relaxed atomic load; ordering is provided by read_op (Acquire).
862    ///
863    /// # Panics
864    /// Panics if `index` is out of bounds.
865    pub fn read_cache_u16(&self, index: usize) -> u16 {
866        let units = unsafe { &*self.units.get() };
867        assert!(index < units.len(), "read_cache_u16: index out of bounds");
868        let ptr = units.as_ptr().wrapping_add(index) as *const AtomicU16;
869        unsafe { &*ptr }.load(Ordering::Relaxed)
870    }
871
872    /// Write a u32 value across two consecutive CACHE code units starting at `index`.
873    ///
874    /// # Safety
875    /// Same requirements as `write_cache_u16`.
876    pub unsafe fn write_cache_u32(&self, index: usize, value: u32) {
877        unsafe {
878            self.write_cache_u16(index, value as u16);
879            self.write_cache_u16(index + 1, (value >> 16) as u16);
880        }
881    }
882
883    /// Read a u32 value from two consecutive CACHE code units starting at `index`.
884    ///
885    /// # Panics
886    /// Panics if `index + 1` is out of bounds.
887    pub fn read_cache_u32(&self, index: usize) -> u32 {
888        let lo = self.read_cache_u16(index) as u32;
889        let hi = self.read_cache_u16(index + 1) as u32;
890        lo | (hi << 16)
891    }
892
893    /// Store a pointer-sized value atomically in the pointer cache at `index`.
894    ///
895    /// Uses a single `AtomicUsize` store to prevent torn writes when
896    /// multiple threads specialize the same instruction concurrently. The
897    /// tear-free width also makes this the right slot for non-pointer guard
898    /// values (e.g. dict keys-version stamps) that must never be observed
899    /// half-written.
900    ///
901    /// # Safety
902    /// - `index` must be in bounds.
903    /// - When the slot holds a `*const PyObject` encoded as `usize` (or `0`),
904    ///   callers must follow the cache invalidation/upgrade protocol:
905    ///   invalidate the version guard before writing and publish the new
906    ///   version after writing.
907    pub unsafe fn write_cache_ptr(&self, index: usize, value: usize) {
908        self.pointer_cache[index].store(value, Ordering::Relaxed);
909    }
910
911    /// Load a pointer-sized value atomically from the pointer cache at `index`.
912    ///
913    /// Uses a single `AtomicUsize` load to prevent torn reads.
914    ///
915    /// # Panics
916    /// Panics if `index` is out of bounds.
917    pub fn read_cache_ptr(&self, index: usize) -> usize {
918        self.pointer_cache[index].load(Ordering::Relaxed)
919    }
920
921    /// Read adaptive counter bits for instruction at `index`.
922    /// Uses Relaxed atomic load.
923    pub fn read_adaptive_counter(&self, index: usize) -> u16 {
924        self.adaptive_counters[index].load(Ordering::Relaxed)
925    }
926
927    /// Write adaptive counter bits for instruction at `index`.
928    /// Uses Relaxed atomic store.
929    ///
930    /// # Safety
931    /// - `index` must be in bounds.
932    pub unsafe fn write_adaptive_counter(&self, index: usize, value: u16) {
933        self.adaptive_counters[index].store(value, Ordering::Relaxed);
934    }
935
936    /// Produce a clean copy of the bytecode suitable for serialization
937    /// (marshal) and `co_code`. Specialized opcodes are mapped back to their
938    /// base variants via `deoptimize()` and all CACHE entries are zeroed.
939    pub fn original_bytes(&self) -> Vec<u8> {
940        let len = self.len();
941        let mut out = Vec::with_capacity(len * 2);
942        let mut i = 0;
943        while i < len {
944            let op = self.read_op(i).deoptimize();
945            let arg = self.read_arg(i);
946            let caches = op.cache_entries();
947            out.push(u8::from(op));
948            out.push(u8::from(arg));
949            // Zero-fill all CACHE entries (counter + cached data)
950            for _ in 0..caches {
951                i += 1;
952                out.push(0); // op = Cache = 0
953                out.push(0); // arg = 0
954            }
955            i += 1;
956        }
957        out
958    }
959
960    /// Initialize adaptive warmup counters for all cacheable instructions.
961    /// Called lazily at RESUME (first execution of a code object).
962    /// Counters are stored out-of-line to preserve `op = Instruction::Cache`.
963    /// All writes are atomic (Relaxed) to avoid data races with concurrent readers.
964    pub fn quicken(&self) {
965        let len = self.len();
966        let mut i = 0;
967        while i < len {
968            let op = self.read_op(i);
969            let caches = op.cache_entries();
970            if caches > 0 {
971                // Don't write adaptive counter for instrumented opcodes;
972                // specialization is skipped while monitoring is active.
973                if !op.is_instrumented() {
974                    let cache_base = i + 1;
975                    if cache_base < len {
976                        let initial_counter = if matches!(op, Instruction::JumpBackward { .. }) {
977                            JUMP_BACKWARD_INITIAL_VALUE
978                        } else {
979                            ADAPTIVE_WARMUP_VALUE
980                        };
981                        unsafe {
982                            self.write_adaptive_counter(cache_base, initial_counter);
983                        }
984                    }
985                }
986                i += 1 + caches;
987            } else {
988                i += 1;
989            }
990        }
991    }
992}
993
994/// A Constant (which usually encapsulates data within it)
995///
996/// # Examples
997/// ```
998/// use rustpython_compiler_core::bytecode::ConstantData;
999/// let a = ConstantData::Float {value: 120f64};
1000/// let b = ConstantData::Boolean {value: false};
1001/// assert_ne!(a, b);
1002/// ```
1003#[derive(Debug, Clone)]
1004pub enum ConstantData {
1005    Tuple {
1006        elements: Vec<Self>,
1007    },
1008    Integer {
1009        value: BigInt,
1010    },
1011    Float {
1012        value: f64,
1013    },
1014    Complex {
1015        value: Complex64,
1016    },
1017    Boolean {
1018        value: bool,
1019    },
1020    Str {
1021        value: Wtf8Buf,
1022    },
1023    Bytes {
1024        value: Vec<u8>,
1025    },
1026    Code {
1027        code: Box<CodeObject>,
1028    },
1029    /// Constant slice(start, stop, step)
1030    Slice {
1031        elements: Box<[Self; 3]>,
1032    },
1033    Frozenset {
1034        elements: Vec<Self>,
1035    },
1036    None,
1037    Ellipsis,
1038}
1039
1040impl ConstantData {
1041    /// Whether or not python would return True/False for the given constant data.
1042    ///
1043    /// ```py
1044    /// bool(0) # False
1045    /// bool(1) # True
1046    /// bool([]) # False
1047    /// bool(...) # True
1048    /// ```
1049    #[must_use]
1050    pub fn truthiness(&self) -> bool {
1051        match self {
1052            Self::Tuple { elements } | Self::Frozenset { elements } => !elements.is_empty(),
1053            Self::Integer { value } => !value.is_zero(),
1054            Self::Float { value } => *value != 0.0,
1055            Self::Complex { value } => value.re != 0.0 || value.im != 0.0,
1056            Self::Boolean { value } => *value,
1057            Self::Str { value } => !value.is_empty(),
1058            Self::Bytes { value } => !value.is_empty(),
1059            Self::Code { .. } | Self::Slice { .. } | Self::Ellipsis => true,
1060            Self::None => false,
1061        }
1062    }
1063}
1064
1065impl PartialEq for ConstantData {
1066    fn eq(&self, other: &Self) -> bool {
1067        match (self, other) {
1068            (Self::Integer { value: a }, Self::Integer { value: b }) => a == b,
1069            (Self::Float { value: a }, Self::Float { value: b }) => a.to_bits() == b.to_bits(),
1070            (Self::Complex { value: a }, Self::Complex { value: b }) => {
1071                a.re.to_bits() == b.re.to_bits() && a.im.to_bits() == b.im.to_bits()
1072            }
1073            (Self::Boolean { value: a }, Self::Boolean { value: b }) => a == b,
1074            (Self::Str { value: a }, Self::Str { value: b }) => a == b,
1075            (Self::Bytes { value: a }, Self::Bytes { value: b }) => a == b,
1076            (Self::Code { code: a }, Self::Code { code: b }) => {
1077                core::ptr::eq(a.as_ref(), b.as_ref())
1078            }
1079            (Self::Tuple { elements: a }, Self::Tuple { elements: b }) => a == b,
1080            (Self::Slice { elements: a }, Self::Slice { elements: b }) => a == b,
1081            (Self::Frozenset { elements: a }, Self::Frozenset { elements: b }) => a == b,
1082            (Self::None, Self::None) => true,
1083            (Self::Ellipsis, Self::Ellipsis) => true,
1084            _ => false,
1085        }
1086    }
1087}
1088
1089impl Eq for ConstantData {}
1090
1091impl hash::Hash for ConstantData {
1092    fn hash<H: hash::Hasher>(&self, state: &mut H) {
1093        mem::discriminant(self).hash(state);
1094
1095        match self {
1096            Self::Integer { value } => value.hash(state),
1097            Self::Float { value } => value.to_bits().hash(state),
1098            Self::Complex { value } => {
1099                value.re.to_bits().hash(state);
1100                value.im.to_bits().hash(state);
1101            }
1102            Self::Boolean { value } => value.hash(state),
1103            Self::Str { value } => value.hash(state),
1104            Self::Bytes { value } => value.hash(state),
1105            Self::Code { code } => core::ptr::hash(code.as_ref(), state),
1106            Self::Tuple { elements } => elements.hash(state),
1107            Self::Slice { elements } => elements.hash(state),
1108            Self::Frozenset { elements } => elements.hash(state),
1109            Self::None => {}
1110            Self::Ellipsis => {}
1111        }
1112    }
1113}
1114
1115/// A borrowed Constant
1116pub enum BorrowedConstant<'a, C: Constant> {
1117    Integer { value: &'a BigInt },
1118    Float { value: f64 },
1119    Complex { value: Complex64 },
1120    Boolean { value: bool },
1121    Str { value: &'a Wtf8 },
1122    Bytes { value: &'a [u8] },
1123    Code { code: &'a CodeObject<C> },
1124    Tuple { elements: &'a [C] },
1125    Slice { elements: &'a [C; 3] },
1126    Frozenset { elements: &'a [C] },
1127    None,
1128    Ellipsis,
1129}
1130
1131impl<C: Constant> Copy for BorrowedConstant<'_, C> {}
1132
1133impl<C: Constant> Clone for BorrowedConstant<'_, C> {
1134    fn clone(&self) -> Self {
1135        *self
1136    }
1137}
1138
1139impl<C: Constant> BorrowedConstant<'_, C> {
1140    pub fn fmt_display(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1141        match self {
1142            BorrowedConstant::Integer { value } => write!(f, "{value}"),
1143            BorrowedConstant::Float { value } => write!(f, "{value}"),
1144            BorrowedConstant::Complex { value } => write!(f, "{value}"),
1145            BorrowedConstant::Boolean { value } => {
1146                write!(f, "{}", if *value { "True" } else { "False" })
1147            }
1148            BorrowedConstant::Str { value } => write!(f, "{value:?}"),
1149            BorrowedConstant::Bytes { value } => write!(f, r#"b"{}""#, value.escape_ascii()),
1150            BorrowedConstant::Code { code } => write!(f, "{code:?}"),
1151            BorrowedConstant::Tuple { elements } => {
1152                write!(f, "(")?;
1153                let mut first = true;
1154                for c in *elements {
1155                    if first {
1156                        first = false
1157                    } else {
1158                        write!(f, ", ")?;
1159                    }
1160                    c.borrow_constant().fmt_display(f)?;
1161                }
1162                write!(f, ")")
1163            }
1164            BorrowedConstant::Slice { elements } => {
1165                write!(f, "slice(")?;
1166                elements[0].borrow_constant().fmt_display(f)?;
1167                write!(f, ", ")?;
1168                elements[1].borrow_constant().fmt_display(f)?;
1169                write!(f, ", ")?;
1170                elements[2].borrow_constant().fmt_display(f)?;
1171                write!(f, ")")
1172            }
1173            BorrowedConstant::Frozenset { elements } => {
1174                write!(f, "frozenset({{")?;
1175                let mut first = true;
1176                for c in *elements {
1177                    if first {
1178                        first = false
1179                    } else {
1180                        write!(f, ", ")?;
1181                    }
1182                    c.borrow_constant().fmt_display(f)?;
1183                }
1184                write!(f, "}})")
1185            }
1186            BorrowedConstant::None => write!(f, "None"),
1187            BorrowedConstant::Ellipsis => write!(f, "..."),
1188        }
1189    }
1190
1191    #[must_use]
1192    pub fn to_owned(self) -> ConstantData {
1193        match self {
1194            BorrowedConstant::Integer { value } => ConstantData::Integer {
1195                value: value.clone(),
1196            },
1197            BorrowedConstant::Float { value } => ConstantData::Float { value },
1198            BorrowedConstant::Complex { value } => ConstantData::Complex { value },
1199            BorrowedConstant::Boolean { value } => ConstantData::Boolean { value },
1200            BorrowedConstant::Str { value } => ConstantData::Str {
1201                value: value.to_owned(),
1202            },
1203            BorrowedConstant::Bytes { value } => ConstantData::Bytes {
1204                value: value.to_owned(),
1205            },
1206            BorrowedConstant::Code { code } => ConstantData::Code {
1207                code: Box::new(code.map_clone_bag(&BasicBag)),
1208            },
1209            BorrowedConstant::Tuple { elements } => ConstantData::Tuple {
1210                elements: elements
1211                    .iter()
1212                    .map(|c| c.borrow_constant().to_owned())
1213                    .collect(),
1214            },
1215            BorrowedConstant::Slice { elements } => ConstantData::Slice {
1216                elements: Box::new(elements.each_ref().map(|c| c.borrow_constant().to_owned())),
1217            },
1218            BorrowedConstant::Frozenset { elements } => ConstantData::Frozenset {
1219                elements: elements
1220                    .iter()
1221                    .map(|c| c.borrow_constant().to_owned())
1222                    .collect(),
1223            },
1224            BorrowedConstant::None => ConstantData::None,
1225            BorrowedConstant::Ellipsis => ConstantData::Ellipsis,
1226        }
1227    }
1228}
1229
1230/*
1231Maintain a stack of blocks on the VM.
1232pub enum BlockType {
1233    Loop,
1234    Except,
1235}
1236*/
1237
1238/// Argument structure
1239pub struct Arguments<'a, N: AsRef<str>> {
1240    pub posonlyargs: &'a [N],
1241    pub args: &'a [N],
1242    pub vararg: Option<&'a N>,
1243    pub kwonlyargs: &'a [N],
1244    pub varkwarg: Option<&'a N>,
1245}
1246
1247impl<N: AsRef<str>> fmt::Debug for Arguments<'_, N> {
1248    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1249        macro_rules! fmt_slice {
1250            ($x:expr) => {
1251                format_args!("[{}]", $x.iter().map(AsRef::as_ref).format(", "))
1252            };
1253        }
1254        f.debug_struct("Arguments")
1255            .field("posonlyargs", &fmt_slice!(self.posonlyargs))
1256            .field("args", &fmt_slice!(self.posonlyargs))
1257            .field("vararg", &self.vararg.map(N::as_ref))
1258            .field("kwonlyargs", &fmt_slice!(self.kwonlyargs))
1259            .field("varkwarg", &self.varkwarg.map(N::as_ref))
1260            .finish()
1261    }
1262}
1263
1264impl<C: Constant> CodeObject<C> {
1265    /// Get all arguments of the code object
1266    /// like inspect.getargs
1267    pub fn arg_names(&self) -> Arguments<'_, C::Name> {
1268        let nargs = self.arg_count as usize;
1269        let nkwargs = self.kwonlyarg_count as usize;
1270        let mut varargs_pos = nargs + nkwargs;
1271        let posonlyargs = &self.varnames[..self.posonlyarg_count as usize];
1272        let args = &self.varnames[..nargs];
1273        let kwonlyargs = &self.varnames[nargs..varargs_pos];
1274
1275        let vararg = if self.flags.contains(CodeFlags::VARARGS) {
1276            let vararg = &self.varnames[varargs_pos];
1277            varargs_pos += 1;
1278            Some(vararg)
1279        } else {
1280            None
1281        };
1282        let varkwarg = if self.flags.contains(CodeFlags::VARKEYWORDS) {
1283            Some(&self.varnames[varargs_pos])
1284        } else {
1285            None
1286        };
1287
1288        Arguments {
1289            posonlyargs,
1290            args,
1291            vararg,
1292            kwonlyargs,
1293            varkwarg,
1294        }
1295    }
1296
1297    /// Return the labels targeted by the instructions of this CodeObject
1298    pub fn label_targets(&self) -> BTreeSet<Label> {
1299        let mut label_targets = BTreeSet::new();
1300        let mut arg_state = OpArgState::default();
1301        for instruction in &*self.instructions {
1302            let (instruction, arg) = arg_state.get(*instruction);
1303            if let Some(l) = instruction.label_arg() {
1304                label_targets.insert(l.get(arg));
1305            }
1306        }
1307        label_targets
1308    }
1309
1310    /// Map this CodeObject to one that holds a Bag::Constant
1311    pub fn map_bag<Bag: ConstantBag>(self, bag: Bag) -> CodeObject<Bag::Constant> {
1312        let map_names = |names: Box<[C::Name]>| {
1313            names
1314                .iter()
1315                .map(|x| bag.make_name(x.as_ref()))
1316                .collect::<Box<[_]>>()
1317        };
1318        CodeObject {
1319            constants: self
1320                .constants
1321                .iter()
1322                .map(|x| bag.make_constant(x.borrow_constant()))
1323                .collect(),
1324            names: map_names(self.names),
1325            varnames: map_names(self.varnames),
1326            cellvars: map_names(self.cellvars),
1327            freevars: map_names(self.freevars),
1328            source_path: bag.make_name(self.source_path.as_ref()),
1329            obj_name: bag.make_name(self.obj_name.as_ref()),
1330            qualname: bag.make_name(self.qualname.as_ref()),
1331
1332            instructions: self.instructions,
1333            locations: self.locations,
1334            flags: self.flags,
1335            posonlyarg_count: self.posonlyarg_count,
1336            arg_count: self.arg_count,
1337            kwonlyarg_count: self.kwonlyarg_count,
1338            first_line_number: self.first_line_number,
1339            max_stackdepth: self.max_stackdepth,
1340            localspluskinds: self.localspluskinds,
1341            linetable: self.linetable,
1342            exceptiontable: self.exceptiontable,
1343        }
1344    }
1345
1346    /// Same as `map_bag` but clones `self`
1347    pub fn map_clone_bag<Bag: ConstantBag>(&self, bag: &Bag) -> CodeObject<Bag::Constant> {
1348        let map_names =
1349            |names: &[C::Name]| names.iter().map(|x| bag.make_name(x.as_ref())).collect();
1350        CodeObject {
1351            constants: self
1352                .constants
1353                .iter()
1354                .map(|x| bag.make_constant(x.borrow_constant()))
1355                .collect(),
1356            names: map_names(&self.names),
1357            varnames: map_names(&self.varnames),
1358            cellvars: map_names(&self.cellvars),
1359            freevars: map_names(&self.freevars),
1360            source_path: bag.make_name(self.source_path.as_ref()),
1361            obj_name: bag.make_name(self.obj_name.as_ref()),
1362            qualname: bag.make_name(self.qualname.as_ref()),
1363
1364            instructions: self.instructions.clone(),
1365            locations: self.locations.clone(),
1366            flags: self.flags,
1367            posonlyarg_count: self.posonlyarg_count,
1368            arg_count: self.arg_count,
1369            kwonlyarg_count: self.kwonlyarg_count,
1370            first_line_number: self.first_line_number,
1371            max_stackdepth: self.max_stackdepth,
1372            localspluskinds: self.localspluskinds.clone(),
1373            linetable: self.linetable.clone(),
1374            exceptiontable: self.exceptiontable.clone(),
1375        }
1376    }
1377}
1378
1379pub trait InstrDisplayContext {
1380    type Constant: Constant;
1381
1382    fn get_constant(&self, consti: oparg::ConstIdx) -> &Self::Constant;
1383
1384    fn get_name(&self, i: usize) -> &str;
1385
1386    fn get_varname(&self, var_num: oparg::VarNum) -> &str;
1387
1388    /// Get name for a localsplus index (used by DEREF instructions).
1389    fn get_localsplus_name(&self, var_num: oparg::VarNum) -> &str;
1390}
1391
1392impl<C: Constant> InstrDisplayContext for CodeObject<C> {
1393    type Constant = C;
1394
1395    fn get_constant(&self, consti: oparg::ConstIdx) -> &C {
1396        &self.constants[consti]
1397    }
1398
1399    fn get_name(&self, i: usize) -> &str {
1400        self.names[i].as_ref()
1401    }
1402
1403    fn get_varname(&self, var_num: oparg::VarNum) -> &str {
1404        self.varnames[var_num].as_ref()
1405    }
1406
1407    fn get_localsplus_name(&self, var_num: oparg::VarNum) -> &str {
1408        let idx = var_num.as_usize();
1409        let nlocals = self.varnames.len();
1410        if idx < nlocals {
1411            self.varnames[idx].as_ref()
1412        } else {
1413            let cell_idx = idx - nlocals;
1414            self.cellvars
1415                .get(cell_idx)
1416                .unwrap_or_else(|| &self.freevars[cell_idx - self.cellvars.len()])
1417                .as_ref()
1418        }
1419    }
1420}
1421
1422impl fmt::Display for ConstantData {
1423    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1424        self.borrow_constant().fmt_display(f)
1425    }
1426}
1427
1428impl<C: Constant> fmt::Debug for CodeObject<C> {
1429    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1430        write!(
1431            f,
1432            "<code object {} at ??? file {:?}, line {}>",
1433            self.obj_name.as_ref(),
1434            self.source_path.as_ref(),
1435            self.first_line_number.map_or(-1, |x| x.get() as i32)
1436        )
1437    }
1438}
1439
1440#[cfg(test)]
1441mod tests {
1442    use super::*;
1443    use alloc::{vec, vec::Vec};
1444
1445    #[test]
1446    fn exception_table_encode_decode() {
1447        let entries = vec![
1448            ExceptionTableEntry::new(0, 10, 20, 2, false),
1449            ExceptionTableEntry::new(15, 25, 30, 1, true),
1450        ];
1451
1452        let encoded = encode_exception_table(&entries);
1453
1454        // Find handler at offset 5 (in range [0, 10))
1455        let handler = find_exception_handler(&encoded, 5);
1456        assert!(handler.is_some());
1457        let handler = handler.unwrap();
1458        assert_eq!(handler.start, 0);
1459        assert_eq!(handler.end, 10);
1460        assert_eq!(handler.target, 20);
1461        assert_eq!(handler.depth, 2);
1462        assert!(!handler.push_lasti);
1463
1464        // Find handler at offset 20 (in range [15, 25))
1465        let handler = find_exception_handler(&encoded, 20);
1466        assert!(handler.is_some());
1467        let handler = handler.unwrap();
1468        assert_eq!(handler.start, 15);
1469        assert_eq!(handler.end, 25);
1470        assert_eq!(handler.target, 30);
1471        assert_eq!(handler.depth, 1);
1472        assert!(handler.push_lasti);
1473
1474        // No handler at offset 12 (not in any range)
1475        let handler = find_exception_handler(&encoded, 12);
1476        assert!(handler.is_none());
1477
1478        // No handler at offset 30 (past all ranges)
1479        let handler = find_exception_handler(&encoded, 30);
1480        assert!(handler.is_none());
1481    }
1482
1483    #[test]
1484    fn exception_table_empty() {
1485        let entries: Vec<ExceptionTableEntry> = vec![];
1486        let encoded = encode_exception_table(&entries);
1487        assert!(encoded.is_empty());
1488        assert!(find_exception_handler(&encoded, 0).is_none());
1489    }
1490
1491    #[test]
1492    fn exception_table_single_entry() {
1493        let entries = vec![ExceptionTableEntry::new(5, 15, 100, 3, true)];
1494        let encoded = encode_exception_table(&entries);
1495
1496        // Inside range
1497        let handler = find_exception_handler(&encoded, 10);
1498        assert!(handler.is_some());
1499        let handler = handler.unwrap();
1500        assert_eq!(handler.target, 100);
1501        assert_eq!(handler.depth, 3);
1502        assert!(handler.push_lasti);
1503
1504        // At start boundary (inclusive)
1505        assert!(find_exception_handler(&encoded, 5).is_some());
1506
1507        // At end boundary (exclusive)
1508        assert!(find_exception_handler(&encoded, 15).is_none());
1509    }
1510}