From 1900a5c0d60c6857c4db6586e2d60d3516ac30bb Mon Sep 17 00:00:00 2001
From: NightEule5 <24661563+NightEule5@users.noreply.github.com>
Date: Sun, 11 Feb 2024 08:21:59 -0700
Subject: [PATCH] Add SIMD read/write support

---
 orio/Cargo.toml      |  1 +
 orio/src/lib.rs      |  1 +
 orio/src/streams.rs  | 50 ++++++++++++++++++++++++++++++++++++++++++++
 orio/tests/buffer.rs | 37 ++++++++++++++++++++++++++++++++
 4 files changed, 89 insertions(+)

diff --git a/orio/Cargo.toml b/orio/Cargo.toml
index 05e9a08..a8bf4be 100644
--- a/orio/Cargo.toml
+++ b/orio/Cargo.toml
@@ -12,6 +12,7 @@ repository = "https://github.com/NightEule5/orio/"
 [features]
 default = ["secure-hash"]
 bytes = ["dep:bytes"]
+simd = ["bytemuck/nightly_portable_simd"]
 shared-pool = []
 hash = ["dep:digest"]
 secure-hash = ["groestl", "sha2", "sha3", "shabal", "whirlpool"]
diff --git a/orio/src/lib.rs b/orio/src/lib.rs
index 7405288..5276698 100644
--- a/orio/src/lib.rs
+++ b/orio/src/lib.rs
@@ -55,6 +55,7 @@
 	type_alias_impl_trait,
 )]
 #![feature(iter_advance_by)]
+#![cfg_attr(feature = "simd", feature(portable_simd))]
 
 mod buffer;
 mod buffered_wrappers;
diff --git a/orio/src/streams.rs b/orio/src/streams.rs
index cac772e..759bc5c 100644
--- a/orio/src/streams.rs
+++ b/orio/src/streams.rs
@@ -20,6 +20,8 @@ pub use crate::buffered_wrappers::{BufferedSink, BufferedSource};
 use crate::error::Context;
 use crate::pattern::Pattern;
 use crate::StreamContext::{Read, Write};
+#[cfg(feature = "simd")]
+use std::simd::{prelude::*, LaneCount, SimdElement, SupportedLaneCount};
 
 /// An "stream closed" error.
 #[derive(Copy, Clone, Debug, Default, thiserror::Error)]
@@ -445,6 +447,30 @@ pub trait BufSource<'d, const N: usize = SIZE>: BufStream<'d, N> + Source<'d, N>
 		self.read_pod().map(T::to_le)
 	}
 
+	/// Reads multiple big-endian integers into a [`Simd`] vector.
+	#[cfg(feature = "simd")]
+	fn read_vector<T: bytemuck::Pod + PrimInt + SimdElement, const L: usize>(&mut self) -> Result<Simd<T, L>>
+	where LaneCount<L>: SupportedLaneCount {
+		let mut vec: Simd<T, L> = self.read_pod()?;
+		// Loop compiles away, generating the same code as SimdInt::swap_bytes.
+		for i in vec.as_mut_array() {
+			*i = i.to_be();
+		}
+		Ok(vec)
+	}
+
+	/// Reads multiple little-endian integers into a [`Simd`] vector.
+	#[cfg(feature = "simd")]
+	fn read_vector_le<T: bytemuck::Pod + PrimInt + SimdElement, const L: usize>(&mut self) -> Result<Simd<T, L>>
+	where LaneCount<L>: SupportedLaneCount {
+		let mut vec: Simd<T, L> = self.read_pod()?;
+		// Loop compiles away, generating the same code as SimdInt::swap_bytes.
+		for i in vec.as_mut_array() {
+			*i = i.to_le();
+		}
+		Ok(vec)
+	}
+
 	/// Reads an arbitrary [`Pod`] data type.
 	///
 	/// [`Pod`]: bytemuck::Pod
@@ -717,6 +743,28 @@ pub trait BufSink<'d, const N: usize = SIZE>: BufStream<'d, N> + Sink<'d, N> {
 		self.write_pod(value.to_le())
 	}
 
+	/// Writes a [`Simd`] vector of big-endian integers.
+	#[cfg(feature = "simd")]
+	fn write_vector<T: bytemuck::Pod + PrimInt + SimdElement, const L: usize>(&mut self, mut vector: Simd<T, L>) -> Result
+	where LaneCount<L>: SupportedLaneCount {
+		// Loop compiles away, generating the same code as SimdInt::swap_bytes.
+		for i in vector.as_mut_array() {
+			*i = i.to_be();
+		}
+		self.write_pod(vector)
+	}
+
+	/// Writes a [`Simd`] vector of little-endian integers.
+	#[cfg(feature = "simd")]
+	fn write_vector_le<T: bytemuck::Pod + PrimInt + SimdElement, const L: usize>(&mut self, mut vector: Simd<T, L>) -> Result
+	where LaneCount<L>: SupportedLaneCount {
+		// Loop compiles away, generating the same code as SimdInt::swap_bytes.
+		for i in vector.as_mut_array() {
+			*i = i.to_le();
+		}
+		self.write_pod(vector)
+	}
+
 	/// Writes an arbitrary [`Pod`] data type.
 	///
 	/// [`Pod`]: bytemuck::Pod
@@ -726,6 +774,8 @@ pub trait BufSink<'d, const N: usize = SIZE>: BufStream<'d, N> + Sink<'d, N> {
 		Ok(())
 	}
 
+
+
 	/// Writes a UTF-8 string.
 	#[inline]
 	fn write_utf8(&mut self, value: &str) -> Result<usize> {
diff --git a/orio/tests/buffer.rs b/orio/tests/buffer.rs
index d135f64..ae6cfe3 100644
--- a/orio/tests/buffer.rs
+++ b/orio/tests/buffer.rs
@@ -1,5 +1,7 @@
 // SPDX-License-Identifier: Apache-2.0
 
+#![cfg_attr(feature = "simd", feature(portable_simd))]
+
 mod dataset;
 
 macro_rules! qc_assert_ok {
@@ -151,6 +153,41 @@ mod read {
 
 	gen! { u8, i8, u16, i16, u32, i32, u64, i64, u128, i128, usize, isize }
 
+	#[cfg(feature = "simd")]
+	#[derive(Copy, Clone, Debug)]
+	struct SimdData(std::simd::u16x16);
+
+	#[cfg(feature = "simd")]
+	impl Arbitrary for SimdData {
+		fn arbitrary(g: &mut Gen) -> Self {
+			let mut arr = [0; 16];
+			arr.fill_with(|| u16::arbitrary(g));
+			Self(arr.into())
+		}
+	}
+
+	#[cfg(feature = "simd")]
+	#[quickcheck]
+	fn u16x16(SimdData(data): SimdData) {
+		let mut buffer = DefaultBuffer::default();
+		for &i in data.as_array() {
+			buffer.write_int(i).unwrap();
+		}
+
+		assert_eq!(buffer.read_vector().unwrap(), data);
+	}
+
+	#[cfg(feature = "simd")]
+	#[quickcheck]
+	fn u16x16_le(SimdData(data): SimdData) {
+		let mut buffer = DefaultBuffer::default();
+		for &i in data.as_array() {
+			buffer.write_int_le(i).unwrap();
+		}
+
+		assert_eq!(buffer.read_vector_le().unwrap(), data);
+	}
+
 	#[quickcheck]
 	fn vec(vec: Vec<u8>) {
 		let mut buffer = DefaultBuffer::default();