AK: Allocate string data in a dedicated heap partition

Add a string heap partition and route long AK string backing
allocations through it.

Give StringBuilder partition-aware outlined storage so adopted String
and Utf16String buffers are allocated from the string heap. Keep the
string heap thread-local because mimalloc heaps may only allocate from
their creating thread, while cross-thread frees are handled by mimalloc.

AK/ByteStringImpl.cpp

@@ -47,7 +47,7 @@ ByteStringImpl::~ByteStringImpl() = default;
 NonnullRefPtr<ByteStringImpl const> ByteStringImpl::create_uninitialized(size_t length, char*& buffer)
 {
     VERIFY(length);
-    void* slot = kmalloc(allocation_size_for_stringimpl(length));
+    void* slot = kmalloc(HeapPartition::String, allocation_size_for_stringimpl(length));
     VERIFY(slot);
     auto new_stringimpl = adopt_ref(*new (slot) ByteStringImpl(ConstructWithInlineBuffer, length));
     buffer = const_cast<char*>(new_stringimpl->characters());

AK/StringBuilder.cpp

@@ -17,11 +17,241 @@
 #include <AK/Utf16String.h>
 #include <AK/Utf16StringData.h>
 #include <AK/Utf16View.h>
+#include <AK/kmalloc.h>
 
 #include <simdutf.h>
 
 namespace AK {
 
+StringBuilder::Buffer::Buffer(Buffer const& other)
+{
+    MUST(try_resize(other.size()));
+    if (other.size() != 0)
+        __builtin_memcpy(data(), other.data(), other.size());
+}
+
+StringBuilder::Buffer::Buffer(Buffer&& other)
+{
+    move_from(move(other));
+}
+
+StringBuilder::Buffer::~Buffer()
+{
+    clear();
+}
+
+auto StringBuilder::Buffer::operator=(Buffer const& other) -> Buffer&
+{
+    if (this != &other) {
+        if (m_size > other.size())
+            trim(other.size(), true);
+        else
+            MUST(try_resize(other.size()));
+
+        if (other.size() != 0)
+            __builtin_memcpy(data(), other.data(), other.size());
+    }
+
+    return *this;
+}
+
+auto StringBuilder::Buffer::operator=(Buffer&& other) -> Buffer&
+{
+    if (this != &other) {
+        clear();
+        move_from(move(other));
+    }
+
+    return *this;
+}
+
+u8* StringBuilder::Buffer::data()
+{
+    return m_inline ? m_inline_buffer : m_outline_buffer;
+}
+
+u8 const* StringBuilder::Buffer::data() const
+{
+    return m_inline ? m_inline_buffer : m_outline_buffer;
+}
+
+Bytes StringBuilder::Buffer::span()
+{
+    return { data(), size() };
+}
+
+ReadonlyBytes StringBuilder::Buffer::span() const
+{
+    return { data(), size() };
+}
+
+void* StringBuilder::Buffer::end_pointer()
+{
+    return data() + m_size;
+}
+
+void StringBuilder::Buffer::clear()
+{
+    if (!m_inline) {
+        kfree(m_outline_buffer);
+        m_inline = true;
+    }
+    m_size = 0;
+}
+
+void StringBuilder::Buffer::resize(size_t new_size)
+{
+    MUST(try_resize(new_size));
+}
+
+void StringBuilder::Buffer::set_size(size_t new_size)
+{
+    ASSERT(new_size <= capacity());
+    m_size = new_size;
+}
+
+void StringBuilder::Buffer::ensure_capacity(size_t new_capacity)
+{
+    MUST(try_ensure_capacity(new_capacity));
+}
+
+ErrorOr<void> StringBuilder::Buffer::try_resize(size_t new_size)
+{
+    if (new_size <= m_size) {
+        trim(new_size, false);
+        return {};
+    }
+
+    TRY(try_ensure_capacity(new_size));
+    set_size(new_size);
+    return {};
+}
+
+ErrorOr<void> StringBuilder::Buffer::try_ensure_capacity(size_t new_capacity)
+{
+    if (new_capacity <= capacity())
+        return {};
+    return try_ensure_capacity_slowpath(new_capacity);
+}
+
+ErrorOr<void> StringBuilder::Buffer::try_append(char byte)
+{
+    auto old_size = size();
+    Checked<size_t> new_size = old_size;
+    new_size += 1;
+    VERIFY(!new_size.has_overflow());
+
+    TRY(try_resize(new_size.value()));
+    data()[old_size] = static_cast<u8>(byte);
+    return {};
+}
+
+ErrorOr<void> StringBuilder::Buffer::try_append(ReadonlyBytes bytes)
+{
+    return try_append(bytes.data(), bytes.size());
+}
+
+ErrorOr<void> StringBuilder::Buffer::try_append(void const* data, size_t data_size)
+{
+    if (data_size == 0)
+        return {};
+    VERIFY(data != nullptr);
+
+    auto old_size = size();
+    Checked<size_t> new_size = old_size;
+    new_size += data_size;
+    VERIFY(!new_size.has_overflow());
+
+    TRY(try_resize(new_size.value()));
+    __builtin_memcpy(this->data() + old_size, data, data_size);
+    return {};
+}
+
+void StringBuilder::Buffer::append(char byte)
+{
+    MUST(try_append(byte));
+}
+
+void StringBuilder::Buffer::append(void const* data, size_t data_size)
+{
+    MUST(try_append(data, data_size));
+}
+
+auto StringBuilder::Buffer::leak_outline_buffer() -> Optional<OutlineBuffer>
+{
+    if (m_inline)
+        return {};
+
+    auto* outline_buffer = m_outline_buffer;
+    auto size = m_size;
+    auto outline_capacity = m_outline_capacity;
+
+    m_inline = true;
+    m_size = 0;
+
+    return OutlineBuffer { Bytes { outline_buffer, size }, outline_capacity };
+}
+
+void StringBuilder::Buffer::move_from(Buffer&& other)
+{
+    m_size = other.m_size;
+    m_inline = other.m_inline;
+
+    if (other.m_inline) {
+        VERIFY(other.m_size <= inline_capacity);
+        if (other.m_size != 0)
+            __builtin_memcpy(m_inline_buffer, other.m_inline_buffer, other.m_size);
+    } else {
+        m_outline_buffer = other.m_outline_buffer;
+        m_outline_capacity = other.m_outline_capacity;
+    }
+
+    other.m_size = 0;
+    other.m_inline = true;
+}
+
+void StringBuilder::Buffer::trim(size_t size, bool may_discard_existing_data)
+{
+    VERIFY(size <= m_size);
+    if (!m_inline && size <= inline_capacity)
+        shrink_into_inline_buffer(size, may_discard_existing_data);
+    m_size = size;
+}
+
+void StringBuilder::Buffer::shrink_into_inline_buffer(size_t size, bool may_discard_existing_data)
+{
+    auto* outline_buffer = m_outline_buffer;
+    if (!may_discard_existing_data)
+        __builtin_memcpy(m_inline_buffer, outline_buffer, size);
+    kfree(outline_buffer);
+    m_inline = true;
+}
+
+ErrorOr<void> StringBuilder::Buffer::try_ensure_capacity_slowpath(size_t new_capacity)
+{
+    new_capacity = max(new_capacity, (capacity() * 3) / 2);
+    new_capacity = kmalloc_good_size(new_capacity);
+
+    if (m_inline) {
+        auto* new_buffer = static_cast<u8*>(kmalloc(HeapPartition::String, new_capacity));
+        if (!new_buffer)
+            return Error::from_errno(ENOMEM);
+
+        __builtin_memcpy(new_buffer, data(), m_size);
+        m_outline_buffer = new_buffer;
+    } else {
+        auto* new_buffer = static_cast<u8*>(krealloc(HeapPartition::String, m_outline_buffer, new_capacity));
+        if (!new_buffer)
+            return Error::from_errno(ENOMEM);
+
+        m_outline_buffer = new_buffer;
+    }
+
+    m_outline_capacity = new_capacity;
+    m_inline = false;
+    return {};
+}
+
 static constexpr size_t string_builder_prefix_size(StringBuilder::Mode mode)
 {
     switch (mode) {
@@ -545,7 +775,7 @@ ErrorOr<void> StringBuilder::try_append_escaped_for_json(StringView string)
 
 auto StringBuilder::leak_buffer_for_string_construction() -> Optional<Buffer::OutlineBuffer>
 {
-    if (auto buffer = m_buffer.leak_outline_buffer({}); buffer.has_value()) {
+    if (auto buffer = m_buffer.leak_outline_buffer(); buffer.has_value()) {
         clear();
         return buffer;
     }

AK/StringBuilder.h

@@ -24,7 +24,64 @@ public:
     static constexpr auto DEFAULT_MODE = Mode::UTF8;
     static constexpr size_t inline_capacity = 256;
 
-    using Buffer = Detail::ByteBuffer<inline_capacity>;
+    class Buffer {
+    public:
+        Buffer() = default;
+        Buffer(Buffer const&);
+        Buffer(Buffer&&);
+        ~Buffer();
+
+        Buffer& operator=(Buffer const&);
+        Buffer& operator=(Buffer&&);
+
+        [[nodiscard]] u8* data();
+        [[nodiscard]] u8 const* data() const;
+
+        [[nodiscard]] Bytes span() LIFETIME_BOUND;
+        [[nodiscard]] ReadonlyBytes span() const LIFETIME_BOUND;
+
+        [[nodiscard]] void* end_pointer();
+
+        [[nodiscard]] size_t size() const { return m_size; }
+        [[nodiscard]] size_t capacity() const { return m_inline ? inline_capacity : m_outline_capacity; }
+        [[nodiscard]] bool is_inline() const { return m_inline; }
+
+        void clear();
+        void resize(size_t);
+        void set_size(size_t);
+        void ensure_capacity(size_t);
+
+        ErrorOr<void> try_resize(size_t);
+        ErrorOr<void> try_ensure_capacity(size_t);
+        ErrorOr<void> try_append(char);
+        ErrorOr<void> try_append(ReadonlyBytes);
+        ErrorOr<void> try_append(void const*, size_t);
+
+        void append(char);
+        void append(void const*, size_t);
+
+        struct OutlineBuffer {
+            Bytes buffer;
+            size_t capacity { 0 };
+        };
+        Optional<OutlineBuffer> leak_outline_buffer();
+
+    private:
+        void move_from(Buffer&&);
+        void trim(size_t, bool may_discard_existing_data);
+        void shrink_into_inline_buffer(size_t, bool may_discard_existing_data);
+        ErrorOr<void> try_ensure_capacity_slowpath(size_t);
+
+        union {
+            u8 m_inline_buffer[inline_capacity];
+            struct {
+                u8* m_outline_buffer;
+                size_t m_outline_capacity;
+            };
+        };
+        size_t m_size { 0 };
+        bool m_inline { true };
+    };
 
     StringBuilder();
     explicit StringBuilder(size_t initial_capacity);

AK/StringData.h

@@ -27,7 +27,7 @@ public:
         VERIFY(byte_count);
 
         auto capacity = allocation_size_for_string_data(byte_count);
-        void* slot = kmalloc(capacity);
+        void* slot = kmalloc(HeapPartition::String, capacity);
         if (!slot)
             return Error::from_errno(ENOMEM);
 
@@ -53,7 +53,7 @@ public:
         VERIFY(byte_count > MAX_SHORT_STRING_BYTE_COUNT);
 
         auto capacity = sizeof(StringData) + sizeof(StringData::SubstringData);
-        void* slot = kmalloc(capacity);
+        void* slot = kmalloc(HeapPartition::String, capacity);
         if (!slot)
             return Error::from_errno(ENOMEM);
 

AK/Utf16StringData.cpp

@@ -21,7 +21,7 @@ NonnullRefPtr<Utf16StringData> Utf16StringData::create_uninitialized(StorageType
 {
     auto allocation_size = allocation_size_for_string_data(storage_type == Utf16StringData::StorageType::ASCII, code_unit_length);
 
-    void* slot = kmalloc(allocation_size);
+    void* slot = kmalloc(HeapPartition::String, allocation_size);
     VERIFY(slot);
 
     return adopt_ref(*new (slot) Utf16StringData(storage_type, code_unit_length));

AK/kmalloc.cpp

@@ -134,6 +134,9 @@ static mi_heap_t* heap_for_partition(HeapPartition partition)
     case HeapPartition::JSObjectStorage:
         static mi_heap_t* js_object_storage_heap = mi_heap_new();
         return js_object_storage_heap;
+    case HeapPartition::String:
+        static thread_local mi_heap_t* string_heap = mi_heap_new();
+        return string_heap;
     }
     VERIFY_NOT_REACHED();
 }

AK/kmalloc.h

@@ -15,6 +15,7 @@ enum class HeapPartition {
     General,
     ArrayBuffer,
     JSObjectStorage,
+    String,
 };
 
 [[nodiscard]] void* ak_kcalloc(size_t count, size_t size);

Tests/AK/TestString.cpp

@@ -282,6 +282,21 @@ TEST_CASE(string_builder)
     EXPECT_EQ(string.bytes().size(), 8u);
 }
 
+TEST_CASE(string_builder_outlined_to_string)
+{
+    StringBuilder builder;
+    builder.append_repeated("abcd"sv, 100);
+
+    auto string = MUST(builder.to_string());
+    EXPECT_EQ(string.bytes_as_string_view().length(), 400u);
+    EXPECT(string.starts_with_bytes("abcd"sv));
+    EXPECT(string.ends_with_bytes("abcd"sv));
+
+    builder.append("reused"sv);
+    auto reused_string = MUST(builder.to_string());
+    EXPECT_EQ(reused_string, "reused"sv);
+}
+
 TEST_CASE(ak_format)
 {
     auto foo = MUST(String::formatted("Hello {}", "friends"_string));