1 files changed, 54 insertions, 20 deletions

diff --git a/src/stanza/mod.rs b/src/stanza/mod.rs
index 16f3bdd..c29b1a2 100644
--- a/src/stanza/mod.rs
+++ b/src/stanza/mod.rs
@@ -9,12 +9,12 @@ use quick_xml::events::Event;
 use quick_xml::{Reader, Writer};
 use tokio::io::{AsyncBufRead, AsyncWrite};
 
-use crate::Result;
+use crate::JabberError;
 
-#[derive(Debug)]
+#[derive(Clone, Debug)]
 pub struct Element<'e> {
     pub event: Event<'e>,
-    pub content: Option<Vec<Element<'e>>>,
+    pub children: Option<Vec<Element<'e>>>,
 }
 
 impl<'e: 'async_recursion, 'async_recursion> Element<'e> {
@@ -23,7 +23,7 @@ impl<'e: 'async_recursion, 'async_recursion> Element<'e> {
         writer: &'life0 mut Writer<W>,
     ) -> ::core::pin::Pin<
         Box<
-            dyn ::core::future::Future<Output = Result<()>>
+            dyn ::core::future::Future<Output = Result<(), JabberError>>
                 + 'async_recursion
                 + ::core::marker::Send,
         >,
@@ -36,9 +36,9 @@ impl<'e: 'async_recursion, 'async_recursion> Element<'e> {
             match &self.event {
                 Event::Start(e) => {
                     writer.write_event_async(Event::Start(e.clone())).await?;
-                    if let Some(content) = &self.content {
-                        for _e in content {
-                            self.write(writer).await?;
+                    if let Some(children) = &self.children {
+                        for e in children {
+                            e.write(writer).await?;
                         }
                     }
                     writer.write_event_async(Event::End(e.to_end())).await?;
@@ -54,7 +54,7 @@ impl<'e> Element<'e> {
     pub async fn write_start<W: AsyncWrite + Unpin + Send>(
         &self,
         writer: &mut Writer<W>,
-    ) -> Result<()> {
+    ) -> Result<(), JabberError> {
         match self.event.as_ref() {
             Event::Start(e) => Ok(writer.write_event_async(Event::Start(e.clone())).await?),
             e => Err(ElementError::NotAStart(e.clone().into_owned()).into()),
@@ -64,7 +64,7 @@ impl<'e> Element<'e> {
     pub async fn write_end<W: AsyncWrite + Unpin + Send>(
         &self,
         writer: &mut Writer<W>,
-    ) -> Result<()> {
+    ) -> Result<(), JabberError> {
         match self.event.as_ref() {
             Event::Start(e) => Ok(writer
                 .write_event_async(Event::End(e.clone().to_end()))
@@ -76,28 +76,38 @@ impl<'e> Element<'e> {
     #[async_recursion]
     pub async fn read<R: AsyncBufRead + Unpin + Send>(
         reader: &mut Reader<R>,
-    ) -> Result<Option<Self>> {
+    ) -> Result<Self, JabberError> {
+        let element = Self::read_recursive(reader)
+            .await?
+            .ok_or(JabberError::UnexpectedEnd);
+        element
+    }
+
+    #[async_recursion]
+    async fn read_recursive<R: AsyncBufRead + Unpin + Send>(
+        reader: &mut Reader<R>,
+    ) -> Result<Option<Self>, JabberError> {
         let mut buf = Vec::new();
         let event = reader.read_event_into_async(&mut buf).await?;
         match event {
             Event::Start(e) => {
-                let mut content_vec = Vec::new();
-                while let Some(sub_element) = Element::read(reader).await? {
-                    content_vec.push(sub_element)
+                let mut children_vec = Vec::new();
+                while let Some(sub_element) = Element::read_recursive(reader).await? {
+                    children_vec.push(sub_element)
                 }
-                let mut content = None;
-                if !content_vec.is_empty() {
-                    content = Some(content_vec)
+                let mut children = None;
+                if !children_vec.is_empty() {
+                    children = Some(children_vec)
                 }
                 Ok(Some(Self {
                     event: Event::Start(e.into_owned()),
-                    content,
+                    children,
                 }))
             }
             Event::End(_) => Ok(None),
             e => Ok(Some(Self {
                 event: e.into_owned(),
-                content: None,
+                children: None,
             })),
         }
     }
@@ -105,14 +115,14 @@ impl<'e> Element<'e> {
     #[async_recursion]
     pub async fn read_start<R: AsyncBufRead + Unpin + Send>(
         reader: &mut Reader<R>,
-    ) -> Result<Self> {
+    ) -> Result<Self, JabberError> {
         let mut buf = Vec::new();
         let event = reader.read_event_into_async(&mut buf).await?;
         match event {
             Event::Start(e) => {
                 return Ok(Self {
                     event: Event::Start(e.into_owned()),
-                    content: None,
+                    children: None,
                 })
             }
             e => Err(ElementError::NotAStart(e.into_owned()).into()),
@@ -120,7 +130,31 @@ impl<'e> Element<'e> {
     }
 }
 
+/// if there is only one child in the vec of children, will return that element
+pub fn child<'p, 'e>(element: &'p Element<'e>) -> Result<&'p Element<'e>, ElementError<'static>> {
+    if let Some(children) = &element.children {
+        if children.len() == 1 {
+            return Ok(&children[0]);
+        } else {
+            return Err(ElementError::MultipleChildren);
+        }
+    }
+    Err(ElementError::NoChildren)
+}
+
+/// returns reference to children
+pub fn children<'p, 'e>(
+    element: &'p Element<'e>,
+) -> Result<&'p Vec<Element<'e>>, ElementError<'e>> {
+    if let Some(children) = &element.children {
+        return Ok(children);
+    }
+    Err(ElementError::NoChildren)
+}
+
 #[derive(Debug)]
 pub enum ElementError<'e> {
     NotAStart(Event<'e>),
+    NoChildren,
+    MultipleChildren,
 }