refactor

src/data/program.rs

@@ -2,9 +2,8 @@ use std::collections::btree_map::Entry;
 use std::collections::{BTreeMap, BTreeSet};
 use std::fmt::{Debug, Formatter};
 
-use anyhow::{anyhow, bail, ensure, Result};
+use anyhow::{anyhow, bail, Result};
 use either::{Left, Right};
-use itertools::Itertools;
 use smallvec::SmallVec;
 use smartstring::{LazyCompact, SmartString};
 
@@ -15,6 +14,7 @@ use crate::data::expr::Expr;
 use crate::data::id::{EntityId, Validity};
 use crate::data::symb::{Symbol, PROG_ENTRY};
 use crate::data::value::DataValue;
+use crate::parse::query::{ConstRules, QueryOutOptions};
 
 #[derive(Default)]
 pub(crate) struct TempSymbGen {
@@ -129,37 +129,11 @@ impl MagicAlgoRuleArg {
 #[derive(Debug, Clone)]
 pub(crate) struct InputProgram {
     pub(crate) prog: BTreeMap<Symbol, InputRulesOrAlgo>,
+    pub(crate) const_rules: ConstRules,
+    pub(crate) out_opts: QueryOutOptions,
 }
 
 impl InputProgram {
-    pub(crate) fn validate_entry(&self) -> Result<()> {
-        match self
-            .prog
-            .get(&PROG_ENTRY)
-            .ok_or_else(|| anyhow!("program entry point not found"))?
-        {
-            InputRulesOrAlgo::Rules(r) => {
-                ensure!(
-                    r.iter().map(|e| &e.head).all_equal(),
-                    "program entry point must have equal bindings"
-                );
-            }
-            InputRulesOrAlgo::Algo(_) => {}
-        }
-        Ok(())
-    }
-    pub(crate) fn get_entry_arity(&self) -> Result<usize> {
-        Ok(
-            match self
-                .prog
-                .get(&PROG_ENTRY)
-                .ok_or_else(|| anyhow!("program entry point not found"))?
-            {
-                InputRulesOrAlgo::Rules(rules) => rules[0].head.len(),
-                InputRulesOrAlgo::Algo(algo_apply) => algo_apply.arity()?,
-            },
-        )
-    }
     pub(crate) fn get_entry_head(&self) -> Result<&[Symbol]> {
         match self
             .prog

src/parse/mod.rs

@@ -3,3 +3,4 @@ pub(crate) mod pull;
 pub(crate) mod query;
 pub(crate) mod triple;
 pub(crate) mod cozoscript;
+pub(crate) mod script;

src/parse/query.rs

@@ -85,12 +85,52 @@ impl QueryOutOptions {
 
 pub(crate) type ConstRules = BTreeMap<MagicSymbol, Vec<Tuple>>;
 
+fn get_entry_head(prog: &BTreeMap<Symbol, InputRulesOrAlgo>) -> Result<&[Symbol]> {
+    match prog
+        .get(&PROG_ENTRY)
+        .ok_or_else(|| anyhow!("program entry point not found"))?
+    {
+        InputRulesOrAlgo::Rules(rules) => Ok(&rules.last().unwrap().head),
+        InputRulesOrAlgo::Algo(_) => {
+            bail!("algo application does not have named entry head")
+        }
+    }
+}
+
+fn validate_entry(prog: &BTreeMap<Symbol, InputRulesOrAlgo>) -> Result<()> {
+    match prog
+        .get(&PROG_ENTRY)
+        .ok_or_else(|| anyhow!("program entry point not found"))?
+    {
+        InputRulesOrAlgo::Rules(r) => {
+            ensure!(
+                r.iter().map(|e| &e.head).all_equal(),
+                "program entry point must have equal bindings"
+            );
+        }
+        InputRulesOrAlgo::Algo(_) => {}
+    }
+    Ok(())
+}
+
+fn get_entry_arity(prog: &BTreeMap<Symbol, InputRulesOrAlgo>) -> Result<usize> {
+    Ok(
+        match prog
+            .get(&PROG_ENTRY)
+            .ok_or_else(|| anyhow!("program entry point not found"))?
+        {
+            InputRulesOrAlgo::Rules(rules) => rules[0].head.len(),
+            InputRulesOrAlgo::Algo(algo_apply) => algo_apply.arity()?,
+        },
+    )
+}
+
 impl SessionTx {
     pub(crate) fn parse_query(
         &mut self,
         payload: &JsonValue,
         params_pool: &BTreeMap<Symbol, DataValue>,
-    ) -> Result<(InputProgram, QueryOutOptions, ConstRules)> {
+    ) -> Result<InputProgram> {
         let vld = match payload.get("since") {
             None => Validity::current(),
             Some(v) => Validity::try_from(v)?,
@@ -101,10 +141,8 @@ impl SessionTx {
         let rules_payload = q
             .as_array()
             .ok_or_else(|| anyhow!("expect field 'q' to be an array in query {}", payload))?;
-        let mut input_prog = if rules_payload.is_empty() {
-            InputProgram {
-                prog: Default::default(),
-            }
+        let mut prog = if rules_payload.is_empty() {
+            Default::default()
         } else if rules_payload.first().unwrap().is_array() {
             let q = json!([{"rule": "?", "args": rules_payload}]);
             self.parse_input_rule_sets(&q, vld, params_pool)?
@@ -114,7 +152,7 @@ impl SessionTx {
         let out_spec = match payload.get("out") {
             None => None,
             Some(spec) => {
-                let entry_bindings = input_prog.get_entry_head()?;
+                let entry_bindings = get_entry_head(&prog)?;
                 Some(self.parse_query_out_spec(spec, entry_bindings))
             }
         };
@@ -220,7 +258,7 @@ impl SessionTx {
                 if !out_name.is_prog_entry() {
                     out_name.validate_not_reserved()?;
                 }
-                match input_prog.prog.entry(out_name) {
+                match prog.entry(out_name) {
                     Entry::Vacant(v) => {
                         v.insert(InputRulesOrAlgo::Algo(AlgoApply {
                             algo: AlgoHandle::new(name_symbol),
@@ -296,8 +334,8 @@ impl SessionTx {
             })
             .try_collect()?;
         if !sorters.is_empty() {
-            input_prog.validate_entry()?;
-            let entry_head = input_prog.get_entry_head()?;
+            validate_entry(&prog)?;
+            let entry_head = get_entry_head(&prog)?;
             if sorters
                 .iter()
                 .map(|(k, _v)| k)
@@ -342,7 +380,7 @@ impl SessionTx {
                     .ok_or_else(|| anyhow!("view name must be a string"))?;
                 let name = Symbol::from(name);
                 ensure!(!name.is_reserved(), "view name {} is reserved", name);
-                let entry_arity = input_prog.get_entry_arity()?;
+                let entry_arity = get_entry_arity(&prog)?;
 
                 (
                     ViewRelMetadata {
@@ -355,9 +393,10 @@ impl SessionTx {
                 )
             }),
         };
-        Ok((
-            input_prog,
-            QueryOutOptions {
+        Ok(InputProgram {
+            prog,
+            const_rules,
+            out_opts: QueryOutOptions {
                 out_spec,
                 vld,
                 limit,
@@ -366,8 +405,7 @@ impl SessionTx {
                 as_view,
                 timeout,
             },
-            const_rules,
-        ))
+        })
     }
     fn parse_query_out_spec(
         &mut self,
@@ -437,7 +475,7 @@ impl SessionTx {
         payload: &JsonValue,
         default_vld: Validity,
         params_pool: &BTreeMap<Symbol, DataValue>,
-    ) -> Result<InputProgram> {
+    ) -> Result<BTreeMap<Symbol, InputRulesOrAlgo>> {
         let rules = payload
             .as_array()
             .ok_or_else(|| anyhow!("expect array for rules, got {}", payload))?
@@ -470,16 +508,16 @@ impl SessionTx {
             .try_collect()?;
 
         match ret.get(&PROG_ENTRY as &Symbol) {
-            None => Ok(InputProgram { prog: ret }),
+            None => Ok(ret),
             Some(ruleset) => match ruleset {
                 InputRulesOrAlgo::Rules(ruleset) => {
                     if !ruleset.iter().map(|r| &r.head).all_equal() {
                         bail!("all heads for the entry query must be identical");
                     } else {
-                        Ok(InputProgram { prog: ret })
+                        Ok(ret)
                     }
                 }
-                InputRulesOrAlgo::Algo(_) => Ok(InputProgram { prog: ret }),
+                InputRulesOrAlgo::Algo(_) => Ok(ret),
             },
         }
     }

src/parse/script/mod.rs

@@ -0,0 +1,33 @@
+use anyhow::Result;
+use pest::Parser;
+
+use crate::data::program::InputProgram;
+
+#[derive(pest_derive::Parser)]
+#[grammar = "cozoscript.pest"]
+pub(crate) struct CozoScriptParser;
+
+pub(crate) type Pair<'a> = pest::iterators::Pair<'a, Rule>;
+pub(crate) type Pairs<'a> = pest::iterators::Pairs<'a, Rule>;
+
+pub(crate) enum CozoScript {
+    Query(InputProgram),
+}
+
+pub(crate) fn parse_script(src: &str) -> Result<CozoScript> {
+    let parsed = CozoScriptParser::parse(Rule::script, src)?.next().unwrap();
+    Ok(match parsed.as_rule() {
+        Rule::query_script => CozoScript::Query(parse_query(parsed.into_inner())?),
+        Rule::schema_script => todo!(),
+        Rule::tx_script => todo!(),
+        Rule::sys_script => todo!(),
+        _ => unreachable!(),
+    })
+}
+
+fn parse_query(src: Pairs<'_>) -> Result<InputProgram> {
+    // let x = InputProgram {
+    //     prog: Default::default(),
+    // };
+    todo!()
+}

src/runtime/db.rs

@@ -435,9 +435,8 @@ impl Db {
     }
     pub fn run_query(&self, payload: &JsonValue) -> Result<JsonValue> {
         let mut tx = self.transact()?;
-        let (input_program, out_opts, const_rules) =
-            tx.parse_query(payload, &Default::default())?;
-        if let Some((meta, op)) = &out_opts.as_view {
+        let input_program = tx.parse_query(payload, &Default::default())?;
+        if let Some((meta, op)) = &input_program.out_opts.as_view {
             if *op == ViewOp::Create {
                 ensure!(
                     !tx.view_exists(&meta.name)?,
@@ -457,10 +456,11 @@ impl Db {
             .stratify()?
             .magic_sets_rewrite();
         debug!("{:#?}", program);
-        let (compiled, stores) = tx.stratified_magic_compile(&program, &const_rules)?;
+        let (compiled, stores) =
+            tx.stratified_magic_compile(&program, &input_program.const_rules)?;
 
         let poison = Poison::default();
-        if let Some(secs) = out_opts.timeout {
+        if let Some(secs) = input_program.out_opts.timeout {
             poison.set_timeout(secs);
         }
         let id = self.queries_count.fetch_add(1, Ordering::AcqRel);
@@ -477,8 +477,8 @@ impl Db {
         let result = tx.stratified_magic_evaluate(
             &compiled,
             &stores,
-            if out_opts.sorters.is_empty() {
-                out_opts.num_to_take()
+            if input_program.out_opts.sorters.is_empty() {
+                input_program.out_opts.num_to_take()
             } else {
                 None
             },
@@ -488,35 +488,35 @@ impl Db {
             Err(_) => JsonValue::Null,
             Ok(headers) => headers.iter().map(|v| json!(v.0)).collect(),
         };
-        if !out_opts.sorters.is_empty() {
+        if !input_program.out_opts.sorters.is_empty() {
             let entry_head = input_program.get_entry_head()?.to_vec();
-            let sorted_result = tx.sort_and_collect(result, &out_opts.sorters, &entry_head)?;
-            let sorted_iter = if let Some(offset) = out_opts.offset {
+            let sorted_result = tx.sort_and_collect(result, &input_program.out_opts.sorters, &entry_head)?;
+            let sorted_iter = if let Some(offset) = input_program.out_opts.offset {
                 Left(sorted_result.scan_sorted().skip(offset))
             } else {
                 Right(sorted_result.scan_sorted())
             };
-            let sorted_iter = if let Some(limit) = out_opts.limit {
+            let sorted_iter = if let Some(limit) = input_program.out_opts.limit {
                 Left(sorted_iter.take(limit))
             } else {
                 Right(sorted_iter)
             };
-            if let Some((meta, view_op)) = out_opts.as_view {
+            if let Some((meta, view_op)) = input_program.out_opts.as_view {
                 tx.execute_view(sorted_iter, view_op, &meta)?;
                 Ok(json!({"view": "OK"}))
             } else {
                 let ret: Vec<_> = tx
-                    .run_pull_on_query_results(sorted_iter, out_opts)?
+                    .run_pull_on_query_results(sorted_iter, input_program.out_opts)?
                     .try_collect()?;
                 Ok(json!({ "rows": ret, "headers": headers }))
             }
         } else {
-            if let Some((meta, view_op)) = out_opts.as_view {
+            if let Some((meta, view_op)) = input_program.out_opts.as_view {
                 tx.execute_view(result.scan_all(), view_op, &meta)?;
                 Ok(json!({"view": "OK"}))
             } else {
                 let ret: Vec<_> = tx
-                    .run_pull_on_query_results(result.scan_all(), out_opts)?
+                    .run_pull_on_query_results(result.scan_all(), input_program.out_opts)?
                     .try_collect()?;
                 Ok(json!({ "rows": ret, "headers": headers }))
             }
@@ -632,9 +632,12 @@ impl Db {
         CustomIter { it, started: false }
     }
     pub fn list_relations(&self) -> Result<JsonValue> {
-        let lower = Tuple(vec![DataValue::Str(SmartString::from(""))]).encode_as_key(ViewRelId::SYSTEM);
-        let upper = Tuple(vec![DataValue::Str(SmartString::from(String::from(LARGEST_UTF_CHAR)))])
-            .encode_as_key(ViewRelId::SYSTEM);
+        let lower =
+            Tuple(vec![DataValue::Str(SmartString::from(""))]).encode_as_key(ViewRelId::SYSTEM);
+        let upper = Tuple(vec![DataValue::Str(SmartString::from(String::from(
+            LARGEST_UTF_CHAR,
+        )))])
+        .encode_as_key(ViewRelId::SYSTEM);
         let mut it = self
             .view_db
             .transact()