function int_range

cozo-core/src/data/expr.rs

@@ -824,6 +824,7 @@ pub(crate) fn get_op(name: &str) -> Option<&'static Op> {
         "l2_normalize" => &OP_L2_NORMALIZE,
         "ip_dist" => &OP_IP_DIST,
         "cos_dist" => &OP_COS_DIST,
+        "int_range" => &OP_INT_RANGE,
         "rand_float" => &OP_RAND_FLOAT,
         "rand_bernoulli" => &OP_RAND_BERNOULLI,
         "rand_int" => &OP_RAND_INT,

cozo-core/src/data/functions.rs

@@ -1850,6 +1850,54 @@ pub(crate) fn op_cos_dist(args: &[DataValue]) -> Result<DataValue> {
     }
 }
 
+define_op!(OP_INT_RANGE, 1, true);
+pub(crate) fn op_int_range(args: &[DataValue]) -> Result<DataValue> {
+    let [start, end] = match args.len() {
+        1 => {
+            let end = args[0]
+                .get_int()
+                .ok_or_else(|| miette!("'int_range' requires integer argument for end"))?;
+            [0, end]
+        }
+        2 => {
+            let start = args[0]
+                .get_int()
+                .ok_or_else(|| miette!("'int_range' requires integer argument for start"))?;
+            let end = args[1]
+                .get_int()
+                .ok_or_else(|| miette!("'int_range' requires integer argument for end"))?;
+            [start, end]
+        }
+        3 => {
+            let start = args[0]
+                .get_int()
+                .ok_or_else(|| miette!("'int_range' requires integer argument for start"))?;
+            let end = args[1]
+                .get_int()
+                .ok_or_else(|| miette!("'int_range' requires integer argument for end"))?;
+            let step = args[2]
+                .get_int()
+                .ok_or_else(|| miette!("'int_range' requires integer argument for step"))?;
+            let mut current = start;
+            let mut result = vec![];
+            if step > 0 {
+                while current < end {
+                    result.push(DataValue::from(current));
+                    current += step;
+                }
+            } else {
+                while current > end {
+                    result.push(DataValue::from(current));
+                    current += step;
+                }
+            }
+            return Ok(DataValue::List(result));
+        }
+        _ => bail!("'int_range' requires 1 to 3 argument"),
+    };
+    Ok(DataValue::List((start..end).map(DataValue::from).collect()))
+}
+
 define_op!(OP_RAND_FLOAT, 0, false);
 pub(crate) fn op_rand_float(_args: &[DataValue]) -> Result<DataValue> {
     Ok(thread_rng().gen::<f64>().into())

cozo-core/src/data/tests/functions.rs

@@ -9,6 +9,7 @@
 use approx::AbsDiffEq;
 use num_traits::FloatConst;
 use regex::Regex;
+use serde_json::json;
 
 use crate::data::functions::*;
 use crate::data::value::{DataValue, RegexWrapper};
@@ -1456,3 +1457,23 @@ fn test_coalesce() {
         .rows;
     assert_eq!(res[0][0], DataValue::from(2));
 }
+
+#[test]
+fn test_range() {
+    let db = new_cozo_mem().unwrap();
+    let res = db
+        .run_script("?[a] := a = int_range(1, 5)", Default::default())
+        .unwrap()
+        .into_json();
+    assert_eq!(res["rows"][0][0], json!([1, 2, 3, 4]));
+    let res = db
+        .run_script("?[a] := a = int_range(5)", Default::default())
+        .unwrap()
+        .into_json();
+    assert_eq!(res["rows"][0][0], json!([0, 1, 2, 3, 4]));
+    let res = db
+        .run_script("?[a] := a = int_range(15, 3, -2)", Default::default())
+        .unwrap()
+        .into_json();
+    assert_eq!(res["rows"][0][0], json!([15, 13, 11, 9, 7, 5]));
+}

cozo-core/src/runtime/relation.rs

@@ -723,6 +723,7 @@ impl<'a> SessionTx<'a> {
 
         // Build key columns definitions
         let mut idx_keys: Vec<ColumnDef> = vec![ColumnDef {
+            // layer -1 stores the self-loops
             name: SmartString::from("layer"),
             typing: NullableColType {
                 coltype: ColType::Int,
@@ -730,6 +731,7 @@ impl<'a> SessionTx<'a> {
             },
             default_gen: None,
         }];
+        // for self-loops, fr and to are identical
         for prefix in ["fr", "to"] {
             for col in rel_handle.metadata.keys.iter() {
                 let mut col = col.clone();
@@ -756,6 +758,7 @@ impl<'a> SessionTx<'a> {
 
         // Build non-key columns definitions
         let non_idx_keys = vec![
+            // For self-loops, stores the number of neighbours
             ColumnDef {
                 name: SmartString::from("dist"),
                 typing: NullableColType {
@@ -764,6 +767,7 @@ impl<'a> SessionTx<'a> {
                 },
                 default_gen: None,
             },
+            // For self-loops, stores a hash of the neighbours, for conflict detection
             ColumnDef {
                 name: SmartString::from("tags"),
                 typing: NullableColType {