Improve literal types (#4712)

Co-authored-by: Raphael Darley <raphael.darley@surrealdb.com>
Co-authored-by: Tobie Morgan Hitchcock <tobie@surrealdb.com>

core/src/doc/clean.rs

@@ -23,7 +23,7 @@ impl Document {
 			// Loop through all field statements
 			for fd in self.fd(ctx, opt).await?.iter() {
 				// Is this a schemaless field?
-				match fd.flex {
+				match fd.flex || fd.kind.as_ref().is_some_and(|k| k.is_literal_nested()) {
 					false => {
 						// Loop over this field in the document
 						for k in self.current.doc.as_ref().each(&fd.name).into_iter() {

core/src/sql/kind.rs

@@ -57,6 +57,93 @@ impl Kind {
 		matches!(self, Kind::Record(_))
 	}
 
+	/// Returns true if this type is optional
+	pub(crate) fn can_be_none(&self) -> bool {
+		matches!(self, Kind::Option(_) | Kind::Any)
+	}
+
+	/// Returns the kind in case of a literal, otherwise returns the kind itself
+	fn to_kind(&self) -> Self {
+		match self {
+			Kind::Literal(l) => l.to_kind(),
+			k => k.to_owned(),
+		}
+	}
+
+	/// Returns true if this type is a literal, or contains a literal
+	pub(crate) fn is_literal_nested(&self) -> bool {
+		if matches!(self, Kind::Literal(_)) {
+			return true;
+		}
+
+		if let Kind::Option(x) = self {
+			return x.is_literal_nested();
+		}
+
+		if let Kind::Either(x) = self {
+			return x.iter().any(|x| x.is_literal_nested());
+		}
+
+		false
+	}
+
+	/// Returns Some if this type can be converted into a discriminated object, None otherwise
+	pub(crate) fn to_discriminated(&self) -> Option<Kind> {
+		match self {
+			Kind::Either(nested) => {
+				if let Some(nested) = nested
+					.iter()
+					.map(|k| match k {
+						Kind::Literal(Literal::Object(o)) => Some(o),
+						_ => None,
+					})
+					.collect::<Option<Vec<&BTreeMap<String, Kind>>>>()
+				{
+					if let Some(first) = nested.first() {
+						let mut key: Option<String> = None;
+
+						'key: for (k, v) in first.iter() {
+							let mut kinds: Vec<Kind> = vec![v.to_owned()];
+							for item in nested[1..].iter() {
+								if let Some(kind) = item.get(k) {
+									match kind {
+										Kind::Literal(l)
+											if kinds.contains(&l.to_kind())
+												|| kinds.contains(&Kind::Literal(l.to_owned())) =>
+										{
+											continue 'key;
+										}
+										kind if kinds.iter().any(|k| *kind == k.to_kind()) => {
+											continue 'key;
+										}
+										kind => {
+											kinds.push(kind.to_owned());
+										}
+									}
+								} else {
+									continue 'key;
+								}
+							}
+
+							key = Some(k.clone());
+							break;
+						}
+
+						if let Some(key) = key {
+							return Some(Kind::Literal(Literal::DiscriminatedObject(
+								key.clone(),
+								nested.into_iter().map(|o| o.to_owned()).collect(),
+							)));
+						}
+					}
+				}
+
+				None
+			}
+			_ => None,
+		}
+	}
+
 	// return the kind of the contained value.
 	//
 	// For example: for `array<number>` or `set<number>` this returns `number`.
@@ -169,6 +256,7 @@ pub enum Literal {
 	Duration(Duration),
 	Array(Vec<Kind>),
 	Object(BTreeMap<String, Kind>),
+	DiscriminatedObject(String, Vec<BTreeMap<String, Kind>>),
 }
 
 impl Literal {
@@ -187,6 +275,7 @@ impl Literal {
 				Kind::Array(Box::new(Kind::Any), None)
 			}
 			Self::Object(_) => Kind::Object,
+			Self::DiscriminatedObject(_, _) => Kind::Object,
 		}
 	}
 
@@ -226,7 +315,7 @@ impl Literal {
 			},
 			Self::Object(o) => match value {
 				Value::Object(x) => {
-					if o.len() != x.len() {
+					if o.len() < x.len() {
 						return false;
 					}
 
@@ -235,7 +324,7 @@ impl Literal {
 							if value.to_owned().coerce_to(v).is_err() {
 								return false;
 							}
-						} else {
+						} else if !v.can_be_none() {
 							return false;
 						}
 					}
@@ -244,6 +333,34 @@ impl Literal {
 				}
 				_ => false,
 			},
+			Self::DiscriminatedObject(key, discriminants) => match value {
+				Value::Object(x) => {
+					let value = x.get(key).unwrap_or(&Value::None);
+					if let Some(o) = discriminants
+						.iter()
+						.find(|o| value.to_owned().coerce_to(o.get(key).unwrap()).is_ok())
+					{
+						if o.len() < x.len() {
+							return false;
+						}
+
+						for (k, v) in o.iter() {
+							if let Some(value) = x.get(k) {
+								if value.to_owned().coerce_to(v).is_err() {
+									return false;
+								}
+							} else if !v.can_be_none() {
+								return false;
+							}
+						}
+
+						true
+					} else {
+						false
+					}
+				}
+				_ => false,
+			},
 		}
 	}
 }
@@ -289,6 +406,40 @@ impl Display for Literal {
 					f.write_str(" }")
 				}
 			}
+			Literal::DiscriminatedObject(_, discriminants) => {
+				let mut f = Pretty::from(f);
+
+				for (i, o) in discriminants.iter().enumerate() {
+					if i > 0 {
+						f.write_str(" | ")?;
+					}
+
+					if is_pretty() {
+						f.write_char('{')?;
+					} else {
+						f.write_str("{ ")?;
+					}
+					if !o.is_empty() {
+						let indent = pretty_indent();
+						write!(
+							f,
+							"{}",
+							Fmt::pretty_comma_separated(o.iter().map(|args| Fmt::new(
+								args,
+								|(k, v), f| write!(f, "{}: {}", escape_key(k), v)
+							)),)
+						)?;
+						drop(indent);
+					}
+					if is_pretty() {
+						f.write_char('}')?;
+					} else {
+						f.write_str(" }")?;
+					}
+				}
+
+				Ok(())
+			}
 		}
 	}
 }

core/src/syn/parser/kind.rs

@@ -34,7 +34,9 @@ impl Parser<'_> {
 					while self.eat(t!("|")) {
 						kind.push(ctx.run(|ctx| self.parse_concrete_kind(ctx)).await?);
 					}
-					Ok(Kind::Either(kind))
+					let kind = Kind::Either(kind);
+					let kind = kind.to_discriminated().unwrap_or(kind);
+					Ok(kind)
 				} else {
 					Ok(first)
 				}
@@ -59,7 +61,9 @@ impl Parser<'_> {
 					while self.eat(t!("|")) {
 						kind.push(ctx.run(|ctx| self.parse_concrete_kind(ctx)).await?);
 					}
-					first = Kind::Either(kind);
+
+					let kind = Kind::Either(kind);
+					first = kind.to_discriminated().unwrap_or(kind);
 				}
 				self.expect_closing_delimiter(t!(">"), delim)?;
 				Ok(Kind::Option(Box::new(first)))
@@ -483,4 +487,55 @@ mod tests {
 		assert_eq!("set<float, 10>", format!("{}", out));
 		assert_eq!(out, Kind::Set(Box::new(Kind::Float), Some(10)));
 	}
+
+	#[test]
+	fn kind_discriminated_object() {
+		let sql = "{ status: 'ok', data: object } | { status: 'error', message: string }";
+		let res = kind(sql);
+		let out = res.unwrap();
+		assert_eq!(
+			"{ data: object, status: 'ok' } | { message: string, status: 'error' }",
+			format!("{}", out)
+		);
+		assert_eq!(
+			out,
+			Kind::Literal(Literal::DiscriminatedObject(
+				"status".to_string(),
+				vec![
+					map! {
+						"status".to_string() => Kind::Literal(Literal::String("ok".into())),
+						"data".to_string() => Kind::Object,
+					},
+					map! {
+						"status".to_string() => Kind::Literal(Literal::String("error".into())),
+						"message".to_string() => Kind::String,
+					},
+				]
+			))
+		);
+	}
+
+	#[test]
+	fn kind_union_literal_object() {
+		let sql = "{ status: 'ok', data: object } | { status: string, message: string }";
+		let res = kind(sql);
+		let out = res.unwrap();
+		assert_eq!(
+			"{ data: object, status: 'ok' } | { message: string, status: string }",
+			format!("{}", out)
+		);
+		assert_eq!(
+			out,
+			Kind::Either(vec![
+				Kind::Literal(Literal::Object(map! {
+					"status".to_string() => Kind::Literal(Literal::String("ok".into())),
+					"data".to_string() => Kind::Object,
+				})),
+				Kind::Literal(Literal::Object(map! {
+					"status".to_string() => Kind::String,
+					"message".to_string() => Kind::String,
+				})),
+			])
+		);
+	}
 }

sdk/tests/typing.rs

@@ -282,6 +282,41 @@ async fn strict_typing_none_null() -> Result<(), Error> {
 	Ok(())
 }
 
+#[tokio::test]
+async fn literal_typing() -> Result<(), Error> {
+	let sql = "
+		DEFINE TABLE test SCHEMAFULL;
+		DEFINE FIELD obj ON test TYPE {
+		    a: int,
+		    b: option<string>
+		};
+
+		CREATE ONLY test:1 SET obj = { a: 1 };
+		CREATE ONLY test:2 SET obj = { a: 2, b: 'foo' };
+		CREATE ONLY test:3 SET obj = { a: 3, b: 'bar', c: 'forbidden' };
+	";
+	let mut t = Test::new(sql).await?;
+	//
+	t.skip_ok(2)?;
+	t.expect_val(
+		"{
+			id: test:1,
+			obj: { a: 1 },
+		}",
+	)?;
+	t.expect_val(
+		"{
+			id: test:2,
+			obj: { a: 2, b: 'foo' },
+		}",
+	)?;
+	t.expect_error(
+		"Found { a: 3, b: 'bar', c: 'forbidden' } for field `obj`, with record `test:3`, but expected a { a: int, b: option<string> }",
+	)?;
+	//
+	Ok(())
+}
+
 #[tokio::test]
 async fn strict_typing_optional_object() -> Result<(), Error> {
 	let sql = "