simplify imperative scripts

cozo-core/src/cozoscript.pest

@@ -202,21 +202,22 @@ tuple_type = {"(" ~ (col_type ~ ",")* ~ col_type? ~ ")"}
 
 imperative_stmt = _{
     break_stmt | continue_stmt | return_stmt | debug_stmt |
-    query_script_inner | ignore_error_script | if_chain | while_block | do_while_block | temp_swap | remove_stmt
+    query_script_inner | ignore_error_script | if_chain | if_not_chain | loop_block | temp_swap
 }
 imperative_condition = _{underscore_ident | query_script_inner}
 if_chain = {"%if" ~ imperative_condition
-          ~ "%then" ~ imperative_block
+          ~ "%then"? ~ imperative_block
+          ~ ("%else" ~ imperative_block)? ~ "%end" }
+if_not_chain = {"%if_not" ~ imperative_condition
+                         ~ "%then"? ~ imperative_block
                          ~ ("%else" ~ imperative_block)? ~ "%end" }
 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)?}
-while_block = {("%mark" ~ ident)? ~ "%while" ~ imperative_condition ~ "%loop" ~ imperative_block ~ "%end"}
-do_while_block = {("%mark" ~ ident)? ~ "%loop" ~ imperative_block ~ "%while" ~ imperative_condition ~ "%end"}
+loop_block = {("%mark" ~ ident)? ~ "%loop" ~ imperative_block ~ "%end"}
 temp_swap = {"%swap" ~ underscore_ident ~ underscore_ident}
-remove_stmt = {"%remove" ~ underscore_ident}
 debug_stmt = {"%debug" ~ (ident | underscore_ident)}
 
 /*

cozo-core/src/parse/imperative.rs

@@ -91,7 +91,8 @@ fn parse_imperative_stmt(
                 },
             }
         }
-        Rule::if_chain => {
+        Rule::if_chain | Rule::if_not_chain => {
+            let negated = pair.as_rule() == Rule::if_not_chain;
             let span = pair.extract_span();
             let mut inner = pair.into_inner();
             let condition = inner.next().unwrap();
@@ -122,11 +123,11 @@ fn parse_imperative_stmt(
                 condition: cond,
                 then_branch: body,
                 else_branch: else_body,
+                negated,
                 span,
             }
         }
-        Rule::while_block => {
-            let span = pair.extract_span();
+        Rule::loop_block => {
             let mut inner = pair.into_inner();
             let mut mark = None;
             let mut nxt = inner.next().unwrap();
@@ -134,60 +135,8 @@ fn parse_imperative_stmt(
                 mark = Some(SmartString::from(nxt.as_str()));
                 nxt = inner.next().unwrap();
             }
-            let cond = match nxt.as_rule() {
-                Rule::underscore_ident => Left(SmartString::from(nxt.as_str())),
-                Rule::query_script_inner => Right(parse_query(
-                    nxt.into_inner(),
-                    param_pool,
-                    fixed_rules,
-                    cur_vld,
-                )?),
-                _ => unreachable!(),
-            };
-            let body = parse_imperative_block(
-                inner.next().unwrap(),
-                param_pool,
-                fixed_rules,
-                cur_vld,
-            )?;
-            ImperativeStmt::While {
-                label: mark,
-                condition: cond,
-                body,
-                span,
-            }
-        }
-        Rule::do_while_block => {
-            let span = pair.extract_span();
-            let mut inner = pair.into_inner();
-            let mut mark = None;
-            let mut nxt = inner.next().unwrap();
-            if nxt.as_rule() == Rule::ident {
-                mark = Some(SmartString::from(nxt.as_str()));
-                nxt = inner.next().unwrap();
-            }
-            let body = parse_imperative_block(
-                inner.next().unwrap(),
-                param_pool,
-                fixed_rules,
-                cur_vld,
-            )?;
-            let cond = match nxt.as_rule() {
-                Rule::underscore_ident => Left(SmartString::from(nxt.as_str())),
-                Rule::query_script_inner => Right(parse_query(
-                    nxt.into_inner(),
-                    param_pool,
-                    fixed_rules,
-                    cur_vld,
-                )?),
-                _ => unreachable!(),
-            };
-            ImperativeStmt::DoWhile {
-                label: mark,
-                body,
-                condition: cond,
-                span,
-            }
+            let body = parse_imperative_block(nxt, param_pool, fixed_rules, cur_vld)?;
+            ImperativeStmt::Loop { label: mark, body }
         }
         Rule::temp_swap => {
             // let span = pair.extract_span();
@@ -202,15 +151,6 @@ fn parse_imperative_stmt(
                 right: SmartString::from(right_name),
             }
         }
-        Rule::remove_stmt => {
-            // let span = pair.extract_span();
-            let name_p = pair.into_inner().next().unwrap();
-            let name = name_p.as_str();
-
-            ImperativeStmt::TempRemove {
-                temp: SmartString::from(name),
-            }
-        }
         Rule::debug_stmt => {
             // let span = pair.extract_span();
             let name_p = pair.into_inner().next().unwrap();

cozo-core/src/parse/mod.rs

@@ -9,6 +9,7 @@
 use either::Either;
 use std::cmp::{max, min};
 use std::collections::BTreeMap;
+use std::fmt::{Display, Formatter};
 use std::sync::Arc;
 
 use miette::{bail, Diagnostic, IntoDiagnostic, Result};
@@ -73,29 +74,18 @@ pub(crate) enum ImperativeStmt {
         condition: ImperativeCondition,
         then_branch: ImperativeProgram,
         else_branch: ImperativeProgram,
+        negated: bool,
         span: SourceSpan,
     },
-    While {
-        label: Option<SmartString<LazyCompact>>,
-        condition: ImperativeCondition,
-        body: ImperativeProgram,
-        span: SourceSpan,
-    },
-    DoWhile {
+    Loop {
         label: Option<SmartString<LazyCompact>>,
         body: ImperativeProgram,
-        condition: ImperativeCondition,
-        span: SourceSpan,
     },
     TempSwap {
         left: SmartString<LazyCompact>,
         right: SmartString<LazyCompact>,
         // span: SourceSpan,
     },
-    TempRemove {
-        temp: SmartString<LazyCompact>,
-        // span: SourceSpan,
-    },
     TempDebug {
         temp: SmartString<LazyCompact>,
     },
@@ -123,24 +113,13 @@ impl ImperativeStmt {
                 }) || then_branch.iter().any(|p| p.needs_write_tx())
                     || else_branch.iter().any(|p| p.needs_write_tx())
             }
-            ImperativeStmt::While {
-                condition, body, ..
-            }
-            | ImperativeStmt::DoWhile {
-                body, condition, ..
-            } => {
-                (match condition {
-                    ImperativeCondition::Left(_) => false,
-                    ImperativeCondition::Right(prog) => prog.needs_write_tx(),
-                }) || body.iter().any(|p| p.needs_write_tx())
-            }
+            ImperativeStmt::Loop { body, .. } => body.iter().any(|p| p.needs_write_tx()),
             ImperativeStmt::TempDebug { .. }
             | ImperativeStmt::ReturnTemp { .. }
             | ImperativeStmt::Break { .. }
             | ImperativeStmt::Continue { .. }
             | ImperativeStmt::ReturnNil { .. }
-            | ImperativeStmt::TempSwap { .. }
-            | ImperativeStmt::TempRemove { .. } => false,
+            | ImperativeStmt::TempSwap { .. } => false,
         }
     }
 }
@@ -165,6 +144,12 @@ impl CozoScript {
 )]
 pub struct SourceSpan(pub(crate) usize, pub(crate) usize);
 
+impl Display for SourceSpan {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}..{}", self.0, self.0 + self.1)
+    }
+}
+
 impl SourceSpan {
     pub(crate) fn merge(self, other: Self) -> Self {
         let s1 = self.0;
@@ -190,7 +175,7 @@ impl From<SourceSpan> for miette::SourceSpan {
 }
 
 #[derive(thiserror::Error, Diagnostic, Debug)]
-#[error("The query parser has encountered unexpected input / end of input")]
+#[error("The query parser has encountered unexpected input / end of input at {span}")]
 #[diagnostic(code(parser::pest))]
 pub(crate) struct ParseError {
     #[label]

cozo-core/src/runtime/db.rs

@@ -681,9 +681,11 @@ impl<'s, S: Storage<'s>> Db<S> {
                     then_branch,
                     else_branch,
                     span,
+                    negated,
                 } => {
                     let cond_val =
                         self.execute_imperative_condition(condition, tx, cleanups, cur_vld, *span)?;
+                    let cond_val = if *negated { !cond_val } else { cond_val };
                     let to_execute = if cond_val { then_branch } else { else_branch };
                     match self.execute_imperative_stmts(to_execute, tx, cleanups, cur_vld)? {
                         Left(rows) => {
@@ -692,57 +694,7 @@ impl<'s, S: Storage<'s>> Db<S> {
                         Right(ctrl) => return Ok(Right(ctrl)),
                     }
                 }
-                ImperativeStmt::While {
-                    label,
-                    condition,
-                    body,
-                    span,
-                } => {
-                    ret = Default::default();
-                    loop {
-                        let cond_val = self.execute_imperative_condition(
-                            condition, tx, cleanups, cur_vld, *span,
-                        )?;
-                        if cond_val {
-                            match self.execute_imperative_stmts(body, tx, cleanups, cur_vld)? {
-                                Left(_) => {}
-                                Right(ctrl) => match ctrl {
-                                    ControlCode::Termination(ret) => {
-                                        return Ok(Right(ControlCode::Termination(ret)))
-                                    }
-                                    ControlCode::Break(break_label, span) => {
-                                        if break_label.is_none() || break_label == *label {
-                                            break;
-                                        } else {
-                                            return Ok(Right(ControlCode::Break(
-                                                break_label,
-                                                span,
-                                            )));
-                                        }
-                                    }
-                                    ControlCode::Continue(cont_label, span) => {
-                                        if cont_label.is_none() || cont_label == *label {
-                                            continue;
-                                        } else {
-                                            return Ok(Right(ControlCode::Continue(
-                                                cont_label, span,
-                                            )));
-                                        }
-                                    }
-                                },
-                            }
-                        } else {
-                            ret = NamedRows::default();
-                            break;
-                        }
-                    }
-                }
-                ImperativeStmt::DoWhile {
-                    label,
-                    body,
-                    condition,
-                    span,
-                } => {
+                ImperativeStmt::Loop { label, body, .. } => {
                     ret = Default::default();
                     loop {
                         match self.execute_imperative_stmts(body, tx, cleanups, cur_vld)? {
@@ -768,33 +720,23 @@ impl<'s, S: Storage<'s>> Db<S> {
                             },
                         }
                     }
-                    let cond_val =
-                        self.execute_imperative_condition(condition, tx, cleanups, cur_vld, *span)?;
-                    if !cond_val {
-                        ret = NamedRows::default();
-                        break;
-                    }
                 }
                 ImperativeStmt::TempSwap { left, right, .. } => {
                     tx.rename_temp_relation(
                         Symbol::new(left.clone(), Default::default()),
-                        Symbol::new(SmartString::from("*temp*"), Default::default()),
+                        Symbol::new(SmartString::from("_*temp*"), Default::default()),
                     )?;
                     tx.rename_temp_relation(
                         Symbol::new(right.clone(), Default::default()),
                         Symbol::new(left.clone(), Default::default()),
                     )?;
                     tx.rename_temp_relation(
-                        Symbol::new(SmartString::from("*temp*"), Default::default()),
+                        Symbol::new(SmartString::from("_*temp*"), Default::default()),
                         Symbol::new(right.clone(), Default::default()),
                     )?;
                     ret = NamedRows::default();
                     break;
                 }
-                ImperativeStmt::TempRemove { temp, .. } => {
-                    tx.destroy_temp_relation(temp)?;
-                    ret = NamedRows::default();
-                }
             }
         }
         return Ok(Left(ret));

cozo-core/src/runtime/relation.rs

@@ -544,16 +544,6 @@ impl<'a> SessionTx<'a> {
         let upper_bound = Tuple::default().encode_as_key(store.id.next());
         Ok((lower_bound, upper_bound))
     }
-    pub(crate) fn destroy_temp_relation(&mut self, name: &str) -> Result<()> {
-        self.get_relation(name, true)?;
-        let key = DataValue::from(name);
-        let encoded = vec![key].encode_as_key(RelationId::SYSTEM);
-        self.temp_store_tx.del(&encoded)?;
-        // TODO: at the moment this doesn't actually remove anything from the store
-        // let lower_bound = Tuple::default().encode_as_key(store.id);
-        // let upper_bound = Tuple::default().encode_as_key(store.id.next());
-        Ok(())
-    }
     pub(crate) fn set_access_level(&mut self, rel: Symbol, level: AccessLevel) -> Result<()> {
         let mut meta = self.get_relation(&rel, true)?;
         meta.access_level = level;

cozo-core/src/runtime/tests.rs

@@ -268,15 +268,61 @@ fn imperative_script() {
     let res = db.run_script(r#"
         {:create _test {a}}
 
-        %while { len[count(x)] := *_test[x]; ?[x] := len[z], x = z < 10 }
         %loop
+            %if { len[count(x)] := *_test[x]; ?[x] := len[z], x = z >= 10 }
+                %then %return _test
+            %end
             { ?[a] := a = rand_uuid_v1(); :put _test {a} }
             %debug _test
         %end
+    "#, Default::default()).unwrap();
+    assert_eq!(res.rows.len(), 10);
+
+    let res = db.run_script(r#"
+        {?[a] <- [[1], [2], [3]]
+         :replace _test {a}}
+
+        %loop
+            { ?[a] := *_test[a]; :limit 1; :rm _test {a} }
+            %debug _test
+
+            %if_not _test
+            %then %break
+            %end
+        %end
 
         %return _test
     "#, Default::default()).unwrap();
-    assert_eq!(res.rows.len(), 10);
+    assert_eq!(res.rows.len(), 0);
+
+    let res = db.run_script(r#"
+        {:create _test {a}}
+
+        %loop
+            { ?[a] := a = rand_uuid_v1(); :put _test {a} }
+
+            %if { len[count(x)] := *_test[x]; ?[x] := len[z], x = z < 10 }
+                %continue
+            %end
+
+            %return _test
+            %debug _test
+        %end
+    "#, Default::default());
+    if let Err(err) = &res {
+        eprintln!("{err:?}");
+    }
+    assert_eq!(res.unwrap().rows.len(), 10);
+
+    let res = db.run_script(r#"
+        {?[a] <- [[1], [2], [3]]
+         :replace _test {a}}
+        {?[a] <- []
+         :replace _test2 {a}}
+        %swap _test _test2
+        %return _test
+    "#, Default::default()).unwrap();
+    assert_eq!(res.rows.len(), 0);
 }
 
 #[test]