Borrow the input of native functions

dissect/util/_native.src/src/buffer.rs

@@ -0,0 +1,22 @@
+use pyo3::prelude::*;
+use pyo3::pybacked::PyBackedBytes;
+
+/// A bytes-like argument. A ``bytes`` is borrowed and a ``bytearray`` copied in one go. Extracting
+/// either as a ``Vec<u8>`` goes through it as a sequence, an int at a time, which costs ~10 ns a
+/// byte: most of the time of a call. Other sequences of ints are still accepted that way.
+#[derive(FromPyObject)]
+pub enum Buffer {
+    Bytes(PyBackedBytes),
+    Sequence(Vec<u8>),
+}
+
+impl std::ops::Deref for Buffer {
+    type Target = [u8];
+
+    fn deref(&self) -> &[u8] {
+        match self {
+            Buffer::Bytes(bytes) => bytes,
+            Buffer::Sequence(sequence) => sequence,
+        }
+    }
+}

dissect/util/_native.src/src/compression/lz4.rs

@@ -3,6 +3,8 @@ use pyo3::prelude::*;
 use pyo3::pybacked::PyBackedBytes;
 use pyo3::types::{PyByteArray, PyBytes};
 
+use crate::buffer::Buffer;
+
 const MAX_DISCOVER_OUTPUT_SIZE: usize = 1024 * 1024 * 1024;
 
 fn decompress_block(
@@ -64,7 +66,7 @@ fn decompress_to_unknown_size(src: &[u8], dictionary: &[u8]) -> Result<Vec<u8>,
 #[pyo3(signature = (src, uncompressed_size=-1, return_bytearray=false, dictionary=None))]
 fn decompress(
     py: Python<'_>,
-    src: Vec<u8>,
+    src: Buffer,
     uncompressed_size: isize,
     return_bytearray: bool,
     dictionary: Option<PyBackedBytes>,

dissect/util/_native.src/src/compression/lzo.rs

@@ -2,6 +2,8 @@ use pyo3::exceptions::PyValueError;
 use pyo3::prelude::*;
 use pyo3::types::PyBytes;
 
+use crate::buffer::Buffer;
+
 /// LZO decompress from bytes. Assumes no header.
 ///
 /// Arguments are largely compatible with python-lzo API.
@@ -20,7 +22,7 @@ use pyo3::types::PyBytes;
 #[pyo3(signature = (src, header=true, buflen=-1))]
 fn decompress(
     py: Python<'_>,
-    src: Vec<u8>,
+    src: Buffer,
     header: bool,
     buflen: isize,
 ) -> PyResult<Bound<'_, PyBytes>> {
@@ -31,13 +33,12 @@ fn decompress(
             return Err(PyValueError::new_err("Invalid header value"));
         }
         let len = u32::from_le_bytes([src[1], src[2], src[3], src[4]]) as usize;
-        (src[5..].to_vec(), Some(len))
+        (&src[5..], Some(len))
     } else {
-        (src, (buflen >= 0).then_some(buflen as usize))
+        (&src[..], (buflen >= 0).then_some(buflen as usize))
     };
 
-    let mut cursor = std::io::Cursor::new(body);
-    lzokay_native::decompress(&mut cursor, out_len)
+    lzokay_native::decompress_all(body, out_len)
         .map_err(|e| PyErr::new::<PyValueError, _>(e.to_string()))
         .map(|result| PyBytes::new(py, &result))
 }

dissect/util/_native.src/src/hash/crc32c.rs

@@ -1,5 +1,7 @@
 use pyo3::prelude::*;
 
+use crate::buffer::Buffer;
+
 const _TABLE: [u32; 256] = [
     0x00000000, 0xf26b8303, 0xe13b70f7, 0x1350f3f4, 0xc79a971f, 0x35f1141c, 0x26a1e7e8, 0xd4ca64eb,
     0x8ad958cf, 0x78b2dbcc, 0x6be22838, 0x9989ab3b, 0x4d43cfd0, 0xbf284cd3, 0xac78bf27, 0x5e133c24,
@@ -46,9 +48,9 @@ const _TABLE: [u32; 256] = [
 ///
 #[pyfunction]
 #[pyo3(signature = (crc, data))]
-fn update(py: Python<'_>, crc: u32, data: Vec<u8>) -> PyResult<Py<PyAny>> {
+fn update(py: Python<'_>, crc: u32, data: Buffer) -> PyResult<Py<PyAny>> {
     let mut crc = crc ^ 0xFFFFFFFF;
-    for &b in &data {
+    for &b in data.iter() {
         crc = _TABLE[((crc ^ b as u32) & 0xFF) as usize] ^ ((crc >> 8) & 0xFFFFFFFF);
     }
     Ok((crc ^ 0xFFFFFFFF).into_pyobject(py)?.into())
@@ -65,7 +67,7 @@ fn update(py: Python<'_>, crc: u32, data: Vec<u8>) -> PyResult<Py<PyAny>> {
 ///
 #[pyfunction]
 #[pyo3(signature = (data, value=0))]
-fn crc32c(py: Python<'_>, data: Vec<u8>, value: u32) -> PyResult<Py<PyAny>> {
+fn crc32c(py: Python<'_>, data: Buffer, value: u32) -> PyResult<Py<PyAny>> {
     update(py, value, data)
 }
 

dissect/util/_native.src/src/hash/crc64c.rs

@@ -1,5 +1,7 @@
 use pyo3::prelude::*;
 
+use crate::buffer::Buffer;
+
 const _TABLE: [u64; 256] = [
     0x0000000000000000,
     0x7f6ef0c830358979,
@@ -270,9 +272,9 @@ const _TABLE: [u64; 256] = [
 ///
 #[pyfunction]
 #[pyo3(signature = (crc, data))]
-fn update(py: Python<'_>, crc: u64, data: Vec<u8>) -> PyResult<Py<PyAny>> {
+fn update(py: Python<'_>, crc: u64, data: Buffer) -> PyResult<Py<PyAny>> {
     let mut crc = crc ^ 0xFFFFFFFFFFFFFFFF;
-    for &b in &data {
+    for &b in data.iter() {
         crc = _TABLE[((crc ^ b as u64) & 0xFF) as usize] ^ (crc >> 8);
     }
     Ok((crc ^ 0xFFFFFFFFFFFFFFFF).into_pyobject(py)?.into())
@@ -289,7 +291,7 @@ fn update(py: Python<'_>, crc: u64, data: Vec<u8>) -> PyResult<Py<PyAny>> {
 ///
 #[pyfunction]
 #[pyo3(signature = (data, value=0))]
-fn crc64c(py: Python<'_>, data: Vec<u8>, value: u64) -> PyResult<Py<PyAny>> {
+fn crc64c(py: Python<'_>, data: Buffer, value: u64) -> PyResult<Py<PyAny>> {
     update(py, value, data)
 }
 

dissect/util/_native.src/src/lib.rs

@@ -1,5 +1,6 @@
 use pyo3::prelude::*;
 
+mod buffer;
 mod compression;
 mod hash;
 

tests/compression/test_lz4.py

@@ -128,6 +128,12 @@ def test_lz4_decompress_chain(lz4: ModuleType) -> None:
     )
 
 
+@pytest.mark.parametrize("source_type", [bytes, bytearray, memoryview])
+def test_lz4_decompress_source_types(lz4: ModuleType, source_type: type) -> None:
+    data, digest = PARAMS[1][1].values
+    assert hashlib.sha256(lz4.decompress(source_type(bytes.fromhex(data)))).hexdigest() == digest
+
+
 @pytest.mark.benchmark
 @pytest.mark.parametrize(*PARAMS)
 def test_benchmark_lz4_decompress(lz4: ModuleType, data: str, digest: str, benchmark: BenchmarkFixture) -> None:

tests/compression/test_lzo.py

@@ -204,6 +204,14 @@ def test_lzo_decompress(lzo: ModuleType, data: str, header: bool, buflen: int, d
     assert hashlib.sha256(lzo.decompress(bytes.fromhex(data), header, buflen)).hexdigest() == digest
 
 
+@pytest.mark.parametrize("source_type", [bytes, bytearray, memoryview])
+@pytest.mark.parametrize(*PARAMS)
+def test_lzo_decompress_source_types(
+    lzo: ModuleType, data: str, header: bool, buflen: int, digest: str, source_type: type
+) -> None:
+    assert hashlib.sha256(lzo.decompress(source_type(bytes.fromhex(data)), header, buflen)).hexdigest() == digest
+
+
 def test_lzo_decompress_invalid_distance() -> None:
     # A literal followed by a match that refers back to before the start of the output.
     # Only the pure Python implementation, as lzokay-native panics on this input instead of raising an error.

tests/test_hash.py

@@ -40,11 +40,23 @@ def test_crc32c(crc32c: ModuleType, data: bytes, value: int, expected: int) -> N
     assert crc32c.crc32c(data, value) == expected
 
 
+@pytest.mark.parametrize("source_type", [bytes, bytearray, memoryview])
+def test_crc32c_source_types(crc32c: ModuleType, source_type: type) -> None:
+    assert crc32c.crc32c(source_type(bytes(range(32))), 0) == 0x46DD794E
+    assert crc32c.update(0x12345678, source_type(b"hello, world!")) == 0x30663976
+
+
 @pytest.mark.benchmark
 def test_crc32c_benchmark(crc32c: ModuleType, benchmark: BenchmarkFixture) -> None:
     benchmark(crc32c.crc32c, b"hello, world!", 0)
 
 
+@pytest.mark.benchmark
+def test_crc32c_benchmark_page(crc32c: ModuleType, benchmark: BenchmarkFixture) -> None:
+    # A 4 KiB page, as in the TIBX archives of dissect.archive
+    assert benchmark(crc32c.crc32c, bytes(range(256)) * 16, 0) == 0x9C71FE32
+
+
 @pytest.mark.parametrize(
     ("data", "value", "expected"),
     [
@@ -58,6 +70,12 @@ def test_crc64c(crc64c: ModuleType, data: bytes, value: int, expected: int) -> N
     assert crc64c.crc64c(data, value) == expected
 
 
+@pytest.mark.parametrize("source_type", [bytes, bytearray, memoryview])
+def test_crc64c_source_types(crc64c: ModuleType, source_type: type) -> None:
+    assert crc64c.crc64c(source_type(b"123456789"), 0) == 0xAE8B14860A799888
+    assert crc64c.update(0x123456789ABCDEF0, source_type(b"hello, world!")) == 0x96270A5EF700596F
+
+
 @pytest.mark.benchmark
 def test_crc64c_benchmark(crc64c: ModuleType, benchmark: BenchmarkFixture) -> None:
     benchmark(crc64c.crc64c, b"hello, world!", 0)