From 16c82f1ec18c221b0d98dfcfb805ed3642354f5b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Kat=20March=C3=A1n?= <kzm@zkat.tech>
Date: Fri, 22 Apr 2022 12:00:32 -0700
Subject: [PATCH] feat(types): add type annotation support (#31)

---
 src/document.rs |  2 +-
 src/entry.rs    |  4 ++--
 src/node.rs     | 16 +++++++++++++++-
 src/parser.rs   | 13 +++++++++++++
 4 files changed, 31 insertions(+), 4 deletions(-)

diff --git a/src/document.rs b/src/document.rs
index 285d3cf..4888c2a 100644
--- a/src/document.rs
+++ b/src/document.rs
@@ -254,7 +254,7 @@ foo 1 2 \"three\" null true bar=\"baz\" {
     - 1
     - 2
     - \"three\"
-    something \"else\"\r
+    (mytype)something (\"name\")\"else\"\r
 }
 
 null_id null_prop=null
diff --git a/src/entry.rs b/src/entry.rs
index 28c74a7..0f05870 100644
--- a/src/entry.rs
+++ b/src/entry.rs
@@ -11,7 +11,7 @@ use crate::{KdlError, KdlErrorKind, KdlIdentifier, KdlValue};
 #[derive(Debug, Clone, PartialEq)]
 pub struct KdlEntry {
     pub(crate) leading: Option<String>,
-    pub(crate) ty: Option<String>,
+    pub(crate) ty: Option<KdlIdentifier>,
     pub(crate) value: KdlValue,
     pub(crate) value_repr: Option<String>,
     pub(crate) name: Option<KdlIdentifier>,
@@ -109,7 +109,7 @@ impl Display for KdlEntry {
             write!(f, "{}", leading)?;
         }
         if let Some(ty) = &self.ty {
-            write!(f, "{}", ty)?;
+            write!(f, "({})", ty)?;
         }
         if let Some(name) = &self.name {
             write!(f, "{}=", name)?;
diff --git a/src/node.rs b/src/node.rs
index 971a615..1f10415 100644
--- a/src/node.rs
+++ b/src/node.rs
@@ -14,7 +14,7 @@ use crate::{KdlDocument, KdlEntry, KdlError, KdlErrorKind, KdlIdentifier, KdlVal
 #[derive(Debug, Clone, PartialEq)]
 pub struct KdlNode {
     pub(crate) leading: Option<String>,
-    pub(crate) ty: Option<String>,
+    pub(crate) ty: Option<KdlIdentifier>,
     pub(crate) name: KdlIdentifier,
     // TODO: consider using `hashlink` for this instead, later.
     pub(crate) entries: Vec<KdlEntry>,
@@ -52,6 +52,20 @@ impl KdlNode {
         self.name = name.into();
     }
 
+    /// Gets the node's type, if any.
+    pub fn ty(&self) -> Option<&KdlIdentifier> {
+        self.ty.as_ref()
+    }
+
+    /// Gets a mutable reference to the node's type.
+    pub fn ty_mut(&mut self) -> &mut Option<KdlIdentifier> {
+        &mut self.ty
+    }
+
+    pub fn set_ty(&mut self, ty: impl Into<KdlIdentifier>) {
+        self.ty = Some(ty.into());
+    }
+
     /// Returns a reference to this node's entries (arguments and properties).
     pub fn entries(&self) -> &[KdlEntry] {
         &self.entries
diff --git a/src/parser.rs b/src/parser.rs
index bbac117..7238ecb 100644
--- a/src/parser.rs
+++ b/src/parser.rs
@@ -23,6 +23,7 @@ pub(crate) fn document(input: &str) -> IResult<&str, KdlDocument, KdlParseError<
 
 pub(crate) fn node(input: &str) -> IResult<&str, KdlNode, KdlParseError<&str>> {
     let (input, leading) = all_whitespace(input)?;
+    let (input, ty) = opt(annotation)(input)?;
     let (input, name) = identifier(input)?;
     let (input, entries) = many0(entry)(input)?;
     let (input, children) = opt(children)(input)?;
@@ -33,6 +34,7 @@ pub(crate) fn node(input: &str) -> IResult<&str, KdlNode, KdlParseError<&str>> {
     let mut node = KdlNode::new(name);
     node.set_leading(leading);
     node.set_trailing(trailing);
+    node.ty = ty;
     let ents = node.entries_mut();
     *ents = entries;
     if let Some((before, children)) = children {
@@ -70,10 +72,12 @@ pub(crate) fn entry(input: &str) -> IResult<&str, KdlEntry, KdlParseError<&str>>
 
 fn property(input: &str) -> IResult<&str, KdlEntry, KdlParseError<&str>> {
     let (input, leading) = recognize(many0(node_space))(input)?;
+    let (input, ty) = opt(annotation)(input)?;
     let (input, name) = identifier(input)?;
     let (input, _) = tag("=")(input)?;
     let (input, (raw, value)) = value(input)?;
     let mut entry = KdlEntry::new_prop(name, value);
+    entry.ty = ty;
     entry.set_leading(if leading.is_empty() { " " } else { leading });
     entry.set_value_repr(raw);
     Ok((input, entry))
@@ -81,8 +85,10 @@ fn property(input: &str) -> IResult<&str, KdlEntry, KdlParseError<&str>> {
 
 fn argument(input: &str) -> IResult<&str, KdlEntry, KdlParseError<&str>> {
     let (input, leading) = recognize(many0(node_space))(input)?;
+    let (input, ty) = opt(annotation)(input)?;
     let (input, (raw, value)) = value(input)?;
     let mut entry = KdlEntry::new(value);
+    entry.ty = ty;
     entry.set_leading(if leading.is_empty() { " " } else { leading });
     entry.set_value_repr(raw);
     Ok((input, entry))
@@ -110,6 +116,13 @@ fn children(input: &str) -> IResult<&str, (&str, KdlDocument), KdlParseError<&st
     Ok((input, (before, children)))
 }
 
+fn annotation(input: &str) -> IResult<&str, KdlIdentifier, KdlParseError<&str>> {
+    let (input, _) = tag("(")(input)?;
+    let (input, ty) = identifier(input)?;
+    let (input, _) = tag(")")(input)?;
+    Ok((input, ty))
+}
+
 fn all_whitespace(input: &str) -> IResult<&str, &str, KdlParseError<&str>> {
     recognize(many0(alt((comment, unicode_space, newline))))(input)
 }