From b3413522e8a1fb71f74f6eca68e929897d0809dc Mon Sep 17 00:00:00 2001
From: kmaasrud <km@aasrud.com>
Date: Wed, 15 Mar 2023 09:44:02 +0100
Subject: [PATCH] feat: add custom CowStr type

Related issue: #20
---
 src/attr.rs   |   1 +
 src/lib.rs    |   4 +-
 src/string.rs | 136 ++++++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 140 insertions(+), 1 deletion(-)
 create mode 100644 src/string.rs

diff --git a/src/attr.rs b/src/attr.rs
index 12a0b595..9836c677 100644
--- a/src/attr.rs
+++ b/src/attr.rs
@@ -57,6 +57,7 @@ impl<'s> AttributeValue<'s> {
                 prev.push(' ');
                 prev.push_str(s);
             }
+            CowStr::Inlined(_, _) => todo!(),
         }
     }
 }
diff --git a/src/lib.rs b/src/lib.rs
index aa7623ab..03db2acb 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -60,13 +60,15 @@ mod block;
 mod inline;
 mod lex;
 mod span;
+mod string;
 mod tree;
 
 use span::Span;
 
 pub use attr::{AttributeValue, AttributeValueParts, Attributes};
+pub use string::CowStr;
 
-type CowStr<'s> = std::borrow::Cow<'s, str>;
+// type CowStr<'s> = std::borrow::Cow<'s, str>;
 
 /// A trait for rendering [`Event`]s to an output format.
 ///
diff --git a/src/string.rs b/src/string.rs
new file mode 100644
index 00000000..3faf5bcb
--- /dev/null
+++ b/src/string.rs
@@ -0,0 +1,136 @@
+use std::{borrow::Borrow, cmp::Ordering, fmt::Display, ops::Deref, str::from_utf8};
+
+// Largest CowStr variant is Owned(String). A String uses 3 words of memory, but a fourth word is
+// needed to hold the tag (the tag takes a byte, but a full word is used for alignment reasons.)
+// This means that the available space we have for an inline string is 4 words - 2 bytes for the
+// tag and length.
+const MAX_INLINE_STR_LEN: usize = 4 * std::mem::size_of::<usize>() - 2;
+
+#[derive(Debug)]
+pub enum CowStr<'s> {
+    Owned(String),
+    Borrowed(&'s str),
+    Inlined([u8; MAX_INLINE_STR_LEN], u8),
+}
+
+impl<'s> CowStr<'s> {
+    pub fn take_borrowed(self) -> &'s str {
+        match self {
+            CowStr::Borrowed(s) => s,
+            CowStr::Owned(_) | CowStr::Inlined(..) => panic!(),
+        }
+    }
+
+    pub fn replace(self, from: &str, to: &str) -> Self {
+        if from.is_empty() {
+            return self;
+        }
+
+        match self {
+            CowStr::Owned(s) => CowStr::Owned(s.replace(from, to)),
+            CowStr::Inlined(mut inner, mut len) => {
+                for (start, _) in self.deref().match_indices(from) {
+                    match from.len().cmp(&to.len()) {
+                        Ordering::Less => {
+                            len += (to.len() - from.len()) as u8;
+                            if len as usize > MAX_INLINE_STR_LEN {
+                                return CowStr::Owned(self.deref().replace(from, to));
+                            }
+                            inner[start + from.len()..].rotate_right(to.len() - from.len());
+                        }
+                        Ordering::Greater => {
+                            inner[start..].rotate_left(from.len() - to.len());
+                            len -= (from.len() - to.len()) as u8;
+                        }
+                        Ordering::Equal => {}
+                    }
+
+                    inner[start..start + to.len()].copy_from_slice(to.as_bytes());
+                }
+
+                CowStr::Inlined(inner, len)
+            }
+            CowStr::Borrowed(s) => CowStr::Owned(s.replace(from, to)),
+        }
+    }
+}
+
+impl<'s> Deref for CowStr<'s> {
+    type Target = str;
+
+    fn deref(&self) -> &Self::Target {
+        match *self {
+            Self::Owned(ref s) => s.borrow(),
+            Self::Borrowed(s) => s,
+            // NOTE: Inlined strings can only be constructed from strings or chars, which means they
+            // are guaranteed to be valid UTF-8. We could consider unchecked conversion as well, but
+            // a benchmark should be done before introducing unsafes.
+            Self::Inlined(ref inner, len) => from_utf8(&inner[..len as usize]).unwrap(),
+        }
+    }
+}
+
+impl<'s> AsRef<str> for CowStr<'s> {
+    fn as_ref(&self) -> &str {
+        self.deref()
+    }
+}
+
+impl<'s> From<char> for CowStr<'s> {
+    fn from(value: char) -> Self {
+        let mut inner = [0u8; MAX_INLINE_STR_LEN];
+        value.encode_utf8(&mut inner);
+        CowStr::Inlined(inner, value.len_utf8() as u8)
+    }
+}
+
+impl<'s> From<&'s str> for CowStr<'s> {
+    fn from(value: &'s str) -> Self {
+        CowStr::Borrowed(value)
+    }
+}
+
+impl<'s> From<String> for CowStr<'s> {
+    fn from(value: String) -> Self {
+        CowStr::Owned(value)
+    }
+}
+
+impl<'s> Clone for CowStr<'s> {
+    fn clone(&self) -> Self {
+        match self {
+            CowStr::Owned(s) => {
+                let len = s.len();
+                if len > MAX_INLINE_STR_LEN {
+                    CowStr::Owned(s.clone())
+                } else {
+                    let mut inner = [0u8; MAX_INLINE_STR_LEN];
+                    inner[..len].copy_from_slice(s.as_bytes());
+                    CowStr::Inlined(inner, len as u8)
+                }
+            }
+            CowStr::Borrowed(s) => CowStr::Borrowed(s),
+            CowStr::Inlined(inner, len) => CowStr::Inlined(*inner, *len),
+        }
+    }
+}
+
+impl<'s> PartialEq for CowStr<'s> {
+    fn eq(&self, other: &Self) -> bool {
+        self.deref() == other.deref()
+    }
+}
+
+impl<'s> Eq for CowStr<'s> {}
+
+impl<'s> Display for CowStr<'s> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.write_str(self.deref())
+    }
+}
+
+impl<'s, 'a> FromIterator<&'a str> for CowStr<'s> {
+    fn from_iter<T: IntoIterator<Item = &'a str>>(iter: T) -> Self {
+        CowStr::Owned(FromIterator::from_iter(iter))
+    }
+}