From 5db7ec74657b893ad6223e067ecc6071b2aed68b Mon Sep 17 00:00:00 2001
From: PoiScript <poiscript@gmail.com>
Date: Sun, 11 Aug 2019 11:40:52 +0800
Subject: [PATCH] feat(node): more headline operations

---
 src/elements/mod.rs         |  26 ++++--
 src/elements/title.rs       |  14 ++++
 src/export/html.rs          |   4 +-
 src/export/org.rs           |   4 +-
 src/node.rs                 | 160 ++++++++++++++++++++++++++++++++++--
 src/org.rs                  |  55 ++++++++++---
 src/parsers.rs              |  72 +++++++++++-----
 tests/node.rs               |  24 ++++++
 tests/{html.rs => parse.rs} |   0
 9 files changed, 311 insertions(+), 48 deletions(-)
 create mode 100644 tests/node.rs
 rename tests/{html.rs => parse.rs} (100%)

diff --git a/src/elements/mod.rs b/src/elements/mod.rs
index 2639237..a059107 100644
--- a/src/elements/mod.rs
+++ b/src/elements/mod.rs
@@ -79,7 +79,7 @@ pub enum Element<'a> {
     DynBlock(DynBlock<'a>),
     FnDef(FnDef<'a>),
     FnRef(FnRef<'a>),
-    Headline,
+    Headline { level: usize },
     InlineCall(InlineCall<'a>),
     InlineSrc(InlineSrc<'a>),
     Keyword(Keyword<'a>),
@@ -112,9 +112,25 @@ impl Element<'_> {
         use Element::*;
 
         match self {
-            SpecialBlock(_) | QuoteBlock(_) | CenterBlock(_) | VerseBlock(_) | Bold | Document
-            | DynBlock(_) | Headline | Italic | List(_) | ListItem(_) | Paragraph | Section
-            | Strike | Underline | Title(_) | Table(_) | TableRow(_) | TableCell => true,
+            SpecialBlock(_)
+            | QuoteBlock(_)
+            | CenterBlock(_)
+            | VerseBlock(_)
+            | Bold
+            | Document
+            | DynBlock(_)
+            | Headline { .. }
+            | Italic
+            | List(_)
+            | ListItem(_)
+            | Paragraph
+            | Section
+            | Strike
+            | Underline
+            | Title(_)
+            | Table(_)
+            | TableRow(_)
+            | TableCell => true,
             _ => false,
         }
     }
@@ -141,7 +157,7 @@ impl Element<'_> {
             DynBlock(e) => DynBlock(e.into_owned()),
             FnDef(e) => FnDef(e.into_owned()),
             FnRef(e) => FnRef(e.into_owned()),
-            Headline => Headline,
+            Headline { level } => Headline { level },
             InlineCall(e) => InlineCall(e.into_owned()),
             InlineSrc(e) => InlineSrc(e.into_owned()),
             Keyword(e) => Keyword(e.into_owned()),
diff --git a/src/elements/title.rs b/src/elements/title.rs
index d6a5231..9dec1a3 100644
--- a/src/elements/title.rs
+++ b/src/elements/title.rs
@@ -93,6 +93,20 @@ impl Title<'_> {
     }
 }
 
+impl Default for Title<'_> {
+    fn default() -> Title<'static> {
+        Title {
+            level: 1,
+            priority: None,
+            tags: Vec::new(),
+            keyword: None,
+            raw: Cow::Borrowed(""),
+            planning: None,
+            properties: HashMap::new(),
+        }
+    }
+}
+
 fn parse_headline<'a>(
     input: &'a str,
     config: &ParseConfig,
diff --git a/src/export/html.rs b/src/export/html.rs
index 60f654b..002db83 100644
--- a/src/export/html.rs
+++ b/src/export/html.rs
@@ -42,7 +42,7 @@ pub trait HtmlHandler<E: From<Error>> {
             Bold => write!(w, "<b>")?,
             Document => write!(w, "<main>")?,
             DynBlock(_dyn_block) => (),
-            Headline => (),
+            Headline { .. } => (),
             List(list) => {
                 if list.ordered {
                     write!(w, "<ol>")?;
@@ -166,7 +166,7 @@ pub trait HtmlHandler<E: From<Error>> {
             Bold => write!(w, "</b>")?,
             Document => write!(w, "</main>")?,
             DynBlock(_dyn_block) => (),
-            Headline => (),
+            Headline { .. } => (),
             List(list) => {
                 if list.ordered {
                     write!(w, "</ol>")?;
diff --git a/src/export/org.rs b/src/export/org.rs
index 7131df0..e20041c 100644
--- a/src/export/org.rs
+++ b/src/export/org.rs
@@ -21,7 +21,7 @@ pub trait OrgHandler<E: From<Error>> {
                 }
                 writeln!(&mut w)?;
             }
-            Headline => (),
+            Headline { .. } => (),
             List(_list) => (),
             Italic => write!(w, "/")?,
             ListItem(list_item) => write!(w, "{}", list_item.bullet)?,
@@ -154,7 +154,7 @@ pub trait OrgHandler<E: From<Error>> {
             Bold => write!(w, "*")?,
             Document => (),
             DynBlock(_dyn_block) => writeln!(w, "#+END:")?,
-            Headline => (),
+            Headline { .. } => (),
             List(_list) => (),
             Italic => write!(w, "/")?,
             ListItem(_) => (),
diff --git a/src/node.rs b/src/node.rs
index d11ce3c..55afd2e 100644
--- a/src/node.rs
+++ b/src/node.rs
@@ -1,27 +1,169 @@
 use indextree::NodeId;
+use std::borrow::Cow;
 
+use crate::config::ParseConfig;
 use crate::elements::{Element, Title};
+use crate::parsers::{parse_container, Container, OwnedArena};
 use crate::Org;
 
-#[derive(Copy, Clone)]
-pub struct HeadlineNode(pub(crate) NodeId);
+#[derive(Copy, Clone, Debug)]
+pub struct HeadlineNode {
+    pub(crate) node: NodeId,
+    pub(crate) level: usize,
+    pub(crate) title_node: NodeId,
+    pub(crate) section_node: Option<NodeId>,
+}
 
-impl HeadlineNode {
-    pub fn get_title<'a: 'b, 'b>(self, org: &'b Org<'a>) -> &'b Title<'a> {
-        let title_node = org.arena[self.0].first_child().unwrap();
-        if let Element::Title(title) = org.arena[title_node].get() {
+impl<'a: 'b, 'b> HeadlineNode {
+    pub(crate) fn new(node: NodeId, level: usize, org: &Org<'_>) -> HeadlineNode {
+        let title_node = org.arena[node].first_child().unwrap();
+        let section_node = if let Some(node) = org.arena[title_node].next_sibling() {
+            if let Element::Section = org.arena[node].get() {
+                Some(node)
+            } else {
+                None
+            }
+        } else {
+            None
+        };
+        HeadlineNode {
+            node,
+            level,
+            title_node,
+            section_node,
+        }
+    }
+
+    pub fn level(self) -> usize {
+        self.level
+    }
+
+    pub fn title(self, org: &'b Org<'a>) -> &'b Title<'a> {
+        if let Element::Title(title) = org.arena[self.title_node].get() {
             title
         } else {
             unreachable!()
         }
     }
 
-    pub fn get_title_mut<'a: 'b, 'b>(self, org: &'b mut Org<'a>) -> &'b mut Title<'a> {
-        let title_node = org.arena[self.0].first_child().unwrap();
-        if let Element::Title(title) = org.arena[title_node].get_mut() {
+    pub fn title_mut(self, org: &'b mut Org<'a>) -> &'b mut Title<'a> {
+        if let Element::Title(title) = org.arena[self.title_node].get_mut() {
             title
         } else {
             unreachable!()
         }
     }
+
+    pub fn set_title_content<S: Into<Cow<'a, str>>>(self, content: S, org: &mut Org<'a>) {
+        let content = content.into();
+
+        let children: Vec<_> = self.title_node.children(&org.arena).collect();
+        for child in children {
+            child.detach(&mut org.arena);
+        }
+
+        match &content {
+            Cow::Borrowed(content) => parse_container(
+                &mut org.arena,
+                Container::Inline {
+                    node: self.title_node,
+                    content,
+                },
+                &ParseConfig::default(),
+            ),
+            Cow::Owned(ref content) => parse_container(
+                &mut OwnedArena::new(&mut org.arena),
+                Container::Inline {
+                    node: self.title_node,
+                    content,
+                },
+                &ParseConfig::default(),
+            ),
+        }
+
+        self.title_mut(org).raw = content;
+    }
+
+    pub fn set_section_content<S: Into<Cow<'a, str>>>(self, content: S, org: &mut Org<'a>) {
+        let node = if let Some(node) = self.section_node {
+            let children: Vec<_> = node.children(&org.arena).collect();
+            for child in children {
+                child.detach(&mut org.arena);
+            }
+            node
+        } else {
+            let node = org.arena.new_node(Element::Section);
+            self.node.append(node, &mut org.arena);
+            node
+        };
+
+        match content.into() {
+            Cow::Borrowed(content) => parse_container(
+                &mut org.arena,
+                Container::Block { node, content },
+                &ParseConfig::default(),
+            ),
+            Cow::Owned(ref content) => parse_container(
+                &mut OwnedArena::new(&mut org.arena),
+                Container::Block { node, content },
+                &ParseConfig::default(),
+            ),
+        }
+    }
+
+    pub fn parent(self, org: &Org<'_>) -> Option<HeadlineNode> {
+        org.arena[self.node].parent().map(|node| {
+            if let &Element::Headline { level } = org.arena[node].get() {
+                HeadlineNode::new(node, level, org)
+            } else {
+                unreachable!()
+            }
+        })
+    }
+
+    pub fn detach(self, org: &mut Org<'_>) {
+        self.node.detach(&mut org.arena);
+    }
+
+    pub fn is_detached(self, org: &Org<'_>) -> bool {
+        self.parent(&org).is_none()
+    }
+
+    pub fn append(self, headline: &HeadlineNode, org: &mut Org<'_>) {
+        if self.is_detached(org) || headline.level <= self.level {
+            // TODO: return an error
+            return;
+        } else {
+            self.node.append(headline.node, &mut org.arena);
+        }
+    }
+
+    pub fn prepend(self, headline: &HeadlineNode, org: &mut Org<'_>) {
+        if self.is_detached(org) || headline.level <= self.level {
+            // TODO: return an error
+            return;
+        } else if let Some(node) = self.section_node {
+            node.insert_after(headline.node, &mut org.arena);
+        } else {
+            self.title_node.insert_after(headline.node, &mut org.arena);
+        }
+    }
+
+    pub fn insert_before(self, headline: &HeadlineNode, org: &mut Org<'_>) {
+        if self.is_detached(org) || headline.level < self.level {
+            // TODO: return an error
+            return;
+        } else {
+            self.node.insert_after(headline.node, &mut org.arena);
+        }
+    }
+
+    pub fn insert_after(self, headline: &HeadlineNode, org: &mut Org<'_>) {
+        if self.is_detached(org) || headline.level < self.level {
+            // TODO: return an error
+            return;
+        } else {
+            self.node.insert_after(headline.node, &mut org.arena);
+        }
+    }
 }
diff --git a/src/org.rs b/src/org.rs
index 471b39c..6d7841d 100644
--- a/src/org.rs
+++ b/src/org.rs
@@ -2,7 +2,7 @@ use indextree::{Arena, NodeEdge, NodeId};
 use std::io::{Error, Write};
 
 use crate::config::ParseConfig;
-use crate::elements::Element;
+use crate::elements::{Element, Title};
 use crate::export::*;
 use crate::node::HeadlineNode;
 use crate::parsers::{parse_container, Container};
@@ -19,17 +19,15 @@ pub enum Event<'a> {
 }
 
 impl Org<'_> {
-    pub fn parse(text: &str) -> Org<'_> {
-        Org::parse_with_config(text, &ParseConfig::default())
+    pub fn new() -> Org<'static> {
+        let mut arena = Arena::new();
+        let root = arena.new_node(Element::Document);
+
+        Org { arena, root }
     }
 
-    pub fn parse_with_config<'a>(content: &'a str, config: &ParseConfig) -> Org<'a> {
-        let mut arena = Arena::new();
-        let node = arena.new_node(Element::Document);
-
-        parse_container(&mut arena, Container::Document { content, node }, config);
-
-        Org { arena, root: node }
+    pub fn parse(text: &str) -> Org<'_> {
+        Org::parse_with_config(text, &ParseConfig::default())
     }
 
     pub fn iter(&self) -> impl Iterator<Item = Event<'_>> + '_ {
@@ -44,7 +42,7 @@ impl Org<'_> {
             .descendants(&self.arena)
             .skip(1)
             .filter_map(move |node| match self.arena[node].get() {
-                Element::Headline => Some(HeadlineNode(node)),
+                &Element::Headline { level } => Some(HeadlineNode::new(node, level, self)),
                 _ => None,
             })
     }
@@ -90,6 +88,41 @@ impl Org<'_> {
     }
 }
 
+impl<'a> Org<'a> {
+    pub fn parse_with_config(content: &'a str, config: &ParseConfig) -> Org<'a> {
+        let mut org = Org::new();
+
+        parse_container(
+            &mut org.arena,
+            Container::Document {
+                content,
+                node: org.root,
+            },
+            config,
+        );
+
+        org
+    }
+
+    pub fn new_headline(&mut self, title: Title<'a>) -> HeadlineNode {
+        let title_level = title.level;
+        let title_raw = title.raw.clone();
+        let headline_node = self
+            .arena
+            .new_node(Element::Headline { level: title_level });
+        let title_node = self.arena.new_node(Element::Title(title));
+        headline_node.append(title_node, &mut self.arena);
+        let headline_node = HeadlineNode {
+            node: headline_node,
+            level: title_level,
+            title_node,
+            section_node: None,
+        };
+        headline_node.set_title_content(title_raw, self);
+        headline_node
+    }
+}
+
 #[cfg(feature = "ser")]
 use serde::{ser::Serializer, Serialize};
 
diff --git a/src/parsers.rs b/src/parsers.rs
index 287321b..b34991a 100644
--- a/src/parsers.rs
+++ b/src/parsers.rs
@@ -9,7 +9,7 @@ use memchr::{memchr, memchr_iter};
 use nom::{
     bytes::complete::take_while1,
     character::complete::{line_ending, not_line_ending},
-    combinator::{map, opt, recognize, verify},
+    combinator::{opt, recognize, verify},
     error::ErrorKind,
     error_position,
     multi::{many0_count, many1_count},
@@ -51,6 +51,42 @@ impl<'a> ElementArena<'a> for Arena<Element<'a>> {
     }
 }
 
+pub struct OwnedArena<'a, 'b, 'c> {
+    arena: &'b mut Arena<Element<'c>>,
+    phantom: PhantomData<&'a ()>,
+}
+
+impl<'a, 'b, 'c> OwnedArena<'a, 'b, 'c> {
+    pub fn new(arena: &'b mut Arena<Element<'c>>) -> OwnedArena<'a, 'b, 'c> {
+        OwnedArena {
+            arena,
+            phantom: PhantomData,
+        }
+    }
+}
+
+impl<'a> ElementArena<'a> for OwnedArena<'a, '_, '_> {
+    fn push_element<T: Into<Element<'a>>>(&mut self, element: T, parent: NodeId) -> NodeId {
+        let node = self.arena.new_node(element.into().into_owned());
+        parent.append(node, self.arena);
+        node
+    }
+
+    fn insert_before_last_child<T: Into<Element<'a>>>(
+        &mut self,
+        element: T,
+        parent: NodeId,
+    ) -> NodeId {
+        if let Some(child) = self.arena[parent].last_child() {
+            let node = self.arena.new_node(element.into().into_owned());
+            child.insert_before(node, self.arena);
+            node
+        } else {
+            self.push_element(element, parent)
+        }
+    }
+}
+
 #[derive(Debug)]
 pub enum Container<'a> {
     // List
@@ -139,22 +175,22 @@ pub fn parse_section_and_headlines<'a, T: ElementArena<'a>>(
 
     let mut last_end = 0;
     for i in memchr_iter(b'\n', content.as_bytes()) {
-        if let Ok((mut tail, headline_content)) = parse_headline(&content[last_end..]) {
+        if let Ok((mut tail, (headline_content, level))) = parse_headline(&content[last_end..]) {
             if last_end != 0 {
                 let node = arena.push_element(Element::Section, parent);
                 let content = &content[0..last_end];
                 containers.push(Container::Block { content, node });
             }
 
-            let node = arena.push_element(Element::Headline, parent);
+            let node = arena.push_element(Element::Headline { level }, parent);
             containers.push(Container::Headline {
                 content: headline_content,
                 node,
             });
 
-            while let Ok((new_tail, content)) = parse_headline(tail) {
+            while let Ok((new_tail, (content, level))) = parse_headline(tail) {
                 debug_assert_ne!(tail, new_tail);
-                let node = arena.push_element(Element::Headline, parent);
+                let node = arena.push_element(Element::Headline { level }, parent);
                 containers.push(Container::Headline { content, node });
                 tail = new_tail;
             }
@@ -719,24 +755,22 @@ pub fn skip_empty_lines(input: &str) -> &str {
         .unwrap_or(input)
 }
 
-pub fn parse_headline(input: &str) -> IResult<&str, &str> {
-    let (input_, level) = get_headline_level(input)?;
-    map(
-        take_lines_while(move |line| {
-            if let Ok((_, l)) = get_headline_level(line) {
-                l.len() > level.len()
-            } else {
-                true
-            }
-        }),
-        move |s: &str| &input[0..level.len() + s.len()],
-    )(input_)
+pub fn parse_headline(input: &str) -> IResult<&str, (&str, usize)> {
+    let (input_, level) = parse_headline_level(input)?;
+    let (input_, content) = take_lines_while(move |line| {
+        if let Ok((_, l)) = parse_headline_level(line) {
+            l > level
+        } else {
+            true
+        }
+    })(input_)?;
+    Ok((input_, (&input[0..level + content.len()], level)))
 }
 
-pub fn get_headline_level(input: &str) -> IResult<&str, &str> {
+pub fn parse_headline_level(input: &str) -> IResult<&str, usize> {
     let (input, stars) = take_while1(|c: char| c == '*')(input)?;
     if input.is_empty() || input.starts_with(' ') || input.starts_with('\n') {
-        Ok((input, stars))
+        Ok((input, stars.len()))
     } else {
         Err(Err::Error(error_position!(input, ErrorKind::Tag)))
     }
diff --git a/tests/node.rs b/tests/node.rs
new file mode 100644
index 0000000..20c7d2f
--- /dev/null
+++ b/tests/node.rs
@@ -0,0 +1,24 @@
+use orgize::Org;
+use pretty_assertions::assert_eq;
+
+#[test]
+fn set_content() {
+    let mut org = Org::parse(
+        r#"* title 1
+section 1
+** title 2
+"#,
+    );
+    let headlines: Vec<_> = org.headlines().collect();
+    for headline in headlines {
+        headline.set_title_content(String::from("a *bold* title"), &mut org);
+        headline.set_section_content("and a _underline_ section", &mut org);
+    }
+    let mut writer = Vec::new();
+    org.html(&mut writer).unwrap();
+    assert_eq!(
+        String::from_utf8(writer).unwrap(),
+        "<main><h1>a <b>bold</b> title</h1><section><p>and a <u>underline</u> section</p></section>\
+         <h2>a <b>bold</b> title</h2><section><p>and a <u>underline</u> section</p></section></main>"
+    );
+}
diff --git a/tests/html.rs b/tests/parse.rs
similarity index 100%
rename from tests/html.rs
rename to tests/parse.rs