From 233acd290fecf1c8b0672d31c2fcd76090618185 Mon Sep 17 00:00:00 2001
From: Ziyang Hu <hu.ziyang@cantab.net>
Date: Wed, 7 Dec 2022 12:04:31 +0800
Subject: [PATCH] to_unity function

---
 cozo-core/src/data/expr.rs            | 16 ++++++--
 cozo-core/src/data/functions.rs       | 21 +++++++++++
 cozo-core/src/data/tests/functions.rs | 53 +++++++++++++++++++++++++++
 3 files changed, 87 insertions(+), 3 deletions(-)

diff --git a/cozo-core/src/data/expr.rs b/cozo-core/src/data/expr.rs
index d6b1f8a3..45170975 100644
--- a/cozo-core/src/data/expr.rs
+++ b/cozo-core/src/data/expr.rs
@@ -326,17 +326,26 @@ impl Expr {
                     #[diagnostic(code(eval::unbound))]
                     struct TupleTooShortError(String, usize, usize, #[label] SourceSpan);
 
-                    Ok(bindings.as_ref()
+                    Ok(bindings
+                        .as_ref()
                         .get(*i)
                         .ok_or_else(|| {
-                            TupleTooShortError(var.name.to_string(), *i, bindings.as_ref().len(), var.span)
+                            TupleTooShortError(
+                                var.name.to_string(),
+                                *i,
+                                bindings.as_ref().len(),
+                                var.span,
+                            )
                         })?
                         .clone())
                 }
             },
             Expr::Const { val, .. } => Ok(val.clone()),
             Expr::Apply { op, args, .. } => {
-                let args: Box<[DataValue]> = args.iter().map(|v| v.eval(bindings.as_ref())).try_collect()?;
+                let args: Box<[DataValue]> = args
+                    .iter()
+                    .map(|v| v.eval(bindings.as_ref()))
+                    .try_collect()?;
                 Ok((op.inner)(&args)
                     .map_err(|err| EvalRaisedError(self.span(), err.to_string()))?)
             }
@@ -701,6 +710,7 @@ pub(crate) fn get_op(name: &str) -> Option<&'static Op> {
         "difference" => &OP_DIFFERENCE,
         "to_uuid" => &OP_TO_UUID,
         "to_bool" => &OP_TO_BOOL,
+        "to_unity" => &OP_TO_UNITY,
         "rand_uuid_v1" => &OP_RAND_UUID_V1,
         "rand_uuid_v4" => &OP_RAND_UUID_V4,
         "uuid_timestamp" => &OP_UUID_TIMESTAMP,
diff --git a/cozo-core/src/data/functions.rs b/cozo-core/src/data/functions.rs
index 98c8beae..aecacd53 100644
--- a/cozo-core/src/data/functions.rs
+++ b/cozo-core/src/data/functions.rs
@@ -1277,10 +1277,31 @@ pub(crate) fn op_to_bool(args: &[DataValue]) -> Result<DataValue> {
     }))
 }
 
+
+define_op!(OP_TO_UNITY, 1, false);
+pub(crate) fn op_to_unity(args: &[DataValue]) -> Result<DataValue> {
+    Ok(DataValue::from(match &args[0] {
+        DataValue::Null => 0,
+        DataValue::Bool(b) => if *b {1} else {0},
+        DataValue::Num(n) => if n.get_float() != 0. {1} else {0},
+        DataValue::Str(s) => if s.is_empty() {0} else { 1},
+        DataValue::Bytes(b) => if b.is_empty() {0} else { 1},
+        DataValue::Uuid(u) => if u.0.is_nil() {0} else { 1 },
+        DataValue::Regex(r) => if r.0.as_str().is_empty() {0 } else { 1},
+        DataValue::List(l) => if l.is_empty() {0} else {1},
+        DataValue::Set(s) => if s.is_empty() {0} else {1},
+        DataValue::Guard => 0,
+        DataValue::Bot => 0,
+    }))
+}
+
+
 define_op!(OP_TO_FLOAT, 1, false);
 pub(crate) fn op_to_float(args: &[DataValue]) -> Result<DataValue> {
     Ok(match &args[0] {
         DataValue::Num(n) => n.get_float().into(),
+        DataValue::Null => DataValue::from(0.0),
+        DataValue::Bool(b) => DataValue::from(if *b { 1.0 } else { 0.0 }),
         DataValue::Str(t) => match t as &str {
             "PI" => f64::PI().into(),
             "E" => f64::E().into(),
diff --git a/cozo-core/src/data/tests/functions.rs b/cozo-core/src/data/tests/functions.rs
index 571a04f2..a6327f71 100644
--- a/cozo-core/src/data/tests/functions.rs
+++ b/cozo-core/src/data/tests/functions.rs
@@ -1217,8 +1217,61 @@ fn test_to_string() {
     );
 }
 
+#[test]
+fn test_to_unity() {
+    assert_eq!(op_to_unity(&[DataValue::Null]).unwrap(), DataValue::from(0));
+    assert_eq!(
+        op_to_unity(&[DataValue::Bool(false)]).unwrap(),
+        DataValue::from(0)
+    );
+    assert_eq!(
+        op_to_unity(&[DataValue::Bool(true)]).unwrap(),
+        DataValue::from(1)
+    );
+    assert_eq!(
+        op_to_unity(&[DataValue::from(10)]).unwrap(),
+        DataValue::from(1)
+    );
+    assert_eq!(
+        op_to_unity(&[DataValue::from(1.0)]).unwrap(),
+        DataValue::from(1)
+    );
+    assert_eq!(
+        op_to_unity(&[DataValue::from(f64::NAN)]).unwrap(),
+        DataValue::from(1)
+    );
+    assert_eq!(
+        op_to_unity(&[DataValue::Str("0".into())]).unwrap(),
+        DataValue::from(1)
+    );
+    assert_eq!(
+        op_to_unity(&[DataValue::Str("".into())]).unwrap(),
+        DataValue::from(0)
+    );
+    assert_eq!(
+        op_to_unity(&[DataValue::List(vec![])]).unwrap(),
+        DataValue::from(0)
+    );
+    assert_eq!(
+        op_to_unity(&[DataValue::List(vec![DataValue::Null])]).unwrap(),
+        DataValue::from(1)
+    );
+}
+
 #[test]
 fn test_to_float() {
+    assert_eq!(
+        op_to_float(&[DataValue::Null]).unwrap(),
+        DataValue::from(0.0)
+    );
+    assert_eq!(
+        op_to_float(&[DataValue::Bool(false)]).unwrap(),
+        DataValue::from(0.0)
+    );
+    assert_eq!(
+        op_to_float(&[DataValue::Bool(true)]).unwrap(),
+        DataValue::from(1.0)
+    );
     assert_eq!(
         op_to_float(&[DataValue::from(1)]).unwrap(),
         DataValue::from(1.0)