things compile again, add a few more caution signs around InstructionTextBuffer

5 files changed, 76 insertions, 9 deletions

diff --git a/src/display.rs b/src/display.rs
index e495aee..9b72cb3 100644
--- a/src/display.rs
+++ b/src/display.rs
@@ -211,6 +211,11 @@ impl<'buf> fmt::Write for InstructionTextSink<'buf> {
         self.buf.write_str(s)
     }
     fn write_char(&mut self, c: char) -> Result<(), core::fmt::Error> {
+        if cfg!(debug_assertions) {
+            if self.buf.capacity() < self.buf.len() + 1 {
+                panic!("InstructionTextSink::write_char would overflow output");
+            }
+        }
         // SAFETY: `buf` is assumed to be long enough to hold all input, `buf` at `underlying.len()`
         // is valid for writing, but may be uninitialized.
         //
@@ -640,6 +645,12 @@ impl DisplaySink for alloc::string::String {
 impl<'buf> DisplaySink for InstructionTextSink<'buf> {
     #[inline(always)]
     fn write_fixed_size(&mut self, s: &str) -> Result<(), core::fmt::Error> {
+        if cfg!(debug_assertions) {
+            if self.buf.capacity() < self.buf.len() + s.len() {
+                panic!("InstructionTextSink::write_fixed_size would overflow output");
+            }
+        }
+
         let buf = unsafe { self.buf.as_mut_vec() };
         let new_bytes = s.as_bytes();
 
@@ -671,6 +682,12 @@ impl<'buf> DisplaySink for InstructionTextSink<'buf> {
         Ok(())
     }
     unsafe fn write_lt_32(&mut self, s: &str) -> Result<(), fmt::Error> {
+        if cfg!(debug_assertions) {
+            if self.buf.capacity() < self.buf.len() + s.len() {
+                panic!("InstructionTextSink::write_lt_32 would overflow output");
+            }
+        }
+
         // SAFETY: todo
         let buf = unsafe { self.buf.as_mut_vec() };
         let new_bytes = s.as_bytes();
@@ -745,6 +762,12 @@ impl<'buf> DisplaySink for InstructionTextSink<'buf> {
         Ok(())
     }
     unsafe fn write_lt_16(&mut self, s: &str) -> Result<(), fmt::Error> {
+        if cfg!(debug_assertions) {
+            if self.buf.capacity() < self.buf.len() + s.len() {
+                panic!("InstructionTextSink::write_lt_16 would overflow output");
+            }
+        }
+
         // SAFETY: todo
         let buf = unsafe { self.buf.as_mut_vec() };
         let new_bytes = s.as_bytes();
@@ -810,6 +833,12 @@ impl<'buf> DisplaySink for InstructionTextSink<'buf> {
         Ok(())
     }
     unsafe fn write_lt_8(&mut self, s: &str) -> Result<(), fmt::Error> {
+        if cfg!(debug_assertions) {
+            if self.buf.capacity() < self.buf.len() + s.len() {
+                panic!("InstructionTextSink::write_lt_8 would overflow output");
+            }
+        }
+
         // SAFETY: todo
         let buf = unsafe { self.buf.as_mut_vec() };
         let new_bytes = s.as_bytes();
@@ -881,6 +910,12 @@ impl<'buf> DisplaySink for InstructionTextSink<'buf> {
         // means we can write directly into the correct offsets of the output string.
         let printed_size = ((8 - v.leading_zeros() + 3) >> 2) as usize;
 
+        if cfg!(debug_assertions) {
+            if self.buf.capacity() < self.buf.len() + printed_size {
+                panic!("InstructionTextSink::write_u8 would overflow output");
+            }
+        }
+
         let buf = unsafe { self.buf.as_mut_vec() };
         let new_len = buf.len() + printed_size;
 
@@ -914,10 +949,17 @@ impl<'buf> DisplaySink for InstructionTextSink<'buf> {
         if v == 0 {
             return self.write_fixed_size("0");
         }
+
         // we can fairly easily predict the size of a formatted string here with lzcnt, which also
         // means we can write directly into the correct offsets of the output string.
         let printed_size = ((16 - v.leading_zeros() + 3) >> 2) as usize;
 
+        if cfg!(debug_assertions) {
+            if self.buf.capacity() < self.buf.len() + printed_size {
+                panic!("InstructionTextSink::write_u16 would overflow output");
+            }
+        }
+
         let buf = unsafe { self.buf.as_mut_vec() };
         let new_len = buf.len() + printed_size;
 
@@ -951,10 +993,17 @@ impl<'buf> DisplaySink for InstructionTextSink<'buf> {
         if v == 0 {
             return self.write_fixed_size("0");
         }
+
         // we can fairly easily predict the size of a formatted string here with lzcnt, which also
         // means we can write directly into the correct offsets of the output string.
         let printed_size = ((32 - v.leading_zeros() + 3) >> 2) as usize;
 
+        if cfg!(debug_assertions) {
+            if self.buf.capacity() < self.buf.len() + printed_size {
+                panic!("InstructionTextSink::write_u32 would overflow output");
+            }
+        }
+
         let buf = unsafe { self.buf.as_mut_vec() };
         let new_len = buf.len() + printed_size;
 
@@ -988,10 +1037,17 @@ impl<'buf> DisplaySink for InstructionTextSink<'buf> {
         if v == 0 {
             return self.write_fixed_size("0");
         }
+
         // we can fairly easily predict the size of a formatted string here with lzcnt, which also
         // means we can write directly into the correct offsets of the output string.
         let printed_size = ((64 - v.leading_zeros() + 3) >> 2) as usize;
 
+        if cfg!(debug_assertions) {
+            if self.buf.capacity() < self.buf.len() + printed_size {
+                panic!("InstructionTextSink::write_u64 would overflow output");
+            }
+        }
+
         let buf = unsafe { self.buf.as_mut_vec() };
         let new_len = buf.len() + printed_size;
 
diff --git a/src/lib.rs b/src/lib.rs
index 709563b..491b6f0 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -140,7 +140,7 @@ pub use real_mode::Arch as x86_16;
 
 mod safer_unchecked;
 #[cfg(feature = "fmt")]
-mod display;
+pub mod display;
 
 const MEM_SIZE_STRINGS: [&'static str; 65] = [
     "BUG",
diff --git a/src/long_mode/display.rs b/src/long_mode/display.rs
index d9799e1..3615538 100644
--- a/src/long_mode/display.rs
+++ b/src/long_mode/display.rs
@@ -4303,6 +4303,7 @@ impl <T: fmt::Write, Y: YaxColors> ShowContextual<u64, [Option<alloc::string::St
     }
 }
 
+// TODO: should include CALL
 static RELATIVE_BRANCHES: [Opcode; 21] = [
     Opcode::JMP, Opcode::JRCXZ,
     Opcode::LOOP, Opcode::LOOPZ, Opcode::LOOPNZ,
@@ -4474,12 +4475,13 @@ impl<'a, F: DisplaySink> crate::long_mode::OperandVisitor for RelativeBranchPrin
 /// ```
 /// use yaxpeax_x86::long_mode::InstDecoder;
 /// use yaxpeax_x86::long_mode::InstructionTextBuffer;
+/// use yaxpeax_x86::long_mode::DisplayStyle;
 ///
 /// let bytes = &[0x33, 0xc0];
 /// let inst = InstDecoder::default().decode_slice(bytes).expect("can decode");
 /// let mut text_buf = InstructionTextBuffer::new();
 /// assert_eq!(
-///     text_buf.format_inst(&inst).expect("can format"),
+///     text_buf.format_inst(&inst.display_with(DisplayStyle::Intel)).expect("can format"),
 ///     "xor eax, eax"
 /// );
 ///
@@ -4511,7 +4513,9 @@ impl InstructionTextBuffer {
     /// this clears and reuses an internal buffer; if an instruction had been previously formatted
     /// through this buffer, it will be overwritten.
     pub fn format_inst<'buf, 'instr>(&'buf mut self, display: &InstructionDisplayer<'instr>) -> Result<&'buf str, fmt::Error> {
-        let mut handle = self.write_handle();
+        // Safety: this sink is used to format exactly one instruction and then dropped. it can
+        // never escape `format_inst`.
+        let mut handle = unsafe { self.write_handle() };
 
         match display.style {
             DisplayStyle::Intel => {
@@ -4531,8 +4535,15 @@ impl InstructionTextBuffer {
         self.content.as_str()
     }
 
-    fn write_handle(&mut self) -> crate::display::InstructionTextSink {
+    /// do the necessary bookkeeping and provide an `InstructionTextSink` to write an instruction
+    /// into.
+    ///
+    /// SAFETY: callers must print at most one instruction into this handle.
+    unsafe fn write_handle(&mut self) -> crate::display::InstructionTextSink {
         self.content.clear();
+        // Safety: `content` was just cleared, so writing begins at the start of the buffer.
+        // `content`is large enough to hold a fully-formatted instruction (see
+        // `InstructionTextBuffer::new`).
         crate::display::InstructionTextSink::new(&mut self.content)
     }
 }
diff --git a/src/long_mode/mod.rs b/src/long_mode/mod.rs
index 418d57f..9aeacdc 100644
--- a/src/long_mode/mod.rs
+++ b/src/long_mode/mod.rs
@@ -7,7 +7,7 @@ pub mod uarch;
 pub use crate::MemoryAccessSize;
 
 #[cfg(feature = "fmt")]
-pub use self::display::{DisplayStyle, InstructionDisplayer};
+pub use self::display::{DisplayStyle, InstructionDisplayer, InstructionTextBuffer};
 
 use core::cmp::PartialEq;
 use crate::safer_unchecked::unreachable_kinda_unchecked as unreachable_unchecked;
diff --git a/test/long_mode/mod.rs b/test/long_mode/mod.rs
index dcc9aad..7742496 100644
--- a/test/long_mode/mod.rs
+++ b/test/long_mode/mod.rs
@@ -79,7 +79,7 @@ fn test_display_under(decoder: &InstDecoder, data: &[u8], expected: &'static str
                     );
 
                     let mut text2 = String::new();
-                    let mut out = yaxpeax_x86::long_mode::NoColorsSink {
+                    let mut out = yaxpeax_x86::display::NoColorsSink {
                         out: &mut text2,
                     };
                     instr.write_to(&mut out).expect("printing succeeds");
@@ -94,12 +94,12 @@ fn test_display_under(decoder: &InstDecoder, data: &[u8], expected: &'static str
                         text,
                     );
 
-                    let mut formatter = yaxpeax_x86::long_mode::InstructionFormatter::new();
-                    let text3 = formatter.format_inst(&instr).expect("printing succeeds");
+                    let mut formatter = yaxpeax_x86::long_mode::InstructionTextBuffer::new();
+                    let text3 = formatter.format_inst(&instr.display_with(yaxpeax_x86::long_mode::DisplayStyle::Intel)).expect("printing succeeds");
 
                     assert!(
                         text3 == text,
-                        "display error through InstructionFormatter for {}:\n  decoded: {:?} under decoder {}\n displayed: {}\n expected: {}\n",
+                        "display error through InstructionTextBuffer for {}:\n  decoded: {:?} under decoder {}\n displayed: {}\n expected: {}\n",
                         hex,
                         instr,
                         decoder,