multiple returns

cozo-core/src/cozoscript.pest

@@ -218,7 +218,7 @@ imperative_block = {imperative_stmt+}
 break_stmt = {"%break" ~ ident?}
 ignore_error_script = {"%ignore_error" ~ query_script_inner}
 continue_stmt = {"%continue" ~ ident?}
-return_stmt = {"%return" ~ (ident | underscore_ident | query_script_inner)?}
+return_stmt = {"%return" ~ (ident | underscore_ident | query_script_inner)*}
 loop_block = {("%mark" ~ ident)? ~ "%loop" ~ imperative_block ~ "%end"}
 temp_swap = {"%swap" ~ underscore_ident ~ underscore_ident}
 debug_stmt = {"%debug" ~ (ident | underscore_ident)}

cozo-core/src/parse/imperative.rs

@@ -7,17 +7,19 @@
  *
  */
 
-use crate::parse::query::parse_query;
-use crate::parse::{ExtractSpan, ImperativeProgram, ImperativeStmt, Pair, Rule, SourceSpan};
-use crate::{DataValue, FixedRule, ValidityTs};
+use std::collections::BTreeMap;
+use std::sync::Arc;
+
 use either::{Left, Right};
 use itertools::Itertools;
 use miette::{Diagnostic, Result};
 use smartstring::SmartString;
-use std::collections::BTreeMap;
-use std::sync::Arc;
 use thiserror::Error;
 
+use crate::parse::query::parse_query;
+use crate::parse::{ExtractSpan, ImperativeProgram, ImperativeStmt, Pair, Rule, SourceSpan};
+use crate::{DataValue, FixedRule, ValidityTs};
+
 pub(crate) fn parse_imperative_block(
     src: Pair<'_>,
     param_pool: &BTreeMap<String, DataValue>,
@@ -76,20 +78,21 @@ fn parse_imperative_stmt(
         }
         Rule::return_stmt => {
             // let span = pair.extract_span();
-            match pair.into_inner().next() {
-                None => ImperativeStmt::ReturnNil,
-                Some(p) => match p.as_rule() {
+            let mut rets = vec![];
+            for p in pair.into_inner() {
+                match p.as_rule() {
                     Rule::ident | Rule::underscore_ident => {
                         let rel = SmartString::from(p.as_str());
-                        ImperativeStmt::ReturnTemp { rel }
+                        rets.push(Right(rel));
                     }
                     Rule::query_script_inner => {
                         let prog = parse_query(p.into_inner(), param_pool, fixed_rules, cur_vld)?;
-                        ImperativeStmt::ReturnProgram { prog }
+                        rets.push(Left(prog))
                     }
                     _ => unreachable!(),
-                },
+                }
             }
+            ImperativeStmt::Return { returns: rets }
         }
         Rule::if_chain | Rule::if_not_chain => {
             let negated = pair.as_rule() == Rule::if_not_chain;

cozo-core/src/parse/mod.rs

@@ -11,7 +11,7 @@ use std::collections::{BTreeMap, BTreeSet};
 use std::fmt::{Display, Formatter};
 use std::sync::Arc;
 
-use either::Either;
+use either::{Either, Left};
 use miette::{bail, Diagnostic, IntoDiagnostic, Result};
 use pest::error::InputLocation;
 use pest::Parser;
@@ -56,13 +56,8 @@ pub(crate) enum ImperativeStmt {
         target: Option<SmartString<LazyCompact>>,
         span: SourceSpan,
     },
-    ReturnNil,
-    ReturnProgram {
-        prog: InputProgram,
-        // span: SourceSpan,
-    },
-    ReturnTemp {
-        rel: SmartString<LazyCompact>,
+    Return {
+        returns: Vec<Either<InputProgram, SmartString<LazyCompact>>>,
     },
     Program {
         prog: InputProgram,
@@ -98,13 +93,21 @@ pub(crate) type ImperativeProgram = Vec<ImperativeStmt>;
 impl ImperativeStmt {
     pub(crate) fn needs_write_locks(&self, collector: &mut BTreeSet<SmartString<LazyCompact>>) {
         match self {
-            ImperativeStmt::ReturnProgram { prog, .. }
-            | ImperativeStmt::Program { prog, .. }
+            ImperativeStmt::Program { prog, .. }
             | ImperativeStmt::IgnoreErrorProgram { prog, .. } => {
                 if let Some(name) = prog.needs_write_lock() {
                     collector.insert(name);
                 }
             }
+            ImperativeStmt::Return { returns, .. } => {
+                for ret in returns {
+                    if let Left(prog) = ret {
+                        if let Some(name) = prog.needs_write_lock() {
+                            collector.insert(name);
+                        }
+                    }
+                }
+            }
             ImperativeStmt::If {
                 condition,
                 then_branch,
@@ -126,10 +129,8 @@ impl ImperativeStmt {
                 }
             }
             ImperativeStmt::TempDebug { .. }
-            | ImperativeStmt::ReturnTemp { .. }
             | ImperativeStmt::Break { .. }
             | ImperativeStmt::Continue { .. }
-            | ImperativeStmt::ReturnNil { .. }
             | ImperativeStmt::TempSwap { .. } => {}
         }
     }

cozo-core/src/query/stored.rs

@@ -248,32 +248,32 @@ impl<'a> SessionTx<'a> {
                             .or_default();
                         target_collector.push((
                             CallbackOp::Rm,
-                            NamedRows {
-                                headers: k_bindings
+                            NamedRows::new(
+                                k_bindings
                                     .into_iter()
                                     .map(|k| k.name.to_string())
                                     .collect_vec(),
-                                rows: new_tuples
+                                new_tuples
                                     .into_iter()
                                     .map(|v| match v {
                                         DataValue::List(l) => l,
                                         _ => unreachable!(),
                                     })
                                     .collect_vec(),
-                            },
-                            NamedRows {
-                                headers: kv_bindings
+                            ),
+                            NamedRows::new(
+                                kv_bindings
                                     .into_iter()
                                     .map(|k| k.name.to_string())
                                     .collect_vec(),
-                                rows: old_tuples
+                                old_tuples
                                     .into_iter()
                                     .map(|v| match v {
                                         DataValue::List(l) => l,
                                         _ => unreachable!(),
                                     })
                                     .collect_vec(),
-                            },
+                            ),
                         ))
                     }
                 }
@@ -530,26 +530,26 @@ impl<'a> SessionTx<'a> {
                             .collect_vec();
                         target_collector.push((
                             CallbackOp::Put,
-                            NamedRows {
-                                headers: headers.clone(),
-                                rows: new_tuples
+                            NamedRows::new(
+                                headers.clone(),
+                                new_tuples
                                     .into_iter()
                                     .map(|v| match v {
                                         DataValue::List(l) => l,
                                         _ => unreachable!(),
                                     })
                                     .collect_vec(),
-                            },
-                            NamedRows {
+                            ),
+                            NamedRows::new(
                                 headers,
-                                rows: old_tuples
+                                old_tuples
                                     .into_iter()
                                     .map(|v| match v {
                                         DataValue::List(l) => l,
                                         _ => unreachable!(),
                                     })
                                     .collect_vec(),
-                            },
+                            ),
                         ))
                     }
                 }

cozo-core/src/runtime/db.rs

@@ -123,11 +123,46 @@ pub struct NamedRows {
     pub headers: Vec<String>,
     /// The rows
     pub rows: Vec<Tuple>,
+    pub(crate) next: Option<Box<NamedRows>>,
 }
 
 impl NamedRows {
+    /// create a named rows with the given headers and rows
+    pub fn new(headers: Vec<String>, rows: Vec<Tuple>) -> Self {
+        Self {
+            headers,
+            rows,
+            next: None,
+        }
+    }
+
+    /// If there are more named rows after the current one
+    pub fn has_more(&self) -> bool {
+        self.next.is_some()
+    }
+
+    /// convert a chain of named rows to individual named rows
+    pub fn flatten(self) -> Vec<Self> {
+        let mut collected = vec![];
+        let mut current = self;
+        loop {
+            let nxt = current.next.take();
+            collected.push(current);
+            if let Some(n) = nxt {
+                current = *n;
+            } else {
+                break;
+            }
+        }
+        collected
+    }
+
     /// Convert to a JSON object
     pub fn into_json(self) -> JsonValue {
+        let nxt = match self.next {
+            None => json!(null),
+            Some(more) => more.into_json(),
+        };
         let rows = self
             .rows
             .into_iter()
@@ -135,7 +170,8 @@ impl NamedRows {
             .collect::<JsonValue>();
         json!({
             "headers": self.headers,
-            "rows": rows
+            "rows": rows,
+            "next": nxt,
         })
     }
 }
@@ -245,7 +281,7 @@ impl<'s, S: Storage<'s>> Db<S> {
                 rows.push(tuple);
             }
             let headers = cols.iter().map(|col| col.to_string()).collect_vec();
-            ret.insert(rel.to_string(), NamedRows { headers, rows });
+            ret.insert(rel.to_string(), NamedRows::new(headers, rows));
         }
         Ok(ret)
     }
@@ -924,7 +960,7 @@ impl<'s, S: Storage<'s>> Db<S> {
             })
             .collect_vec();
 
-        Ok(NamedRows { headers, rows })
+        Ok(NamedRows::new(headers, rows))
     }
     fn run_sys_op(&'s self, op: SysOp) -> Result<NamedRows> {
         match op {
@@ -939,10 +975,10 @@ impl<'s, S: Storage<'s>> Db<S> {
             }
             SysOp::Compact => {
                 self.compact_relation()?;
-                Ok(NamedRows {
-                    headers: vec![STATUS_STR.to_string()],
-                    rows: vec![vec![DataValue::from(OK_STR)]],
-                })
+                Ok(NamedRows::new(
+                    vec![STATUS_STR.to_string()],
+                    vec![vec![DataValue::from(OK_STR)]],
+                ))
             }
             SysOp::ListRelations => self.list_relations(),
             SysOp::RemoveRelation(rel_names) => {
@@ -961,10 +997,10 @@ impl<'s, S: Storage<'s>> Db<S> {
                 for (lower, upper) in bounds {
                     self.db.del_range(&lower, &upper)?;
                 }
-                Ok(NamedRows {
-                    headers: vec![STATUS_STR.to_string()],
-                    rows: vec![vec![DataValue::from(OK_STR)]],
-                })
+                Ok(NamedRows::new(
+                    vec![STATUS_STR.to_string()],
+                    vec![vec![DataValue::from(OK_STR)]],
+                ))
             }
             SysOp::CreateIndex(rel_name, idx_name, cols) => {
                 let lock = self
@@ -975,10 +1011,10 @@ impl<'s, S: Storage<'s>> Db<S> {
                 let mut tx = self.transact_write()?;
                 tx.create_index(&rel_name, &idx_name, cols)?;
                 tx.commit_tx()?;
-                Ok(NamedRows {
-                    headers: vec![STATUS_STR.to_string()],
-                    rows: vec![vec![DataValue::from(OK_STR)]],
-                })
+                Ok(NamedRows::new(
+                    vec![STATUS_STR.to_string()],
+                    vec![vec![DataValue::from(OK_STR)]],
+                ))
             }
             SysOp::RemoveIndex(rel_name, idx_name) => {
                 let lock = self
@@ -989,10 +1025,10 @@ impl<'s, S: Storage<'s>> Db<S> {
                 let mut tx = self.transact_write()?;
                 tx.remove_index(&rel_name, &idx_name)?;
                 tx.commit_tx()?;
-                Ok(NamedRows {
-                    headers: vec![STATUS_STR.to_string()],
-                    rows: vec![vec![DataValue::from(OK_STR)]],
-                })
+                Ok(NamedRows::new(
+                    vec![STATUS_STR.to_string()],
+                    vec![vec![DataValue::from(OK_STR)]],
+                ))
             }
             SysOp::ListRelation(rs) => self.list_relation(&rs),
             SysOp::RenameRelation(rename_pairs) => {
@@ -1004,25 +1040,25 @@ impl<'s, S: Storage<'s>> Db<S> {
                     tx.rename_relation(old, new)?;
                 }
                 tx.commit_tx()?;
-                Ok(NamedRows {
-                    headers: vec![STATUS_STR.to_string()],
-                    rows: vec![vec![DataValue::from(OK_STR)]],
-                })
+                Ok(NamedRows::new(
+                    vec![STATUS_STR.to_string()],
+                    vec![vec![DataValue::from(OK_STR)]],
+                ))
             }
             SysOp::ListRunning => self.list_running(),
             SysOp::KillRunning(id) => {
                 let queries = self.running_queries.lock().unwrap();
                 Ok(match queries.get(&id) {
-                    None => NamedRows {
-                        headers: vec![STATUS_STR.to_string()],
-                        rows: vec![vec![DataValue::from("NOT_FOUND")]],
-                    },
+                    None => NamedRows::new(
+                        vec![STATUS_STR.to_string()],
+                        vec![vec![DataValue::from("NOT_FOUND")]],
+                    ),
                     Some(handle) => {
                         handle.poison.0.store(true, Ordering::Relaxed);
-                        NamedRows {
-                            headers: vec![STATUS_STR.to_string()],
-                            rows: vec![vec![DataValue::from("KILLING")]],
-                        }
+                        NamedRows::new(
+                            vec![STATUS_STR.to_string()],
+                            vec![vec![DataValue::from("KILLING")]],
+                        )
                     }
                 })
             }
@@ -1044,19 +1080,19 @@ impl<'s, S: Storage<'s>> Db<S> {
                     .map(|row| row.into_iter().map(DataValue::from).collect_vec())
                     .collect_vec();
                 tx.commit_tx()?;
-                Ok(NamedRows {
-                    headers: vec!["type".to_string(), "idx".to_string(), "trigger".to_string()],
+                Ok(NamedRows::new(
+                    vec!["type".to_string(), "idx".to_string(), "trigger".to_string()],
                     rows,
-                })
+                ))
             }
             SysOp::SetTriggers(name, puts, rms, replaces) => {
                 let mut tx = self.transact_write()?;
                 tx.set_relation_triggers(name, puts, rms, replaces)?;
                 tx.commit_tx()?;
-                Ok(NamedRows {
-                    headers: vec![STATUS_STR.to_string()],
-                    rows: vec![vec![DataValue::from(OK_STR)]],
-                })
+                Ok(NamedRows::new(
+                    vec![STATUS_STR.to_string()],
+                    vec![vec![DataValue::from(OK_STR)]],
+                ))
             }
             SysOp::SetAccessLevel(names, level) => {
                 let mut tx = self.transact_write()?;
@@ -1064,10 +1100,10 @@ impl<'s, S: Storage<'s>> Db<S> {
                     tx.set_access_level(name, level)?;
                 }
                 tx.commit_tx()?;
-                Ok(NamedRows {
-                    headers: vec![STATUS_STR.to_string()],
-                    rows: vec![vec![DataValue::from(OK_STR)]],
-                })
+                Ok(NamedRows::new(
+                    vec![STATUS_STR.to_string()],
+                    vec![vec![DataValue::from(OK_STR)]],
+                ))
             }
         }
     }
@@ -1229,23 +1265,23 @@ impl<'s, S: Storage<'s>> Db<S> {
                     .wrap_err_with(|| format!("when executing against relation '{}'", meta.name))?;
                 clean_ups.extend(to_clear);
                 Ok((
-                    NamedRows {
-                        headers: vec![STATUS_STR.to_string()],
-                        rows: vec![vec![DataValue::from(OK_STR)]],
-                    },
+                    NamedRows::new(
+                        vec![STATUS_STR.to_string()],
+                        vec![vec![DataValue::from(OK_STR)]],
+                    ),
                     clean_ups,
                 ))
             } else {
                 // not sorting outputs
                 let rows: Vec<Tuple> = sorted_iter.collect_vec();
                 Ok((
-                    NamedRows {
-                        headers: entry_head_or_default
+                    NamedRows::new(
+                        entry_head_or_default
                             .iter()
                             .map(|s| s.to_string())
                             .collect_vec(),
                         rows,
-                    },
+                    ),
                     clean_ups,
                 ))
             }
@@ -1284,23 +1320,23 @@ impl<'s, S: Storage<'s>> Db<S> {
                     .wrap_err_with(|| format!("when executing against relation '{}'", meta.name))?;
                 clean_ups.extend(to_clear);
                 Ok((
-                    NamedRows {
-                        headers: vec![STATUS_STR.to_string()],
-                        rows: vec![vec![DataValue::from(OK_STR)]],
-                    },
+                    NamedRows::new(
+                        vec![STATUS_STR.to_string()],
+                        vec![vec![DataValue::from(OK_STR)]],
+                    ),
                     clean_ups,
                 ))
             } else {
                 let rows: Vec<Tuple> = scan.collect_vec();
 
                 Ok((
-                    NamedRows {
-                        headers: entry_head_or_default
+                    NamedRows::new(
+                        entry_head_or_default
                             .iter()
                             .map(|s| s.to_string())
                             .collect_vec(),
                         rows,
-                    },
+                    ),
                     clean_ups,
                 ))
             }
@@ -1319,10 +1355,10 @@ impl<'s, S: Storage<'s>> Db<S> {
                 ]
             })
             .collect_vec();
-        Ok(NamedRows {
-            headers: vec!["id".to_string(), "started_at".to_string()],
+        Ok(NamedRows::new(
+            vec!["id".to_string(), "started_at".to_string()],
             rows,
-        })
+        ))
     }
     fn list_relation(&'s self, name: &str) -> Result<NamedRows> {
         let mut tx = self.transact()?;
@@ -1354,8 +1390,8 @@ impl<'s, S: Storage<'s>> Db<S> {
             .into_iter()
             .map(|row| row.into_iter().map(DataValue::from).collect_vec())
             .collect_vec();
-        Ok(NamedRows {
-            headers: vec![
+        Ok(NamedRows::new(
+            vec![
                 "column".to_string(),
                 "is_key".to_string(),
                 "index".to_string(),
@@ -1363,7 +1399,7 @@ impl<'s, S: Storage<'s>> Db<S> {
                 "has_default".to_string(),
             ],
             rows,
-        })
+        ))
     }
     fn list_relations(&'s self) -> Result<NamedRows> {
         let lower = vec![DataValue::from("")].encode_as_key(RelationId::SYSTEM);
@@ -1401,8 +1437,8 @@ impl<'s, S: Storage<'s>> Db<S> {
             .into_iter()
             .map(|row| row.into_iter().map(DataValue::from).collect_vec())
             .collect_vec();
-        Ok(NamedRows {
-            headers: vec![
+        Ok(NamedRows::new(
+            vec![
                 "name".to_string(),
                 "arity".to_string(),
                 "access_level".to_string(),
@@ -1413,7 +1449,7 @@ impl<'s, S: Storage<'s>> Db<S> {
                 "n_replace_triggers".to_string(),
             ],
             rows,
-        })
+        ))
     }
 }
 

cozo-core/src/runtime/imperative.rs

@@ -7,19 +7,20 @@
  */
 
 use std::collections::{BTreeMap, BTreeSet};
+
 use either::{Either, Left, Right};
 use itertools::Itertools;
+use miette::{bail, Diagnostic, Report, Result};
 use smartstring::{LazyCompact, SmartString};
+use thiserror::Error;
 
-use crate::parse::{ImperativeCondition, ImperativeProgram, ImperativeStmt, SourceSpan};
-use crate::{DataValue, Db, NamedRows, Storage, ValidityTs};
 use crate::data::expr::PredicateTypeError;
 use crate::data::functions::op_to_bool;
+use crate::data::symb::Symbol;
+use crate::parse::{ImperativeCondition, ImperativeProgram, ImperativeStmt, SourceSpan};
 use crate::runtime::callback::CallbackCollector;
 use crate::runtime::transact::SessionTx;
-use miette::{bail, Report, Result, Diagnostic};
-use crate::data::symb::Symbol;
-use thiserror::Error;
+use crate::{DataValue, Db, NamedRows, Storage, ValidityTs};
 
 enum ControlCode {
     Termination(NamedRows),
@@ -84,23 +85,32 @@ impl<'s, S: Storage<'s>> Db<S> {
                 ImperativeStmt::Continue { target, span, .. } => {
                     return Ok(Right(ControlCode::Continue(target.clone(), *span)));
                 }
-                ImperativeStmt::ReturnNil { .. } => {
-                    return Ok(Right(ControlCode::Termination(NamedRows::default())))
-                }
-                ImperativeStmt::ReturnProgram { prog, .. } => {
-                    ret = self.execute_single_program(
-                        prog.clone(),
-                        tx,
-                        cleanups,
-                        cur_vld,
-                        callback_targets,
-                        callback_collector,
-                    )?;
-                    return Ok(Right(ControlCode::Termination(ret)));
-                }
-                ImperativeStmt::ReturnTemp { rel, .. } => {
-                    let relation = tx.get_relation(rel, false)?;
-                    return Ok(Right(ControlCode::Termination(relation.as_named_rows(tx)?)));
+                ImperativeStmt::Return {returns} => {
+                    if returns.is_empty() {
+                        return Ok(Right(ControlCode::Termination(NamedRows::default())))
+                    }
+                    let mut current = None;
+                    for nxt in returns.iter().rev() {
+                        let mut nr = match nxt {
+                            Left(prog) => {
+                                self.execute_single_program(
+                                    prog.clone(),
+                                    tx,
+                                    cleanups,
+                                    cur_vld,
+                                    callback_targets,
+                                    callback_collector,
+                                )?
+                            }
+                            Right(rel) => {
+                                let relation = tx.get_relation(rel, false)?;
+                                 relation.as_named_rows(tx)?
+                            }
+                        };
+                        nr.next = current;
+                        current = Some(Box::new(nr))
+                    }
+                    return Ok(Right(ControlCode::Termination(*current.unwrap())))
                 }
                 ImperativeStmt::TempDebug { temp, .. } => {
                     let relation = tx.get_relation(temp, false)?;
@@ -128,10 +138,10 @@ impl<'s, S: Storage<'s>> Db<S> {
                     ) {
                         Ok(res) => ret = res,
                         Err(_) => {
-                            ret = NamedRows {
-                                headers: vec!["status".to_string()],
-                                rows: vec![vec![DataValue::from("FAILED")]],
-                            }
+                            ret = NamedRows::new(
+                                vec!["status".to_string()],
+                                vec![vec![DataValue::from("FAILED")]],
+                            )
                         }
                     }
                 }
@@ -221,7 +231,11 @@ impl<'s, S: Storage<'s>> Db<S> {
         }
         Ok(Left(ret))
     }
-    pub(crate) fn execute_imperative(&'s self, cur_vld: ValidityTs, ps: &ImperativeProgram) -> Result<NamedRows, Report> {
+    pub(crate) fn execute_imperative(
+        &'s self,
+        cur_vld: ValidityTs,
+        ps: &ImperativeProgram,
+    ) -> Result<NamedRows, Report> {
         let mut callback_collector = BTreeMap::new();
         let mut write_lock_names = BTreeSet::new();
         for p in ps {

cozo-core/src/runtime/relation.rs

@@ -145,7 +145,7 @@ impl RelationHandle {
                 .iter()
                 .map(|col| col.name.to_string()),
         );
-        Ok(NamedRows { headers, rows })
+        Ok(NamedRows::new(headers, rows))
     }
     #[allow(dead_code)]
     pub(crate) fn amend_key_prefix(&self, data: &mut [u8]) {