Skip to content

Commit

Permalink
refactor interface_py.rs
Browse files Browse the repository at this point in the history
  • Loading branch information
cdump committed Feb 16, 2025
1 parent 3f0f3c3 commit 15799d5
Show file tree
Hide file tree
Showing 2 changed files with 140 additions and 158 deletions.
5 changes: 4 additions & 1 deletion python/test_python.py
Original file line number Diff line number Diff line change
@@ -1,15 +1,18 @@
from evmole import contract_info
from evmole import Contract, Function, StorageRecord

code = '6080604052348015600e575f80fd5b50600436106026575f3560e01c8063fae7ab8214602a575b5f80fd5b603960353660046062565b6052565b60405163ffffffff909116815260200160405180910390f35b5f605c826001608a565b92915050565b5f602082840312156071575f80fd5b813563ffffffff811681146083575f80fd5b9392505050565b63ffffffff8181168382160190811115605c57634e487b7160e01b5f52601160045260245ffd'

info = contract_info(code, selectors=True, arguments=True, state_mutability=True, disassemble=True)
assert isinstance(info, Contract)
assert info.functions is not None
assert len(info.functions) == 1
assert isinstance(info.functions[0], Function)
assert info.functions[0].selector == 'fae7ab82'
assert info.functions[0].arguments == 'uint32'
assert info.functions[0].state_mutability == 'pure'

assert info.disassembled is not None
assert info.disassembled[0] == (0, 'PUSH1 80')

print(f'Success, {info}')
print(f'Success #1, {info}')
293 changes: 136 additions & 157 deletions src/interface_py.rs
Original file line number Diff line number Diff line change
Expand Up @@ -22,173 +22,152 @@ fn input_to_bytes<'a>(code: &'a Bound<'a, PyAny>) -> PyResult<Cow<'a, [u8]>> {
}
}

#[pyclass(name = "Function")]
#[derive(Clone)]
struct PyFunction {
#[pyo3(get)]
selector: String,

#[pyo3(get)]
bytecode_offset: usize,

#[pyo3(get)]
arguments: Option<String>,

#[pyo3(get)]
state_mutability: Option<String>,
}

impl PyFunction {
fn repr(&self) -> String {
format!(
"Function(selector={},bytecode_offset={},arguments={},state_mutability={})",
self.selector,
self.bytecode_offset,
self.arguments.as_deref().unwrap_or("None"),
self.state_mutability.as_deref().unwrap_or("None"),
)
#[pymodule]
mod evmole {
use super::*;

#[pyclass(name = "Function", get_all)]
#[derive(Clone)]
struct PyFunction {
selector: String,
bytecode_offset: usize,
arguments: Option<String>,
state_mutability: Option<String>,
}
}

#[pymethods]
impl PyFunction {
fn __repr__(slf: &Bound<'_, Self>) -> PyResult<String> {
Ok(PyFunction::repr(&slf.borrow()))
#[pymethods]
impl PyFunction {
fn __repr__(&self) -> String {
format!(
"Function(selector={:?}, bytecode_offset={}, arguments={}, state_mutability={})",
self.selector,
self.bytecode_offset,
self.arguments
.as_deref()
.map_or_else(|| "None".to_string(), |v| format!("\"{v}\"")),
self.state_mutability
.as_deref()
.map_or_else(|| "None".to_string(), |v| format!("\"{v}\"")),
)
}
}
}

#[pyclass(name = "StorageRecord")]
#[derive(Clone)]
struct PyStorageRecord {
slot: String,
offset: u8,
r#type: String,
reads: Vec<String>,
writes: Vec<String>,
}

impl PyStorageRecord {
fn repr(&self) -> String {
format!(
"StorageRecord(slot={},offset={},type={},reads={:?},writes={:?})",
self.slot, self.offset, self.r#type, self.reads, self.writes
)
#[pyclass(name = "StorageRecord")]
#[derive(Clone)]
struct PyStorageRecord {
slot: String,
offset: u8,
r#type: String,
reads: Vec<String>,
writes: Vec<String>,
}
}

#[pymethods]
impl PyStorageRecord {
fn __repr__(slf: &Bound<'_, Self>) -> PyResult<String> {
Ok(PyStorageRecord::repr(&slf.borrow()))
#[pymethods]
impl PyStorageRecord {
fn __repr__(&self) -> String {
format!(
"StorageRecord(slot=\"{}\", offset={}, type=\"{}\", reads={:?}, writes={:?})",
self.slot, self.offset, self.r#type, self.reads, self.writes
)
}
}
}

#[pyclass(name = "Contract")]
struct PyContract {
#[pyo3(get)]
functions: Option<Vec<PyFunction>>,

#[pyo3(get)]
storage: Option<Vec<PyStorageRecord>>,

#[pyo3(get)]
disassembled: Option<Vec<(usize, String)>>,
}

#[pymethods]
impl PyContract {
fn __repr__(slf: &Bound<'_, Self>) -> PyResult<String> {
Ok(format!(
"Contract(functions={},storage={})",
if let Some(ref v) = slf.borrow().functions {
format!(
"[{}]",
v.iter().map(|v| v.repr()).collect::<Vec<_>>().join(",")
)
} else {
"None".to_string()
},
if let Some(ref v) = slf.borrow().storage {
format!(
"[{}]",
v.iter().map(|v| v.repr()).collect::<Vec<_>>().join(",")
)
} else {
"None".to_string()
},
))
#[pyclass(name = "Contract", get_all)]
struct PyContract {
functions: Option<Vec<PyFunction>>,
storage: Option<Vec<PyStorageRecord>>,
disassembled: Option<Vec<(usize, String)>>,
}
}

#[pyfunction]
#[pyo3(signature = (code, *, selectors=false, arguments=false, state_mutability=false, storage=false, disassemble=false))]
fn contract_info(
code: &Bound<'_, PyAny>,
selectors: bool,
arguments: bool,
state_mutability: bool,
storage: bool,
disassemble: bool,
) -> PyResult<PyContract> {
let code_bytes = input_to_bytes(code)?;
let mut args = crate::ContractInfoArgs::new(&code_bytes);

if selectors {
args = args.with_selectors();
}
if arguments {
args = args.with_arguments();
}
if state_mutability {
args = args.with_state_mutability();
}
if storage {
args = args.with_storage();
}
if disassemble {
args = args.with_disassemble();
#[pymethods]
impl PyContract {
fn __repr__(&self) -> String {
format!(
"Contract(functions={}, storage={}, disassembled={})",
self.functions.as_ref().map_or_else(|| "None".to_string(), |v|
format!(
"[{}]",
v.iter().map(|v| v.__repr__()).collect::<Vec<_>>().join(", ")
)
),
self.storage.as_ref().map_or_else(|| "None".to_string(), |v|
format!(
"[{}]",
v.iter().map(|v| v.__repr__()).collect::<Vec<_>>().join(", ")
)
),
self.disassembled.as_ref().map_or_else(|| "None".to_string(), |v|
format!("{:?}", v)
),
)
}
}

let info = crate::contract_info(args);

let functions = info.functions.map(|fns| {
fns.into_iter()
.map(|f| PyFunction {
selector: hex::encode(f.selector),
bytecode_offset: f.bytecode_offset,
arguments: f.arguments.map(|fargs| {
fargs
.into_iter()
.map(|t| t.sol_type_name().to_string())
.collect::<Vec<String>>()
.join(",")
}),
state_mutability: f.state_mutability.map(|sm| sm.as_json_str().to_string()),
})
.collect()
});

let storage = info.storage.map(|st| {
st.into_iter()
.map(|v| PyStorageRecord {
slot: hex::encode(v.slot),
offset: v.offset,
r#type: v.r#type,
reads: v.reads.into_iter().map(hex::encode).collect(),
writes: v.writes.into_iter().map(hex::encode).collect(),
})
.collect()
});

Ok(PyContract {
functions,
storage,
disassembled: info.disassembled,
})
}

#[pymodule]
fn evmole(m: &Bound<'_, PyModule>) -> PyResult<()> {
m.add_function(wrap_pyfunction!(contract_info, m)?)?;
Ok(())
#[pyfunction]
#[pyo3(signature = (code, *, selectors=false, arguments=false, state_mutability=false, storage=false, disassemble=false))]
#[allow(clippy::too_many_arguments)]
fn contract_info(
code: &Bound<'_, PyAny>,
selectors: bool,
arguments: bool,
state_mutability: bool,
storage: bool,
disassemble: bool,
) -> PyResult<PyContract> {
let code_bytes = input_to_bytes(code)?;
let mut args = crate::ContractInfoArgs::new(&code_bytes);

if selectors {
args = args.with_selectors();
}
if arguments {
args = args.with_arguments();
}
if state_mutability {
args = args.with_state_mutability();
}
if storage {
args = args.with_storage();
}
if disassemble {
args = args.with_disassemble();
}

let info = crate::contract_info(args);

let functions = info.functions.map(|fns| {
fns.into_iter()
.map(|f| PyFunction {
selector: hex::encode(f.selector),
bytecode_offset: f.bytecode_offset,
arguments: f.arguments.map(|fargs| {
fargs
.into_iter()
.map(|t| t.sol_type_name().to_string())
.collect::<Vec<String>>()
.join(", ")
}),
state_mutability: f.state_mutability.map(|sm| sm.as_json_str().to_string()),
})
.collect()
});

let storage = info.storage.map(|st| {
st.into_iter()
.map(|v| PyStorageRecord {
slot: hex::encode(v.slot),
offset: v.offset,
r#type: v.r#type,
reads: v.reads.into_iter().map(hex::encode).collect(),
writes: v.writes.into_iter().map(hex::encode).collect(),
})
.collect()
});

Ok(PyContract {
functions,
storage,
disassembled: info.disassembled,
})
}
}

0 comments on commit 15799d5

Please sign in to comment.