From c1e6f7e918eacaad9c8b1a4b27fcd4d6245aaf68 Mon Sep 17 00:00:00 2001
From: cel 🌸 <cel@bunny.garden>
Date: Wed, 20 Nov 2024 16:43:34 +0000
Subject: implement element writing

---
 src/element.rs |  27 ++++++++++
 src/reader.rs  |  20 ++++----
 src/writer.rs  | 160 +++++++++++++++++++++++++++++++++++++++++++++++++++++----
 src/xml/mod.rs |   2 +-
 4 files changed, 188 insertions(+), 21 deletions(-)

(limited to 'src')

diff --git a/src/element.rs b/src/element.rs
index 5b5f048..04f2e5e 100644
--- a/src/element.rs
+++ b/src/element.rs
@@ -54,6 +54,33 @@ pub struct Element {
     pub content: Vec<Content>,
 }
 
+pub fn escape_str(s: &str) -> String {
+    let mut string = String::new();
+    for str in s.split_inclusive(|c| c == '<' || c == '&' || c == '>') {
+        if let Some(str) = str.strip_suffix('<') {
+            if !str.is_empty() {
+                string.push_str(str)
+            }
+            string.push_str("&lt;");
+        } else if let Some(str) = str.strip_suffix('&') {
+            if !str.is_empty() {
+                string.push_str(str)
+            }
+            string.push_str("&amp;");
+        } else if let Some(str) = str.strip_suffix('>') {
+            if !str.is_empty() {
+                string.push_str(str)
+            }
+            string.push_str("&gt;");
+        } else {
+            if !str.is_empty() {
+                string.push_str(str)
+            }
+        }
+    }
+    string
+}
+
 // impl<'s> TryFrom<xml::Element<'s>> for Element<'s> {
 //     type Error = Error;
 
diff --git a/src/reader.rs b/src/reader.rs
index 8387373..f1f3744 100644
--- a/src/reader.rs
+++ b/src/reader.rs
@@ -48,7 +48,7 @@ where
         Ok(self.inner.read_buf(&mut self.buffer).await?)
     }
 
-    async fn read_prolog<'s>(&'s mut self) -> Result<()> {
+    pub async fn read_prolog<'s>(&'s mut self) -> Result<()> {
         loop {
             self.read_buf().await?;
             let input = str::from_utf8(self.buffer.data())?;
@@ -68,7 +68,7 @@ where
         }
     }
 
-    async fn read_start_tag<'s>(&'s mut self) -> Result<Element> {
+    pub async fn read_start_tag<'s>(&'s mut self) -> Result<Element> {
         loop {
             self.read_buf().await?;
             let input = str::from_utf8(self.buffer.data())?;
@@ -93,7 +93,7 @@ where
         }
     }
 
-    async fn read_end_tag<'s>(&'s mut self) -> Result<()> {
+    pub async fn read_end_tag<'s>(&'s mut self) -> Result<()> {
         loop {
             self.read_buf().await?;
             let input = str::from_utf8(self.buffer.data())?;
@@ -118,7 +118,7 @@ where
         }
     }
 
-    async fn read_element<'s>(&'s mut self) -> Result<Element> {
+    pub async fn read_element<'s>(&'s mut self) -> Result<Element> {
         loop {
             self.read_buf().await?;
             let input = str::from_utf8(self.buffer.data())?;
@@ -140,7 +140,7 @@ where
         }
     }
 
-    async fn read_content<'s>(&'s mut self) -> Result<Content> {
+    pub async fn read_content<'s>(&'s mut self) -> Result<Content> {
         let mut last_char = false;
         let mut text = String::new();
         loop {
@@ -674,19 +674,21 @@ impl<R: AsyncRead + Unpin> Stream for Reader<R> {
 }
 
 #[cfg(test)]
-mod test {
+pub(crate) mod test {
     use futures::{sink::Buffer, StreamExt};
     use tokio::io::AsyncRead;
 
+    use crate::element::Element;
+
     use super::Reader;
 
-    struct MockAsyncReader<'s> {
+    pub struct MockAsyncReader<'s> {
         put: bool,
         data: &'s str,
     }
 
     impl<'s> MockAsyncReader<'s> {
-        fn new(data: &'s str) -> Self {
+        pub fn new(data: &'s str) -> Self {
             Self { put: false, data }
         }
     }
@@ -705,7 +707,7 @@ mod test {
         }
     }
 
-    const TEST_DOC: &'static str = "<xs:schema
+    pub const TEST_DOC: &'static str = "<xs:schema
        xmlns:xs='http://www.w3.org/2001/XMLSchema'
        targetNamespace='http://etherx.jabber.org/streams'
        xmlns='http://etherx.jabber.org/streams'
diff --git a/src/writer.rs b/src/writer.rs
index 4881f57..dc5b48a 100644
--- a/src/writer.rs
+++ b/src/writer.rs
@@ -1,10 +1,11 @@
 use std::{collections::HashSet, str::FromStr};
 
+use async_recursion::async_recursion;
 use futures::Sink;
-use tokio::io::AsyncWrite;
+use tokio::io::{AsyncWrite, AsyncWriteExt};
 
 use crate::{
-    element::{Element, Name, NamespaceDeclaration},
+    element::{escape_str, Content, Element, Name, NamespaceDeclaration},
     error::Error,
     xml::{self, composers::Composer, parsers_complete::Parser, ETag},
     Result,
@@ -17,12 +18,32 @@ pub struct Writer<W> {
     namespace_declarations: Vec<HashSet<NamespaceDeclaration>>,
 }
 
-impl<W: AsyncWrite + Unpin> Writer<W> {
-    pub async fn write(&mut self, element: Element) -> Result<()> {
-        todo!()
+impl<W> Writer<W> {
+    pub fn new(writer: W) -> Self {
+        Self {
+            inner: writer,
+            depth: Vec::new(),
+            namespace_declarations: Vec::new(),
+        }
     }
+}
 
-    pub async fn write_start(&mut self, element: Element) -> Result<()> {
+impl<W: AsyncWrite + Unpin + Send> Writer<W> {
+    #[async_recursion]
+    pub async fn write_element(&mut self, element: &Element) -> Result<()> {
+        if element.content.is_empty() {
+            self.write_empty(element).await?;
+        } else {
+            self.write_start(element).await?;
+            for content in &element.content {
+                self.write_content(content).await?;
+            }
+            self.write_end().await?;
+        }
+        Ok(())
+    }
+
+    pub async fn write_empty(&mut self, element: &Element) -> Result<()> {
         let namespace_declarations_stack: Vec<_> = self
             .namespace_declarations
             .iter()
@@ -60,9 +81,97 @@ impl<W: AsyncWrite + Unpin> Writer<W> {
                 .prefix
                 .as_ref()
                 .map(|prefix| -> Result<_> {
-                    Ok(xml::NSAttName::PrefixedAttName(
-                        xml::PrefixedAttName::parse_full(&prefix)?,
-                    ))
+                    Ok(xml::NSAttName::PrefixedAttName(xml::PrefixedAttName(
+                        xml::NCName::parse_full(&prefix)?,
+                    )))
+                })
+                .unwrap_or(Ok(xml::NSAttName::DefaultAttName))?;
+            let value = xml::AttValue::from(namespace_declaration.namespace.as_str());
+            let xml_attribute = xml::Attribute::NamespaceDeclaration { ns_name, value };
+            attributes.push(xml_attribute);
+        }
+
+        for (name, value) in &element.attributes {
+            let prefix;
+            if let Some(namespace) = &name.namespace {
+                let name_namespace_declaration = namespace_declarations_stack
+                    .iter()
+                    .rfind(|namespace_declaration| namespace_declaration.namespace == *namespace)
+                    .ok_or(Error::UndeclaredNamespace(namespace.clone()))?;
+                prefix = name_namespace_declaration.prefix.as_ref();
+            } else {
+                prefix = None
+            }
+
+            let att_name;
+            if let Some(prefix) = &prefix {
+                att_name = xml::QName::PrefixedName(xml::PrefixedName {
+                    prefix: xml::Prefix::parse_full(prefix)?,
+                    local_part: xml::LocalPart::parse_full(&element.name.local_name)?,
+                })
+            } else {
+                att_name = xml::QName::UnprefixedName(xml::UnprefixedName::parse_full(
+                    &element.name.local_name,
+                )?)
+            }
+
+            let value = xml::AttValue::from(value.as_str());
+
+            let xml_attribute = xml::Attribute::Attribute {
+                name: att_name,
+                value,
+            };
+            attributes.push(xml_attribute);
+        }
+
+        let tag = xml::EmptyElemTag { name, attributes };
+
+        tag.write(&mut self.inner).await?;
+
+        Ok(())
+    }
+
+    pub async fn write_start(&mut self, element: &Element) -> Result<()> {
+        let namespace_declarations_stack: Vec<_> = self
+            .namespace_declarations
+            .iter()
+            .flatten()
+            .chain(&element.namespace_declarations)
+            .collect();
+
+        let prefix;
+        if let Some(namespace) = &element.name.namespace {
+            let name_namespace_declaration = namespace_declarations_stack
+                .iter()
+                .rfind(|namespace_declaration| namespace_declaration.namespace == *namespace)
+                .ok_or(Error::UndeclaredNamespace(namespace.clone()))?;
+            prefix = name_namespace_declaration.prefix.as_ref();
+        } else {
+            prefix = None
+        }
+
+        let name;
+        if let Some(prefix) = &prefix {
+            name = xml::QName::PrefixedName(xml::PrefixedName {
+                prefix: xml::Prefix::parse_full(prefix)?,
+                local_part: xml::LocalPart::parse_full(&element.name.local_name)?,
+            })
+        } else {
+            name = xml::QName::UnprefixedName(xml::UnprefixedName::parse_full(
+                &element.name.local_name,
+            )?)
+        }
+
+        let mut attributes = Vec::new();
+
+        for namespace_declaration in &element.namespace_declarations {
+            let ns_name = namespace_declaration
+                .prefix
+                .as_ref()
+                .map(|prefix| -> Result<_> {
+                    Ok(xml::NSAttName::PrefixedAttName(xml::PrefixedAttName(
+                        xml::NCName::parse_full(&prefix)?,
+                    )))
                 })
                 .unwrap_or(Ok(xml::NSAttName::DefaultAttName))?;
             let value = xml::AttValue::from(namespace_declaration.namespace.as_str());
@@ -107,9 +216,20 @@ impl<W: AsyncWrite + Unpin> Writer<W> {
 
         s_tag.write(&mut self.inner).await?;
 
-        self.depth.push(element.name);
+        self.depth.push(element.name.clone());
         self.namespace_declarations
-            .push(element.namespace_declarations);
+            .push(element.namespace_declarations.clone());
+        Ok(())
+    }
+
+    pub async fn write_content(&mut self, content: &Content) -> Result<()> {
+        match content {
+            Content::Element(element) => self.write_element(element).await?,
+            Content::Text(text) => self.inner.write_all(escape_str(text).as_bytes()).await?,
+            // TODO: comments and PI
+            Content::PI => {}
+            Content::Comment(_) => {}
+        }
         Ok(())
     }
 
@@ -181,3 +301,21 @@ impl<W: AsyncWrite, E: Into<Element>> Sink<E> for Writer<W> {
         todo!()
     }
 }
+
+#[cfg(test)]
+mod test {
+    use crate::{
+        reader::{test::*, Reader},
+        writer::Writer,
+    };
+
+    #[tokio::test]
+    async fn test_element_write() {
+        let mock = MockAsyncReader::new(TEST_DOC);
+        let mut reader = Reader::new(mock);
+        let element = reader.read_element().await.unwrap();
+        let stdout = tokio::io::stdout();
+        let mut writer = Writer::new(stdout);
+        writer.write_element(&element).await.unwrap();
+    }
+}
diff --git a/src/xml/mod.rs b/src/xml/mod.rs
index 8fb5419..3150df0 100644
--- a/src/xml/mod.rs
+++ b/src/xml/mod.rs
@@ -17,7 +17,7 @@ pub enum NSAttName<'s> {
 
 /// [2]   	PrefixedAttName	   ::=   	'xmlns:' NCName
 #[derive(Clone, Debug)]
-pub struct PrefixedAttName<'s>(NCName<'s>);
+pub struct PrefixedAttName<'s>(pub NCName<'s>);
 
 impl<'s> Deref for PrefixedAttName<'s> {
     type Target = NCName<'s>;
-- 
cgit