jit_assembler/common/
mod.rs

1//! Common types and traits shared across all target architectures.
2
3use core::fmt;
4
5/// Register usage tracking functionality
6#[cfg(feature = "register-tracking")]
7pub mod register_usage;
8
9#[cfg(feature = "register-tracking")]
10use register_usage::RegisterUsageInfo;
11
12#[cfg(feature = "std")]
13use std::vec::Vec;
14#[cfg(not(feature = "std"))]
15use alloc::vec::Vec;
16
17/// A collection of instructions with associated register usage information.
18/// 
19/// This struct combines an `InstructionCollection` with `RegisterUsageInfo`,
20/// allowing you to merge both instructions and their usage statistics together.
21/// This is useful when combining code sequences from multiple builders while
22/// preserving register usage information.
23/// 
24/// # Example
25/// 
26/// ```rust,ignore
27/// // Create three builders for different parts of the code
28/// let mut prologue = Riscv64InstructionBuilder::new();
29/// prologue.sd(reg::SP, reg::S0, -8);  // Save S0
30/// 
31/// let mut main = Riscv64InstructionBuilder::new();
32/// main.add(reg::S0, reg::A0, reg::A1);  // Main computation
33/// 
34/// let mut epilogue = Riscv64InstructionBuilder::new();
35/// epilogue.ld(reg::S0, reg::SP, -8);  // Restore S0
36/// 
37/// // Create tracked collections
38/// let prologue_tracked = InstructionCollectionWithUsage::new(
39///     prologue.instructions(),
40///     prologue.register_usage().clone()
41/// );
42/// let main_tracked = InstructionCollectionWithUsage::new(
43///     main.instructions(),
44///     main.register_usage().clone()
45/// );
46/// let epilogue_tracked = InstructionCollectionWithUsage::new(
47///     epilogue.instructions(),
48///     epilogue.register_usage().clone()
49/// );
50/// 
51/// // Merge them: prologue + main + epilogue
52/// let combined = prologue_tracked + main_tracked + epilogue_tracked;
53/// 
54/// // Now we have the complete function with accurate register usage
55/// let instructions = combined.instructions();
56/// let usage = combined.register_usage();
57/// ```
58#[cfg(feature = "register-tracking")]
59#[derive(Debug, Clone)]
60pub struct InstructionCollectionWithUsage<I: Instruction, R: Register> {
61    instructions: InstructionCollection<I>,
62    register_usage: RegisterUsageInfo<R>,
63}
64
65#[cfg(feature = "register-tracking")]
66impl<I: Instruction, R: Register> InstructionCollectionWithUsage<I, R> {
67    /// Create a new tracked instruction collection.
68    pub fn new(instructions: InstructionCollection<I>, register_usage: RegisterUsageInfo<R>) -> Self {
69        Self {
70            instructions,
71            register_usage,
72        }
73    }
74    
75    /// Create from raw parts.
76    pub fn from_parts(instructions: InstructionCollection<I>, register_usage: RegisterUsageInfo<R>) -> Self {
77        Self::new(instructions, register_usage)
78    }
79    
80    /// Get a reference to the instructions.
81    pub fn instructions(&self) -> &InstructionCollection<I> {
82        &self.instructions
83    }
84    
85    /// Get a mutable reference to the instructions.
86    pub fn instructions_mut(&mut self) -> &mut InstructionCollection<I> {
87        &mut self.instructions
88    }
89    
90    /// Get a reference to the register usage information.
91    pub fn register_usage(&self) -> &RegisterUsageInfo<R> {
92        &self.register_usage
93    }
94    
95    /// Get a mutable reference to the register usage information.
96    pub fn register_usage_mut(&mut self) -> &mut RegisterUsageInfo<R> {
97        &mut self.register_usage
98    }
99    
100    /// Consume this collection and return the instructions and register usage.
101    pub fn into_parts(self) -> (InstructionCollection<I>, RegisterUsageInfo<R>) {
102        (self.instructions, self.register_usage)
103    }
104    
105    /// Consume this collection and return just the instructions.
106    pub fn into_instructions(self) -> InstructionCollection<I> {
107        self.instructions
108    }
109    
110    /// Merge another tracked collection into this one.
111    /// 
112    /// This appends the instructions and merges the register usage information.
113    pub fn append(&mut self, other: InstructionCollectionWithUsage<I, R>) {
114        self.instructions.append(other.instructions);
115        self.register_usage.merge(&other.register_usage);
116    }
117    
118    /// Extend this collection with instructions and register usage from another.
119    /// 
120    /// This clones instructions and merges the register usage information.
121    pub fn extend_from(&mut self, other: &InstructionCollectionWithUsage<I, R>) {
122        self.instructions.extend_from_collection(&other.instructions);
123        self.register_usage.merge(&other.register_usage);
124    }
125    
126    /// Concatenate two tracked collections, consuming both.
127    pub fn concat(mut self, other: InstructionCollectionWithUsage<I, R>) -> Self {
128        self.instructions.append(other.instructions);
129        self.register_usage.merge(&other.register_usage);
130        self
131    }
132}
133
134#[cfg(feature = "register-tracking")]
135impl<I: Instruction, R: Register> core::ops::Add for InstructionCollectionWithUsage<I, R> {
136    type Output = InstructionCollectionWithUsage<I, R>;
137    
138    /// Concatenate two tracked instruction collections using the `+` operator.
139    fn add(self, other: InstructionCollectionWithUsage<I, R>) -> InstructionCollectionWithUsage<I, R> {
140        self.concat(other)
141    }
142}
143
144#[cfg(feature = "register-tracking")]
145impl<I: Instruction, R: Register> core::ops::AddAssign for InstructionCollectionWithUsage<I, R> {
146    /// Append another tracked instruction collection using the `+=` operator.
147    fn add_assign(&mut self, other: InstructionCollectionWithUsage<I, R>) {
148        self.append(other);
149    }
150}
151
152/// A machine instruction that can be encoded to bytes
153pub trait Instruction: Copy + Clone + fmt::Debug + fmt::Display {
154    /// Get the instruction as a 32-bit or 64-bit value
155    fn value(&self) -> u64;
156    
157    /// Get the instruction as bytes (little-endian)
158    fn bytes(&self) -> Vec<u8>;
159    
160    /// Get the size of this instruction in bytes
161    fn size(&self) -> usize;
162}
163    
164/// ABI classification for registers based on preservation requirements.
165/// 
166/// This simplified classification focuses on whether registers need to be
167/// preserved across function calls, which is the most critical information
168/// for JIT compilation and register allocation.
169#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
170pub enum AbiClass {
171    /// Caller-saved registers that don't need to be preserved across calls.
172    /// 
173    /// These registers can be freely used by JIT-compiled functions without
174    /// saving/restoring their values. Examples: argument registers, temp registers.
175    CallerSaved,
176    
177    /// Callee-saved registers that must be preserved across calls.
178    /// 
179    /// If a JIT-compiled function uses these registers, it must save their
180    /// values on entry and restore them before returning. Examples: saved registers.
181    CalleeSaved,
182    
183    /// Special-purpose registers with specific ABI requirements.
184    /// 
185    /// These registers have specific roles (stack pointer, frame pointer, zero register, etc.)
186    /// and require careful handling. Generally should be avoided in JIT code
187    /// unless specifically needed for their intended purpose.
188    Special,
189}
190
191impl fmt::Display for AbiClass {
192    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
193        match self {
194            AbiClass::CallerSaved => write!(f, "caller-saved"),
195            AbiClass::CalleeSaved => write!(f, "callee-saved"),
196            AbiClass::Special => write!(f, "special"),
197        }
198    }
199}
200
201/// A register identifier for a target architecture
202pub trait Register: Copy + Clone + fmt::Debug + core::hash::Hash + Eq {
203    /// Get the register number/identifier
204    fn id(&self) -> u32;
205    
206    /// Get the ABI classification for this register.
207    /// 
208    /// This method should return the appropriate `AbiClass` based on the
209    /// target architecture's calling convention.
210    /// 
211    /// # Example
212    /// 
213    /// ```rust,ignore
214    /// // For RISC-V
215    /// match self {
216    ///     Register::T0 | Register::T1 => AbiClass::CallerSaved,
217    ///     Register::S0 | Register::S1 => AbiClass::CalleeSaved,
218    ///     Register::SP | Register::FP => AbiClass::Other,
219    ///     // ...
220    /// }
221    /// ```
222    fn abi_class(&self) -> AbiClass;
223    
224    /// Check if this register is caller-saved.
225    /// 
226    /// Convenience method equivalent to `self.abi_class() == AbiClass::CallerSaved`.
227    fn is_caller_saved(&self) -> bool {
228        self.abi_class() == AbiClass::CallerSaved
229    }
230    
231    /// Check if this register is callee-saved.
232    /// 
233    /// Convenience method equivalent to `self.abi_class() == AbiClass::CalleeSaved`.
234    fn is_callee_saved(&self) -> bool {
235        self.abi_class() == AbiClass::CalleeSaved
236    }
237    
238    /// Check if this register is special-purpose.
239    /// 
240    /// Convenience method equivalent to `self.abi_class() == AbiClass::Special`.
241    fn is_special(&self) -> bool {
242        self.abi_class() == AbiClass::Special
243    }
244}
245
246/// An instruction builder for a specific architecture
247pub trait InstructionBuilder<I: Instruction> {
248    /// The register type used by this architecture
249    type Register: Register;
250    
251    /// Create a new instruction builder
252    fn new() -> Self;
253    
254    /// Get the generated instructions
255    fn instructions(&self) -> InstructionCollection<I>;
256    
257    /// Add an instruction to the builder
258    fn push(&mut self, instr: I);
259    
260    /// Clear all instructions
261    fn clear(&mut self);
262    
263    /// Get register usage information (register-tracking feature only)
264    /// 
265    /// This method returns information about which registers have been used
266    /// by the instructions in this builder, enabling register allocation
267    /// analysis and ABI compliance checking.
268    #[cfg(feature = "register-tracking")]
269    fn register_usage(&self) -> &crate::common::register_usage::RegisterUsageInfo<Self::Register>;
270    
271    /// Get mutable register usage information (register-tracking feature only)
272    /// 
273    /// This allows direct manipulation of the usage tracking, which can be
274    /// useful for advanced use cases or manual register tracking.
275    #[cfg(feature = "register-tracking")]
276    fn register_usage_mut(&mut self) -> &mut crate::common::register_usage::RegisterUsageInfo<Self::Register>;
277    
278    /// Clear register usage information (register-tracking feature only)
279    /// 
280    /// This resets the usage tracking to an empty state, which can be
281    /// useful when reusing a builder for multiple functions.
282    #[cfg(feature = "register-tracking")]
283    fn clear_register_usage(&mut self) {
284        self.register_usage_mut().clear();
285    }
286    
287    /// Create a JIT-compiled function from the assembled instructions (std-only)
288    /// 
289    /// This method converts the assembled instructions into executable machine code
290    /// that can be called directly as a function. The generic type parameter `F`
291    /// specifies the function signature.
292    /// 
293    /// # ABI Compatibility
294    /// 
295    /// While you specify a Rust function type like `fn() -> u64`, the actual JIT
296    /// code uses C ABI internally for stability across Rust versions. This conversion
297    /// is handled transparently by the `call()` methods.
298    /// 
299    /// # Limitations
300    /// 
301    /// Currently supports function signatures with up to 7 arguments. For functions
302    /// with more arguments or complex calling conventions, use manual function pointer
303    /// conversion.
304    /// 
305    /// # Safety
306    /// 
307    /// This function is unsafe because:
308    /// - It allocates executable memory
309    /// - It assumes the assembled code follows the correct C ABI
310    /// - The caller must ensure the function signature matches the actual code
311    /// - The assembled code must be valid for the target architecture
312    /// 
313    #[cfg(feature = "std")]
314    unsafe fn function<F>(&self) -> Result<crate::common::jit::CallableJitFunction<F>, crate::common::jit::JitError>;
315
316    /// Create a raw JIT-compiled function for manual type conversion (std-only)
317    /// 
318    /// This method is similar to `function()` but returns a type-erased function
319    /// that allows manual conversion to any function signature. Use this for
320    /// function signatures with more than 7 arguments or custom calling conventions.
321    /// 
322    /// # Safety
323    /// 
324    /// This function is unsafe because:
325    /// - It allocates executable memory  
326    /// - The caller must manually ensure type safety when calling `as_fn()`
327    /// - The assembled code must be valid for the target architecture
328    /// 
329    #[cfg(feature = "std")]
330    unsafe fn raw_function(&self) -> Result<crate::common::jit::RawCallableJitFunction, crate::common::jit::JitError>;
331}
332
333/// Convenience functions for instruction collections
334/// These functions work with any collection of instructions
335pub fn instructions_to_bytes<I: Instruction>(instructions: &[I]) -> Vec<u8> {
336    let mut result = Vec::new();
337    for instr in instructions {
338        result.extend_from_slice(&instr.bytes());
339    }
340    result
341}
342
343/// Get the total size in bytes of a collection of instructions
344pub fn instructions_total_size<I: Instruction>(instructions: &[I]) -> usize {
345    instructions.iter().map(|i| i.size()).sum()
346}
347
348/// A collection of instructions with convenient methods for byte manipulation
349#[derive(Debug, Clone, PartialEq, Eq)]
350pub struct InstructionCollection<I: Instruction> {
351    instructions: Vec<I>,
352}
353
354impl<I: Instruction> InstructionCollection<I> {
355    /// Create a new empty instruction collection
356    pub fn new() -> Self {
357        Self {
358            instructions: Vec::new(),
359        }
360    }
361    
362    /// Create from a vector of instructions
363    pub fn from_vec(instructions: Vec<I>) -> Self {
364        Self { instructions }
365    }
366    
367    /// Create from a slice of instructions
368    pub fn from_slice(instructions: &[I]) -> Self {
369        Self {
370            instructions: instructions.to_vec(),
371        }
372    }
373    
374    /// Get the instructions as a slice
375    pub fn as_slice(&self) -> &[I] {
376        &self.instructions
377    }
378    
379    /// Get the instructions as a mutable slice
380    pub fn as_mut_slice(&mut self) -> &mut [I] {
381        &mut self.instructions
382    }
383    
384    /// Convert instructions to a single byte vector
385    pub fn to_bytes(&self) -> Vec<u8> {
386        instructions_to_bytes(&self.instructions)
387    }
388    
389    /// Get the total size in bytes of all instructions
390    pub fn total_size(&self) -> usize {
391        instructions_total_size(&self.instructions)
392    }
393    
394    /// Get the number of instructions
395    pub fn len(&self) -> usize {
396        self.instructions.len()
397    }
398    
399    /// Check if the collection is empty
400    pub fn is_empty(&self) -> bool {
401        self.instructions.is_empty()
402    }
403    
404    /// Add an instruction to the collection
405    pub fn push(&mut self, instruction: I) {
406        self.instructions.push(instruction);
407    }
408    
409    /// Remove all instructions
410    pub fn clear(&mut self) {
411        self.instructions.clear();
412    }
413    
414    /// Get an iterator over the instructions
415    pub fn iter(&self) -> core::slice::Iter<'_, I> {
416        self.instructions.iter()
417    }
418    
419    /// Get a mutable iterator over the instructions
420    pub fn iter_mut(&mut self) -> core::slice::IterMut<'_, I> {
421        self.instructions.iter_mut()
422    }
423    
424    /// Convert to a Vec (consumes self)
425    pub fn into_vec(self) -> Vec<I> {
426        self.instructions
427    }
428    
429    /// Create a Vec by cloning the instructions
430    pub fn to_vec(&self) -> Vec<I> {
431        self.instructions.clone()
432    }
433    
434    /// Get a reference to the instruction at the given index
435    pub fn get(&self, index: usize) -> Option<&I> {
436        self.instructions.get(index)
437    }
438    
439    /// Get a mutable reference to the instruction at the given index
440    pub fn get_mut(&mut self, index: usize) -> Option<&mut I> {
441        self.instructions.get_mut(index)
442    }
443    
444    /// Append another instruction collection to this one (consumes other).
445    /// 
446    /// This moves all instructions from `other` into this collection.
447    /// 
448    /// # Example
449    /// 
450    /// ```rust,ignore
451    /// let mut collection1 = InstructionCollection::from_slice(&[instr1, instr2]);
452    /// let collection2 = InstructionCollection::from_slice(&[instr3, instr4]);
453    /// collection1.append(collection2);
454    /// // collection1 now contains [instr1, instr2, instr3, instr4]
455    /// ```
456    pub fn append(&mut self, mut other: InstructionCollection<I>) {
457        self.instructions.append(&mut other.instructions);
458    }
459    
460    /// Extend this collection with instructions from another collection.
461    /// 
462    /// This clones instructions from `other` into this collection.
463    /// 
464    /// # Example
465    /// 
466    /// ```rust,ignore
467    /// let mut collection1 = InstructionCollection::from_slice(&[instr1, instr2]);
468    /// let collection2 = InstructionCollection::from_slice(&[instr3, instr4]);
469    /// collection1.extend_from_collection(&collection2);
470    /// // collection1 now contains [instr1, instr2, instr3, instr4]
471    /// // collection2 is still valid
472    /// ```
473    pub fn extend_from_collection(&mut self, other: &InstructionCollection<I>) {
474        self.instructions.extend_from_slice(&other.instructions);
475    }
476    
477    /// Concatenate two instruction collections, consuming both.
478    /// 
479    /// Creates a new collection containing all instructions from `self` 
480    /// followed by all instructions from `other`.
481    /// 
482    /// # Example
483    /// 
484    /// ```rust,ignore
485    /// let collection1 = InstructionCollection::from_slice(&[instr1, instr2]);
486    /// let collection2 = InstructionCollection::from_slice(&[instr3, instr4]);
487    /// let combined = collection1.concat(collection2);
488    /// // combined contains [instr1, instr2, instr3, instr4]
489    /// ```
490    pub fn concat(mut self, mut other: InstructionCollection<I>) -> Self {
491        self.instructions.append(&mut other.instructions);
492        self
493    }
494}
495
496impl<I: Instruction> Default for InstructionCollection<I> {
497    fn default() -> Self {
498        Self::new()
499    }
500}
501
502impl<I: Instruction> From<Vec<I>> for InstructionCollection<I> {
503    fn from(instructions: Vec<I>) -> Self {
504        Self::from_vec(instructions)
505    }
506}
507
508impl<I: Instruction> From<&[I]> for InstructionCollection<I> {
509    fn from(instructions: &[I]) -> Self {
510        Self::from_slice(instructions)
511    }
512}
513
514impl<I: Instruction> AsRef<[I]> for InstructionCollection<I> {
515    fn as_ref(&self) -> &[I] {
516        &self.instructions
517    }
518}
519
520impl<I: Instruction> AsMut<[I]> for InstructionCollection<I> {
521    fn as_mut(&mut self) -> &mut [I] {
522        &mut self.instructions
523    }
524}
525
526impl<I: Instruction> core::ops::Index<usize> for InstructionCollection<I> {
527    type Output = I;
528    
529    fn index(&self, index: usize) -> &Self::Output {
530        &self.instructions[index]
531    }
532}
533
534impl<I: Instruction> core::ops::IndexMut<usize> for InstructionCollection<I> {
535    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
536        &mut self.instructions[index]
537    }
538}
539
540impl<I: Instruction> IntoIterator for InstructionCollection<I> {
541    type Item = I;
542    type IntoIter = <Vec<I> as IntoIterator>::IntoIter;
543    
544    fn into_iter(self) -> Self::IntoIter {
545        self.instructions.into_iter()
546    }
547}
548
549impl<'a, I: Instruction> IntoIterator for &'a InstructionCollection<I> {
550    type Item = &'a I;
551    type IntoIter = core::slice::Iter<'a, I>;
552    
553    fn into_iter(self) -> Self::IntoIter {
554        self.instructions.iter()
555    }
556}
557
558impl<'a, I: Instruction> IntoIterator for &'a mut InstructionCollection<I> {
559    type Item = &'a mut I;
560    type IntoIter = core::slice::IterMut<'a, I>;
561    
562    fn into_iter(self) -> Self::IntoIter {
563        self.instructions.iter_mut()
564    }
565}
566
567impl<I: Instruction> core::ops::Deref for InstructionCollection<I> {
568    type Target = [I];
569    
570    fn deref(&self) -> &Self::Target {
571        &self.instructions
572    }
573}
574
575impl<I: Instruction> core::ops::DerefMut for InstructionCollection<I> {
576    fn deref_mut(&mut self) -> &mut Self::Target {
577        &mut self.instructions
578    }
579}
580
581impl<I: Instruction> core::ops::Add for InstructionCollection<I> {
582    type Output = InstructionCollection<I>;
583    
584    /// Concatenate two instruction collections using the `+` operator.
585    /// 
586    /// # Example
587    /// 
588    /// ```rust,ignore
589    /// let collection1 = InstructionCollection::from_slice(&[instr1, instr2]);
590    /// let collection2 = InstructionCollection::from_slice(&[instr3, instr4]);
591    /// let combined = collection1 + collection2;
592    /// // combined contains [instr1, instr2, instr3, instr4]
593    /// ```
594    fn add(self, other: InstructionCollection<I>) -> InstructionCollection<I> {
595        self.concat(other)
596    }
597}
598
599impl<I: Instruction> core::ops::AddAssign for InstructionCollection<I> {
600    /// Append another instruction collection to this one using the `+=` operator.
601    /// 
602    /// # Example
603    /// 
604    /// ```rust,ignore
605    /// let mut collection1 = InstructionCollection::from_slice(&[instr1, instr2]);
606    /// let collection2 = InstructionCollection::from_slice(&[instr3, instr4]);
607    /// collection1 += collection2;
608    /// // collection1 now contains [instr1, instr2, instr3, instr4]
609    /// ```
610    fn add_assign(&mut self, other: InstructionCollection<I>) {
611        self.append(other);
612    }
613}
614
615/// Trait extension for instruction collections
616/// This allows you to call `.to_bytes()` and `.total_size()` directly on slices and vectors
617pub trait InstructionCollectionExt<I: Instruction> {
618    /// Convert instructions to a single byte vector
619    fn to_bytes(&self) -> Vec<u8>;
620    
621    /// Get the total size in bytes of all instructions
622    fn total_size(&self) -> usize;
623}
624
625impl<I: Instruction> InstructionCollectionExt<I> for [I] {
626    fn to_bytes(&self) -> Vec<u8> {
627        instructions_to_bytes(self)
628    }
629    
630    fn total_size(&self) -> usize {
631        instructions_total_size(self)
632    }
633}
634
635impl<I: Instruction> InstructionCollectionExt<I> for Vec<I> {
636    fn to_bytes(&self) -> Vec<u8> {
637        instructions_to_bytes(self)
638    }
639    
640    fn total_size(&self) -> usize {
641        instructions_total_size(self)
642    }
643}
644
645/// Architecture-specific encoding functions
646pub trait ArchitectureEncoder<I: Instruction> {
647    /// Encode an instruction with specific opcode and operands
648    fn encode(&self, opcode: u32, operands: &[u32]) -> I;
649}
650
651/// Common result type for instruction building
652pub type BuildResult<I> = Result<I, BuildError>;
653
654/// Errors that can occur during instruction building
655#[derive(Debug, Clone, Copy, PartialEq, Eq)]
656pub enum BuildError {
657    /// Invalid register identifier
658    InvalidRegister(u32),
659    /// Invalid immediate value
660    InvalidImmediate(i64),
661    /// Unsupported instruction
662    UnsupportedInstruction,
663    /// Invalid operand combination
664    InvalidOperands,
665}
666
667impl fmt::Display for BuildError {
668    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
669        match self {
670            BuildError::InvalidRegister(id) => write!(f, "Invalid register ID: {}", id),
671            BuildError::InvalidImmediate(val) => write!(f, "Invalid immediate value: {}", val),
672            BuildError::UnsupportedInstruction => write!(f, "Unsupported instruction"),
673            BuildError::InvalidOperands => write!(f, "Invalid operand combination"),
674        }
675    }
676}
677
678#[cfg(feature = "std")]
679impl std::error::Error for BuildError {}
680
681/// JIT execution functionality (std-only)
682#[cfg(feature = "std")]
683pub mod jit {
684    use std::marker::PhantomData;
685    use jit_allocator2::JitAllocator;
686
687    /// A raw JIT-compiled function for manual type conversion
688    /// 
689    /// This is a type-erased version of `CallableJitFunction` that allows
690    /// manual function pointer conversion for cases not covered by the
691    /// predefined `call()` methods.
692    /// 
693    /// # Usage
694    /// 
695    /// ```rust,no_run
696    /// # use jit_assembler::common::jit::*;
697    /// # let code = &[0u8; 4];
698    /// let raw_func = RawCallableJitFunction::new(code)?;
699    /// 
700    /// // Manual conversion to any function type
701    /// let func: extern "C" fn(u64, u64, u64, u64, u64, u64, u64, u64) -> u64 = 
702    ///     unsafe { raw_func.as_fn() };
703    /// let result = func(1, 2, 3, 4, 5, 6, 7, 8);
704    /// # Ok::<(), JitError>(())
705    /// ```
706    pub struct RawCallableJitFunction {
707        _allocator: Box<JitAllocator>,
708        exec_ptr: *const u8,
709    }
710
711    impl RawCallableJitFunction {
712        /// Create a new raw callable JIT function from instruction bytes
713        pub fn new(code: &[u8]) -> Result<Self, JitError> {
714            let mut allocator = JitAllocator::new(Default::default());
715            let (exec_ptr, mut_ptr) = allocator.alloc(code.len()).map_err(JitError::AllocationFailed)?;
716            
717            unsafe {
718                std::ptr::copy_nonoverlapping(code.as_ptr(), mut_ptr, code.len());
719            }
720
721            Ok(RawCallableJitFunction {
722                _allocator: allocator,
723                exec_ptr,
724            })
725        }
726
727        /// Convert to a function pointer of the specified type
728        /// 
729        /// # Safety
730        /// 
731        /// The caller must ensure that:
732        /// - The function signature `F` matches the actual assembled code
733        /// - The assembled code follows the expected calling convention (usually C ABI)
734        /// - The function pointer is called with valid arguments
735        /// 
736        /// # Example
737        /// 
738        /// ```rust,no_run
739        /// # use jit_assembler::common::jit::*;
740        /// # let code = &[0u8; 4];
741        /// let raw_func = RawCallableJitFunction::new(code)?;
742        /// 
743        /// // For complex signatures not covered by CallableJitFunction
744        /// let func: extern "C" fn(u64, u64, u64, u64, u64, u64, u64, u64) -> u64 = 
745        ///     unsafe { raw_func.as_fn() };
746        /// let result = func(1, 2, 3, 4, 5, 6, 7, 8);
747        /// # Ok::<(), JitError>(())
748        /// ```
749        pub unsafe fn as_fn<F>(&self) -> F {
750            std::mem::transmute_copy(&self.exec_ptr)
751        }
752    }
753
754    /// A JIT-compiled function that can be called directly
755    /// 
756    /// This structure wraps executable machine code and provides type-safe
757    /// calling methods. While the type parameter `F` represents a Rust function
758    /// signature, the actual execution uses C ABI for stability.
759    /// 
760    /// # Supported Signatures
761    /// 
762    /// Currently supports function signatures with 0-7 arguments:
763    /// - `fn() -> R`
764    /// - `fn(A1) -> R`  
765    /// - `fn(A1, A2) -> R`
766    /// - ... up to `fn(A1, A2, A3, A4, A5, A6, A7) -> R`
767    /// 
768    /// For unsupported signatures, use `RawCallableJitFunction` instead.
769    pub struct CallableJitFunction<F> {
770        _allocator: Box<JitAllocator>,
771        exec_ptr: *const u8,
772        _phantom: PhantomData<F>,
773    }
774
775    impl<F> CallableJitFunction<F> {
776        /// Create a new callable JIT function from instruction bytes
777        pub fn new(code: &[u8]) -> Result<Self, JitError> {
778            let mut allocator = JitAllocator::new(Default::default());
779            let (exec_ptr, mut_ptr) = allocator.alloc(code.len()).map_err(JitError::AllocationFailed)?;
780            
781            unsafe {
782                std::ptr::copy_nonoverlapping(code.as_ptr(), mut_ptr, code.len());
783            }
784
785            Ok(CallableJitFunction {
786                _allocator: allocator,
787                exec_ptr,
788                _phantom: PhantomData,
789            })
790        }
791    }
792
793    /// Direct call methods based on function signature
794    /// 
795    /// These methods provide type-safe calling with automatic ABI conversion from
796    /// Rust function types to C ABI. Currently supports 0-7 arguments.
797    /// 
798    /// The conversion from `fn(...)` to `extern "C" fn(...)` is handled internally
799    /// for ABI stability across Rust versions.
800    
801    impl<R> CallableJitFunction<fn() -> R> {
802        /// Call with no arguments - natural syntax: func.call()
803        pub fn call(&self) -> R {
804            let func: extern "C" fn() -> R = unsafe { std::mem::transmute_copy(&self.exec_ptr) };
805            func()
806        }
807    }
808
809    impl<A1, R> CallableJitFunction<fn(A1) -> R> {
810        /// Call with one argument - natural syntax: func.call(arg)
811        pub fn call(&self, arg1: A1) -> R {
812            let func: extern "C" fn(A1) -> R = unsafe { std::mem::transmute_copy(&self.exec_ptr) };
813            func(arg1)
814        }
815    }
816
817    impl<A1, A2, R> CallableJitFunction<fn(A1, A2) -> R> {
818        /// Call with two arguments - natural syntax: func.call(arg1, arg2)
819        pub fn call(&self, arg1: A1, arg2: A2) -> R {
820            let func: extern "C" fn(A1, A2) -> R = unsafe { std::mem::transmute_copy(&self.exec_ptr) };
821            func(arg1, arg2)
822        }
823    }
824
825    impl<A1, A2, A3, R> CallableJitFunction<fn(A1, A2, A3) -> R> {
826        /// Call with three arguments - natural syntax: func.call(arg1, arg2, arg3)
827        pub fn call(&self, arg1: A1, arg2: A2, arg3: A3) -> R {
828            let func: extern "C" fn(A1, A2, A3) -> R = unsafe { std::mem::transmute_copy(&self.exec_ptr) };
829            func(arg1, arg2, arg3)
830        }
831    }
832
833    impl<A1, A2, A3, A4, R> CallableJitFunction<fn(A1, A2, A3, A4) -> R> {
834        /// Call with four arguments - natural syntax: func.call(arg1, arg2, arg3, arg4)
835        pub fn call(&self, arg1: A1, arg2: A2, arg3: A3, arg4: A4) -> R {
836            let func: extern "C" fn(A1, A2, A3, A4) -> R = unsafe { std::mem::transmute_copy(&self.exec_ptr) };
837            func(arg1, arg2, arg3, arg4)
838        }
839    }
840
841    impl<A1, A2, A3, A4, A5, R> CallableJitFunction<fn(A1, A2, A3, A4, A5) -> R> {
842        /// Call with five arguments
843        pub fn call(&self, arg1: A1, arg2: A2, arg3: A3, arg4: A4, arg5: A5) -> R {
844            let func: extern "C" fn(A1, A2, A3, A4, A5) -> R = unsafe { std::mem::transmute_copy(&self.exec_ptr) };
845            func(arg1, arg2, arg3, arg4, arg5)
846        }
847    }
848
849    impl<A1, A2, A3, A4, A5, A6, R> CallableJitFunction<fn(A1, A2, A3, A4, A5, A6) -> R> {
850        /// Call with six arguments
851        pub fn call(&self, arg1: A1, arg2: A2, arg3: A3, arg4: A4, arg5: A5, arg6: A6) -> R {
852            let func: extern "C" fn(A1, A2, A3, A4, A5, A6) -> R = unsafe { std::mem::transmute_copy(&self.exec_ptr) };
853            func(arg1, arg2, arg3, arg4, arg5, arg6)
854        }
855    }
856
857    impl<A1, A2, A3, A4, A5, A6, A7, R> CallableJitFunction<fn(A1, A2, A3, A4, A5, A6, A7) -> R> {
858        /// Call with seven arguments
859        pub fn call(&self, arg1: A1, arg2: A2, arg3: A3, arg4: A4, arg5: A5, arg6: A6, arg7: A7) -> R {
860            let func: extern "C" fn(A1, A2, A3, A4, A5, A6, A7) -> R = unsafe { std::mem::transmute_copy(&self.exec_ptr) };
861            func(arg1, arg2, arg3, arg4, arg5, arg6, arg7)
862        }
863    }
864
865    // Note: For void return functions, we don't generate them here 
866    // as they would need special handling with unit type ()
867
868    /// Errors that can occur during JIT execution
869    #[derive(Debug)]
870    pub enum JitError {
871        AllocationFailed(jit_allocator2::Error),
872    }
873
874    impl std::fmt::Display for JitError {
875        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
876            match self {
877                JitError::AllocationFailed(e) => write!(f, "Failed to allocate JIT memory: {:?}", e),
878            }
879        }
880    }
881
882    impl std::error::Error for JitError {}
883}
884
885#[cfg(test)]
886mod tests {
887    use super::*;
888    
889    // Test instruction type for unit tests
890    #[derive(Debug, Clone, Copy, PartialEq, Eq)]
891    struct TestInstruction(u32);
892    
893    impl fmt::Display for TestInstruction {
894        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
895            write!(f, "TestInstruction(0x{:08x})", self.0)
896        }
897    }
898    
899    impl Instruction for TestInstruction {
900        fn value(&self) -> u64 {
901            self.0 as u64
902        }
903        
904        fn bytes(&self) -> Vec<u8> {
905            self.0.to_le_bytes().to_vec()
906        }
907        
908        fn size(&self) -> usize {
909            4
910        }
911    }
912    
913    #[test]
914    fn test_instruction_collection_append() {
915        let mut collection1 = InstructionCollection::from_slice(&[
916            TestInstruction(1),
917            TestInstruction(2),
918        ]);
919        let collection2 = InstructionCollection::from_slice(&[
920            TestInstruction(3),
921            TestInstruction(4),
922        ]);
923        
924        collection1.append(collection2);
925        
926        assert_eq!(collection1.len(), 4);
927        assert_eq!(collection1[0], TestInstruction(1));
928        assert_eq!(collection1[1], TestInstruction(2));
929        assert_eq!(collection1[2], TestInstruction(3));
930        assert_eq!(collection1[3], TestInstruction(4));
931    }
932    
933    #[test]
934    fn test_instruction_collection_extend_from_collection() {
935        let mut collection1 = InstructionCollection::from_slice(&[
936            TestInstruction(1),
937            TestInstruction(2),
938        ]);
939        let collection2 = InstructionCollection::from_slice(&[
940            TestInstruction(3),
941            TestInstruction(4),
942        ]);
943        
944        collection1.extend_from_collection(&collection2);
945        
946        // collection1 should have all 4 instructions
947        assert_eq!(collection1.len(), 4);
948        assert_eq!(collection1[0], TestInstruction(1));
949        assert_eq!(collection1[1], TestInstruction(2));
950        assert_eq!(collection1[2], TestInstruction(3));
951        assert_eq!(collection1[3], TestInstruction(4));
952        
953        // collection2 should still be valid
954        assert_eq!(collection2.len(), 2);
955        assert_eq!(collection2[0], TestInstruction(3));
956        assert_eq!(collection2[1], TestInstruction(4));
957    }
958    
959    #[test]
960    fn test_instruction_collection_concat() {
961        let collection1 = InstructionCollection::from_slice(&[
962            TestInstruction(1),
963            TestInstruction(2),
964        ]);
965        let collection2 = InstructionCollection::from_slice(&[
966            TestInstruction(3),
967            TestInstruction(4),
968        ]);
969        
970        let combined = collection1.concat(collection2);
971        
972        assert_eq!(combined.len(), 4);
973        assert_eq!(combined[0], TestInstruction(1));
974        assert_eq!(combined[1], TestInstruction(2));
975        assert_eq!(combined[2], TestInstruction(3));
976        assert_eq!(combined[3], TestInstruction(4));
977    }
978    
979    #[test]
980    fn test_instruction_collection_add_operator() {
981        let collection1 = InstructionCollection::from_slice(&[
982            TestInstruction(1),
983            TestInstruction(2),
984        ]);
985        let collection2 = InstructionCollection::from_slice(&[
986            TestInstruction(3),
987            TestInstruction(4),
988        ]);
989        
990        let combined = collection1 + collection2;
991        
992        assert_eq!(combined.len(), 4);
993        assert_eq!(combined[0], TestInstruction(1));
994        assert_eq!(combined[1], TestInstruction(2));
995        assert_eq!(combined[2], TestInstruction(3));
996        assert_eq!(combined[3], TestInstruction(4));
997    }
998    
999    #[test]
1000    fn test_instruction_collection_add_assign_operator() {
1001        let mut collection1 = InstructionCollection::from_slice(&[
1002            TestInstruction(1),
1003            TestInstruction(2),
1004        ]);
1005        let collection2 = InstructionCollection::from_slice(&[
1006            TestInstruction(3),
1007            TestInstruction(4),
1008        ]);
1009        
1010        collection1 += collection2;
1011        
1012        assert_eq!(collection1.len(), 4);
1013        assert_eq!(collection1[0], TestInstruction(1));
1014        assert_eq!(collection1[1], TestInstruction(2));
1015        assert_eq!(collection1[2], TestInstruction(3));
1016        assert_eq!(collection1[3], TestInstruction(4));
1017    }
1018    
1019    #[test]
1020    fn test_instruction_collection_multiple_merge() {
1021        let collection1 = InstructionCollection::from_slice(&[TestInstruction(1)]);
1022        let collection2 = InstructionCollection::from_slice(&[TestInstruction(2)]);
1023        let collection3 = InstructionCollection::from_slice(&[TestInstruction(3)]);
1024        
1025        // Test chaining with + operator
1026        let combined = collection1 + collection2 + collection3;
1027        
1028        assert_eq!(combined.len(), 3);
1029        assert_eq!(combined[0], TestInstruction(1));
1030        assert_eq!(combined[1], TestInstruction(2));
1031        assert_eq!(combined[2], TestInstruction(3));
1032    }
1033    
1034    #[test]
1035    fn test_instruction_collection_merge_empty() {
1036        let mut collection1 = InstructionCollection::from_slice(&[
1037            TestInstruction(1),
1038            TestInstruction(2),
1039        ]);
1040        let collection2 = InstructionCollection::<TestInstruction>::new();
1041        
1042        collection1.append(collection2);
1043        
1044        assert_eq!(collection1.len(), 2);
1045        assert_eq!(collection1[0], TestInstruction(1));
1046        assert_eq!(collection1[1], TestInstruction(2));
1047    }
1048    
1049    #[test]
1050    fn test_instruction_collection_merge_into_empty() {
1051        let mut collection1 = InstructionCollection::<TestInstruction>::new();
1052        let collection2 = InstructionCollection::from_slice(&[
1053            TestInstruction(3),
1054            TestInstruction(4),
1055        ]);
1056        
1057        collection1.append(collection2);
1058        
1059        assert_eq!(collection1.len(), 2);
1060        assert_eq!(collection1[0], TestInstruction(3));
1061        assert_eq!(collection1[1], TestInstruction(4));
1062    }
1063}
1064
1065#[cfg(all(test, feature = "register-tracking"))]
1066mod register_tracking_tests {
1067    use super::*;
1068    
1069    // Test instruction and register types for unit tests
1070    #[derive(Debug, Clone, Copy, PartialEq, Eq)]
1071    struct TestInstruction(u32);
1072    
1073    impl fmt::Display for TestInstruction {
1074        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1075            write!(f, "TestInstruction(0x{:08x})", self.0)
1076        }
1077    }
1078    
1079    impl Instruction for TestInstruction {
1080        fn value(&self) -> u64 {
1081            self.0 as u64
1082        }
1083        
1084        fn bytes(&self) -> Vec<u8> {
1085            self.0.to_le_bytes().to_vec()
1086        }
1087        
1088        fn size(&self) -> usize {
1089            4
1090        }
1091    }
1092    
1093    #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
1094    enum TestRegister {
1095        T0, T1,    // Caller-saved
1096        S0, S1,    // Callee-saved
1097        SP, FP,    // Special
1098    }
1099    
1100    impl Register for TestRegister {
1101        fn id(&self) -> u32 {
1102            match self {
1103                TestRegister::T0 => 0,
1104                TestRegister::T1 => 1,
1105                TestRegister::S0 => 2,
1106                TestRegister::S1 => 3,
1107                TestRegister::SP => 4,
1108                TestRegister::FP => 5,
1109            }
1110        }
1111        
1112        fn abi_class(&self) -> AbiClass {
1113            match self {
1114                TestRegister::T0 | TestRegister::T1 => AbiClass::CallerSaved,
1115                TestRegister::S0 | TestRegister::S1 => AbiClass::CalleeSaved,
1116                TestRegister::SP | TestRegister::FP => AbiClass::Special,
1117            }
1118        }
1119    }
1120    
1121    #[test]
1122    fn test_instruction_collection_with_usage_new() {
1123        let instructions = InstructionCollection::from_slice(&[
1124            TestInstruction(1),
1125            TestInstruction(2),
1126        ]);
1127        let mut usage = register_usage::RegisterUsageInfo::new();
1128        usage.add_written_register(TestRegister::T0);
1129        usage.add_read_register(TestRegister::T1);
1130        
1131        let tracked = InstructionCollectionWithUsage::new(instructions.clone(), usage.clone());
1132        
1133        assert_eq!(tracked.instructions().len(), 2);
1134        assert_eq!(tracked.register_usage().register_count(), 2);
1135    }
1136    
1137    #[test]
1138    fn test_instruction_collection_with_usage_append() {
1139        let instructions1 = InstructionCollection::from_slice(&[TestInstruction(1)]);
1140        let mut usage1 = register_usage::RegisterUsageInfo::new();
1141        usage1.add_written_register(TestRegister::T0);
1142        
1143        let instructions2 = InstructionCollection::from_slice(&[TestInstruction(2)]);
1144        let mut usage2 = register_usage::RegisterUsageInfo::new();
1145        usage2.add_read_register(TestRegister::S0);
1146        
1147        let mut tracked1 = InstructionCollectionWithUsage::new(instructions1, usage1);
1148        let tracked2 = InstructionCollectionWithUsage::new(instructions2, usage2);
1149        
1150        tracked1.append(tracked2);
1151        
1152        assert_eq!(tracked1.instructions().len(), 2);
1153        assert_eq!(tracked1.register_usage().register_count(), 2);
1154        assert!(tracked1.register_usage().contains_written_register(&TestRegister::T0));
1155        assert!(tracked1.register_usage().contains_read_register(&TestRegister::S0));
1156    }
1157    
1158    #[test]
1159    fn test_instruction_collection_with_usage_concat() {
1160        let instructions1 = InstructionCollection::from_slice(&[TestInstruction(1)]);
1161        let mut usage1 = register_usage::RegisterUsageInfo::new();
1162        usage1.add_written_register(TestRegister::T0);
1163        
1164        let instructions2 = InstructionCollection::from_slice(&[TestInstruction(2)]);
1165        let mut usage2 = register_usage::RegisterUsageInfo::new();
1166        usage2.add_read_register(TestRegister::S0);
1167        
1168        let tracked1 = InstructionCollectionWithUsage::new(instructions1, usage1);
1169        let tracked2 = InstructionCollectionWithUsage::new(instructions2, usage2);
1170        
1171        let combined = tracked1.concat(tracked2);
1172        
1173        assert_eq!(combined.instructions().len(), 2);
1174        assert_eq!(combined.register_usage().register_count(), 2);
1175        assert!(combined.register_usage().contains_written_register(&TestRegister::T0));
1176        assert!(combined.register_usage().contains_read_register(&TestRegister::S0));
1177    }
1178    
1179    #[test]
1180    fn test_instruction_collection_with_usage_add_operator() {
1181        let instructions1 = InstructionCollection::from_slice(&[TestInstruction(1)]);
1182        let mut usage1 = register_usage::RegisterUsageInfo::new();
1183        usage1.add_written_register(TestRegister::T0);
1184        
1185        let instructions2 = InstructionCollection::from_slice(&[TestInstruction(2)]);
1186        let mut usage2 = register_usage::RegisterUsageInfo::new();
1187        usage2.add_read_register(TestRegister::S0);
1188        
1189        let tracked1 = InstructionCollectionWithUsage::new(instructions1, usage1);
1190        let tracked2 = InstructionCollectionWithUsage::new(instructions2, usage2);
1191        
1192        let combined = tracked1 + tracked2;
1193        
1194        assert_eq!(combined.instructions().len(), 2);
1195        assert_eq!(combined.register_usage().register_count(), 2);
1196        assert!(combined.register_usage().contains_written_register(&TestRegister::T0));
1197        assert!(combined.register_usage().contains_read_register(&TestRegister::S0));
1198    }
1199    
1200    #[test]
1201    fn test_instruction_collection_with_usage_multiple_merge() {
1202        // Simulate the use case from the issue: prologue + main + epilogue
1203        
1204        // Prologue: save registers
1205        let prologue_instrs = InstructionCollection::from_slice(&[TestInstruction(0x10)]);
1206        let mut prologue_usage = register_usage::RegisterUsageInfo::new();
1207        prologue_usage.add_read_register(TestRegister::SP);
1208        prologue_usage.add_read_register(TestRegister::S0);
1209        
1210        // Main: actual computation
1211        let main_instrs = InstructionCollection::from_slice(&[
1212            TestInstruction(0x20),
1213            TestInstruction(0x21),
1214        ]);
1215        let mut main_usage = register_usage::RegisterUsageInfo::new();
1216        main_usage.add_written_register(TestRegister::T0);
1217        main_usage.add_written_register(TestRegister::S0);
1218        main_usage.add_read_register(TestRegister::T1);
1219        
1220        // Epilogue: restore registers
1221        let epilogue_instrs = InstructionCollection::from_slice(&[TestInstruction(0x30)]);
1222        let mut epilogue_usage = register_usage::RegisterUsageInfo::new();
1223        epilogue_usage.add_read_register(TestRegister::SP);
1224        epilogue_usage.add_written_register(TestRegister::S0);
1225        
1226        let prologue = InstructionCollectionWithUsage::new(prologue_instrs, prologue_usage);
1227        let main = InstructionCollectionWithUsage::new(main_instrs, main_usage);
1228        let epilogue = InstructionCollectionWithUsage::new(epilogue_instrs, epilogue_usage);
1229        
1230        // Merge: prologue + main + epilogue
1231        let combined = prologue + main + epilogue;
1232        
1233        // Check instructions
1234        assert_eq!(combined.instructions().len(), 4);
1235        assert_eq!(combined.instructions()[0], TestInstruction(0x10));
1236        assert_eq!(combined.instructions()[1], TestInstruction(0x20));
1237        assert_eq!(combined.instructions()[2], TestInstruction(0x21));
1238        assert_eq!(combined.instructions()[3], TestInstruction(0x30));
1239        
1240        // Check register usage is properly merged
1241        let usage = combined.register_usage();
1242        assert_eq!(usage.register_count(), 4);  // T0, T1, S0, SP
1243        assert!(usage.contains_register(&TestRegister::T0));
1244        assert!(usage.contains_register(&TestRegister::T1));
1245        assert!(usage.contains_register(&TestRegister::S0));
1246        assert!(usage.contains_register(&TestRegister::SP));
1247        
1248        // Check that S0 appears as both written and read
1249        assert!(usage.contains_written_register(&TestRegister::S0));
1250        assert!(usage.contains_read_register(&TestRegister::S0));
1251    }
1252    
1253    #[test]
1254    fn test_instruction_collection_with_usage_into_parts() {
1255        let instructions = InstructionCollection::from_slice(&[TestInstruction(1)]);
1256        let mut usage = register_usage::RegisterUsageInfo::new();
1257        usage.add_written_register(TestRegister::T0);
1258        
1259        let tracked = InstructionCollectionWithUsage::new(instructions.clone(), usage.clone());
1260        let (instrs, reg_usage) = tracked.into_parts();
1261        
1262        assert_eq!(instrs.len(), 1);
1263        assert_eq!(reg_usage.register_count(), 1);
1264    }
1265    
1266    #[test]
1267    fn test_instruction_collection_with_usage_extend_from() {
1268        let instructions1 = InstructionCollection::from_slice(&[TestInstruction(1)]);
1269        let mut usage1 = register_usage::RegisterUsageInfo::new();
1270        usage1.add_written_register(TestRegister::T0);
1271        
1272        let instructions2 = InstructionCollection::from_slice(&[TestInstruction(2)]);
1273        let mut usage2 = register_usage::RegisterUsageInfo::new();
1274        usage2.add_read_register(TestRegister::S0);
1275        
1276        let mut tracked1 = InstructionCollectionWithUsage::new(instructions1, usage1);
1277        let tracked2 = InstructionCollectionWithUsage::new(instructions2, usage2);
1278        
1279        tracked1.extend_from(&tracked2);
1280        
1281        // tracked1 should be extended
1282        assert_eq!(tracked1.instructions().len(), 2);
1283        assert_eq!(tracked1.register_usage().register_count(), 2);
1284        
1285        // tracked2 should still exist
1286        assert_eq!(tracked2.instructions().len(), 1);
1287        assert_eq!(tracked2.register_usage().register_count(), 1);
1288    }
1289}