From 0819561fe72109193a8e5fe05845dd7119e47da9 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 2 Oct 2025 00:25:18 +0000 Subject: [PATCH 1/3] Initial plan From 94b3c9fceeb3cbfd03e44770a4383375a133d7a2 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 2 Oct 2025 00:32:23 +0000 Subject: [PATCH 2/3] Add InstructionCollection merge functionality with register tracking support Co-authored-by: petitstrawberry <14030185+petitstrawberry@users.noreply.github.com> --- src/common/mod.rs | 629 ++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 629 insertions(+) diff --git a/src/common/mod.rs b/src/common/mod.rs index 9e0443a..9685b7c 100644 --- a/src/common/mod.rs +++ b/src/common/mod.rs @@ -6,11 +6,149 @@ use core::fmt; #[cfg(feature = "register-tracking")] pub mod register_usage; +#[cfg(feature = "register-tracking")] +use register_usage::RegisterUsageInfo; + #[cfg(feature = "std")] use std::vec::Vec; #[cfg(not(feature = "std"))] use alloc::vec::Vec; +/// A collection of instructions with associated register usage information. +/// +/// This struct combines an `InstructionCollection` with `RegisterUsageInfo`, +/// allowing you to merge both instructions and their usage statistics together. +/// This is useful when combining code sequences from multiple builders while +/// preserving register usage information. +/// +/// # Example +/// +/// ```rust,ignore +/// // Create three builders for different parts of the code +/// let mut prologue = Riscv64InstructionBuilder::new(); +/// prologue.sd(reg::SP, reg::S0, -8); // Save S0 +/// +/// let mut main = Riscv64InstructionBuilder::new(); +/// main.add(reg::S0, reg::A0, reg::A1); // Main computation +/// +/// let mut epilogue = Riscv64InstructionBuilder::new(); +/// epilogue.ld(reg::S0, reg::SP, -8); // Restore S0 +/// +/// // Create tracked collections +/// let prologue_tracked = InstructionCollectionWithUsage::new( +/// prologue.instructions(), +/// prologue.register_usage().clone() +/// ); +/// let main_tracked = InstructionCollectionWithUsage::new( +/// main.instructions(), +/// main.register_usage().clone() +/// ); +/// let epilogue_tracked = InstructionCollectionWithUsage::new( +/// epilogue.instructions(), +/// epilogue.register_usage().clone() +/// ); +/// +/// // Merge them: prologue + main + epilogue +/// let combined = prologue_tracked + main_tracked + epilogue_tracked; +/// +/// // Now we have the complete function with accurate register usage +/// let instructions = combined.instructions(); +/// let usage = combined.register_usage(); +/// ``` +#[cfg(feature = "register-tracking")] +#[derive(Debug, Clone)] +pub struct InstructionCollectionWithUsage { + instructions: InstructionCollection, + register_usage: RegisterUsageInfo, +} + +#[cfg(feature = "register-tracking")] +impl InstructionCollectionWithUsage { + /// Create a new tracked instruction collection. + pub fn new(instructions: InstructionCollection, register_usage: RegisterUsageInfo) -> Self { + Self { + instructions, + register_usage, + } + } + + /// Create from raw parts. + pub fn from_parts(instructions: InstructionCollection, register_usage: RegisterUsageInfo) -> Self { + Self::new(instructions, register_usage) + } + + /// Get a reference to the instructions. + pub fn instructions(&self) -> &InstructionCollection { + &self.instructions + } + + /// Get a mutable reference to the instructions. + pub fn instructions_mut(&mut self) -> &mut InstructionCollection { + &mut self.instructions + } + + /// Get a reference to the register usage information. + pub fn register_usage(&self) -> &RegisterUsageInfo { + &self.register_usage + } + + /// Get a mutable reference to the register usage information. + pub fn register_usage_mut(&mut self) -> &mut RegisterUsageInfo { + &mut self.register_usage + } + + /// Consume this collection and return the instructions and register usage. + pub fn into_parts(self) -> (InstructionCollection, RegisterUsageInfo) { + (self.instructions, self.register_usage) + } + + /// Consume this collection and return just the instructions. + pub fn into_instructions(self) -> InstructionCollection { + self.instructions + } + + /// Merge another tracked collection into this one. + /// + /// This appends the instructions and merges the register usage information. + pub fn append(&mut self, other: InstructionCollectionWithUsage) { + self.instructions.append(other.instructions); + self.register_usage.merge(&other.register_usage); + } + + /// Extend this collection with instructions and register usage from another. + /// + /// This clones instructions and merges the register usage information. + pub fn extend_from(&mut self, other: &InstructionCollectionWithUsage) { + self.instructions.extend_from_collection(&other.instructions); + self.register_usage.merge(&other.register_usage); + } + + /// Concatenate two tracked collections, consuming both. + pub fn concat(mut self, other: InstructionCollectionWithUsage) -> Self { + self.instructions.append(other.instructions); + self.register_usage.merge(&other.register_usage); + self + } +} + +#[cfg(feature = "register-tracking")] +impl core::ops::Add for InstructionCollectionWithUsage { + type Output = InstructionCollectionWithUsage; + + /// Concatenate two tracked instruction collections using the `+` operator. + fn add(self, other: InstructionCollectionWithUsage) -> InstructionCollectionWithUsage { + self.concat(other) + } +} + +#[cfg(feature = "register-tracking")] +impl core::ops::AddAssign for InstructionCollectionWithUsage { + /// Append another tracked instruction collection using the `+=` operator. + fn add_assign(&mut self, other: InstructionCollectionWithUsage) { + self.append(other); + } +} + /// A machine instruction that can be encoded to bytes pub trait Instruction: Copy + Clone + fmt::Debug + fmt::Display { /// Get the instruction as a 32-bit or 64-bit value @@ -302,6 +440,57 @@ impl InstructionCollection { pub fn get_mut(&mut self, index: usize) -> Option<&mut I> { self.instructions.get_mut(index) } + + /// Append another instruction collection to this one (consumes other). + /// + /// This moves all instructions from `other` into this collection. + /// + /// # Example + /// + /// ```rust,ignore + /// let mut collection1 = InstructionCollection::from_slice(&[instr1, instr2]); + /// let collection2 = InstructionCollection::from_slice(&[instr3, instr4]); + /// collection1.append(collection2); + /// // collection1 now contains [instr1, instr2, instr3, instr4] + /// ``` + pub fn append(&mut self, mut other: InstructionCollection) { + self.instructions.append(&mut other.instructions); + } + + /// Extend this collection with instructions from another collection. + /// + /// This clones instructions from `other` into this collection. + /// + /// # Example + /// + /// ```rust,ignore + /// let mut collection1 = InstructionCollection::from_slice(&[instr1, instr2]); + /// let collection2 = InstructionCollection::from_slice(&[instr3, instr4]); + /// collection1.extend_from_collection(&collection2); + /// // collection1 now contains [instr1, instr2, instr3, instr4] + /// // collection2 is still valid + /// ``` + pub fn extend_from_collection(&mut self, other: &InstructionCollection) { + self.instructions.extend_from_slice(&other.instructions); + } + + /// Concatenate two instruction collections, consuming both. + /// + /// Creates a new collection containing all instructions from `self` + /// followed by all instructions from `other`. + /// + /// # Example + /// + /// ```rust,ignore + /// let collection1 = InstructionCollection::from_slice(&[instr1, instr2]); + /// let collection2 = InstructionCollection::from_slice(&[instr3, instr4]); + /// let combined = collection1.concat(collection2); + /// // combined contains [instr1, instr2, instr3, instr4] + /// ``` + pub fn concat(mut self, mut other: InstructionCollection) -> Self { + self.instructions.append(&mut other.instructions); + self + } } impl Default for InstructionCollection { @@ -389,6 +578,40 @@ impl core::ops::DerefMut for InstructionCollection { } } +impl core::ops::Add for InstructionCollection { + type Output = InstructionCollection; + + /// Concatenate two instruction collections using the `+` operator. + /// + /// # Example + /// + /// ```rust,ignore + /// let collection1 = InstructionCollection::from_slice(&[instr1, instr2]); + /// let collection2 = InstructionCollection::from_slice(&[instr3, instr4]); + /// let combined = collection1 + collection2; + /// // combined contains [instr1, instr2, instr3, instr4] + /// ``` + fn add(self, other: InstructionCollection) -> InstructionCollection { + self.concat(other) + } +} + +impl core::ops::AddAssign for InstructionCollection { + /// Append another instruction collection to this one using the `+=` operator. + /// + /// # Example + /// + /// ```rust,ignore + /// let mut collection1 = InstructionCollection::from_slice(&[instr1, instr2]); + /// let collection2 = InstructionCollection::from_slice(&[instr3, instr4]); + /// collection1 += collection2; + /// // collection1 now contains [instr1, instr2, instr3, instr4] + /// ``` + fn add_assign(&mut self, other: InstructionCollection) { + self.append(other); + } +} + /// Trait extension for instruction collections /// This allows you to call `.to_bytes()` and `.total_size()` directly on slices and vectors pub trait InstructionCollectionExt { @@ -657,4 +880,410 @@ pub mod jit { } impl std::error::Error for JitError {} +} + +#[cfg(test)] +mod tests { + use super::*; + + // Test instruction type for unit tests + #[derive(Debug, Clone, Copy, PartialEq, Eq)] + struct TestInstruction(u32); + + impl fmt::Display for TestInstruction { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "TestInstruction(0x{:08x})", self.0) + } + } + + impl Instruction for TestInstruction { + fn value(&self) -> u64 { + self.0 as u64 + } + + fn bytes(&self) -> Vec { + self.0.to_le_bytes().to_vec() + } + + fn size(&self) -> usize { + 4 + } + } + + #[test] + fn test_instruction_collection_append() { + let mut collection1 = InstructionCollection::from_slice(&[ + TestInstruction(1), + TestInstruction(2), + ]); + let collection2 = InstructionCollection::from_slice(&[ + TestInstruction(3), + TestInstruction(4), + ]); + + collection1.append(collection2); + + assert_eq!(collection1.len(), 4); + assert_eq!(collection1[0], TestInstruction(1)); + assert_eq!(collection1[1], TestInstruction(2)); + assert_eq!(collection1[2], TestInstruction(3)); + assert_eq!(collection1[3], TestInstruction(4)); + } + + #[test] + fn test_instruction_collection_extend_from_collection() { + let mut collection1 = InstructionCollection::from_slice(&[ + TestInstruction(1), + TestInstruction(2), + ]); + let collection2 = InstructionCollection::from_slice(&[ + TestInstruction(3), + TestInstruction(4), + ]); + + collection1.extend_from_collection(&collection2); + + // collection1 should have all 4 instructions + assert_eq!(collection1.len(), 4); + assert_eq!(collection1[0], TestInstruction(1)); + assert_eq!(collection1[1], TestInstruction(2)); + assert_eq!(collection1[2], TestInstruction(3)); + assert_eq!(collection1[3], TestInstruction(4)); + + // collection2 should still be valid + assert_eq!(collection2.len(), 2); + assert_eq!(collection2[0], TestInstruction(3)); + assert_eq!(collection2[1], TestInstruction(4)); + } + + #[test] + fn test_instruction_collection_concat() { + let collection1 = InstructionCollection::from_slice(&[ + TestInstruction(1), + TestInstruction(2), + ]); + let collection2 = InstructionCollection::from_slice(&[ + TestInstruction(3), + TestInstruction(4), + ]); + + let combined = collection1.concat(collection2); + + assert_eq!(combined.len(), 4); + assert_eq!(combined[0], TestInstruction(1)); + assert_eq!(combined[1], TestInstruction(2)); + assert_eq!(combined[2], TestInstruction(3)); + assert_eq!(combined[3], TestInstruction(4)); + } + + #[test] + fn test_instruction_collection_add_operator() { + let collection1 = InstructionCollection::from_slice(&[ + TestInstruction(1), + TestInstruction(2), + ]); + let collection2 = InstructionCollection::from_slice(&[ + TestInstruction(3), + TestInstruction(4), + ]); + + let combined = collection1 + collection2; + + assert_eq!(combined.len(), 4); + assert_eq!(combined[0], TestInstruction(1)); + assert_eq!(combined[1], TestInstruction(2)); + assert_eq!(combined[2], TestInstruction(3)); + assert_eq!(combined[3], TestInstruction(4)); + } + + #[test] + fn test_instruction_collection_add_assign_operator() { + let mut collection1 = InstructionCollection::from_slice(&[ + TestInstruction(1), + TestInstruction(2), + ]); + let collection2 = InstructionCollection::from_slice(&[ + TestInstruction(3), + TestInstruction(4), + ]); + + collection1 += collection2; + + assert_eq!(collection1.len(), 4); + assert_eq!(collection1[0], TestInstruction(1)); + assert_eq!(collection1[1], TestInstruction(2)); + assert_eq!(collection1[2], TestInstruction(3)); + assert_eq!(collection1[3], TestInstruction(4)); + } + + #[test] + fn test_instruction_collection_multiple_merge() { + let collection1 = InstructionCollection::from_slice(&[TestInstruction(1)]); + let collection2 = InstructionCollection::from_slice(&[TestInstruction(2)]); + let collection3 = InstructionCollection::from_slice(&[TestInstruction(3)]); + + // Test chaining with + operator + let combined = collection1 + collection2 + collection3; + + assert_eq!(combined.len(), 3); + assert_eq!(combined[0], TestInstruction(1)); + assert_eq!(combined[1], TestInstruction(2)); + assert_eq!(combined[2], TestInstruction(3)); + } + + #[test] + fn test_instruction_collection_merge_empty() { + let mut collection1 = InstructionCollection::from_slice(&[ + TestInstruction(1), + TestInstruction(2), + ]); + let collection2 = InstructionCollection::::new(); + + collection1.append(collection2); + + assert_eq!(collection1.len(), 2); + assert_eq!(collection1[0], TestInstruction(1)); + assert_eq!(collection1[1], TestInstruction(2)); + } + + #[test] + fn test_instruction_collection_merge_into_empty() { + let mut collection1 = InstructionCollection::::new(); + let collection2 = InstructionCollection::from_slice(&[ + TestInstruction(3), + TestInstruction(4), + ]); + + collection1.append(collection2); + + assert_eq!(collection1.len(), 2); + assert_eq!(collection1[0], TestInstruction(3)); + assert_eq!(collection1[1], TestInstruction(4)); + } +} + +#[cfg(all(test, feature = "register-tracking"))] +mod register_tracking_tests { + use super::*; + + // Test instruction and register types for unit tests + #[derive(Debug, Clone, Copy, PartialEq, Eq)] + struct TestInstruction(u32); + + impl fmt::Display for TestInstruction { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "TestInstruction(0x{:08x})", self.0) + } + } + + impl Instruction for TestInstruction { + fn value(&self) -> u64 { + self.0 as u64 + } + + fn bytes(&self) -> Vec { + self.0.to_le_bytes().to_vec() + } + + fn size(&self) -> usize { + 4 + } + } + + #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] + enum TestRegister { + T0, T1, // Caller-saved + S0, S1, // Callee-saved + SP, FP, // Special + } + + impl Register for TestRegister { + fn id(&self) -> u32 { + match self { + TestRegister::T0 => 0, + TestRegister::T1 => 1, + TestRegister::S0 => 2, + TestRegister::S1 => 3, + TestRegister::SP => 4, + TestRegister::FP => 5, + } + } + + fn abi_class(&self) -> AbiClass { + match self { + TestRegister::T0 | TestRegister::T1 => AbiClass::CallerSaved, + TestRegister::S0 | TestRegister::S1 => AbiClass::CalleeSaved, + TestRegister::SP | TestRegister::FP => AbiClass::Special, + } + } + } + + #[test] + fn test_instruction_collection_with_usage_new() { + let instructions = InstructionCollection::from_slice(&[ + TestInstruction(1), + TestInstruction(2), + ]); + let mut usage = register_usage::RegisterUsageInfo::new(); + usage.add_written_register(TestRegister::T0); + usage.add_read_register(TestRegister::T1); + + let tracked = InstructionCollectionWithUsage::new(instructions.clone(), usage.clone()); + + assert_eq!(tracked.instructions().len(), 2); + assert_eq!(tracked.register_usage().register_count(), 2); + } + + #[test] + fn test_instruction_collection_with_usage_append() { + let instructions1 = InstructionCollection::from_slice(&[TestInstruction(1)]); + let mut usage1 = register_usage::RegisterUsageInfo::new(); + usage1.add_written_register(TestRegister::T0); + + let instructions2 = InstructionCollection::from_slice(&[TestInstruction(2)]); + let mut usage2 = register_usage::RegisterUsageInfo::new(); + usage2.add_read_register(TestRegister::S0); + + let mut tracked1 = InstructionCollectionWithUsage::new(instructions1, usage1); + let tracked2 = InstructionCollectionWithUsage::new(instructions2, usage2); + + tracked1.append(tracked2); + + assert_eq!(tracked1.instructions().len(), 2); + assert_eq!(tracked1.register_usage().register_count(), 2); + assert!(tracked1.register_usage().contains_written_register(&TestRegister::T0)); + assert!(tracked1.register_usage().contains_read_register(&TestRegister::S0)); + } + + #[test] + fn test_instruction_collection_with_usage_concat() { + let instructions1 = InstructionCollection::from_slice(&[TestInstruction(1)]); + let mut usage1 = register_usage::RegisterUsageInfo::new(); + usage1.add_written_register(TestRegister::T0); + + let instructions2 = InstructionCollection::from_slice(&[TestInstruction(2)]); + let mut usage2 = register_usage::RegisterUsageInfo::new(); + usage2.add_read_register(TestRegister::S0); + + let tracked1 = InstructionCollectionWithUsage::new(instructions1, usage1); + let tracked2 = InstructionCollectionWithUsage::new(instructions2, usage2); + + let combined = tracked1.concat(tracked2); + + assert_eq!(combined.instructions().len(), 2); + assert_eq!(combined.register_usage().register_count(), 2); + assert!(combined.register_usage().contains_written_register(&TestRegister::T0)); + assert!(combined.register_usage().contains_read_register(&TestRegister::S0)); + } + + #[test] + fn test_instruction_collection_with_usage_add_operator() { + let instructions1 = InstructionCollection::from_slice(&[TestInstruction(1)]); + let mut usage1 = register_usage::RegisterUsageInfo::new(); + usage1.add_written_register(TestRegister::T0); + + let instructions2 = InstructionCollection::from_slice(&[TestInstruction(2)]); + let mut usage2 = register_usage::RegisterUsageInfo::new(); + usage2.add_read_register(TestRegister::S0); + + let tracked1 = InstructionCollectionWithUsage::new(instructions1, usage1); + let tracked2 = InstructionCollectionWithUsage::new(instructions2, usage2); + + let combined = tracked1 + tracked2; + + assert_eq!(combined.instructions().len(), 2); + assert_eq!(combined.register_usage().register_count(), 2); + assert!(combined.register_usage().contains_written_register(&TestRegister::T0)); + assert!(combined.register_usage().contains_read_register(&TestRegister::S0)); + } + + #[test] + fn test_instruction_collection_with_usage_multiple_merge() { + // Simulate the use case from the issue: prologue + main + epilogue + + // Prologue: save registers + let prologue_instrs = InstructionCollection::from_slice(&[TestInstruction(0x10)]); + let mut prologue_usage = register_usage::RegisterUsageInfo::new(); + prologue_usage.add_read_register(TestRegister::SP); + prologue_usage.add_read_register(TestRegister::S0); + + // Main: actual computation + let main_instrs = InstructionCollection::from_slice(&[ + TestInstruction(0x20), + TestInstruction(0x21), + ]); + let mut main_usage = register_usage::RegisterUsageInfo::new(); + main_usage.add_written_register(TestRegister::T0); + main_usage.add_written_register(TestRegister::S0); + main_usage.add_read_register(TestRegister::T1); + + // Epilogue: restore registers + let epilogue_instrs = InstructionCollection::from_slice(&[TestInstruction(0x30)]); + let mut epilogue_usage = register_usage::RegisterUsageInfo::new(); + epilogue_usage.add_read_register(TestRegister::SP); + epilogue_usage.add_written_register(TestRegister::S0); + + let prologue = InstructionCollectionWithUsage::new(prologue_instrs, prologue_usage); + let main = InstructionCollectionWithUsage::new(main_instrs, main_usage); + let epilogue = InstructionCollectionWithUsage::new(epilogue_instrs, epilogue_usage); + + // Merge: prologue + main + epilogue + let combined = prologue + main + epilogue; + + // Check instructions + assert_eq!(combined.instructions().len(), 4); + assert_eq!(combined.instructions()[0], TestInstruction(0x10)); + assert_eq!(combined.instructions()[1], TestInstruction(0x20)); + assert_eq!(combined.instructions()[2], TestInstruction(0x21)); + assert_eq!(combined.instructions()[3], TestInstruction(0x30)); + + // Check register usage is properly merged + let usage = combined.register_usage(); + assert_eq!(usage.register_count(), 4); // T0, T1, S0, SP + assert!(usage.contains_register(&TestRegister::T0)); + assert!(usage.contains_register(&TestRegister::T1)); + assert!(usage.contains_register(&TestRegister::S0)); + assert!(usage.contains_register(&TestRegister::SP)); + + // Check that S0 appears as both written and read + assert!(usage.contains_written_register(&TestRegister::S0)); + assert!(usage.contains_read_register(&TestRegister::S0)); + } + + #[test] + fn test_instruction_collection_with_usage_into_parts() { + let instructions = InstructionCollection::from_slice(&[TestInstruction(1)]); + let mut usage = register_usage::RegisterUsageInfo::new(); + usage.add_written_register(TestRegister::T0); + + let tracked = InstructionCollectionWithUsage::new(instructions.clone(), usage.clone()); + let (instrs, reg_usage) = tracked.into_parts(); + + assert_eq!(instrs.len(), 1); + assert_eq!(reg_usage.register_count(), 1); + } + + #[test] + fn test_instruction_collection_with_usage_extend_from() { + let instructions1 = InstructionCollection::from_slice(&[TestInstruction(1)]); + let mut usage1 = register_usage::RegisterUsageInfo::new(); + usage1.add_written_register(TestRegister::T0); + + let instructions2 = InstructionCollection::from_slice(&[TestInstruction(2)]); + let mut usage2 = register_usage::RegisterUsageInfo::new(); + usage2.add_read_register(TestRegister::S0); + + let mut tracked1 = InstructionCollectionWithUsage::new(instructions1, usage1); + let tracked2 = InstructionCollectionWithUsage::new(instructions2, usage2); + + tracked1.extend_from(&tracked2); + + // tracked1 should be extended + assert_eq!(tracked1.instructions().len(), 2); + assert_eq!(tracked1.register_usage().register_count(), 2); + + // tracked2 should still exist + assert_eq!(tracked2.instructions().len(), 1); + assert_eq!(tracked2.register_usage().register_count(), 1); + } } \ No newline at end of file From 3e6d9bcd42f1e6ff66084cde36e9a5b23d39ca3a Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 2 Oct 2025 00:34:41 +0000 Subject: [PATCH 3/3] Add example and documentation for instruction collection merging Co-authored-by: petitstrawberry <14030185+petitstrawberry@users.noreply.github.com> --- README.md | 84 ++++++++++++++ examples/instruction_collection_merge.rs | 135 +++++++++++++++++++++++ 2 files changed, 219 insertions(+) create mode 100644 examples/instruction_collection_merge.rs diff --git a/README.md b/README.md index 2a8d12b..7e591f6 100644 --- a/README.md +++ b/README.md @@ -339,6 +339,90 @@ This information is invaluable for: - **Performance optimization**: Minimize unnecessary register saves/restores - **Code analysis**: Understand register pressure and usage patterns +## Merging Instruction Collections + +When building complex functions, you may need to combine multiple instruction sequences in arbitrary order. For example, you might want to build the main function body first, then add prologue and epilogue code after determining which registers need to be saved. + +### Basic Merging + +```rust +use jit_assembler::riscv64::{reg, Riscv64InstructionBuilder}; +use jit_assembler::common::InstructionBuilder; + +// Build different parts separately +let mut prologue = Riscv64InstructionBuilder::new(); +prologue + .addi(reg::SP, reg::SP, -16) + .sd(reg::SP, reg::RA, 8); + +let mut main_code = Riscv64InstructionBuilder::new(); +main_code + .add(reg::A0, reg::A1, reg::A2) + .mul(reg::A0, reg::A0, reg::A3); + +let mut epilogue = Riscv64InstructionBuilder::new(); +epilogue + .ld(reg::RA, reg::SP, 8) + .addi(reg::SP, reg::SP, 16) + .ret(); + +// Combine them using + operator: prologue + main + epilogue +let combined = prologue.instructions() + main_code.instructions() + epilogue.instructions(); + +// Or use method chaining +let mut combined = prologue.instructions(); +combined += main_code.instructions(); +combined += epilogue.instructions(); + +// Convert to executable code +let bytes = combined.to_bytes(); +``` + +### Merging with Register Tracking + +When the `register-tracking` feature is enabled, you can merge instruction collections while preserving register usage information: + +```rust +use jit_assembler::riscv64::{reg, Riscv64InstructionBuilder}; +use jit_assembler::common::{InstructionBuilder, InstructionCollectionWithUsage}; + +// Build code in separate builders +let mut prologue = Riscv64InstructionBuilder::new(); +prologue.sd(reg::SP, reg::S0, -8); // Save S0 + +let mut main_code = Riscv64InstructionBuilder::new(); +main_code.add(reg::S0, reg::A0, reg::A1); // Use S0 + +let mut epilogue = Riscv64InstructionBuilder::new(); +epilogue.ld(reg::S0, reg::SP, -8); // Restore S0 + +// Create tracked collections +let prologue_tracked = InstructionCollectionWithUsage::new( + prologue.instructions(), + prologue.register_usage().clone() +); +let main_tracked = InstructionCollectionWithUsage::new( + main_code.instructions(), + main_code.register_usage().clone() +); +let epilogue_tracked = InstructionCollectionWithUsage::new( + epilogue.instructions(), + epilogue.register_usage().clone() +); + +// Merge with register usage tracking +let combined = prologue_tracked + main_tracked + epilogue_tracked; + +// Access both instructions and register usage +let instructions = combined.instructions(); +let usage = combined.register_usage(); + +println!("Combined {} instructions", instructions.len()); +println!("Register usage: {}", usage); +``` + +See `examples/instruction_collection_merge.rs` for a complete working example. + ## Contributing Contributions are welcome! Please feel free to submit a Pull Request. diff --git a/examples/instruction_collection_merge.rs b/examples/instruction_collection_merge.rs new file mode 100644 index 0000000..5998907 --- /dev/null +++ b/examples/instruction_collection_merge.rs @@ -0,0 +1,135 @@ +//! Example demonstrating InstructionCollection merge functionality +//! +//! This example shows how to combine multiple instruction sequences, +//! such as function prologue, main body, and epilogue, in arbitrary order. +//! When the register-tracking feature is enabled, register usage information +//! is also properly merged. + +use jit_assembler::riscv64::{Riscv64InstructionBuilder, reg}; +use jit_assembler::common::InstructionBuilder; + +#[cfg(feature = "register-tracking")] +use jit_assembler::common::InstructionCollectionWithUsage; + +fn main() { + println!("=== Instruction Collection Merge Demo ===\n"); + + // Create three separate builders for different parts of the code + + // Builder 1: Main computation + println!("1. Building main computation..."); + let mut main_builder = Riscv64InstructionBuilder::new(); + main_builder + .add(reg::A0, reg::A1, reg::A2) // a0 = a1 + a2 + .mul(reg::A0, reg::A0, reg::A3) // a0 = a0 * a3 + .addi(reg::A0, reg::A0, 100); // a0 = a0 + 100 + + // Builder 2: Function prologue (register save) + // This is typically built AFTER the main code to know which registers to save + println!("2. Building prologue (register save)..."); + let mut prologue_builder = Riscv64InstructionBuilder::new(); + prologue_builder + .addi(reg::SP, reg::SP, -16) // Allocate stack space + .sd(reg::SP, reg::RA, 8) // Save return address + .sd(reg::SP, reg::S0, 0); // Save frame pointer + + // Builder 3: Function epilogue (register restore) + println!("3. Building epilogue (register restore)..."); + let mut epilogue_builder = Riscv64InstructionBuilder::new(); + epilogue_builder + .ld(reg::RA, reg::SP, 8) // Restore return address + .ld(reg::S0, reg::SP, 0) // Restore frame pointer + .addi(reg::SP, reg::SP, 16) // Deallocate stack space + .ret(); // Return + + // Get instruction collections from each builder + let prologue = prologue_builder.instructions(); + let main_code = main_builder.instructions(); + let epilogue = epilogue_builder.instructions(); + + println!("\nInstruction counts:"); + println!(" Prologue: {} instructions", prologue.len()); + println!(" Main: {} instructions", main_code.len()); + println!(" Epilogue: {} instructions", epilogue.len()); + + // Combine them in the desired order: prologue + main + epilogue + println!("\n4. Combining instruction collections..."); + let combined = prologue + main_code + epilogue; + + println!(" Combined: {} instructions", combined.len()); + + // Display the combined instructions + println!("\nCombined instructions:"); + for (i, instr) in combined.iter().enumerate() { + println!(" [{:2}] {}", i, instr); + } + + // Convert to bytes for execution + let bytes = combined.to_bytes(); + println!("\nGenerated {} bytes of machine code", bytes.len()); + + // Demonstrate with register tracking (when feature is enabled) + #[cfg(feature = "register-tracking")] + demonstrate_with_register_tracking(); +} + +#[cfg(feature = "register-tracking")] +fn demonstrate_with_register_tracking() { + println!("\n\n=== With Register Tracking ===\n"); + + // Create the same three builders + let mut main_builder = Riscv64InstructionBuilder::new(); + main_builder + .add(reg::A0, reg::A1, reg::A2) + .mul(reg::A0, reg::A0, reg::A3) + .addi(reg::A0, reg::A0, 100); + + let mut prologue_builder = Riscv64InstructionBuilder::new(); + prologue_builder + .addi(reg::SP, reg::SP, -16) + .sd(reg::SP, reg::RA, 8) + .sd(reg::SP, reg::S0, 0); + + let mut epilogue_builder = Riscv64InstructionBuilder::new(); + epilogue_builder + .ld(reg::RA, reg::SP, 8) + .ld(reg::S0, reg::SP, 0) + .addi(reg::SP, reg::SP, 16) + .ret(); + + // Create tracked collections that include register usage info + let prologue = InstructionCollectionWithUsage::new( + prologue_builder.instructions(), + prologue_builder.register_usage().clone() + ); + let main_code = InstructionCollectionWithUsage::new( + main_builder.instructions(), + main_builder.register_usage().clone() + ); + let epilogue = InstructionCollectionWithUsage::new( + epilogue_builder.instructions(), + epilogue_builder.register_usage().clone() + ); + + println!("Register usage before merge:"); + println!(" Prologue: {}", prologue.register_usage()); + println!(" Main: {}", main_code.register_usage()); + println!(" Epilogue: {}", epilogue.register_usage()); + + // Merge with register usage tracking + let combined = prologue + main_code + epilogue; + + println!("\nRegister usage after merge:"); + println!(" Combined: {}", combined.register_usage()); + + let usage = combined.register_usage(); + println!("\nDetailed register analysis:"); + println!(" Total registers used: {}", usage.register_count()); + println!(" Caller-saved: {:?}", usage.caller_saved_registers()); + println!(" Callee-saved: {:?}", usage.callee_saved_registers()); + println!(" Special: {:?}", usage.special_registers()); + println!(" Needs stack frame: {}", usage.needs_stack_frame()); + + // You can still access the instructions + println!("\nFinal instruction count: {}", combined.instructions().len()); +}