Rework py api

src/py/api.rs

@@ -1,24 +1,104 @@
+use std::cell::RefMut;
+
 use pyo3::prelude::*;
+use pyo3::types::{PyBytes, PyDict};
 use pyo3::wrap_pyfunction;
 
+use crate::allocator::Allocator;
+use crate::core_ops::*;
+use crate::err_utils::err;
+use crate::more_ops::*;
+use crate::node::Node;
+use crate::serialize::node_to_bytes;
+
+use super::arena::Arena;
+use super::dialect::{Dialect, PyMultiOpFn};
+use super::f_table::OpFn;
 use super::lazy_node::LazyNode;
+use super::native_op::NativeOp;
 use super::run_program::{
     __pyo3_get_function_deserialize_and_run_program,
-    __pyo3_get_function_deserialize_and_run_program2,
-    __pyo3_get_function_serialize_and_run_program, __pyo3_get_function_serialized_length,
+    __pyo3_get_function_deserialize_and_run_program2, __pyo3_get_function_serialized_length,
     STRICT_MODE,
 };
 
+#[pyfunction]
+pub fn native_opcodes_dict(py: Python) -> PyResult<PyObject> {
+    let opcode_lookup: [(OpFn, &str); 30] = [
+        (op_if, "op_if"),
+        (op_cons, "op_cons"),
+        (op_first, "op_first"),
+        (op_rest, "op_rest"),
+        (op_listp, "op_listp"),
+        (op_raise, "op_raise"),
+        (op_eq, "op_eq"),
+        (op_sha256, "op_sha256"),
+        (op_add, "op_add"),
+        (op_subtract, "op_subtract"),
+        (op_multiply, "op_multiply"),
+        (op_divmod, "op_divmod"),
+        (op_substr, "op_substr"),
+        (op_strlen, "op_strlen"),
+        (op_point_add, "op_point_add"),
+        (op_pubkey_for_exp, "op_pubkey_for_exp"),
+        (op_concat, "op_concat"),
+        (op_gr, "op_gr"),
+        (op_gr_bytes, "op_gr_bytes"),
+        (op_logand, "op_logand"),
+        (op_logior, "op_logior"),
+        (op_logxor, "op_logxor"),
+        (op_lognot, "op_lognot"),
+        (op_ash, "op_ash"),
+        (op_lsh, "op_lsh"),
+        (op_not, "op_not"),
+        (op_any, "op_any"),
+        (op_all, "op_all"),
+        (op_softfork, "op_softfork"),
+        (op_div, "op_div"),
+    ];
+    let r = PyDict::new(py);
+    for (f, name) in opcode_lookup.iter() {
+        r.set_item(name, PyCell::new(py, NativeOp::new(*f))?)?;
+    }
+    Ok(r.to_object(py))
+}
+
+#[pyfunction]
+fn serialize_to_bytes<'p>(py: Python<'p>, sexp: &PyAny) -> PyResult<&'p PyBytes> {
+    let arena_cell = Arena::new_cell(py)?;
+    let arena = arena_cell.borrow();
+    let mut allocator_refcell: RefMut<Allocator> = arena.allocator();
+    let allocator: &mut Allocator = &mut allocator_refcell as &mut Allocator;
+
+    let ptr = arena.as_native(py, allocator, sexp)?;
+
+    let node = Node::new(allocator, ptr);
+    let s: Vec<u8> = node_to_bytes(&node)?;
+    Ok(PyBytes::new(py, &s))
+}
+
 /// This module is a python module implemented in Rust.
 #[pymodule]
 fn clvm_rs(_py: Python, m: &PyModule) -> PyResult<()> {
-    m.add_function(wrap_pyfunction!(serialize_and_run_program, m)?)?;
-    m.add_function(wrap_pyfunction!(deserialize_and_run_program, m)?)?;
-    m.add_function(wrap_pyfunction!(deserialize_and_run_program2, m)?)?;
-    m.add("STRICT_MODE", STRICT_MODE)?;
+    m.add_class::<Arena>()?;
+    m.add_class::<Dialect>()?;
     m.add_class::<LazyNode>()?;
 
+    m.add_function(wrap_pyfunction!(native_opcodes_dict, m)?)?;
     m.add_function(wrap_pyfunction!(serialized_length, m)?)?;
 
+    m.add(
+        "NATIVE_OP_UNKNOWN_STRICT",
+        PyMultiOpFn::new(|_a, b, _op, _d| err(b, "unimplemented operator")),
+    )?;
+
+    m.add("NATIVE_OP_UNKNOWN_NON_STRICT", PyMultiOpFn::new(op_unknown))?;
+
+    m.add_function(wrap_pyfunction!(serialize_to_bytes, m)?)?;
+
+    m.add_function(wrap_pyfunction!(deserialize_and_run_program, m)?)?;
+    m.add_function(wrap_pyfunction!(deserialize_and_run_program2, m)?)?;
+    m.add("STRICT_MODE", STRICT_MODE)?;
+
     Ok(())
 }

src/py/arena.rs

@@ -0,0 +1,86 @@
+/// An `Arena` is a collection of objects representing a program and
+/// its arguments, and intermediate values reached while running
+/// a program. Objects can be created in an `Arena` but are never
+/// dropped until the `Arena` is dropped.
+use std::cell::{RefCell, RefMut};
+
+use pyo3::prelude::pyclass;
+use pyo3::prelude::*;
+
+use super::bridge_cache::BridgeCache;
+use crate::allocator::{Allocator, NodePtr};
+use crate::serialize::node_from_bytes;
+
+#[pyclass(subclass, unsendable)]
+pub struct Arena {
+    arena: RefCell<Allocator>,
+    cache: BridgeCache,
+}
+
+#[pymethods]
+impl Arena {
+    #[new]
+    pub fn new(py: Python, new_obj_f: PyObject) -> PyResult<Self> {
+        Ok(Arena {
+            arena: RefCell::new(Allocator::default()),
+            cache: BridgeCache::new(py, new_obj_f)?,
+        })
+    }
+
+    /// deserialize `bytes` into an object in this `Arena`
+    pub fn deserialize<'p>(&self, py: Python<'p>, blob: &[u8]) -> PyResult<&'p PyAny> {
+        let allocator: &mut Allocator = &mut self.allocator() as &mut Allocator;
+        let ptr = node_from_bytes(allocator, blob)?;
+        self.as_python(py, allocator, ptr)
+    }
+
+    /// copy this python object into this `Arena` if it's not yet in the cache
+    /// (otherwise it returns the previously cached object)
+    pub fn include<'p>(&self, py: Python<'p>, obj: &'p PyAny) -> PyResult<&'p PyAny> {
+        let allocator = &mut self.allocator();
+        let ptr = self.as_native(py, allocator, obj)?;
+        self.as_python(py, allocator, ptr)
+    }
+
+    /// copy this python object into this `Arena` if it's not yet in the cache
+    /// (otherwise it returns the previously cached object)
+    pub fn ptr_for_obj(&self, py: Python, obj: &PyAny) -> PyResult<i32> {
+        self.as_native(py, &mut self.allocator(), obj)
+    }
+}
+
+impl Arena {
+    pub fn new_cell_obj(py: Python, new_obj_f: PyObject) -> PyResult<&PyCell<Self>> {
+        PyCell::new(py, Arena::new(py, new_obj_f)?)
+    }
+
+    pub fn new_cell(py: Python) -> PyResult<&PyCell<Self>> {
+        Self::new_cell_obj(py, py.eval("lambda sexp: sexp", None, None)?.to_object(py))
+    }
+
+    pub fn obj_for_ptr<'p>(&self, py: Python<'p>, ptr: i32) -> PyResult<&'p PyAny> {
+        self.as_python(py, &mut self.allocator(), ptr)
+    }
+
+    pub fn allocator(&self) -> RefMut<Allocator> {
+        self.arena.borrow_mut()
+    }
+
+    pub fn as_native(
+        &self,
+        py: Python,
+        allocator: &mut Allocator,
+        obj: &PyAny,
+    ) -> PyResult<NodePtr> {
+        self.cache.as_native(py, allocator, obj)
+    }
+
+    pub fn as_python<'p>(
+        &self,
+        py: Python<'p>,
+        allocator: &mut Allocator,
+        ptr: NodePtr,
+    ) -> PyResult<&'p PyAny> {
+        self.cache.as_python(py, allocator, ptr)
+    }
+}