From b9b3475a146e493dfaca61c32247d90f0e2b3448 Mon Sep 17 00:00:00 2001
From: Ziyang Hu <hu.ziyang@cantab.net>
Date: Wed, 14 Jun 2023 00:33:30 +0800
Subject: [PATCH] mutation shorthands

---
 cozo-core/src/parse/query.rs   | 36 ++++++++++++++++++++++++++++++++--
 cozo-core/src/parse/schema.rs  | 11 -----------
 cozo-core/src/runtime/tests.rs |  9 +++++++++
 3 files changed, 43 insertions(+), 13 deletions(-)

diff --git a/cozo-core/src/parse/query.rs b/cozo-core/src/parse/query.rs
index 3392eda8..a47da3fc 100644
--- a/cozo-core/src/parse/query.rs
+++ b/cozo-core/src/parse/query.rs
@@ -449,14 +449,15 @@ pub(crate) fn parse_query(
         }
     }
 
+    let head_args = prog.get_entry_out_head()?;
+
+
     if !prog.out_opts.sorters.is_empty() {
         #[derive(Debug, Error, Diagnostic)]
         #[error("Sort key '{0}' not found")]
         #[diagnostic(code(parser::sort_key_not_found))]
         struct SortKeyNotFound(String, #[label] SourceSpan);
 
-        let head_args = prog.get_entry_out_head()?;
-
         for (sorter, _) in &prog.out_opts.sorters {
             ensure!(
                 head_args.contains(sorter),
@@ -465,6 +466,37 @@ pub(crate) fn parse_query(
         }
     }
 
+    #[derive(Debug, Error, Diagnostic)]
+    #[error("Input relation '{0}' has no keys")]
+    #[diagnostic(code(parser::relation_has_no_keys))]
+    struct RelationHasNoKeys(String, #[label] SourceSpan);
+
+
+    if let Some((handle, _, _)) = &mut prog.out_opts.store_relation {
+        if handle.key_bindings.is_empty() {
+            if handle.dep_bindings.is_empty() {
+                if head_args.is_empty() {
+                    bail!(RelationHasNoKeys(handle.name.to_string(), handle.span));
+                } else {
+                    handle.key_bindings = head_args.clone();
+                    handle.metadata.keys = head_args
+                        .iter()
+                        .map(|s| ColumnDef {
+                            name: s.name.clone(),
+                            typing: NullableColType {
+                                coltype: ColType::Any,
+                                nullable: true,
+                            },
+                            default_gen: None,
+                        })
+                        .collect();
+                }
+            } else {
+                bail!(RelationHasNoKeys(handle.name.to_string(), handle.span));
+            }
+        }
+    }
+
     Ok(prog)
 }
 
diff --git a/cozo-core/src/parse/schema.rs b/cozo-core/src/parse/schema.rs
index 17df2030..812dec98 100644
--- a/cozo-core/src/parse/schema.rs
+++ b/cozo-core/src/parse/schema.rs
@@ -22,9 +22,6 @@ use crate::parse::{ExtractSpan, Pair, Rule, SourceSpan};
 pub(crate) fn parse_schema(
     pair: Pair<'_>,
 ) -> Result<(StoredRelationMetadata, Vec<Symbol>, Vec<Symbol>)> {
-    // assert_eq!(pair.as_rule(), Rule::table_schema);
-    let span = pair.extract_span();
-
     let mut src = pair.into_inner();
     let mut keys = vec![];
     let mut dependents = vec![];
@@ -57,14 +54,6 @@ pub(crate) fn parse_schema(
         }
     }
 
-    if seen_names.is_empty() {
-        #[derive(Debug, Error, Diagnostic)]
-        #[error("Empty schema is not allowed here")]
-        #[diagnostic(code(parser::empty_rel_schema))]
-        struct EmptySchema(#[label] SourceSpan);
-        bail!(EmptySchema(span))
-    }
-
     Ok((
         StoredRelationMetadata {
             keys,
diff --git a/cozo-core/src/runtime/tests.rs b/cozo-core/src/runtime/tests.rs
index 9c450333..6f1e814c 100644
--- a/cozo-core/src/runtime/tests.rs
+++ b/cozo-core/src/runtime/tests.rs
@@ -1315,4 +1315,13 @@ fn puts() {
                 seg_vecs = [], seg_pos = [[0, 10]]
         :put cm_txt {tid, aid, tag, text, info_amount, seg_vecs, seg_pos, dup_for}
     ").unwrap();
+}
+
+#[test]
+fn short_hand() {
+    let db = DbInstance::default();
+    db.run_default(r":create x {x => y, z}").unwrap();
+    db.run_default(r"?[x, y, z] <- [[1, 2, 3]] :put x {}").unwrap();
+    let r = db.run_default(r"?[x, y, z] := *x {x, y, z}").unwrap();
+    assert_eq!(r.into_json()["rows"], json!([[1, 2, 3]]));
 }
\ No newline at end of file