python bindings for binutils assemble and disassemble

.github/workflows/build-test.yml

@@ -152,7 +152,11 @@ jobs:
         pip install --no-index --find-links target/wheels/ clvm_tools_rs
         pip install clvm_rs
         pip install clvm_tools
-        cd resources/tests && python test_clvm_step.py && python mandelbrot-cldb.py && python test_compile_from_string.py
+        cd resources/tests
+        python test_clvm_step.py
+        python mandelbrot-cldb.py
+        python test_compile_from_string.py
+        python test_binutils_api.py
 
     - name: "Test step run with mandelbrot, setting print only"
       run: |

resources/tests/test_binutils_api.py

@@ -0,0 +1,82 @@
+from clvm_tools_rs import binutils
+from clvm_tools.binutils import assemble, disassemble
+from clvm_rs.program import Program
+from random import randint
+
+test_cases = [
+    "(q)",
+    "()",
+    "(1 . 2)",
+    "(() . ())",
+    "(1 2 3 . ())",
+    "(1 2 3)",
+    "((4 5 6) (7 8 9))",
+    "(i (q . 1) (q . 2) (q . 3))",
+    "((0x33 . i))"
+]
+
+def extend_atom():
+    n_choice = randint(0,14)
+    if n_choice > 10:
+        return ""
+    return f"{chr(n_choice+97)}@"
+
+def pad(n,v,s):
+    if len(s) < n:
+        return "".join(([v]*(n - len(s)))) + s
+    else:
+        return s
+
+def extend_hex():
+    n_choice = randint(0,400)
+    if n_choice > 255:
+        return ""
+    return f"{pad(2,'0',hex(n_choice)[2:])}^"
+
+def any_item():
+    keys = list(filter(lambda k: k not in "&^", production_rules.keys()))
+    key_choice = randint(0,len(keys) - 1)
+    return keys[key_choice]
+
+production_rules = {
+    "!": lambda: "()",
+    "~": lambda: "(*&",
+    "&": lambda: [")", " *&", " . *)"][randint(0,2)],
+    "#": lambda: str(randint(0,2**48) - (2**47)),
+    "@": lambda: f"{chr(randint(0,10)+97)}{extend_atom()}",
+    "$": lambda: f"0x{pad(2,'0',hex(randint(0,255))[2:])}^",
+    "^": extend_hex,
+    "*": any_item
+}
+production_keys = "".join(production_rules.keys())
+
+def do_replacement(x):
+    if x in production_rules:
+        return production_rules[x]()
+    else:
+        return x
+
+def generate_random_sexp(s):
+    def generate_round(s):
+        return "".join([do_replacement(x) for x in s])
+
+    s = generate_round(s)
+    while any([x in production_keys for x in s]):
+        s = generate_round(s)
+
+    return s
+
+for i in range(20):
+    test_cases.append(generate_random_sexp("*"))
+
+for t in test_cases:
+    print(t)
+    expected = assemble(t)
+    assembled = binutils.assemble_generic(lambda x,y: Program.to((x,y)), Program.to, t)
+    print(expected, assembled)
+
+    assert expected == assembled
+
+    disassembled = binutils.disassemble_generic(bytes(assembled))
+    expected_dis = disassemble(expected)
+    assert expected_dis == disassembled

src/py/api.rs

@@ -42,6 +42,7 @@ use crate::util::{gentle_overwrite, version};
 
 use crate::py::pyval::{clvm_value_to_python, python_value_to_clvm};
 
+use super::binutils::create_binutils_module;
 use super::cmds::create_cmds_module;
 
 create_exception!(mymodule, CldbError, PyException);
@@ -495,6 +496,7 @@ pub fn compose_run_function(
 #[pymodule]
 fn clvm_tools_rs(py: Python, m: &PyModule) -> PyResult<()> {
     m.add_submodule(create_cmds_module(py)?)?;
+    m.add_submodule(create_binutils_module(py)?)?;
 
     m.add("CldbError", py.get_type::<CldbError>())?;
     m.add("CompError", py.get_type::<CompError>())?;

src/py/binutils.rs

@@ -0,0 +1,105 @@
+use std::collections::HashMap;
+
+use crate::classic::clvm::__type_compatibility__::{Bytes, BytesFromType, Stream};
+use crate::classic::clvm::serialize::{sexp_from_stream, SimpleCreateCLVMObject};
+use clvm_rs::allocator::{Allocator, NodePtr, SExp};
+
+use pyo3::exceptions::PyException;
+use pyo3::prelude::*;
+use pyo3::types::{PyBytes, PyTuple};
+
+use crate::classic::clvm_tools::binutils;
+
+create_exception!(mymodule, ConvError, PyException);
+
+fn convert_to_external(
+    allocator: &Allocator,
+    cons: &PyAny,
+    from_bytes: &PyAny,
+    root_node: NodePtr,
+) -> PyResult<PyObject> {
+    let mut stack: Vec<NodePtr> = vec![root_node];
+    let mut finished = HashMap::new();
+
+    while let Some(node) = stack.last() {
+        let node = *node; // To avoid borrowing issues with the stack
+        match allocator.sexp(node) {
+            SExp::Pair(left, right) => {
+                let left_finished = finished.contains_key(&left);
+                let right_finished = finished.contains_key(&right);
+
+                if left_finished && right_finished {
+                    stack.pop();
+
+                    let result: PyObject = Python::with_gil(|py| {
+                        let a = finished.get(&left).unwrap();
+                        let b = finished.get(&right).unwrap();
+                        let args = PyTuple::new(py, &[a, b]);
+                        cons.call1(args)
+                    })?
+                    .into();
+
+                    finished.insert(node, result);
+                } else {
+                    if !left_finished {
+                        stack.push(left);
+                    }
+                    if !right_finished {
+                        stack.push(right);
+                    }
+                }
+            }
+            SExp::Atom => {
+                stack.pop();
+
+                if !finished.contains_key(&node) {
+                    let converted: PyObject = Python::with_gil(|py| {
+                        let bytes = PyBytes::new(py, allocator.atom(node));
+                        let args = PyTuple::new(py, &[bytes]);
+                        from_bytes.call1(args)
+                    })?
+                    .into();
+                    finished.insert(node, converted);
+                }
+            }
+        }
+    }
+
+    finished
+        .get(&root_node)
+        .cloned()
+        .ok_or_else(|| ConvError::new_err("error converting assembled value"))
+}
+
+#[pyfunction]
+pub fn assemble_generic(cons: &PyAny, from_bytes: &PyAny, args: String) -> PyResult<PyObject> {
+    let mut allocator = Allocator::new();
+    let assembled =
+        binutils::assemble(&mut allocator, &args).map_err(|e| ConvError::new_err(e.to_string()))?;
+    convert_to_external(&allocator, cons, from_bytes, assembled)
+}
+
+#[pyfunction]
+pub fn disassemble_generic(program_bytes: &PyBytes) -> PyResult<String> {
+    let mut allocator = Allocator::new();
+    let mut stream = Stream::new(Some(Bytes::new(Some(BytesFromType::Raw(
+        program_bytes.as_bytes().to_vec(),
+    )))));
+
+    let sexp = sexp_from_stream(
+        &mut allocator,
+        &mut stream,
+        Box::new(SimpleCreateCLVMObject {}),
+    )
+    .map_err(|e| ConvError::new_err(e.to_string()))?;
+
+    let disassembled = binutils::disassemble(&allocator, sexp.1, Some(0));
+    Ok(disassembled)
+}
+
+pub fn create_binutils_module(py: Python) -> PyResult<&'_ PyModule> {
+    let m = PyModule::new(py, "binutils")?;
+    m.add_function(wrap_pyfunction!(assemble_generic, m)?)?;
+    m.add_function(wrap_pyfunction!(disassemble_generic, m)?)?;
+    Ok(m)
+}

src/py/mod.rs

@@ -1,3 +1,4 @@
 pub mod api;
+mod binutils;
 mod cmds;
 pub mod pyval;

src/tests/classic/smoke.rs

@@ -13,8 +13,9 @@ use crate::classic::clvm::__type_compatibility__::{
     pybytes_repr, t, Bytes, Stream, UnvalidatedBytesFromType,
 };
 use crate::classic::clvm::serialize::{sexp_from_stream, SimpleCreateCLVMObject};
-use crate::classic::clvm::sexp::{First, NodeSel, Rest, SelectNode, ThisNode};
+use crate::classic::clvm::sexp::{sexp_as_bin, First, NodeSel, Rest, SelectNode, ThisNode};
 use crate::classic::clvm::syntax_error::SyntaxErr;
+use crate::classic::clvm_tools::binutils;
 use crate::classic::clvm_tools::cmds::{launch_tool, OpcConversion, OpdConversion, TConversion};
 
 use crate::classic::clvm_tools::binutils::{assemble, assemble_from_ir, disassemble};
@@ -258,6 +259,17 @@ fn compile_program<'a>(
     return res.map(|x| disassemble(allocator, x.1, Some(0)));
 }
 
+#[test]
+fn binutils_assemble() {
+    let mut allocator = Allocator::new();
+    let src = "(q)".to_string();
+    let assembled = binutils::assemble(&mut allocator, &src)
+        .map_err(|e| e.to_string())
+        .map(|sexp| t(sexp, sexp_as_bin(&mut allocator, sexp).hex()))
+        .unwrap();
+    assert_eq!(assembled.rest(), "ff0180");
+}
+
 #[test]
 fn quoted_negative() {
     let mut s = Stream::new(None);