use alloc::boxed::Box;
use alloc::rc::Rc;
use alloc::string::String;
use alloc::vec::Vec;
use crate::wtf8;
pub const MAX_STRING_LEN: usize = crate::limits::DEFAULT_MAX_STRING_LEN;
#[derive(Clone)]
pub struct Rope(Rc<Node>);
enum Node {
Leaf(Box<[u8]>),
Concat { left: Rope, right: Rope, len: usize },
}
impl Rope {
#[must_use]
pub fn new() -> Self {
Self::leaf("")
}
#[must_use]
pub fn leaf(s: &str) -> Self {
Rope(Rc::new(Node::Leaf(Box::from(s.as_bytes()))))
}
#[must_use]
pub fn from_wtf8(bytes: Vec<u8>) -> Self {
Rope(Rc::new(Node::Leaf(bytes.into_boxed_slice())))
}
#[must_use]
pub fn from_bytes(bytes: &[u8]) -> Self {
Rope(Rc::new(Node::Leaf(Box::from(bytes))))
}
#[must_use]
pub fn len(&self) -> usize {
match &*self.0 {
Node::Leaf(s) => s.len(),
Node::Concat { len, .. } => *len,
}
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.len() == 0
}
#[must_use]
pub fn try_concat(&self, other: &Rope) -> Option<Rope> {
if self.is_empty() {
return Some(other.clone());
}
if other.is_empty() {
return Some(self.clone());
}
let len = self.len().checked_add(other.len())?;
if len > MAX_STRING_LEN {
return None;
}
Some(Rope(Rc::new(Node::Concat {
left: self.clone(),
right: other.clone(),
len,
})))
}
#[must_use]
pub fn concat(&self, other: &Rope) -> Rope {
self.try_concat(other).unwrap_or_else(|| self.clone())
}
#[must_use]
pub fn push_str(&self, s: &str) -> Rope {
self.concat(&Rope::leaf(s))
}
#[must_use]
pub fn as_leaf_bytes(&self) -> Option<&[u8]> {
match &*self.0 {
Node::Leaf(s) => Some(s),
Node::Concat { .. } => None,
}
}
#[must_use]
pub fn materialize_bytes(&self) -> Vec<u8> {
let mut out = Vec::with_capacity(self.len());
let mut stack: Vec<&Rope> = alloc::vec![self];
while let Some(node) = stack.pop() {
match &*node.0 {
Node::Leaf(s) => out.extend_from_slice(s),
Node::Concat { left, right, .. } => {
stack.push(right);
stack.push(left);
}
}
}
out
}
#[must_use]
pub fn materialize(&self) -> String {
wtf8::to_string_lossy(&self.materialize_bytes())
}
}
impl Drop for Node {
fn drop(&mut self) {
let sentinel = || Rc::new(Node::Leaf(Box::from(&[][..])));
let Node::Concat { left, right, .. } = self else {
return; };
let mut stack = alloc::vec![
core::mem::replace(&mut left.0, sentinel()),
core::mem::replace(&mut right.0, sentinel()),
];
while let Some(rc) = stack.pop() {
if let Ok(mut node) = Rc::try_unwrap(rc)
&& let Node::Concat { left, right, .. } = &mut node
{
stack.push(core::mem::replace(&mut left.0, sentinel()));
stack.push(core::mem::replace(&mut right.0, sentinel()));
}
}
}
}
impl Default for Rope {
fn default() -> Self {
Self::new()
}
}
impl From<&str> for Rope {
fn from(s: &str) -> Self {
Rope::leaf(s)
}
}
impl core::fmt::Display for Rope {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
let mut stack: Vec<&Rope> = alloc::vec![self];
while let Some(node) = stack.pop() {
match &*node.0 {
Node::Leaf(s) => match wtf8::as_str(s) {
Some(valid) => f.write_str(valid)?,
None => f.write_str(&wtf8::to_string_lossy(s))?,
},
Node::Concat { left, right, .. } => {
stack.push(right);
stack.push(left);
}
}
}
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
use alloc::string::ToString;
#[test]
fn leaf_basics() {
let r = Rope::leaf("hello");
assert_eq!(r.len(), 5);
assert!(!r.is_empty());
assert_eq!(r.materialize(), "hello");
assert!(Rope::new().is_empty());
assert_eq!(Rope::new().materialize(), "");
}
#[test]
fn concat_is_o1_and_flattens() {
let a = Rope::leaf("foo");
let b = Rope::leaf("bar");
let ab = a.concat(&b);
assert_eq!(ab.len(), 6);
assert_eq!(ab.materialize(), "foobar");
assert_eq!(a.materialize(), "foo");
assert_eq!(b.materialize(), "bar");
}
#[test]
fn concat_with_empty_shares_the_other_side() {
let a = Rope::leaf("x");
let empty = Rope::new();
assert_eq!(a.concat(&empty).materialize(), "x");
assert_eq!(empty.concat(&a).materialize(), "x");
}
#[test]
fn repeated_append_builds_correctly() {
let mut r = Rope::new();
let mut expected = String::new();
for i in 0..50 {
let part = i.to_string();
r = r.push_str(&part);
expected.push_str(&part);
}
assert_eq!(r.len(), expected.len());
assert_eq!(r.materialize(), expected);
}
#[test]
fn deeply_nested_rope_flattens_without_recursion() {
let mut r = Rope::leaf("a");
for _ in 0..10_000 {
r = r.push_str("b");
}
let s = r.materialize();
assert_eq!(s.len(), 10_001);
assert!(s.starts_with("ab"));
assert!(s.ends_with("bb"));
}
#[cfg(feature = "std")]
#[test]
fn deeply_nested_rope_drops_without_overflow() {
std::thread::Builder::new()
.stack_size(256 * 1024)
.spawn(|| {
let mut r = Rope::leaf("a");
for _ in 0..100_000 {
r = r.push_str("b");
}
assert_eq!(r.materialize().len(), 100_001);
})
.expect("spawn")
.join()
.expect("the deep rope dropped without overflowing the stack");
}
#[test]
fn display_matches_to_string() {
let r = Rope::leaf("a").push_str("b").concat(&Rope::leaf("cd"));
assert_eq!(alloc::format!("{r}"), r.materialize());
assert_eq!(r.materialize(), "abcd");
}
#[test]
fn non_surrogate_string_is_byte_identical() {
for s in ["", "hello", "héllo 中 😀 mix"] {
let r = Rope::from(s);
assert_eq!(r.materialize_bytes(), s.as_bytes());
assert_eq!(r.materialize(), s);
assert_eq!(r.len(), s.len());
}
}
#[test]
fn lone_surrogate_round_trips_through_rope() {
let bytes = crate::wtf8::from_utf16(&[0x61, 0xD800, 0x62]);
let r = Rope::from_wtf8(bytes.clone());
assert_eq!(r.materialize_bytes(), bytes);
assert_eq!(r.materialize(), "a\u{FFFD}b");
assert_eq!(alloc::format!("{r}"), "a\u{FFFD}b");
assert_eq!(r.len(), bytes.len()); }
#[test]
fn concat_preserves_surrogate_bytes() {
let left = Rope::from("x");
let right = Rope::from_wtf8(crate::wtf8::from_utf16(&[0xD800]));
let joined = left.concat(&right);
let mut expected = alloc::vec![b'x'];
expected.extend_from_slice(&crate::wtf8::from_utf16(&[0xD800]));
assert_eq!(joined.materialize_bytes(), expected);
}
#[test]
fn from_bytes_matches_from_wtf8() {
let bytes = crate::wtf8::from_utf16(&[0xDC00, 0x41]);
assert_eq!(
Rope::from_bytes(&bytes).materialize_bytes(),
Rope::from_wtf8(bytes.clone()).materialize_bytes()
);
}
#[test]
fn as_leaf_bytes_borrows_unconcatenated_leaf() {
let leaf = Rope::leaf("hello");
assert_eq!(leaf.as_leaf_bytes(), Some(&b"hello"[..]));
let bytes = crate::wtf8::from_utf16(&[0xD800]);
let surr = Rope::from_wtf8(bytes.clone());
assert_eq!(surr.as_leaf_bytes(), Some(&bytes[..]));
let joined = Rope::leaf("foo").concat(&Rope::leaf("bar"));
assert_eq!(joined.as_leaf_bytes(), None);
assert_eq!(joined.materialize_bytes(), b"foobar");
}
#[test]
fn clone_is_cheap_and_shares() {
let a = Rope::leaf("shared");
let b = a.clone();
assert!(Rc::ptr_eq(&a.0, &b.0));
assert_eq!(b.materialize(), "shared");
}
}