1use 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#[derive(Clone, Copy, Debug, PartialEq, Eq)]
44pub struct ExceptionTableEntry {
45 pub start: u32,
47 pub end: u32,
49 pub target: u32,
51 pub depth: u16,
53 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#[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#[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#[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#[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 match code {
165 0..=9 => pos += 1, 10..=12 => pos += 2, 13 => {
168 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 while pos < linetable.len() {
181 let b = linetable[pos];
182 pos += 1;
183 if b & 0x40 == 0 {
184 break;
185 }
186 }
187 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 => {} _ => {}
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#[derive(Copy, Clone, Debug, PartialEq, Eq)]
215#[repr(u8)]
216pub enum PyCodeLocationInfoKind {
217 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 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 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
314pub 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
412impl<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
421impl<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 #[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#[derive(Clone)]
466pub struct CodeObject<C: Constant = ConstantData> {
467 pub instructions: CodeUnits,
468 pub locations: Box<[(SourceLocation, SourceLocation)]>,
469 pub flags: CodeFlags,
470 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 pub obj_name: C::Name,
479 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 pub localspluskinds: Box<[u8]>,
490 pub linetable: Box<[u8]>,
492 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 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 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#[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
554pub const ADAPTIVE_WARMUP_VALUE: u16 = adaptive_counter_bits(1, 1);
558pub const ADAPTIVE_COOLDOWN_VALUE: u16 = adaptive_counter_bits(52, 0);
562pub 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#[must_use]
571pub const fn adaptive_counter_bits(value: u16, backoff: u16) -> u16 {
572 (value << BACKOFF_BITS) | backoff
573}
574
575#[inline]
577#[must_use]
578pub const fn adaptive_counter_triggers(counter: u16) -> bool {
579 counter < UNREACHABLE_BACKOFF
580}
581
582#[inline]
584#[must_use]
585pub const fn advance_adaptive_counter(counter: u16) -> u16 {
586 counter.wrapping_sub(1 << BACKOFF_BITS)
587}
588
589#[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_cache: Box<[AtomicUsize]>,
634}
635
636unsafe impl Sync for CodeUnits {}
644
645impl Clone for CodeUnits {
646 fn clone(&self) -> Self {
647 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 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 unsafe { &*self.units.get() }
733 }
734}
735
736impl CodeUnits {
737 pub fn disable_specialization(&self) {
740 for counter in &self.adaptive_counters {
741 counter.store(UNREACHABLE_BACKOFF, Ordering::Relaxed);
742 }
743 }
744
745 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 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 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 unsafe { mem::transmute::<u8, Instruction>(byte) }
789 }
790
791 #[inline(always)]
798 #[must_use]
799 pub fn units_ptr(&self) -> *const CodeUnit {
800 unsafe { &*self.units.get() }.as_ptr()
801 }
802
803 #[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 CodeUnit {
817 op: unsafe { mem::transmute::<u8, Instruction>(op) },
818 arg: OpArgByte::from(arg),
819 }
820 }
821
822 #[inline(always)]
834 #[must_use]
835 pub fn read_unit(&self, index: usize) -> CodeUnit {
836 unsafe { Self::read_unit_from(self.units_ptr(), index) }
838 }
839
840 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 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 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 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 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 pub unsafe fn write_cache_ptr(&self, index: usize, value: usize) {
908 self.pointer_cache[index].store(value, Ordering::Relaxed);
909 }
910
911 pub fn read_cache_ptr(&self, index: usize) -> usize {
918 self.pointer_cache[index].load(Ordering::Relaxed)
919 }
920
921 pub fn read_adaptive_counter(&self, index: usize) -> u16 {
924 self.adaptive_counters[index].load(Ordering::Relaxed)
925 }
926
927 pub unsafe fn write_adaptive_counter(&self, index: usize, value: u16) {
933 self.adaptive_counters[index].store(value, Ordering::Relaxed);
934 }
935
936 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 for _ in 0..caches {
951 i += 1;
952 out.push(0); out.push(0); }
955 i += 1;
956 }
957 out
958 }
959
960 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 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#[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 Slice {
1031 elements: Box<[Self; 3]>,
1032 },
1033 Frozenset {
1034 elements: Vec<Self>,
1035 },
1036 None,
1037 Ellipsis,
1038}
1039
1040impl ConstantData {
1041 #[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
1115pub 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
1230pub 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 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 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 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 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 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 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 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 let handler = find_exception_handler(&encoded, 12);
1476 assert!(handler.is_none());
1477
1478 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 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 assert!(find_exception_handler(&encoded, 5).is_some());
1506
1507 assert!(find_exception_handler(&encoded, 15).is_none());
1509 }
1510}