1 files changed, 166 insertions, 0 deletions
diff --git a/src/testkit/display.rs b/src/testkit/display.rs
new file mode 100644
index 0000000..a745b5c
--- /dev/null
+++ b/src/testkit/display.rs
@@ -0,0 +1,166 @@
+//! tools to test the correctness of `yaxpeax-arch` trait implementations.
+
+use core::fmt;
+use core::fmt::Write;
+
+use crate::display::DisplaySink;
+
+/// `DisplaySinkValidator` is a `DisplaySink` that panics if invariants required of
+/// `DisplaySink`-writing functions are not upheld.
+///
+/// there are two categories of invariants that `DisplaySinkValidator` validates.
+///
+/// first, this panics if spans are not `span_end_*`-ed in first-in-last-out order with
+/// corresponding `span_start_*. second, this panics if `write_lt_*` functions are ever provided
+/// inputs longer than the corresponding maximum length.
+///
+/// functions that write to a `DisplaySink` are strongly encouraged to come with fuzzing that for
+/// all inputs `DisplaySinkValidator` does not panic.
+pub struct DisplaySinkValidator {
+    spans: alloc::vec::Vec<&'static str>,
+}
+
+impl fmt::Write for DisplaySinkValidator {
+    fn write_str(&mut self, _s: &str) -> Result<(), fmt::Error> {
+        Ok(())
+    }
+    fn write_char(&mut self, _c: char) -> Result<(), fmt::Error> {
+        Ok(())
+    }
+}
+
+impl DisplaySink for DisplaySinkValidator {
+    unsafe fn write_lt_32(&mut self, s: &str) -> Result<(), fmt::Error> {
+        if s.len() >= 32 {
+            panic!("DisplaySinkValidator::write_lt_32 was given a string longer than the maximum permitted length");
+        }
+
+        self.write_str(s)
+    }
+    unsafe fn write_lt_16(&mut self, s: &str) -> Result<(), fmt::Error> {
+        if s.len() >= 16 {
+            panic!("DisplaySinkValidator::write_lt_16 was given a string longer than the maximum permitted length");
+        }
+
+        self.write_str(s)
+    }
+    unsafe fn write_lt_8(&mut self, s: &str) -> Result<(), fmt::Error> {
+        if s.len() >= 8 {
+            panic!("DisplaySinkValidator::write_lt_8 was given a string longer than the maximum permitted length");
+        }
+
+        self.write_str(s)
+    }
+
+    fn span_start_immediate(&mut self) {
+        self.spans.push("immediate");
+    }
+
+    fn span_end_immediate(&mut self) {
+        let last = self.spans.pop().expect("item to pop");
+        assert_eq!(last, "immediate");
+    }
+
+    fn span_start_register(&mut self) {
+        self.spans.push("register");
+    }
+
+    fn span_end_register(&mut self) {
+        let last = self.spans.pop().expect("item to pop");
+        assert_eq!(last, "register");
+    }
+
+    fn span_start_opcode(&mut self) {
+        self.spans.push("opcode");
+    }
+
+    fn span_end_opcode(&mut self) {
+        let last = self.spans.pop().expect("item to pop");
+        assert_eq!(last, "opcode");
+    }
+
+    fn span_start_program_counter(&mut self) {
+        self.spans.push("program counter");
+    }
+
+    fn span_end_program_counter(&mut self) {
+        let last = self.spans.pop().expect("item to pop");
+        assert_eq!(last, "program counter");
+    }
+
+    fn span_start_number(&mut self) {
+        self.spans.push("number");
+    }
+
+    fn span_end_number(&mut self) {
+        let last = self.spans.pop().expect("item to pop");
+        assert_eq!(last, "number");
+    }
+
+    fn span_start_address(&mut self) {
+        self.spans.push("address");
+    }
+
+    fn span_end_address(&mut self) {
+        let last = self.spans.pop().expect("item to pop");
+        assert_eq!(last, "address");
+    }
+
+    fn span_start_function_expr(&mut self) {
+        self.spans.push("function expr");
+    }
+
+    fn span_end_function_expr(&mut self) {
+        let last = self.spans.pop().expect("item to pop");
+        assert_eq!(last, "function expr");
+    }
+}
+
+/// `DisplaySinkWriteComparator` helps test that two `DisplaySink` implementations which should
+/// produce the same output actually do.
+///
+/// this is most useful for cases like testing specialized `write_lt_*` functions, which ought to
+/// behave the same as if `write_str()` were called instead and so can be used as a very simple
+/// oracle.
+///
+/// this is somewhat less useful when the sinks are expected to produce unequal text, such as when
+/// one sink writes ANSI color sequences and the other does not.
+pub struct DisplaySinkWriteComparator<'sinks, T: DisplaySink, U: DisplaySink> {
+    sink1: &'sinks mut T,
+    sink1_check: fn(&T) -> &str,
+    sink2: &'sinks mut U,
+    sink2_check: fn(&U) -> &str,
+}
+
+impl<'sinks, T: DisplaySink, U: DisplaySink> DisplaySinkWriteComparator<'sinks, T, U> {
+    fn compare_sinks(&self) {
+        let sink1_text = (self.sink1_check)(self.sink1);
+        let sink2_text = (self.sink2_check)(self.sink2);
+
+        if sink1_text != sink2_text {
+            panic!("sinks produced different output: {} != {}", sink1_text, sink2_text);
+        }
+    }
+}
+
+impl<'sinks, T: DisplaySink, U: DisplaySink> DisplaySink for DisplaySinkWriteComparator<'sinks, T, U> {
+    fn write_u8(&mut self, v: u8) -> Result<(), fmt::Error> {
+        self.sink1.write_u8(v).expect("write to sink1 succeeds");
+        self.sink2.write_u8(v).expect("write to sink2 succeeds");
+        self.compare_sinks();
+        Ok(())
+    }
+}
+
+impl<'sinks, T: DisplaySink, U: DisplaySink> fmt::Write for DisplaySinkWriteComparator<'sinks, T, U> {
+    fn write_str(&mut self, s: &str) -> Result<(), fmt::Error> {
+        self.sink1.write_str(s).expect("write to sink1 succeeds");
+        self.sink2.write_str(s).expect("write to sink2 succeeds");
+        Ok(())
+    }
+    fn write_char(&mut self, c: char) -> Result<(), fmt::Error> {
+        self.sink1.write_char(c).expect("write to sink1 succeeds");
+        self.sink2.write_char(c).expect("write to sink2 succeeds");
+        Ok(())
+    }
+}