Allows Python Function Registration

Cargo.toml

@@ -8,13 +8,17 @@ description = "A library for defining domain specific languages in a polymorphic
 repository = "https://github.com/mlb2251/lambdas"
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+[features]
+python = ["dep:pyo3"]
+default = []
 
 [dependencies]
 serde_json = {version = "1.0", features = ["preserve_order"]}
 serde = {version = "1.0", features = ["derive"]}
 once_cell = "1.16.0"
 string_cache = "0.8.4"
 rustc-hash = "1.1.0"
+pyo3 = { version = "0.25.1", features = ["auto-initialize"], optional=true}
 
 # [profile.release]
 # debug = true # for flamegraphs

src/domains/mod.rs

@@ -2,3 +2,9 @@
 // just register each domain here by including it with `pub mod domain_name;`
 pub mod simple;
 pub mod prim_lists;
+
+#[cfg(feature = "python")]
+pub mod simple_python;
+
+#[cfg(feature = "python")]
+pub mod py;

src/domains/py.rs

@@ -0,0 +1,120 @@
+use crate::*;
+use std::collections::HashSet;
+use std::sync::Arc;
+
+
+use pyo3::prelude::*;
+use pyo3::types::PyAny;
+
+#[derive(Clone,Debug, PartialEq, Eq, Hash)]
+pub enum PyVal {
+    Int(i32),
+    List(Vec<Val>),
+}
+
+#[derive(Clone,Debug, PartialEq, Eq, Hash)]
+pub enum PyType {
+    TInt,
+    TList
+}
+
+
+type Val = crate::eval::Val<PyVal>;
+
+
+#[cfg(feature = "python")]
+pub fn create_python_production<D: Domain>(
+    name: Symbol,
+    tp: SlowType,
+    lazy_args: Option<&[usize]>,
+    pyfunc: Py<PyAny>,
+) -> Production<D> {
+    let arity = tp.arity();
+    let lazy: HashSet<usize> = lazy_args
+        .map(|xs| xs.iter().copied().collect())
+        .unwrap_or_default();
+
+    use crate::eval::{CurriedFn, Val};
+
+    Production {
+        name: name.clone(),
+        val: Val::PrimFun(CurriedFn::<D>::new(name.clone(), arity)),
+        tp,
+        arity,
+        lazy_args: lazy,
+        fn_ptr: Some(FnPtr::Python(Arc::new(pyfunc)))
+    }
+}
+
+// From<Val> impls are needed for unwrapping values. We can assume the program
+// has been type checked so it's okay to panic if the type is wrong. Each val variant
+// must map to exactly one unwrapped type (though it doesnt need to be one to one in the
+// other direction)
+impl FromVal<PyVal> for i32 {
+    fn from_val(v: Val) -> Result<Self, VError> {
+        match v {
+            Dom(PyVal::Int(i)) => Ok(i),
+            _ => Err("from_val_to_i32: not an int".into())
+        }
+    }
+}
+impl<T: FromVal<PyVal>> FromVal<PyVal> for Vec<T> {
+    fn from_val(v: Val) -> Result<Self, VError> {
+        match v {
+            Dom(PyVal::List(v)) => v.into_iter().map(|v| T::from_val(v)).collect(),
+            _ => Err("from_val_to_vec: not a list".into())
+        }
+    }
+}
+
+impl From<i32> for Val {
+    fn from(i: i32) -> Val {
+        Dom(PyVal::Int(i))
+    }
+}
+impl<T: Into<Val>> From<Vec<T>> for Val {
+    fn from(vec: Vec<T>) -> Val {
+        Dom(PyVal::List(vec.into_iter().map(|v| v.into()).collect()))
+    }
+}
+
+impl Domain for PyVal {
+    type Data = ();
+
+    #[cfg(feature = "python")]
+    fn py_val_to_py(py: Python<'_>, v: crate::eval::Val<Self>) -> PyResult<Py<PyAny>> {
+        crate::domains::simple_python::val_to_py(py, v)
+    }
+
+    #[cfg(feature = "python")]
+    fn py_py_to_val(obj: &Bound<'_, PyAny>) -> Result<crate::eval::Val<Self>, String> {
+        crate::domains::simple_python::py_to_val(obj)
+    }
+
+    fn new_dsl() -> DSL<Self> {
+        let prods = vec![];
+        let dsl = DSL::new(prods);
+        dsl
+    }
+
+    fn val_of_prim_fallback(p: &Symbol) -> Option<Val> {
+        None
+    }
+
+    fn type_of_dom_val(&self) -> SlowType {
+        match self {
+            PyVal::Int(_) => SlowType::base(Symbol::from("int")),
+            PyVal::List(xs) => {
+                let elem_tp = if xs.is_empty() {
+                    SlowType::Var(0) // (list t0)
+                } else {
+                    let result = Self::type_of_dom_val(&xs.first().unwrap().clone().dom().unwrap());
+                    assert!(xs.iter().all(|v| result == Self::type_of_dom_val(&v.clone().dom().unwrap())));
+                    result
+                };
+                SlowType::Term("list".into(),vec![elem_tp])
+            },
+        }
+    }
+
+}

src/domains/simple_python.rs

@@ -0,0 +1,43 @@
+#![cfg(feature = "python")]
+
+use pyo3::prelude::*;
+use pyo3::types::PyList;
+use pyo3::conversion::IntoPyObjectExt;
+
+use crate::domains::py::PyVal;
+use crate::eval;
+type Val = eval::Val<PyVal>;
+use PyVal::*;
+
+
+pub fn val_to_py(py: Python<'_>, v: Val) -> PyResult<Py<PyAny>> {
+    match v.dom().expect("Val should be Dom") {
+        Int(i) => {
+            Ok(i.into_py_any(py)?)
+        }
+        List(xs) => {
+            let mut elems: Vec<Py<PyAny>> = Vec::with_capacity(xs.len());
+            for x in xs {
+                elems.push(val_to_py(py, x.clone())?);
+            }
+
+            let list_bound: Bound<'_, PyList> = PyList::new(py, &elems)?;
+
+            Ok(list_bound.into_any().unbind())
+        }
+    }
+}
+
+pub fn py_to_val(obj: &Bound<'_, PyAny>) -> Result<Val, String> {
+    if let Ok(i) = obj.extract::<i32>() {
+        return Ok(Val::from(Int(i)));
+    }
+    if let Ok(list) = obj.downcast::<PyList>() {
+        let mut out: Vec<Val> = Vec::with_capacity(list.len());
+        for item in list.iter() {
+            out.push(py_to_val(&item)?);
+        }
+        return Ok(Val::from(List(out)));
+    }
+    Err("unsupported Python type for PyVal".into())
+}

src/dsl.rs

@@ -3,18 +3,31 @@ use crate::*;
 use std::collections::{HashMap, HashSet};
 use std::fmt::{Debug};
 use std::hash::Hash;
+#[cfg(feature = "python")]
+use std::sync::Arc;
 
+#[cfg(feature = "python")]
+use pyo3::prelude::*;
+#[cfg(feature = "python")]
+use pyo3::types::{PyList, PyTuple};
 
 pub type DSLFn<D> = fn(Env<D>, &Evaluator<D>) -> VResult<D>;
 
+#[derive(Clone)]
+pub enum FnPtr<D: Domain> {
+    Native(DSLFn<D>),
+    #[cfg(feature = "python")]
+    Python(Arc<Py<PyAny>>),
+}
+
 #[derive(Clone)]
 pub struct Production<D: Domain> {
     pub name: Symbol, // eg "map" or "0" or "[1,2,3]"
     pub val: Val<D>,
     pub tp: SlowType,
     pub arity: usize,
     pub lazy_args: HashSet<usize>,
-    pub fn_ptr: Option<DSLFn<D>>,
+    pub fn_ptr: Option<FnPtr<D>>,
 }
 
 impl<D:Domain> Debug for Production<D> {
@@ -23,6 +36,40 @@ impl<D:Domain> Debug for Production<D> {
     }
 }
 
+impl<D: Domain> Production<D> {
+    #[inline]
+    pub fn call(&self, args: Env<D>, handle: &Evaluator<D>) -> VResult<D> {
+        match &self.fn_ptr{
+            Some(FnPtr::Native(f)) => f(args, handle),
+
+            #[cfg(feature = "python")]
+            Some(FnPtr::Python(pyf)) => {
+                Python::with_gil(|py| -> Result<Val<D>, String> {
+                    // Env<D> -> Python list
+                    let mut elems: Vec<Py<PyAny>> = Vec::with_capacity(args.len());
+                    for v in &args.env {
+                        let obj = D::py_val_to_py(py, v.clone())
+                            .map_err(|e| e.to_string())?;
+                        elems.push(obj);
+                    }
+
+                    let tuple_bound:Bound<PyTuple> = PyTuple::new(py, &elems)
+                            .map_err(|e| e.to_string())?;
+
+                    let ret = (**pyf)
+                        .bind(py)
+                        .call1(tuple_bound)
+                        .map_err(|e| e.to_string())?;
+
+                    // Python -> Val<D>
+                    D::py_py_to_val(&ret)
+                })
+            }
+            None => Err("primitive has no function pointer".into()),
+        }
+    }
+}
+
 
 #[derive(Clone, Debug)]
 pub struct DSL<D:Domain> {
@@ -65,7 +112,7 @@ impl<D: Domain> Production<D> {
             tp,
             arity,
             lazy_args,
-            fn_ptr: Some(fn_ptr),
+            fn_ptr: Some(FnPtr::Native(fn_ptr)),
         }
     }
 
@@ -85,6 +132,13 @@ impl<D: Domain> DSL<D> {
         self.productions.insert(entry.name.clone(), entry);
     }
 
+    /// NEW: add a constant (arity 0) to the DSL
+    pub fn add_constant(&mut self, name: Symbol, tp: SlowType, val: Val<D>) {
+        assert_eq!(tp.arity(), 0);
+        let prod = Production::val_raw(name, tp, val);
+        self.add_entry(prod);
+    }
+
     /// given a primitive's symbol return a runtime Val object. For function primitives
     /// this should return a PrimFun(CurriedFn) object.
     pub fn val_of_prim(&self, p: &Symbol) -> Option<Val<D>> {
@@ -110,5 +164,19 @@ pub trait Domain: Clone + Debug + PartialEq + Eq + Hash + Send + Sync {
     fn type_of_dom_val(&self) -> SlowType;
 
     fn new_dsl() -> DSL<Self>;
+
+    #[cfg(feature = "python")]
+    fn py_val_to_py(py: Python<'_>, v: Val<Self>) -> PyResult<Py<PyAny>> {
+        // default: not supported for this domain
+        Err(pyo3::exceptions::PyTypeError::new_err(
+            "Python bridge not implemented for this domain",
+        ))
+    }
+
+    #[cfg(feature = "python")]
+    fn py_py_to_val(_obj: &Bound<'_, PyAny>) -> Result<Val<Self>, String> {
+        Err("Python bridge not implemented for this domain".into())
+    }
+
 }
 

src/eval.rs

@@ -131,7 +131,7 @@ impl<D: Domain> CurriedFn<D> {
     pub fn apply(mut self, arg: Val<D>, handle: &Evaluator<D>) -> VResult<D> {
         self.partial_args.push_back(arg);
         if self.partial_args.len() == self.arity {
-            handle.dsl.productions.get(&self.name).unwrap().fn_ptr.unwrap() (self.partial_args, handle)
+            handle.dsl.productions.get(&self.name).unwrap().call(self.partial_args, handle)
         } else {
             Ok(Val::PrimFun(self))
         }