zero-copy bytes passing

This commit is contained in:
Richard Patel 2021-08-19 01:13:19 +02:00
parent 57a317a6b2
commit 4fc2d32e02
1 changed files with 27 additions and 15 deletions

View File

@ -1,4 +1,5 @@
use pyo3::prelude::*; use pyo3::prelude::*;
use pyo3::types::PyBytes;
use pyo3::{PyIterProtocol, PyObjectProtocol}; use pyo3::{PyIterProtocol, PyObjectProtocol};
#[pyclass] #[pyclass]
@ -95,9 +96,10 @@ impl<'a> PyObjectProtocol<'a> for Ins {
#[pyclass] #[pyclass]
struct DisasmIterator { struct DisasmIterator {
bytes: Vec<u8>, bytes: Py<PyBytes>,
addr: u32, addr: u32,
offset: u32, offset: u32,
left: usize,
} }
#[pyproto] #[pyproto]
@ -106,14 +108,16 @@ impl PyIterProtocol for DisasmIterator {
slf slf
} }
fn __next__(mut slf: PyRefMut<Self>) -> PyResult<Option<Ins>> { fn __next__(mut slf: PyRefMut<Self>) -> PyResult<Option<Ins>> {
if (slf.bytes.len() as u32) - slf.offset < 4 { if slf.left < 4 {
return Ok(None); return Ok(None);
} }
let code = ((slf.bytes[(slf.offset) as usize] as u32) << 24) let bytes = slf.bytes.as_ref(slf.py());
| ((slf.bytes[(slf.offset + 1) as usize] as u32) << 16) let code = ((bytes[(slf.offset) as usize] as u32) << 24)
| ((slf.bytes[(slf.offset + 2) as usize] as u32) << 8) | ((bytes[(slf.offset + 1) as usize] as u32) << 16)
| (slf.bytes[(slf.offset + 3) as usize] as u32); | ((bytes[(slf.offset + 2) as usize] as u32) << 8)
| (bytes[(slf.offset + 3) as usize] as u32);
slf.offset += 4; slf.offset += 4;
slf.left -= 4;
let ins = Ins::new(code, slf.addr); let ins = Ins::new(code, slf.addr);
slf.addr += 4; slf.addr += 4;
Ok(Some(ins)) Ok(Some(ins))
@ -121,17 +125,25 @@ impl PyIterProtocol for DisasmIterator {
} }
#[pyfunction(code, addr, offset = "0", size = "None")] #[pyfunction(code, addr, offset = "0", size = "None")]
fn disasm_iter(code: &[u8], addr: u32, offset: u32, size: Option<u32>) -> PyResult<DisasmIterator> { fn disasm_iter(
let mut bytes = &code[offset as usize..]; code: &PyBytes,
if let Some(size) = size { addr: u32,
if (size as usize) < bytes.len() { offset: u32,
bytes = &bytes[..size as usize]; size: Option<u32>,
} ) -> PyResult<DisasmIterator> {
} let left = match size {
None => code
.as_bytes()
.len()
.checked_sub(offset as usize)
.unwrap_or(0),
Some(v) => v as usize,
};
Ok(DisasmIterator { Ok(DisasmIterator {
bytes: bytes.to_vec(), bytes: code.into(),
addr, addr,
offset: 0, offset,
left,
}) })
} }