impl hashlib like interface for testing

README.md

@@ -26,22 +26,16 @@ maturin develop --release
 
 You should now be able to use the package:
 
-```
-(.env) Jack: xof-py % python3
-Python 3.12.4 (main, Jun  6 2024, 18:26:44) [Clang 15.0.0 (clang-1500.3.9.4)] on darwin
-Type "help", "copyright", "credits" or "license" for more information.
->>> from xof_py import shake_123
->>> shake_123().hex()
-'d6b9bdbda14c3858c36d'
->>> from hashlib import shake_128
->>> shake_128(b"123").hexdigest(10)
-'d6b9bdbda14c3858c36d'
+```py
+>>> from xof_py import pyo3_shake_128
+>>> pyo3_shake_128(b"cryptohack", 100).hex()
+'8d043455562ebedd1b3fcf5b0e0a058091752d161e7eef40364a565aacb3b5d3bbefa804de6087e77c4c211ef57ab83869e3e18627f8421540ae9a8b61da847d0da513c56c5feba397ab2b4a1a2ef67c6f17162c8dfdb41901ad70bca8195fd35bcea259'
 ```
 
 ### Speed Test
 
 ```
-(.env) Jack: xof-py % python3 speed_test.py
-10_000 calls with rust sha3: 0.010552167892456055
-10_000 calls with hashlib: 0.011549949645996094
+(.env) Jack: xof-py % python3 speed_test.py               
+10_000 calls with rust sha3: 0.0849299430847168
+10_000 calls with hashlib: 0.06996297836303711
 ```

speed_test.py

@@ -1,14 +1,26 @@
+import random
 from hashlib import shake_128
-from xof_py import shake_123
-
+from xof_py import pyo3_shake_128
 import time
+import os
+
+for _ in range(100):
+    absorb = os.urandom(32)
+    n = random.randint(1, 1000)
+    a = shake_128(absorb).digest(n)
+    b = pyo3_shake_128(absorb, n)
+    assert a == b
 
+random.seed(0)
 t0 = time.time()
 for _ in range(10_000):
-    a = shake_123()
+    n = random.randint(1, 5000)
+    a = pyo3_shake_128(b"123", n)
 print(f"10_000 calls with rust sha3: {time.time() - t0 }")
 
+random.seed(0)
 t0 = time.time()
 for _ in range(10_000):
-    a = shake_128(b"123").digest(10)
-print(f"10_000 calls with hashlib: {time.time() - t0}")
+    n = random.randint(1, 5000)
+    a = shake_128(b"123").digest(n)
+print(f"10_000 calls with hashlib: {time.time() - t0}")

src/lib.rs

@@ -2,21 +2,21 @@ use pyo3::prelude::*;
 use pyo3::types::PyBytes;
 use sha3::{Shake128, digest::{Update, ExtendableOutput, XofReader}};
 
-/// Returns the first 10 bytes of Shake128 initialized with
-/// b"123"
+/// Returns the first n bytes of Shake128 initialized with
+/// input_bytes
 #[pyfunction]
-fn shake_123(py: Python) -> PyObject {
+fn pyo3_shake_128(py: Python, input_bytes: &[u8], n: usize) -> PyObject {
     let mut hasher = Shake128::default();
-    hasher.update(b"123");
+    hasher.update(input_bytes);
     let mut xof = hasher.finalize_xof();
-    let mut res = [0u8; 10];
+    let mut res = vec![0u8; n];
     xof.read(&mut res);
     PyBytes::new_bound(py, &res).into()
 }
 
 /// A Python module implemented in Rust.
 #[pymodule]
 fn xof_py(m: &Bound<'_, PyModule>) -> PyResult<()> {
-    m.add_function(wrap_pyfunction!(shake_123, m)?)?;
+    m.add_function(wrap_pyfunction!(pyo3_shake_128, m)?)?;
     Ok(())
 }