From 525cfeff5965677c19eda7a01f30747e2cdf6956 Mon Sep 17 00:00:00 2001
From: Ziyang Hu <hu.ziyang@cantab.net>
Date: Sat, 14 Jan 2023 12:27:08 +0800
Subject: [PATCH] Don't go through JSON

---
 cozo-core/src/fixed_rule/mod.rs | 11 +++++------
 cozo-lib-python/src/lib.rs      | 23 ++++++++++++++---------
 2 files changed, 19 insertions(+), 15 deletions(-)

diff --git a/cozo-core/src/fixed_rule/mod.rs b/cozo-core/src/fixed_rule/mod.rs
index d9edca6b..15a68b3c 100644
--- a/cozo-core/src/fixed_rule/mod.rs
+++ b/cozo-core/src/fixed_rule/mod.rs
@@ -23,7 +23,6 @@ use smartstring::{LazyCompact, SmartString};
 use thiserror::Error;
 
 use crate::data::expr::Expr;
-use crate::data::json::JsonValue;
 use crate::data::program::{
     FixedRuleOptionNotFoundError, MagicFixedRuleApply, MagicFixedRuleRuleArg, MagicSymbol,
     WrongFixedRuleOptionError,
@@ -567,7 +566,7 @@ pub trait FixedRule: Send + Sync {
 /// but implementation is simpler.
 pub struct SimpleFixedRule {
     return_arity: usize,
-    rule: Box<dyn Fn(Vec<NamedRows>, JsonValue) -> Result<NamedRows> + Send + Sync + 'static>,
+    rule: Box<dyn Fn(Vec<NamedRows>, BTreeMap<String, DataValue>) -> Result<NamedRows> + Send + Sync + 'static>,
 }
 
 impl SimpleFixedRule {
@@ -581,7 +580,7 @@ impl SimpleFixedRule {
     //    Every row of the returned relation must have length equal to `return_arity`.
     pub fn new<R>(return_arity: usize, rule: R) -> Self
     where
-        R: Fn(Vec<NamedRows>, JsonValue) -> Result<NamedRows> + Send + Sync + 'static,
+        R: Fn(Vec<NamedRows>, BTreeMap<String, DataValue>) -> Result<NamedRows> + Send + Sync + 'static,
     {
         Self {
             return_arity,
@@ -593,7 +592,7 @@ impl SimpleFixedRule {
         return_arity: usize,
     ) -> (
         Self,
-        Receiver<(Vec<NamedRows>, JsonValue, Sender<Result<NamedRows>>)>,
+        Receiver<(Vec<NamedRows>, BTreeMap<String, DataValue>, Sender<Result<NamedRows>>)>,
     ) {
         let (db2app_sender, db2app_receiver) = bounded(0);
         (
@@ -628,13 +627,13 @@ impl FixedRule for SimpleFixedRule {
         out: &'_ mut RegularTempStore,
         _poison: Poison,
     ) -> Result<()> {
-        let options: JsonValue = payload
+        let options: BTreeMap<_, _> = payload
             .manifest
             .options
             .iter()
             .map(|(k, v)| -> Result<_> {
                 let val = v.clone().eval_to_const()?;
-                Ok((k.to_string(), JsonValue::from(val)))
+                Ok((k.to_string(), val))
             })
             .try_collect()?;
         let input_arity = payload.manifest.rule_args.len();
diff --git a/cozo-lib-python/src/lib.rs b/cozo-lib-python/src/lib.rs
index 077b6ef3..b28ccbe1 100644
--- a/cozo-lib-python/src/lib.rs
+++ b/cozo-lib-python/src/lib.rs
@@ -69,6 +69,17 @@ fn convert_params(ob: &PyDict) -> PyResult<BTreeMap<String, DataValue>> {
     Ok(ret)
 }
 
+fn options_to_py(opts: BTreeMap<String, DataValue>, py: Python<'_>) -> PyResult<PyObject> {
+    let ret = PyDict::new(py);
+
+    for (k, v) in opts {
+        let val = value_to_py(v, py);
+        ret.set_item(k, val)?;
+    }
+
+    Ok(ret.into())
+}
+
 fn value_to_py(val: DataValue, py: Python<'_>) -> PyObject {
     match val {
         DataValue::Null => py.None(),
@@ -175,23 +186,17 @@ impl CozoDbPy {
     ) -> PyResult<()> {
         if let Some(db) = &self.db {
             let cb: Py<PyAny> = callback.into();
-            let (rule_impl, receiver, sender) = SimpleFixedRule::rule_with_channel(arity);
+            let (rule_impl, receiver) = SimpleFixedRule::rule_with_channel(arity);
             match db.register_fixed_rule(name, rule_impl) {
                 Ok(_) => {
                     thread::spawn(move || {
-                        for (inputs, options) in receiver {
+                        for (inputs, options, sender) in receiver {
                             let res = Python::with_gil(|py| -> Result<NamedRows> {
-                                let json_convert =
-                                    PyModule::import(py, "json").into_diagnostic()?;
-                                let json_convert: Py<PyAny> =
-                                    json_convert.getattr("loads").into_diagnostic()?.into();
                                 let py_inputs = PyList::new(
                                     py,
                                     inputs.into_iter().map(|nr| rows_to_py_rows(nr.rows, py)),
                                 );
-                                let opts_str = PyString::new(py, &options.to_string());
-                                let args = PyTuple::new(py, vec![opts_str]);
-                                let py_opts = json_convert.call1(py, args).into_diagnostic()?;
+                                let py_opts = options_to_py(options, py).into_diagnostic()?;
                                 let args =
                                     PyTuple::new(py, vec![PyObject::from(py_inputs), py_opts]);
                                 let res = cb.as_ref(py).call1(args).into_diagnostic()?;