pyc_editor/
traits.rs

1use std::{
2    collections::HashMap,
3    fmt::Debug,
4    ops::{Deref, DerefMut},
5};
6
7use crate::{
8    cfg::{CFGIndexRange, ControlFlowGraph},
9    error::Error,
10    utils::{ExceptionTableEntry, StackEffect},
11};
12
13#[cfg(feature = "sir")]
14use crate::sir::{SIR, SIRControlFlowGraph, SIRStatement, StackItem};
15
16pub trait Oparg: Copy + PartialEq + 'static + Debug {
17    fn is_u32() -> bool;
18
19    fn to_u32(self) -> u32;
20
21    fn is_u8() -> bool;
22
23    fn to_u8(self) -> u8;
24}
25
26impl Oparg for u8 {
27    fn is_u32() -> bool {
28        false
29    }
30
31    #[inline]
32    fn to_u32(self) -> u32 {
33        self as u32
34    }
35
36    fn is_u8() -> bool {
37        true
38    }
39
40    #[inline]
41    fn to_u8(self) -> u8 {
42        self
43    }
44}
45impl Oparg for u32 {
46    fn is_u32() -> bool {
47        true
48    }
49
50    #[inline]
51    fn to_u32(self) -> u32 {
52        self
53    }
54
55    fn is_u8() -> bool {
56        false
57    }
58
59    #[inline]
60    fn to_u8(self) -> u8 {
61        self as u8
62    }
63}
64
65pub trait InstructionAccess<OpargType, I>
66where
67    Self: AsRef<[I]> + Deref<Target = [I]>,
68    OpargType: Oparg,
69    I: GenericInstruction + std::fmt::Debug,
70{
71    // type Instruction: GenericInstruction + std::fmt::Debug;
72    type Jump;
73
74    fn get_instructions(&self) -> &[I] {
75        self.as_ref()
76    }
77
78    /// Returns the index and the instruction of the jump target. None if the index is not a valid jump.
79    // TODO: Create a bounded version of this function
80    fn get_jump_target(&self, index: u32) -> Option<(u32, I)>;
81
82    /// Returns a list of all indexes that jump to the given index
83    fn get_jump_xrefs(&self, index: u32) -> Vec<u32> {
84        let jump_map = self.get_jump_map();
85
86        jump_map
87            .iter()
88            .filter(|(_, to)| **to == index)
89            .map(|(from, _)| *from)
90            .collect()
91    }
92
93    /// Returns a hashmap of jump indexes and their jump target
94    fn get_jump_map(&self) -> HashMap<u32, u32> {
95        let mut jump_map: HashMap<u32, u32> = HashMap::new();
96
97        for index in 0..self.as_ref().len() {
98            let jump_target = self.get_jump_target(index as u32);
99
100            if let Some((jump_index, _)) = jump_target {
101                jump_map.insert(index as u32, jump_index);
102            }
103        }
104
105        jump_map
106    }
107
108    fn get_jump_value(&self, index: u32) -> Option<Self::Jump>;
109
110    /// Calculates the full argument for an index (keeping in mind extended args). None if the index is not within bounds.
111    /// NOTE: If there is a jump skipping the extended arg(s) before this instruction, this will return an incorrect value.
112    fn get_full_arg(&self, index: usize) -> Option<u32> {
113        if self.as_ref().len() > index {
114            if OpargType::is_u32() {
115                self.as_ref().get(index).map(|i| i.get_raw_value().to_u32())
116            } else {
117                let mut curr_index = index;
118                let mut extended_args = vec![];
119
120                while curr_index > 0 {
121                    curr_index -= 1;
122
123                    if self.as_ref()[curr_index].is_extended_arg() {
124                        extended_args.push(self.as_ref()[curr_index].get_raw_value());
125                    } else {
126                        break;
127                    }
128                }
129
130                let mut extended_arg = 0;
131
132                for arg in extended_args.iter().rev() {
133                    // We collected them in the reverse order
134                    let arg = (*arg).to_u32() | extended_arg;
135                    extended_arg = arg << 8;
136                }
137
138                Some(self.as_ref()[index].get_raw_value().to_u32() | extended_arg)
139            }
140        } else {
141            None
142        }
143    }
144
145    /// The same as `get_full_arg` but you can specify a lower limit while searching for extended args.
146    /// NOTE: If there is a jump skipping the extended arg(s) before this instruction, this will return an incorrect value.
147    fn get_full_arg_bounded(&self, index: usize, lower_bound: usize) -> Option<u32> {
148        if self.as_ref().len() > index {
149            if OpargType::is_u32() {
150                self.as_ref().get(index).map(|i| i.get_raw_value().to_u32())
151            } else {
152                let mut curr_index = index;
153                let mut extended_args = vec![];
154
155                while curr_index > lower_bound {
156                    curr_index -= 1;
157
158                    if self.as_ref()[curr_index].is_extended_arg() {
159                        extended_args.push(self.as_ref()[curr_index].get_raw_value());
160                    } else {
161                        break;
162                    }
163                }
164
165                let mut extended_arg = 0;
166
167                for arg in extended_args.iter().rev() {
168                    // We collected them in the reverse order
169                    let arg = (*arg).to_u32() | extended_arg;
170                    extended_arg = arg << 8;
171                }
172
173                Some(self.as_ref()[index].get_raw_value().to_u32() | extended_arg)
174            }
175        } else {
176            None
177        }
178    }
179}
180
181pub trait SimpleInstructionAccess<I>
182where
183    Self: InstructionAccess<u8, I> + AsRef<[I]>,
184    I: GenericInstruction,
185{
186    /// This finds jumps that jump to instructions after an extended arg. This is a very unique case.
187    /// This kind of bytecode should never be emitted by the Python compiler but it's possible for custom bytecode to have this.
188    /// Returns a list of indexes of jump instructions that have a jump target like this.
189    fn find_ext_arg_jumps(&self) -> Vec<u32> {
190        let mut jumps: Vec<u32> = vec![];
191
192        for (index, instruction) in self.as_ref().iter().enumerate() {
193            if instruction.is_jump() {
194                let jump_target = self.get_jump_target(index as u32);
195
196                // Jump target is valid
197                if let Some(jump) = jump_target {
198                    // The jump target has a value bigger than 1 byte, this means we skipped an extended arg
199                    if self
200                        .get_full_arg(jump.0 as usize)
201                        .expect("We know the index is valid")
202                        > u8::MAX.into()
203                    {
204                        jumps.push(index as u32);
205                    }
206                }
207            }
208        }
209
210        jumps
211    }
212
213    fn to_bytes(&self) -> Vec<u8> {
214        let mut bytearray = Vec::with_capacity(self.as_ref().len() * 2);
215
216        for instruction in self.as_ref().iter() {
217            bytearray.push(instruction.get_opcode().into());
218            bytearray.push(instruction.get_raw_value().to_u8())
219        }
220
221        bytearray
222    }
223
224    /// Calculates the maximum stack size the instructions will use. `start_stacksize` is only used for generator code objects that start with a stacksize of 1.
225    /// For 3.11+ you should also pass the exception table for correct stack calculation.
226    /// For 3.13 you should put allow_zero to false since it's not allowed to put 0 as a stacksize in 3.13
227    fn max_stack_size(
228        &self,
229        start_stacksize: u32,
230        exception_table: Option<Vec<ExceptionTableEntry>>,
231        allow_zero: bool,
232    ) -> Result<u32, Error> {
233        // Contains starting indexes of blocks that need to be processed along with their starting stack size.
234        let mut block_queue = vec![(start_stacksize, 0usize)];
235        let mut max_stack_size: u32 = 0;
236        let mut visited: Vec<(u32, usize)> = Vec::new();
237
238        if let Some(exception_entries) = exception_table {
239            for exception in exception_entries {
240                block_queue.push((
241                    exception.depth + 1 + if exception.lasti { 1 } else { 0 },
242                    exception.target as usize,
243                ));
244            }
245        }
246
247        while let Some((stack_size, start_index)) = block_queue.pop() {
248            if visited.contains(&(stack_size, start_index)) {
249                // already analyzed
250                continue;
251            }
252
253            visited.push((stack_size, start_index));
254
255            let mut curr_stack_size = stack_size;
256
257            if curr_stack_size >= max_stack_size {
258                max_stack_size = curr_stack_size;
259            }
260
261            for instruction_index in start_index..self.as_ref().len() {
262                let instruction = self.as_ref().get(instruction_index).unwrap();
263                let arg = self
264                    .get_full_arg_bounded(instruction_index, start_index)
265                    .unwrap();
266
267                if instruction.is_jump() || instruction.stops_execution() {
268                    // Not the last instruction
269                    if instruction_index != self.as_ref().len() - 1
270                        && instruction.is_conditional_jump()
271                        && !instruction.stops_execution()
272                    {
273                        // Block for not taking the jump
274                        let stack_effect = instruction.stack_effect(arg, false, false).net_total();
275
276                        let (stack_size, indx) = (
277                            curr_stack_size.checked_add_signed(stack_effect).ok_or(
278                                Error::InvalidStacksize(curr_stack_size as i32 + stack_effect),
279                            )?,
280                            instruction_index + 1,
281                        );
282
283                        block_queue.push((stack_size, indx));
284                    }
285
286                    if let Some((jump_index, _)) = self.get_jump_target(instruction_index as u32) {
287                        // Block for valid jump target
288                        let stack_effect = instruction.stack_effect(arg, true, false).net_total();
289
290                        let (stack_size, indx) = (
291                            curr_stack_size.checked_add_signed(stack_effect).ok_or(
292                                Error::InvalidStacksize(curr_stack_size as i32 + stack_effect),
293                            )?,
294                            jump_index as usize,
295                        );
296
297                        block_queue.push((stack_size, indx));
298                    }
299
300                    // Process new block, as this one has ended
301                    break;
302                } else {
303                    let stack_effect = instruction.stack_effect(arg, false, true).net_total();
304
305                    curr_stack_size = curr_stack_size.checked_add_signed(stack_effect).ok_or(
306                        Error::InvalidStacksize(curr_stack_size as i32 + stack_effect),
307                    )?;
308
309                    if curr_stack_size >= max_stack_size {
310                        max_stack_size = curr_stack_size;
311                    }
312                }
313            }
314        }
315
316        if !allow_zero && max_stack_size == 0 {
317            max_stack_size = 1;
318        }
319
320        Ok(max_stack_size)
321    }
322}
323
324pub trait IsExtInstruction {}
325
326pub trait ToExtInstructions<I>
327where
328    Self: SimpleInstructionAccess<I>,
329    I: GenericInstruction,
330{
331    /// Returns the instructions but with the extended_args resolved
332    /// For 3.11+: Also pass exception table so that it can return the new exception table with updated indexes
333    fn to_resolved(
334        &self,
335        exception_table: Option<&[ExceptionTableEntry]>,
336    ) -> Result<
337        (
338            <I::OtherType as GenericInstruction>::Instructions,
339            Option<Vec<ExceptionTableEntry>>,
340        ),
341        Error,
342    >;
343}
344
345pub trait ExtInstructionAccess<I, ExtI>
346where
347    ExtI: GenericInstruction,
348    I: GenericInstruction,
349{
350    type ExtInstructions: InstructionAccess<u32, ExtI>;
351    type Instructions: SimpleInstructionAccess<I>;
352
353    /// Convert the resolved instructions back into instructions with extended args.
354    fn to_instructions(&self) -> Self::Instructions;
355
356    /// Resolve instructions into extended instructions.
357    /// For 3.11+, also pass the exception_table to update the indexes when removing extended_args
358    fn from_instructions(
359        instructions: &[I],
360        exception_table: Option<&[ExceptionTableEntry]>,
361    ) -> Result<(Self::ExtInstructions, Option<Vec<ExceptionTableEntry>>), Error>;
362
363    fn to_bytes(&self) -> Vec<u8> {
364        self.to_instructions().to_bytes()
365    }
366}
367
368pub trait InstructionsOwned<T>
369where
370    Self: DerefMut<Target = [Self::Instruction]>,
371    T: Copy,
372{
373    type Instruction;
374
375    // TODO: REQUIRE EXCEPTION TABLE TO BE PASSED
376
377    fn push(&mut self, item: T);
378
379    fn get_instructions_mut(&mut self) -> &mut [Self::Instruction] {
380        self.deref_mut()
381    }
382
383    /// Append multiple instructions at the end
384    fn append_instructions(&mut self, instructions: &[T]) {
385        for instruction in instructions {
386            self.push(*instruction);
387        }
388    }
389    /// Append an instruction at the end
390    fn append_instruction(&mut self, instruction: T) {
391        self.push(instruction);
392    }
393}
394
395pub trait ExtInstructionsOwned<T>
396where
397    Self: DerefMut<Target = [Self::Instruction]>,
398    Self::Instruction: Copy,
399{
400    type Instruction;
401
402    /// Delete instruction at index
403    fn delete_instruction(&mut self, index: usize);
404
405    /// Delete instructions in range (ex. 1..10)
406    fn delete_instructions(&mut self, range: std::ops::Range<usize>) {
407        range
408            .into_iter()
409            .for_each(|index| self.delete_instruction(index));
410    }
411
412    /// Insert instruction at a specific index. It automatically fixes jump offsets in other instructions.
413    fn insert_instruction(&mut self, index: usize, instruction: Self::Instruction);
414
415    /// Insert a slice of instructions at an index
416    fn insert_instructions(&mut self, index: usize, instructions: &[Self::Instruction]) {
417        for (idx, instruction) in instructions.iter().enumerate() {
418            self.insert_instruction(index + idx, *instruction);
419        }
420    }
421}
422
423/// Generic opcode functions each version has to implement
424pub trait GenericOpcode: StackEffectTrait + PartialEq + Into<u8> + Debug + Clone {
425    type BranchReason: BranchReasonTrait;
426
427    fn is_jump(&self) -> bool;
428    fn is_absolute_jump(&self) -> bool;
429    fn is_relative_jump(&self) -> bool;
430    fn is_jump_forwards(&self) -> bool;
431    fn is_jump_backwards(&self) -> bool;
432    fn is_conditional_jump(&self) -> bool;
433    fn stops_execution(&self) -> bool;
434    fn is_extended_arg(&self) -> bool;
435    fn is_cache(&self) -> bool;
436    fn get_nop() -> Self;
437}
438
439/// Generic instruction functions used by all versions
440pub trait GenericInstruction: PartialEq + Debug + Clone {
441    type OpargType: Oparg;
442    type Opcode: GenericOpcode;
443
444    /// This is the "list" version of the actual instruction (e.g. ExtInstructions)
445    type Instructions: InstructionAccess<Self::OpargType, Self>;
446
447    /// This is the related type. For example for Instruction this would be ExtInstruction and visa versa.
448    type OtherType: GenericInstruction;
449
450    fn get_opcode(&self) -> Self::Opcode;
451
452    fn get_raw_value(&self) -> Self::OpargType;
453
454    /// Relative or absolute jump
455    fn is_jump(&self) -> bool {
456        self.get_opcode().is_jump()
457    }
458
459    fn is_absolute_jump(&self) -> bool {
460        self.get_opcode().is_absolute_jump()
461    }
462
463    fn is_relative_jump(&self) -> bool {
464        self.get_opcode().is_relative_jump()
465    }
466
467    fn is_jump_forwards(&self) -> bool {
468        self.get_opcode().is_jump_forwards()
469    }
470
471    fn is_jump_backwards(&self) -> bool {
472        self.get_opcode().is_jump_backwards()
473    }
474
475    fn is_conditional_jump(&self) -> bool {
476        self.get_opcode().is_conditional_jump()
477    }
478
479    fn stops_execution(&self) -> bool {
480        self.get_opcode().stops_execution()
481    }
482
483    fn is_extended_arg(&self) -> bool {
484        self.get_opcode().is_extended_arg()
485    }
486
487    fn is_cache(&self) -> bool {
488        self.get_opcode().is_cache()
489    }
490
491    fn get_nop() -> Self;
492
493    /// If the code has a jump target and `jump` is true, `stack_effect()` will return the stack effect of jumping.
494    /// If jump is false, it will return the stack effect of not jumping.
495    /// And if calculate_max is true, it will return the maximal stack effect of both cases.
496    fn stack_effect(&self, oparg: u32, jump: bool, calculate_max: bool) -> StackEffect {
497        self.get_opcode().stack_effect(oparg, jump, calculate_max)
498    }
499}
500
501/// Should be automatically implemented by define_opcodes!()
502pub trait StackEffectTrait {
503    /// If the code has a jump target and `jump` is true, `stack_effect()` will return the stack effect of jumping.
504    /// If jump is false, it will return the stack effect of not jumping.
505    /// And if calculate_max is true, it will return the maximal stack effect of both cases.
506    fn stack_effect(&self, oparg: u32, jump: bool, calculate_max: bool) -> StackEffect;
507}
508
509#[cfg(feature = "sir")]
510pub trait GenericSIRNode: Clone + Debug + PartialEq {
511    type Opcode: GenericOpcode;
512    type SIRException: GenericSIRException;
513
514    fn new(opcode: Self::Opcode, oparg: u32, jump: bool) -> Self;
515
516    fn get_outputs(&self) -> &[StackItem];
517
518    fn get_inputs(&self) -> &[StackItem];
519
520    fn get_net_stack_delta(&self) -> isize;
521}
522
523#[cfg(feature = "sir")]
524pub trait GenericSIRException: Clone + Debug + PartialEq {
525    type Opcode: GenericOpcode;
526
527    fn new(lasti: bool, stack_depth: usize, jump: bool) -> Self;
528
529    fn get_outputs(&self) -> &[StackItem];
530
531    fn get_inputs(&self) -> &[StackItem];
532
533    fn get_net_stack_delta(&self) -> isize;
534
535    fn get_stack_depth(&self) -> usize;
536}
537
538#[cfg(feature = "sir")]
539/// A trait to indicate that the SIR statements are owned.
540pub trait SIROwned<SIRNode: GenericSIRNode>: std::fmt::Display {
541    fn new(statements: Vec<SIRStatement<SIRNode>>) -> Self;
542}
543
544#[cfg(feature = "sir")]
545impl<SIRNode: GenericSIRNode> Deref for SIR<SIRNode> {
546    type Target = [SIRStatement<SIRNode>];
547
548    fn deref(&self) -> &Self::Target {
549        &self.0
550    }
551}
552
553#[cfg(feature = "sir")]
554impl<SIRNode: GenericSIRNode> DerefMut for SIR<SIRNode> {
555    fn deref_mut(&mut self) -> &mut Self::Target {
556        &mut self.0
557    }
558}
559
560#[cfg(feature = "sir")]
561/// A trait for passes that can run on a SIRControlFlowGraph
562pub trait SIRCFGPass<SIRNode: GenericSIRNode> {
563    fn run_on(&self, cfg: &mut SIRControlFlowGraph<SIRNode>);
564}
565
566#[cfg(feature = "sir")]
567pub trait ToSIR<SIRNode: GenericSIRNode> {
568    /// Exception table should be passed for 3.11+
569    fn to_sir(
570        &self,
571        exception_table: Option<&[ExceptionTableEntry]>,
572    ) -> Result<SIRControlFlowGraph<SIRNode>, Error>;
573}
574
575/// Trait to show what the branch reason is (opcode or exception)
576pub trait BranchReasonTrait: Clone + Debug + std::fmt::Display + PartialEq {
577    type Opcode: GenericOpcode;
578
579    fn from_opcode(opcode: Self::Opcode) -> Result<Self, Error>;
580
581    fn from_exception(lasti: bool, stack_depth: usize) -> Result<Self, Error>;
582
583    fn is_opcode(&self) -> bool;
584
585    fn is_exception(&self) -> bool;
586
587    fn get_opcode(&self) -> Option<&Self::Opcode>;
588
589    fn get_lasti(&self) -> Option<bool>;
590
591    fn get_stack_depth(&self) -> Option<usize>;
592}
593
594pub trait BlockSliceExt<I> {
595    fn find_exception_block(&self, index_to_search: usize) -> Option<usize>;
596}
597
598/// This trait allows modification of the CFG after creating it.
599/// This is necessary for versions where there is no exception table and we need to do block stack analysis to find the exception ranges.
600pub(crate) trait FinalizeCFG<I> {
601    fn finalize_cfg(&mut self, map: Option<&mut Vec<CFGIndexRange>>) -> Result<(), Error>
602    where
603        I: GenericInstruction,
604        for<'a> &'a [I]: InstructionAccess<I::OpargType, I>,
605        <I::Opcode as GenericOpcode>::BranchReason: BranchReasonTrait<Opcode = I::Opcode>;
606}
607
608#[allow(private_bounds)]
609pub trait CreateCFG<I> {
610    fn create_cfg(
611        self,
612        exception_table: Option<&[crate::utils::ExceptionTableEntry]>,
613    ) -> Result<ControlFlowGraph<I>, Error>
614    where
615        I: GenericInstruction,
616        ControlFlowGraph<I>: FinalizeCFG<I>,
617        for<'a> &'a [I]: InstructionAccess<I::OpargType, I>,
618        <I::Opcode as GenericOpcode>::BranchReason: BranchReasonTrait<Opcode = I::Opcode>;
619}
620
621#[allow(private_bounds)]
622impl<I> CreateCFG<I> for &[I] {
623    fn create_cfg(
624        self,
625        exception_table: Option<&[crate::utils::ExceptionTableEntry]>,
626    ) -> Result<ControlFlowGraph<I>, Error>
627    where
628        I: GenericInstruction,
629        ControlFlowGraph<I>: FinalizeCFG<I>,
630        for<'a> &'a [I]: InstructionAccess<<I>::OpargType, I>,
631        <<I>::Opcode as GenericOpcode>::BranchReason: BranchReasonTrait<Opcode = <I>::Opcode>,
632    {
633        crate::cfg::create_cfg(self, exception_table)
634    }
635}
636
637#[cfg(all(test, feature = "v311"))]
638mod test {
639    use crate::traits::SimpleInstructionAccess;
640
641    #[test]
642    fn test_invalid_extended_arg_jump() {
643        let instructions = crate::v311::instructions::Instructions::new(vec![
644            crate::v311::instructions::Instruction::JumpForward(1),
645            crate::v311::instructions::Instruction::ExtendedArg(1),
646            crate::v311::instructions::Instruction::Nop(1),
647        ]);
648
649        assert_eq!(instructions.find_ext_arg_jumps().len(), 1)
650    }
651
652    #[test]
653    fn test_stack_size() {
654        let instructions = crate::v311::instructions::Instructions::new(vec![
655            crate::v311::instructions::Instruction::Resume(0),
656            crate::v311::instructions::Instruction::PushNull(0),
657            crate::v311::instructions::Instruction::LoadName(0),
658            crate::v311::instructions::Instruction::LoadConst(0),
659            crate::v311::instructions::Instruction::Precall(1),
660            crate::v311::instructions::Instruction::Cache(0),
661            crate::v311::instructions::Instruction::Call(1),
662            crate::v311::instructions::Instruction::Cache(0),
663            crate::v311::instructions::Instruction::Cache(0),
664            crate::v311::instructions::Instruction::Cache(0),
665            crate::v311::instructions::Instruction::Cache(0),
666            crate::v311::instructions::Instruction::PopTop(0),
667            crate::v311::instructions::Instruction::LoadConst(1),
668            crate::v311::instructions::Instruction::StoreName(1),
669            crate::v311::instructions::Instruction::PushNull(0),
670            crate::v311::instructions::Instruction::LoadName(0),
671            crate::v311::instructions::Instruction::LoadConst(2),
672            crate::v311::instructions::Instruction::LoadName(1),
673            crate::v311::instructions::Instruction::FormatValue(2),
674            crate::v311::instructions::Instruction::BuildString(2),
675            crate::v311::instructions::Instruction::Precall(1),
676            crate::v311::instructions::Instruction::Cache(0),
677            crate::v311::instructions::Instruction::Call(1),
678            crate::v311::instructions::Instruction::Cache(0),
679            crate::v311::instructions::Instruction::Cache(0),
680            crate::v311::instructions::Instruction::Cache(0),
681            crate::v311::instructions::Instruction::Cache(0),
682            crate::v311::instructions::Instruction::PopTop(0),
683            crate::v311::instructions::Instruction::LoadConst(3),
684            crate::v311::instructions::Instruction::ReturnValue(0),
685        ]);
686
687        assert_eq!(instructions.max_stack_size(0, None, true).unwrap(), 4);
688    }
689}
pyc_editor/traits.rs

pyc_editor/
traits.rs
