Wrap unsafe functions to catch errors in debug

Closes https://github.com/iximeow/yaxpeax-x86/issues/16

8 files changed, 42 insertions, 9 deletions

diff --git a/src/lib.rs b/src/lib.rs
index 46bebdb..b8bd4f7 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -121,6 +121,8 @@ pub use protected_mode::Arch as x86_32;
 pub mod real_mode;
 pub use real_mode::Arch as x86_16;
 
+mod safer_unchecked;
+
 const MEM_SIZE_STRINGS: [&'static str; 64] = [
     "byte", "word", "BUG", "dword", "ptr", "far", "BUG", "qword",
     "BUG", "mword", "BUG", "BUG", "BUG", "BUG", "BUG", "xmmword",
diff --git a/src/long_mode/display.rs b/src/long_mode/display.rs
index 35b9c1f..3295ae7 100644
--- a/src/long_mode/display.rs
+++ b/src/long_mode/display.rs
@@ -3,6 +3,7 @@ use core::fmt;
 use yaxpeax_arch::{Colorize, ShowContextual, NoColors, YaxColors};
 use yaxpeax_arch::display::*;
 
+use crate::safer_unchecked::GetSaferUnchecked as _;
 use crate::MEM_SIZE_STRINGS;
 use crate::long_mode::{RegSpec, Opcode, Operand, MergeMode, InstDecoder, Instruction, Segment, PrefixRex, OperandSpec};
 
@@ -127,7 +128,7 @@ const REG_NAMES: &[&'static str] = &[
 ];
 
 pub(crate) fn regspec_label(spec: &RegSpec) -> &'static str {
-    unsafe { REG_NAMES.get_unchecked((spec.num as u16 + ((spec.bank as u16) << 3)) as usize) }
+    unsafe { REG_NAMES.get_kinda_unchecked((spec.num as u16 + ((spec.bank as u16) << 3)) as usize) }
 }
 
 impl fmt::Display for RegSpec {
@@ -1810,7 +1811,7 @@ const MNEMONICS: &[&'static str] = &[
 impl Opcode {
     fn name(&self) -> &'static str {
         unsafe {
-            MNEMONICS.get_unchecked(*self as usize)
+            MNEMONICS.get_kinda_unchecked(*self as usize)
         }
     }
 }
diff --git a/src/long_mode/mod.rs b/src/long_mode/mod.rs
index 44cf992..d2cb2f7 100644
--- a/src/long_mode/mod.rs
+++ b/src/long_mode/mod.rs
@@ -10,7 +10,7 @@ pub use crate::MemoryAccessSize;
 pub use self::display::{DisplayStyle, InstructionDisplayer};
 
 use core::cmp::PartialEq;
-use core::hint::unreachable_unchecked;
+use crate::safer_unchecked::unreachable_kinda_unchecked as unreachable_unchecked;
 
 use yaxpeax_arch::{AddressDiff, Decoder, Reader, LengthedInstruction};
 use yaxpeax_arch::annotation::{AnnotatingDecoder, DescriptionSink, NullSink};
diff --git a/src/protected_mode/display.rs b/src/protected_mode/display.rs
index f545111..c2b52eb 100644
--- a/src/protected_mode/display.rs
+++ b/src/protected_mode/display.rs
@@ -1,4 +1,5 @@
 use core::fmt;
+use crate::safer_unchecked::GetSaferUnchecked as _;
 
 use yaxpeax_arch::{Colorize, ShowContextual, NoColors, YaxColors};
 use yaxpeax_arch::display::*;
@@ -123,7 +124,7 @@ const REG_NAMES: &[&'static str] = &[
 ];
 
 pub(crate) fn regspec_label(spec: &RegSpec) -> &'static str {
-    unsafe { REG_NAMES.get_unchecked((spec.num as u16 + ((spec.bank as u16) << 3)) as usize) }
+    unsafe { REG_NAMES.get_kinda_unchecked((spec.num as u16 + ((spec.bank as u16) << 3)) as usize) }
 }
 
 impl fmt::Display for RegSpec {
@@ -1811,7 +1812,7 @@ const MNEMONICS: &[&'static str] = &[
 impl Opcode {
     fn name(&self) -> &'static str {
         unsafe {
-            MNEMONICS.get_unchecked(*self as usize)
+            MNEMONICS.get_kinda_unchecked(*self as usize)
         }
     }
 }
diff --git a/src/protected_mode/mod.rs b/src/protected_mode/mod.rs
index 040a23e..0a6fcf8 100644
--- a/src/protected_mode/mod.rs
+++ b/src/protected_mode/mod.rs
@@ -10,7 +10,7 @@ pub use crate::MemoryAccessSize;
 pub use self::display::{DisplayStyle, InstructionDisplayer};
 
 use core::cmp::PartialEq;
-use core::hint::unreachable_unchecked;
+use crate::safer_unchecked::unreachable_kinda_unchecked as unreachable_unchecked;
 
 use yaxpeax_arch::{AddressDiff, Decoder, Reader, LengthedInstruction};
 use yaxpeax_arch::annotation::{AnnotatingDecoder, DescriptionSink, NullSink};
diff --git a/src/real_mode/display.rs b/src/real_mode/display.rs
index 258f5d9..0ae0a46 100644
--- a/src/real_mode/display.rs
+++ b/src/real_mode/display.rs
@@ -3,6 +3,7 @@ use core::fmt;
 use yaxpeax_arch::{Colorize, ShowContextual, NoColors, YaxColors};
 use yaxpeax_arch::display::*;
 
+use crate::safer_unchecked::GetSaferUnchecked as _;
 use crate::MEM_SIZE_STRINGS;
 use crate::real_mode::{RegSpec, Opcode, Operand, MergeMode, InstDecoder, Instruction, Segment, PrefixVex, OperandSpec};
 
@@ -123,7 +124,7 @@ const REG_NAMES: &[&'static str] = &[
 ];
 
 pub(crate) fn regspec_label(spec: &RegSpec) -> &'static str {
-    unsafe { REG_NAMES.get_unchecked((spec.num as u16 + ((spec.bank as u16) << 3)) as usize) }
+    unsafe { REG_NAMES.get_kinda_unchecked((spec.num as u16 + ((spec.bank as u16) << 3)) as usize) }
 }
 
 impl fmt::Display for RegSpec {
@@ -1811,7 +1812,7 @@ const MNEMONICS: &[&'static str] = &[
 impl Opcode {
     fn name(&self) -> &'static str {
         unsafe {
-            MNEMONICS.get_unchecked(*self as usize)
+            MNEMONICS.get_kinda_unchecked(*self as usize)
         }
     }
 }
diff --git a/src/real_mode/mod.rs b/src/real_mode/mod.rs
index 0a66ab9..f4cfdd5 100644
--- a/src/real_mode/mod.rs
+++ b/src/real_mode/mod.rs
@@ -10,7 +10,7 @@ pub use crate::MemoryAccessSize;
 pub use self::display::{DisplayStyle, InstructionDisplayer};
 
 use core::cmp::PartialEq;
-use core::hint::unreachable_unchecked;
+use crate::safer_unchecked::unreachable_kinda_unchecked as unreachable_unchecked;
 
 use yaxpeax_arch::{AddressDiff, Decoder, Reader, LengthedInstruction};
 use yaxpeax_arch::annotation::{AnnotatingDecoder, DescriptionSink, NullSink};
diff --git a/src/safer_unchecked.rs b/src/safer_unchecked.rs
new file mode 100644
index 0000000..afd0355
--- /dev/null
+++ b/src/safer_unchecked.rs
@@ -0,0 +1,28 @@
+use core::slice::SliceIndex;
+
+pub trait GetSaferUnchecked<T> {
+    unsafe fn get_kinda_unchecked<I>(&self, index: I) -> &<I as SliceIndex<[T]>>::Output
+    where
+        I: SliceIndex<[T]>;
+}
+
+impl<T> GetSaferUnchecked<T> for [T] {
+    unsafe fn get_kinda_unchecked<I>(&self, index: I) -> &<I as SliceIndex<[T]>>::Output
+    where
+        I: SliceIndex<[T]>,
+    {
+        if cfg!(debug_assertions) {
+            &self[index]
+        } else {
+            self.get_unchecked(index)
+        }
+    }
+}
+
+pub unsafe fn unreachable_kinda_unchecked() -> ! {
+    if cfg!(debug_assertions) {
+        panic!("UB: Unreachable unchecked was executed")
+    } else {
+        core::hint::unreachable_unchecked()
+    }
+}