output string contains aggr name

src/cozoscript.pest

@@ -95,9 +95,10 @@ grouping = { "(" ~ expr ~ ")" }
 
 option = _{(limit_option|offset_option|out_option|sort_option|relation_option|timeout_option|
             assert_none_option|assert_some_option) ~ ";"?}
+out_arg = @{var ~ ("(" ~ var ~ ")")?}
 limit_option = {":limit"  ~ expr}
 offset_option = {":offset" ~ expr}
-out_option = {":pull" ~ var ~ ("@" ~ expr)? ~ out_spec}
+out_option = {":pull" ~ out_arg ~ ("@" ~ expr)? ~ out_spec}
 sort_option = {(":sort" | ":order") ~ (sort_arg ~ ",")* ~ sort_arg }
 relation_option = {":relation" ~ relation_op ~ compound_ident}
 relation_op = _{relation_create | relation_rederive | relation_put | relation_retract}
@@ -106,7 +107,7 @@ relation_rederive = {"rederive"}
 relation_put = {"put"}
 relation_retract = {"retract"}
 timeout_option = {":timeout" ~ expr }
-sort_arg = { sort_dir? ~ var }
+sort_arg = { sort_dir? ~ out_arg }
 sort_dir = _{ sort_asc | sort_desc }
 sort_asc = {"+"}
 sort_desc = {"-"}

src/data/program.rs

@@ -1,9 +1,9 @@
-use std::collections::btree_map::Entry;
 use std::collections::{BTreeMap, BTreeSet};
+use std::collections::btree_map::Entry;
 use std::fmt::{Debug, Formatter};
 
 use either::{Left, Right};
-use miette::{ensure, Diagnostic, Result};
+use miette::{Diagnostic, ensure, Result};
 use smallvec::SmallVec;
 use smartstring::{LazyCompact, SmartString};
 use thiserror::Error;
@@ -13,7 +13,7 @@ use crate::data::aggr::Aggregation;
 use crate::data::attr::Attribute;
 use crate::data::expr::Expr;
 use crate::data::id::Validity;
-use crate::data::symb::{Symbol, PROG_ENTRY};
+use crate::data::symb::{PROG_ENTRY, Symbol};
 use crate::data::tuple::Tuple;
 use crate::data::value::DataValue;
 use crate::parse::pull::OutPullSpec;
@@ -515,6 +515,46 @@ impl InputProgram {
 
         Err(NoEntryError.into())
     }
+    pub(crate) fn get_entry_out_head(&self) -> Result<Vec<Symbol>> {
+        if let Some(entry) = self.prog.get(&Symbol::new(PROG_ENTRY, SourceSpan(0, 0))) {
+            return match entry {
+                InputRulesOrAlgo::Rules { rules } => {
+                    let head = &rules.last().unwrap().head;
+                    let mut ret = Vec::with_capacity(head.len());
+                    let aggrs = &rules.last().unwrap().aggr;
+                    for (symb, aggr) in head.iter().zip(aggrs.iter()) {
+                        if let Some((aggr, _)) = aggr {
+                            ret.push(Symbol::new(
+                                &format!("{}({})", aggr.name.strip_prefix("AGGR_").unwrap().to_ascii_lowercase(),
+                                         symb), symb.span))
+                        } else {
+                            ret.push(symb.clone())
+                        }
+                    }
+                    Ok(ret)
+                }
+                InputRulesOrAlgo::Algo { algo: algo_apply } => {
+                    if algo_apply.head.is_empty() {
+                        Err(EntryHeadNotExplicitlyDefinedError(entry.first_span()).into())
+                    } else {
+                        Ok(algo_apply.head.to_vec())
+                    }
+                }
+            };
+        }
+
+        if let Some(ConstRule { bindings, span, .. }) = self.const_rules.get(&MagicSymbol::Muggle {
+            inner: Symbol::new(PROG_ENTRY, SourceSpan(0, 0)),
+        }) {
+            return if bindings.is_empty() {
+                Err(EntryHeadNotExplicitlyDefinedError(*span).into())
+            } else {
+                Ok(bindings.to_vec())
+            };
+        }
+
+        Err(NoEntryError.into())
+    }
     pub(crate) fn get_entry_head(&self) -> Result<&[Symbol]> {
         if let Some(entry) = self.prog.get(&Symbol::new(PROG_ENTRY, SourceSpan(0, 0))) {
             return match entry {

src/parse/query.rs

@@ -332,7 +332,7 @@ pub(crate) fn parse_query(
                     let mut span = part.extract_span();
                     for a in part.into_inner() {
                         match a.as_rule() {
-                            Rule::var => {
+                            Rule::out_arg => {
                                 var = a.as_str();
                                 span = a.extract_span();
                             }
@@ -438,7 +438,7 @@ pub(crate) fn parse_query(
         #[diagnostic(code(parser::sort_key_not_found))]
         struct SortKeyNotFound(String, #[label] SourceSpan);
 
-        let head_args = prog.get_entry_head().unwrap_or(&[]);
+        let head_args = prog.get_entry_out_head()?;
 
         for (sorter, _) in &prog.out_opts.sorters {
             ensure!(

src/query/pull.rs

@@ -196,7 +196,7 @@ impl SessionTx {
     pub(crate) fn run_pull_on_query_results(
         &self,
         res_iter: impl Iterator<Item=Result<Tuple>>,
-        headers: Option<&[Symbol]>,
+        headers: Option<Vec<Symbol>>,
         out_spec: &BTreeMap<Symbol, (Vec<OutPullSpec>, Option<Validity>)>,
         default_vld: Validity,
     ) -> Result<Vec<JsonValue>> {
@@ -205,7 +205,7 @@ impl SessionTx {
                 .map_ok(|tuple| tuple.0.into_iter().map(JsonValue::from).collect())
                 .try_collect()?)
         } else {
-            let headers = headers.unwrap_or(&[]);
+            let headers = headers.unwrap_or(vec![]);
             let mut idx2pull: Vec<Option<(Vec<_>, _)>> = Vec::with_capacity(headers.len());
             for head in headers.iter() {
                 match out_spec.get(head) {

src/runtime/db.rs

@@ -554,12 +554,12 @@ impl Db {
                 }
             }
         }
-        let json_headers = match input_program.get_entry_head() {
+        let json_headers = match input_program.get_entry_out_head() {
             Err(_) => JsonValue::Null,
-            Ok(headers) => headers.iter().map(|v| json!(v.name)).collect(),
+            Ok(headers) => headers.into_iter().map(|v| json!(v.name)).collect(),
         };
         if !input_program.out_opts.sorters.is_empty() {
-            let entry_head = input_program.get_entry_head()?.to_vec();
+            let entry_head = input_program.get_entry_out_head()?;
             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 {
@@ -581,7 +581,7 @@ impl Db {
             } else {
                 let ret: Vec<_> = tx.run_pull_on_query_results(
                     sorted_iter,
-                    input_program.get_entry_head().ok(),
+                    input_program.get_entry_out_head().ok(),
                     &input_program.out_opts.out_spec,
                     default_vld,
                 )?;
@@ -607,7 +607,7 @@ impl Db {
             } else {
                 let ret: Vec<_> = tx.run_pull_on_query_results(
                     scan,
-                    input_program.get_entry_head().ok(),
+                    input_program.get_entry_out_head().ok(),
                     &input_program.out_opts.out_spec,
                     default_vld,
                 )?;

tests/air_routes.rs

@@ -729,7 +729,7 @@ fn air_routes() -> Result<()> {
                                [r route.dst a2],
                                [a2 airport.country us],
                                [a2 airport.iata code];
-        :order r;
+        :order count(r);
     "#,
         &params, false,
     )?;