prepare for yielding relations

cozo-core/src/cozoscript.pest

@@ -108,10 +108,11 @@ list = { "[" ~ (expr ~ ",")* ~ expr? ~ "]" }
 grouping = { "(" ~ expr ~ ")" }
 
 option = _{(limit_option|offset_option|sort_option|relation_option|timeout_option|sleep_option|
-            assert_none_option|assert_some_option) ~ ";"?}
+            assert_none_option|assert_some_option|yield_option) ~ ";"?}
 out_arg = @{var ~ ("(" ~ var ~ ")")?}
 limit_option = {":limit"  ~ expr}
 offset_option = {":offset" ~ expr}
+yield_option = {":yield" ~ ident}
 sort_option = {(":sort" | ":order") ~ (sort_arg ~ ",")* ~ sort_arg }
 relation_option = {relation_op ~ compound_ident ~ table_schema?}
 relation_op = _{relation_create | relation_replace | relation_put | relation_rm | relation_ensure | relation_ensure_not}

cozo-core/src/data/program.rs

@@ -43,6 +43,7 @@ pub(crate) struct QueryOutOptions {
     pub(crate) sorters: Vec<(Symbol, SortDir)>,
     pub(crate) store_relation: Option<(InputRelationHandle, RelationOp)>,
     pub(crate) assertion: Option<QueryAssertion>,
+    pub(crate) yield_const: Option<Symbol>,
 }
 
 impl Debug for QueryOutOptions {

cozo-core/src/parse/query.rs

@@ -62,6 +62,11 @@ struct MultipleRuleDefinitionError(String, Vec<SourceSpan>);
 #[diagnostic(code(parser::multiple_out_assert))]
 struct DuplicateQueryAssertion(#[label] SourceSpan);
 
+#[derive(Debug, Error, Diagnostic)]
+#[error("Multiple query yields defined")]
+#[diagnostic(code(parser::multiple_yields))]
+struct DuplicateYield(#[label] SourceSpan);
+
 impl Error for MultipleRuleDefinitionError {}
 
 impl Display for MultipleRuleDefinitionError {
@@ -342,6 +347,14 @@ pub(crate) fn parse_query(
                 );
                 out_opts.assertion = Some(QueryAssertion::AssertSome(pair.extract_span()))
             }
+            Rule::yield_option => {
+                ensure!(
+                    out_opts.yield_const.is_none(),
+                    DuplicateYield(pair.extract_span())
+                );
+                let p = pair.into_inner().next().unwrap();
+                out_opts.yield_const = Some(Symbol::new(p.as_str(), p.extract_span()));
+            }
             Rule::EOI => break,
             r => unreachable!("{:?}", r),
         }
@@ -453,7 +466,12 @@ fn parse_rule(
     let mut body_clauses = vec![];
     let mut ignored_counter = 0;
     for atom_src in body.into_inner() {
-        body_clauses.push(parse_disjunction(atom_src, param_pool, cur_vld, &mut ignored_counter)?)
+        body_clauses.push(parse_disjunction(
+            atom_src,
+            param_pool,
+            cur_vld,
+            &mut ignored_counter,
+        )?)
     }
 
     Ok((
@@ -489,7 +507,7 @@ fn parse_atom(
     src: Pair<'_>,
     param_pool: &BTreeMap<String, DataValue>,
     cur_vld: ValidityTs,
-    ignored_counter: &mut u32
+    ignored_counter: &mut u32,
 ) -> Result<InputAtom> {
     Ok(match src.as_rule() {
         Rule::rule_body => {
@@ -506,7 +524,12 @@ fn parse_atom(
         Rule::disjunction => parse_disjunction(src, param_pool, cur_vld, ignored_counter)?,
         Rule::negation => {
             let span = src.extract_span();
-            let inner = parse_atom(src.into_inner().next().unwrap(), param_pool, cur_vld, ignored_counter)?;
+            let inner = parse_atom(
+                src.into_inner().next().unwrap(),
+                param_pool,
+                cur_vld,
+                ignored_counter,
+            )?;
             InputAtom::Negation {
                 inner: inner.into(),
                 span,
@@ -746,7 +769,8 @@ fn parse_fixed_rule(
                         for v in els {
                             let s = v.as_str();
                             if s == "_" {
-                                let symb = Symbol::new(format!("*_*{}", binding_gen_id) ,v.extract_span());
+                                let symb =
+                                    Symbol::new(format!("*_*{}", binding_gen_id), v.extract_span());
                                 binding_gen_id += 1;
                                 bindings.push(symb);
                             } else {
@@ -773,7 +797,10 @@ fn parse_fixed_rule(
                                 Rule::var => {
                                     let s = v.as_str();
                                     if s == "_" {
-                                        let symb = Symbol::new(format!("*_*{}", binding_gen_id) ,v.extract_span());
+                                        let symb = Symbol::new(
+                                            format!("*_*{}", binding_gen_id),
+                                            v.extract_span(),
+                                        );
                                         binding_gen_id += 1;
                                         bindings.push(symb);
                                     } else {

cozo-core/src/runtime/db.rs

@@ -27,7 +27,7 @@ use thiserror::Error;
 
 use crate::data::functions::current_validity;
 use crate::data::json::JsonValue;
-use crate::data::program::{InputProgram, QueryAssertion, RelationOp};
+use crate::data::program::{InputProgram, MagicSymbol, QueryAssertion, RelationOp};
 use crate::data::relation::ColumnDef;
 use crate::data::tuple::{Tuple, TupleT};
 use crate::data::value::{DataValue, ValidityTs, LARGEST_UTF_CHAR};
@@ -100,7 +100,7 @@ impl<S> Debug for Db<S> {
 #[diagnostic(code(db::init))]
 pub(crate) struct BadDbInit(#[help] pub(crate) String);
 
-#[derive(serde_derive::Serialize, serde_derive::Deserialize, Debug)]
+#[derive(serde_derive::Serialize, serde_derive::Deserialize, Debug, Clone)]
 /// Rows in a relation, together with headers for the fields.
 pub struct NamedRows {
     /// The headers
@@ -1462,4 +1462,23 @@ grandparent[gcld, gp] := parent[gcld, p], parent[p, gp]
 
         assert_eq!(res.len(), 1);
     }
+
+    #[test]
+    fn returning_relations() {
+        let db = new_cozo_mem().unwrap();
+        let res = db
+            .run_script(
+                r#"
+        {
+            ?[] <- [[1,2,3]]
+            :yield nxt
+        }
+        {
+            ?[a,b,c] := nxt[a, b, c]
+        }
+        "#,
+                Default::default(),
+            )
+            .unwrap();
+    }
 }