use core::fmt;
use crate::{
bytecode::{CodeUnit, instruction::Instruction},
marshal::MarshalError,
};
pub trait OpArgType: Copy + Into<u32> + TryFrom<u32> {}
#[derive(Copy, Clone, PartialEq, Eq)]
#[repr(transparent)]
pub struct OpArgByte(u8);
impl OpArgByte {
pub const NULL: Self = Self::new(0);
#[must_use]
pub const fn new(value: u8) -> Self {
Self(value)
}
}
impl From<u8> for OpArgByte {
fn from(raw: u8) -> Self {
Self::new(raw)
}
}
impl From<OpArgByte> for u8 {
fn from(value: OpArgByte) -> Self {
value.0
}
}
impl fmt::Debug for OpArgByte {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
self.0.fmt(f)
}
}
#[derive(Copy, Clone, Debug)]
#[repr(transparent)]
pub struct OpArg(u32);
impl OpArg {
pub const NULL: Self = Self::new(0);
#[must_use]
pub const fn new(value: u32) -> Self {
Self(value)
}
#[inline]
pub const fn instr_size(self) -> usize {
(self.0 > 0xff) as usize + (self.0 > 0xff_ff) as usize + (self.0 > 0xff_ff_ff) as usize + 1
}
#[inline(always)]
pub fn split(self) -> (impl ExactSizeIterator<Item = OpArgByte>, OpArgByte) {
let mut it = self
.0
.to_le_bytes()
.map(OpArgByte)
.into_iter()
.take(self.instr_size());
let lo = it.next().unwrap();
(it.rev(), lo)
}
}
impl From<u32> for OpArg {
fn from(raw: u32) -> Self {
Self::new(raw)
}
}
impl From<OpArg> for u32 {
fn from(value: OpArg) -> Self {
value.0
}
}
#[derive(Default, Copy, Clone)]
#[repr(transparent)]
pub struct OpArgState {
state: u32,
}
impl OpArgState {
#[inline(always)]
pub fn get(&mut self, ins: CodeUnit) -> (Instruction, OpArg) {
let arg = self.extend(ins.arg);
if !matches!(ins.op, Instruction::ExtendedArg) {
self.reset();
}
(ins.op, arg)
}
#[inline(always)]
pub fn extend(&mut self, arg: OpArgByte) -> OpArg {
self.state = (self.state << 8) | u32::from(arg.0);
self.state.into()
}
#[inline(always)]
pub const fn reset(&mut self) {
self.state = 0
}
}
macro_rules! oparg_enum {
(
$(#[$enum_meta:meta])*
$vis:vis enum $name:ident {
$(
$(#[$variant_meta:meta])*
$variant:ident = $value:literal $(| $alternatives:expr)*
),* $(,)?
}
) => {
$(#[$enum_meta])*
$vis enum $name {
$(
$(#[$variant_meta])*
$variant, )*
}
impl_oparg_enum!(
enum $name {
$(
$variant = $value $(| $alternatives)*,
)*
}
);
};
}
macro_rules! impl_oparg_enum {
(
enum $name:ident {
$(
$variant:ident = $value:literal $(| $alternatives:expr)*
),* $(,)?
}
) => {
impl TryFrom<u8> for $name {
type Error = $crate::marshal::MarshalError;
fn try_from(value: u8) -> Result<Self, Self::Error> {
Ok(match value {
$(
$value $(| $alternatives)* => Self::$variant,
)*
_ => return Err(Self::Error::InvalidBytecode),
})
}
}
impl TryFrom<u32> for $name {
type Error = $crate::marshal::MarshalError;
fn try_from(value: u32) -> Result<Self, Self::Error> {
u8::try_from(value)
.map_err(|_| Self::Error::InvalidBytecode)
.map(TryInto::try_into)?
}
}
impl From<$name> for u8 {
fn from(value: $name) -> Self {
match value {
$(
$name::$variant => $value,
)*
}
}
}
impl From<$name> for u32 {
fn from(value: $name) -> Self {
Self::from(u8::from(value))
}
}
impl OpArgType for $name {}
};
}
oparg_enum!(
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub enum ConvertValueOparg {
None = 0 | 255,
Str = 1,
Repr = 2,
Ascii = 3,
}
);
impl fmt::Display for ConvertValueOparg {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
let out = match self {
Self::Str => "1 (str)",
Self::Repr => "2 (repr)",
Self::Ascii => "3 (ascii)",
Self::None => "",
};
write!(f, "{out}")
}
}
pub type NameIdx = u32;
impl OpArgType for u32 {}
oparg_enum!(
#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub enum RaiseKind {
BareRaise = 0,
Raise = 1,
RaiseCause = 2,
ReraiseFromStack = 3,
}
);
oparg_enum!(
#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub enum IntrinsicFunction1 {
Print = 1,
ImportStar = 2,
StopIterationError = 3,
AsyncGenWrap = 4,
UnaryPositive = 5,
ListToTuple = 6,
TypeVar = 7,
ParamSpec = 8,
TypeVarTuple = 9,
SubscriptGeneric = 10,
TypeAlias = 11,
}
);
oparg_enum!(
#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub enum IntrinsicFunction2 {
PrepReraiseStar = 1,
TypeVarWithBound = 2,
TypeVarWithConstraint = 3,
SetFunctionTypeParams = 4,
SetTypeparamDefault = 5,
}
);
bitflagset::bitflag! {
#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
#[repr(u8)]
pub enum MakeFunctionFlag {
Defaults = 0,
KwOnlyDefaults = 1,
Annotations = 2,
Closure = 3,
Annotate = 4,
TypeParams = 5,
}
}
bitflagset::bitflagset! {
#[derive(Copy, Clone, PartialEq, Eq)]
pub struct MakeFunctionFlags(u8): MakeFunctionFlag
}
impl TryFrom<u32> for MakeFunctionFlag {
type Error = MarshalError;
fn try_from(value: u32) -> Result<Self, Self::Error> {
match value {
0x01 => Ok(Self::Defaults),
0x02 => Ok(Self::KwOnlyDefaults),
0x04 => Ok(Self::Annotations),
0x08 => Ok(Self::Closure),
0x10 => Ok(Self::Annotate),
0x20 => Ok(Self::TypeParams),
_ => Err(MarshalError::InvalidBytecode),
}
}
}
impl From<MakeFunctionFlag> for u32 {
fn from(flag: MakeFunctionFlag) -> Self {
1u32 << (flag as u32)
}
}
impl OpArgType for MakeFunctionFlag {}
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub enum ComparisonOperator {
Less,
LessOrEqual,
Equal,
NotEqual,
Greater,
GreaterOrEqual,
}
impl TryFrom<u8> for ComparisonOperator {
type Error = MarshalError;
fn try_from(value: u8) -> Result<Self, Self::Error> {
Self::try_from(value as u32)
}
}
impl TryFrom<u32> for ComparisonOperator {
type Error = MarshalError;
fn try_from(value: u32) -> Result<Self, Self::Error> {
match value >> 5 {
0 => Ok(Self::Less),
1 => Ok(Self::LessOrEqual),
2 => Ok(Self::Equal),
3 => Ok(Self::NotEqual),
4 => Ok(Self::Greater),
5 => Ok(Self::GreaterOrEqual),
_ => Err(MarshalError::InvalidBytecode),
}
}
}
impl From<ComparisonOperator> for u8 {
fn from(value: ComparisonOperator) -> Self {
match value {
ComparisonOperator::Less => 0,
ComparisonOperator::LessOrEqual => 1 << 5,
ComparisonOperator::Equal => 2 << 5,
ComparisonOperator::NotEqual => 3 << 5,
ComparisonOperator::Greater => 4 << 5,
ComparisonOperator::GreaterOrEqual => 5 << 5,
}
}
}
impl From<ComparisonOperator> for u32 {
fn from(value: ComparisonOperator) -> Self {
Self::from(u8::from(value))
}
}
impl OpArgType for ComparisonOperator {}
oparg_enum!(
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum BinaryOperator {
Add = 0,
And = 1,
FloorDivide = 2,
Lshift = 3,
MatrixMultiply = 4,
Multiply = 5,
Remainder = 6,
Or = 7,
Power = 8,
Rshift = 9,
Subtract = 10,
TrueDivide = 11,
Xor = 12,
InplaceAdd = 13,
InplaceAnd = 14,
InplaceFloorDivide = 15,
InplaceLshift = 16,
InplaceMatrixMultiply = 17,
InplaceMultiply = 18,
InplaceRemainder = 19,
InplaceOr = 20,
InplacePower = 21,
InplaceRshift = 22,
InplaceSubtract = 23,
InplaceTrueDivide = 24,
InplaceXor = 25,
Subscr = 26,
}
);
impl BinaryOperator {
#[must_use]
pub const fn as_inplace(self) -> Self {
match self {
Self::Add => Self::InplaceAdd,
Self::And => Self::InplaceAnd,
Self::FloorDivide => Self::InplaceFloorDivide,
Self::Lshift => Self::InplaceLshift,
Self::MatrixMultiply => Self::InplaceMatrixMultiply,
Self::Multiply => Self::InplaceMultiply,
Self::Remainder => Self::InplaceRemainder,
Self::Or => Self::InplaceOr,
Self::Power => Self::InplacePower,
Self::Rshift => Self::InplaceRshift,
Self::Subtract => Self::InplaceSubtract,
Self::TrueDivide => Self::InplaceTrueDivide,
Self::Xor => Self::InplaceXor,
_ => self,
}
}
}
impl fmt::Display for BinaryOperator {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
let op = match self {
Self::Add => "+",
Self::And => "&",
Self::FloorDivide => "//",
Self::Lshift => "<<",
Self::MatrixMultiply => "@",
Self::Multiply => "*",
Self::Remainder => "%",
Self::Or => "|",
Self::Power => "**",
Self::Rshift => ">>",
Self::Subtract => "-",
Self::TrueDivide => "/",
Self::Xor => "^",
Self::InplaceAdd => "+=",
Self::InplaceAnd => "&=",
Self::InplaceFloorDivide => "//=",
Self::InplaceLshift => "<<=",
Self::InplaceMatrixMultiply => "@=",
Self::InplaceMultiply => "*=",
Self::InplaceRemainder => "%=",
Self::InplaceOr => "|=",
Self::InplacePower => "**=",
Self::InplaceRshift => ">>=",
Self::InplaceSubtract => "-=",
Self::InplaceTrueDivide => "/=",
Self::InplaceXor => "^=",
Self::Subscr => "[]",
};
write!(f, "{op}")
}
}
oparg_enum!(
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub enum Invert {
No = 0,
Yes = 1,
}
);
oparg_enum!(
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub enum SpecialMethod {
Enter = 0,
Exit = 1,
AEnter = 2,
AExit = 3,
}
);
impl fmt::Display for SpecialMethod {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
let method_name = match self {
Self::Enter => "__enter__",
Self::Exit => "__exit__",
Self::AEnter => "__aenter__",
Self::AExit => "__aexit__",
};
write!(f, "{method_name}")
}
}
oparg_enum!(
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub enum CommonConstant {
AssertionError = 0,
NotImplementedError = 1,
BuiltinTuple = 2,
BuiltinAll = 3,
BuiltinAny = 4,
BuiltinList = 5,
BuiltinSet = 6,
}
);
impl fmt::Display for CommonConstant {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
let name = match self {
Self::AssertionError => "AssertionError",
Self::NotImplementedError => "NotImplementedError",
Self::BuiltinTuple => "tuple",
Self::BuiltinAll => "all",
Self::BuiltinAny => "any",
Self::BuiltinList => "list",
Self::BuiltinSet => "set",
};
write!(f, "{name}")
}
}
oparg_enum!(
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum BuildSliceArgCount {
Two = 2,
Three = 3,
}
);
#[derive(Copy, Clone)]
pub struct UnpackExArgs {
pub before: u8,
pub after: u8,
}
impl From<u32> for UnpackExArgs {
fn from(value: u32) -> Self {
let [before, after, ..] = value.to_le_bytes();
Self { before, after }
}
}
impl From<UnpackExArgs> for u32 {
fn from(value: UnpackExArgs) -> Self {
Self::from_le_bytes([value.before, value.after, 0, 0])
}
}
impl OpArgType for UnpackExArgs {}
impl fmt::Display for UnpackExArgs {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "before: {}, after: {}", self.before, self.after)
}
}
macro_rules! newtype_oparg {
(
$(#[$oparg_meta:meta])*
$vis:vis struct $name:ident(u32)
) => {
$(#[$oparg_meta])*
$vis struct $name(u32);
impl $name {
#[must_use]
pub const fn from_u32(value: u32) -> Self {
Self(value)
}
#[must_use]
pub const fn as_u32(self) -> u32 {
self.0
}
#[must_use]
pub const fn as_usize(self) -> usize {
self.0 as usize
}
}
impl From<u32> for $name {
fn from(value: u32) -> Self {
Self::from_u32(value)
}
}
impl From<$name> for u32 {
fn from(value: $name) -> Self {
value.as_u32()
}
}
impl From<$name> for usize {
fn from(value: $name) -> Self {
value.as_usize()
}
}
impl ::core::fmt::Display for $name {
fn fmt(&self, f: &mut ::core::fmt::Formatter<'_>) -> ::core::fmt::Result {
self.0.fmt(f)
}
}
impl OpArgType for $name {}
}
}
newtype_oparg!(
#[derive(Clone, Copy)]
#[repr(transparent)]
pub struct ConstIdx(u32)
);
newtype_oparg!(
#[derive(Clone, Copy)]
#[repr(transparent)]
pub struct VarNum(u32)
);
newtype_oparg!(
#[derive(Clone, Copy)]
#[repr(transparent)]
pub struct VarNums(u32)
);
newtype_oparg!(
#[derive(Clone, Copy)]
#[repr(transparent)]
pub struct LoadAttr(u32)
);
newtype_oparg!(
#[derive(Clone, Copy)]
#[repr(transparent)]
pub struct LoadSuperAttr(u32)
);
newtype_oparg!(
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Ord, PartialOrd)]
#[repr(transparent)]
pub struct Label(u32)
);
newtype_oparg!(
#[derive(Clone, Copy)]
#[repr(transparent)]
pub struct ResumeContext(u32)
);
impl ResumeContext {
pub const LOCATION_MASK: u32 = 0x3;
pub const DEPTH1_MASK: u32 = 0x4;
#[must_use]
pub const fn new(location: ResumeLocation, is_exception_depth1: bool) -> Self {
let value = if is_exception_depth1 {
Self::DEPTH1_MASK
} else {
0
};
Self::from_u32(location.as_u32() | value)
}
#[must_use]
pub fn location(&self) -> ResumeLocation {
unsafe { ResumeLocation::try_from(self.as_u32() & Self::LOCATION_MASK).unwrap_unchecked() }
}
#[must_use]
pub const fn is_exception_depth1(&self) -> bool {
(self.as_u32() & Self::DEPTH1_MASK) != 0
}
}
#[derive(Copy, Clone)]
pub enum ResumeLocation {
AtFuncStart,
AfterYield,
AfterYieldFrom,
AfterAwait,
}
impl From<ResumeLocation> for ResumeContext {
fn from(location: ResumeLocation) -> Self {
Self::new(location, false)
}
}
impl TryFrom<u32> for ResumeLocation {
type Error = MarshalError;
fn try_from(value: u32) -> Result<Self, Self::Error> {
Ok(match value {
0 => Self::AtFuncStart,
1 => Self::AfterYield,
2 => Self::AfterYieldFrom,
3 => Self::AfterAwait,
_ => return Err(Self::Error::InvalidBytecode),
})
}
}
impl ResumeLocation {
#[must_use]
pub const fn as_u8(&self) -> u8 {
match self {
Self::AtFuncStart => 0,
Self::AfterYield => 1,
Self::AfterYieldFrom => 2,
Self::AfterAwait => 3,
}
}
#[must_use]
pub const fn as_u32(&self) -> u32 {
self.as_u8() as u32
}
}
impl From<ResumeLocation> for u8 {
fn from(location: ResumeLocation) -> Self {
location.as_u8()
}
}
impl From<ResumeLocation> for u32 {
fn from(location: ResumeLocation) -> Self {
location.as_u32()
}
}
impl VarNums {
#[must_use]
pub const fn idx_1(self) -> VarNum {
VarNum::from_u32(self.0 >> 4)
}
#[must_use]
pub const fn idx_2(self) -> VarNum {
VarNum::from_u32(self.0 & 15)
}
#[must_use]
pub const fn indexes(self) -> (VarNum, VarNum) {
(self.idx_1(), self.idx_2())
}
}
impl LoadAttr {
#[must_use]
pub const fn new(name_idx: u32, is_method: bool) -> Self {
Self::from_u32((name_idx << 1) | (is_method as u32))
}
#[must_use]
pub const fn name_idx(self) -> u32 {
self.0 >> 1
}
#[must_use]
pub const fn is_method(self) -> bool {
(self.0 & 1) == 1
}
}
impl LoadSuperAttr {
#[must_use]
pub const fn new(name_idx: u32, is_load_method: bool, has_class: bool) -> Self {
Self::from_u32((name_idx << 2) | (is_load_method as u32) | ((has_class as u32) << 1))
}
#[must_use]
pub const fn name_idx(self) -> u32 {
self.0 >> 2
}
#[must_use]
pub const fn is_load_method(self) -> bool {
(self.0 & 1) == 1
}
#[must_use]
pub const fn has_class(self) -> bool {
(self.0 & 2) == 2
}
}