From 229082a1b2993218d645221dfd601a3c46264c49 Mon Sep 17 00:00:00 2001
From: Bananymous <oskari.alaranta@bananymous.com>
Date: Sun, 29 Oct 2023 02:29:49 +0300
Subject: [PATCH] BAN: Rewrite String with small string optimizations

String now holds a 15 byte sso buffer. I'm not sure what the size
should actually be but 15 will work for now. Maybe the sso buffer
should be contained in an union with one bit flag in size instead of
variant that uses extra 8 bytes for type index.

This patch buffs sizeof(String) from 24 bytes to 32 bytes on 64 bit.
I assume this is much better version than the old which had to make
allocation even for empty strings :D.
---
 BAN/BAN/String.cpp       | 298 ++++++++++++++++++++-------------------
 BAN/include/BAN/String.h |  55 +++++---
 2 files changed, 186 insertions(+), 167 deletions(-)

diff --git a/BAN/BAN/String.cpp b/BAN/BAN/String.cpp
index 96f2693618..5f94a52bf5 100644
--- a/BAN/BAN/String.cpp
+++ b/BAN/BAN/String.cpp
@@ -1,102 +1,108 @@
-#include <BAN/Errors.h>
-#include <BAN/Math.h>
-#include <BAN/Move.h>
-#include <BAN/New.h>
 #include <BAN/String.h>
-#include <BAN/StringView.h>
-
-#include <string.h>
+#include <BAN/New.h>
+#include <BAN/Variant.h>
 
 namespace BAN
 {
 
 	String::String()
 	{
-		MUST(copy_impl(""sv));
 	}
 
 	String::String(const String& other)
 	{
-		MUST(copy_impl(other.sv()));
+		*this = other;
 	}
 
 	String::String(String&& other)
 	{
-		move_impl(move(other));
+		*this = move(other);
 	}
 
 	String::String(StringView other)
 	{
-		MUST(copy_impl(other));
+		*this = other;
 	}
 
 	String::~String()
 	{
-		BAN::deallocator(m_data);
+		clear();
 	}
 
 	String& String::operator=(const String& other)
 	{
-		MUST(copy_impl(other.sv()));
+		clear();
+		if (!other.fits_in_sso())
+			MUST(ensure_capacity(other.size()));
+		memcpy(data(), other.data(), other.size() + 1);
+		m_size = other.size();
 		return *this;
 	}
 
 	String& String::operator=(String&& other)
 	{
-		BAN::deallocator(m_data);
-		move_impl(move(other));
+		clear();
+
+		if (other.fits_in_sso())
+			memcpy(data(), other.data(), other.size() + 1);
+		else
+			m_storage = other.m_storage.get<GeneralStorage>();
+		m_size = other.m_size;
+
+		other.m_size = 0;
+		other.m_storage = SSOStorage();
+
 		return *this;
 	}
 
 	String& String::operator=(StringView other)
 	{
-		MUST(copy_impl(other));
+		clear();
+		if (!fits_in_sso(other.size()))
+			MUST(ensure_capacity(other.size()));
+		memcpy(data(), other.data(), other.size());
+		m_size = other.size();
+		data()[m_size] = '\0';
 		return *this;
 	}
 
-	ErrorOr<void> String::push_back(char ch)
+	ErrorOr<void> String::push_back(char c)
 	{
-		TRY(ensure_capacity(m_size + 2));
-		m_data[m_size] = ch;
+		TRY(ensure_capacity(m_size + 1));
+		data()[m_size] = c;
 		m_size++;
-		m_data[m_size] = '\0';
+		data()[m_size] = '\0';
 		return {};
 	}
 
-	ErrorOr<void> String::insert(char ch, size_type index)
+	ErrorOr<void> String::insert(char c, size_type index)
 	{
 		ASSERT(index <= m_size);
-		TRY(ensure_capacity(m_size + 1 + 1));
-		memmove(m_data + index + 1, m_data + index, m_size - index);
-		m_data[index] = ch;
-		m_size += 1;
-		m_data[m_size] = '\0';
+		TRY(ensure_capacity(m_size + 1));
+		memmove(data() + index + 1, data() + index, m_size - index);
+		data()[index] = c;
+		m_size++;
+		data()[m_size] = '\0';
 		return {};
 	}
 
-	ErrorOr<void> String::insert(StringView other, size_type index)
+	ErrorOr<void> String::insert(StringView str, size_type index)
 	{
 		ASSERT(index <= m_size);
-		TRY(ensure_capacity(m_size + other.size() + 1));
-		memmove(m_data + index + other.size(), m_data + index, m_size - index);
-		memcpy(m_data + index, other.data(), other.size());
-		m_size += other.size();
-		m_data[m_size] = '\0';
+		TRY(ensure_capacity(m_size + str.size()));
+		memmove(data() + index + str.size(), data() + index, m_size - index);
+		memcpy(data() + index, str.data(), str.size());
+		m_size += str.size();
+		data()[m_size] = '\0';
 		return {};
 	}
 
-	ErrorOr<void> String::append(StringView other)
+	ErrorOr<void> String::append(StringView str)
 	{
-		TRY(ensure_capacity(m_size + other.size() + 1));
-		memcpy(m_data + m_size, other.data(), other.size());
-		m_size += other.size();
-		m_data[m_size] = '\0';
-		return {};
-	}
-
-	ErrorOr<void> String::append(const String& string)
-	{
-		TRY(append(string.sv()));
+		TRY(ensure_capacity(m_size + str.size()));
+		memcpy(data() + m_size, str.data(), str.size());
+		m_size += str.size();
+		data()[m_size] = '\0';
 		return {};
 	}
 
@@ -104,159 +110,161 @@ namespace BAN
 	{
 		ASSERT(m_size > 0);
 		m_size--;
-		m_data[m_size] = '\0';
+		data()[m_size] = '\0';
 	}
 
 	void String::remove(size_type index)
 	{
-		erase(index,  1);
-	}
-
-	void String::erase(size_type index, size_type count)
-	{
-		ASSERT(index + count <= m_size);
-		memmove(m_data + index, m_data + index + count, m_size - index - count);
-		m_size -= count;
-		m_data[m_size] = '\0';
+		ASSERT(index < m_size);
+		memcpy(data() + index, data() + index + 1, m_size - index);
+		m_size--;
+		data()[m_size] = '\0';
 	}
 
 	void String::clear()
 	{
+		if (!has_sso())
+		{
+			deallocator(m_storage.get<GeneralStorage>().data);
+			m_storage = SSOStorage();
+		}
 		m_size = 0;
-		m_data[0] = '\0';
+		data()[m_size] = '\0';
 	}
 
-	char String::operator[](size_type index) const
+	bool String::operator==(StringView str) const
 	{
-		ASSERT(index < m_size);
-		return m_data[index];
-	}
-
-	char& String::operator[](size_type index)
-	{
-		ASSERT(index < m_size);
-		return m_data[index];
-	}
-
-	bool String::operator==(const String& other) const
-	{
-		if (m_size != other.m_size)
+		if (size() != str.size())
 			return false;
-		return memcmp(m_data, other.m_data, m_size) == 0;
-	}
-
-	bool String::operator==(StringView other) const
-	{
-		if (m_size != other.size())
-			return false;
-		return memcmp(m_data, other.data(), m_size) == 0;
-	}
-
-	bool String::operator==(const char* other) const
-	{
-		for (size_type i = 0; i <= m_size; i++)
-			if (m_data[i] != other[i])
+		for (size_type i = 0; i < m_size; i++)
+			if (data()[i] != str.data()[i])
 				return false;
 		return true;
 	}
 
-	ErrorOr<void> String::resize(size_type size, char ch)
+	bool String::operator==(const char* cstr) const
 	{
-		if (size < m_size)
+		for (size_type i = 0; i < m_size; i++)
+			if (data()[i] != cstr[i])
+				return false;
+		if (cstr[size()] != '\0')
+			return false;
+		return true;
+	}
+
+	ErrorOr<void> String::resize(size_type new_size, char init_c)
+	{
+		if (m_size == new_size)
+			return {};
+
+		// expanding
+		if (m_size < new_size)
 		{
-			m_data[size] = '\0';
-			m_size = size;
+			TRY(ensure_capacity(new_size));
+			memset(data() + m_size, init_c, new_size - m_size);
+			m_size = new_size;
+			data()[m_size] = '\0';
+			return {};
 		}
-		else if (size > m_size)
+		
+		// shrink general -> sso
+		if (!has_sso() && fits_in_sso(new_size))
 		{
-			TRY(ensure_capacity(size + 1));
-			for (size_type i = m_size; i < size; i++)
-				m_data[i] = ch;
-			m_data[size] = '\0';
-			m_size = size;
+			char* data = m_storage.get<GeneralStorage>().data;
+			m_storage = SSOStorage();
+			memcpy(m_storage.get<SSOStorage>().storage, data, new_size);
+			deallocator(data);
 		}
-		m_size = size;
+
+		m_size = new_size;
+		data()[m_size] = '\0';
 		return {};
 	}
 
-	ErrorOr<void> String::reserve(size_type size)
+	ErrorOr<void> String::reserve(size_type new_size)
 	{
-		TRY(ensure_capacity(size));
+		TRY(ensure_capacity(new_size));
 		return {};
 	}
 
 	ErrorOr<void> String::shrink_to_fit()
 	{
-		size_type temp = m_capacity;
-		m_capacity = 0;
-		auto error_or = ensure_capacity(m_size);
-		if (error_or.is_error())
+		if (has_sso())
+			return {};
+
+		if (fits_in_sso())
 		{
-			m_capacity = temp;
-			return error_or;
+			char* data = m_storage.get<GeneralStorage>().data;
+			m_storage = SSOStorage();
+			memcpy(m_storage.get<SSOStorage>().storage, data, m_size + 1);
+			deallocator(data);
+			return {};
 		}
+
+		GeneralStorage& storage = m_storage.get<GeneralStorage>();
+		if (storage.capacity == m_size)
+			return {};
+
+		char* new_data = (char*)allocator(m_size + 1);
+		if (new_data == nullptr)
+			return BAN::Error::from_errno(ENOMEM);
+
+		memcpy(new_data, storage.data, m_size);
+		deallocator(storage.data);
+
+		storage.capacity = m_size;
+		storage.data = new_data;
+
 		return {};
 	}
 
-	StringView String::sv() const
-	{
-		return StringView(*this);
-	}
-
-	bool String::empty() const
-	{
-		return m_size == 0;
-	}
-
-	String::size_type String::size() const
-	{
-		return m_size;
-	}
-
 	String::size_type String::capacity() const
 	{
-		return m_capacity;
+		if (has_sso())
+			return sso_capacity;
+		return m_storage.get<GeneralStorage>().capacity;
+	}
+
+	char* String::data()
+	{
+		if (has_sso())
+			return m_storage.get<SSOStorage>().storage;
+		return m_storage.get<GeneralStorage>().data;
 	}
 
 	const char* String::data() const
 	{
-		return m_data;
+		if (has_sso())
+			return m_storage.get<SSOStorage>().storage;
+		return m_storage.get<GeneralStorage>().data;
 	}
 
-	ErrorOr<void> String::ensure_capacity(size_type size)
+	ErrorOr<void> String::ensure_capacity(size_type new_size)
 	{
-		if (m_capacity >= size)
+		if (m_size >= new_size || fits_in_sso(new_size))
 			return {};
-		size_type new_cap = BAN::Math::max<size_type>(size, m_capacity * 2);
-		void* new_data = BAN::allocator(new_cap);
+		
+		char* new_data = (char*)allocator(new_size + 1);
 		if (new_data == nullptr)
-			return Error::from_errno(ENOMEM);
-		if (m_data)
-			memcpy(new_data, m_data, m_size + 1);
-		BAN::deallocator(m_data);
-		m_data = (char*)new_data;
-		m_capacity = new_cap;
+			return BAN::Error::from_errno(ENOMEM);
+		
+		memcpy(new_data, data(), m_size + 1);
+
+		if (has_sso())
+			m_storage = GeneralStorage();
+		else
+			deallocator(m_storage.get<GeneralStorage>().data);
+
+		auto& storage = m_storage.get<GeneralStorage>();
+		storage.capacity = new_size;
+		storage.data = new_data;
+
 		return {};
 	}
 
-	ErrorOr<void> String::copy_impl(StringView other)
+	bool String::has_sso() const
 	{
-		TRY(ensure_capacity(other.size() + 1));
-		memcpy(m_data, other.data(), other.size());
-		m_size = other.size();
-		m_data[m_size] = '\0';
-		return {};
-	}
-
-	void String::move_impl(String&& other)
-	{
-		m_data		= other.m_data;
-		m_size		= other.m_size;
-		m_capacity	= other.m_capacity;
-
-		other.m_data = nullptr;
-		other.m_size = 0;
-		other.m_capacity = 0;
+		return m_storage.has<SSOStorage>();
 	}
 
 }
diff --git a/BAN/include/BAN/String.h b/BAN/include/BAN/String.h
index f1c4f3ab66..26685afdc8 100644
--- a/BAN/include/BAN/String.h
+++ b/BAN/include/BAN/String.h
@@ -1,8 +1,8 @@
 #pragma once
 
 #include <BAN/Errors.h>
-#include <BAN/ForwardList.h>
 #include <BAN/Formatter.h>
+#include <BAN/ForwardList.h>
 #include <BAN/Hash.h>
 #include <BAN/Iterators.h>
 
@@ -15,6 +15,7 @@ namespace BAN
 		using size_type = size_t;
 		using iterator = IteratorSimple<char, String>;
 		using const_iterator = ConstIteratorSimple<char, String>;
+		static constexpr size_type sso_capacity = 15;
 
 	public:
 		String();
@@ -34,29 +35,26 @@ namespace BAN
 		ErrorOr<void> insert(char, size_type);
 		ErrorOr<void> insert(StringView, size_type);
 		ErrorOr<void> append(StringView);
-		ErrorOr<void> append(const String&);
 
 		void pop_back();
 		void remove(size_type);
-		void erase(size_type, size_type);
 
 		void clear();
 
-		const_iterator begin() const { return const_iterator(m_data); }
-		iterator begin() { return iterator(m_data); }
-		const_iterator end() const { return const_iterator(m_data + m_size); }
-		iterator end() { return iterator(m_data + m_size); }
+		const_iterator begin() const	{ return const_iterator(data()); }
+		iterator begin()				{ return iterator(data()); }
+		const_iterator end() const		{ return const_iterator(data() + size()); }
+		iterator end()					{ return iterator(data() + size()); }
 
-		char front() const	{ ASSERT(!empty()); return m_data[0]; }
-		char& front()		{ ASSERT(!empty()); return m_data[0]; }
+		char front() const	{ ASSERT(m_size > 0); return data()[0]; }
+		char& front()		{ ASSERT(m_size > 0); return data()[0]; }
 
-		char back() const	{ ASSERT(!empty()); return m_data[m_size - 1]; }
-		char& back()		{ ASSERT(!empty()); return m_data[m_size - 1]; }
+		char back() const	{ ASSERT(m_size > 0); return data()[m_size - 1]; }
+		char& back()		{ ASSERT(m_size > 0); return data()[m_size - 1]; }
 
-		char operator[](size_type) const;
-		char& operator[](size_type);
+		char operator[](size_type index) const	{ ASSERT(index < m_size); return data()[index]; }
+		char& operator[](size_type index)		{ ASSERT(index < m_size); return data()[index]; }
 
-		bool operator==(const String&) const;
 		bool operator==(StringView) const;
 		bool operator==(const char*) const;
 
@@ -64,24 +62,37 @@ namespace BAN
 		ErrorOr<void> reserve(size_type);
 		ErrorOr<void> shrink_to_fit();
 
-		StringView sv() const;
+		StringView sv() const	{ return StringView(data(), size()); }
 
-		bool empty() const;
-		size_type size() const;
+		bool empty() const		{ return m_size == 0; }
+		size_type size() const	{ return m_size; }
 		size_type capacity() const;
 
+		char* data();
 		const char* data() const;
 
 	private:
 		ErrorOr<void> ensure_capacity(size_type);
 
-		ErrorOr<void> copy_impl(StringView);
-		void move_impl(String&&);
+		bool has_sso() const;
+
+		bool fits_in_sso() const { return fits_in_sso(m_size); }
+		static bool fits_in_sso(size_type size) { return size < sso_capacity; }
 
 	private:
-		char*		m_data		= nullptr;
-		size_type	m_capacity	= 0;
-		size_type	m_size		= 0;	
+		struct SSOStorage
+		{
+			char storage[sso_capacity + 1] {};
+		};
+		struct GeneralStorage
+		{
+			size_type capacity { 0 };
+			char* data;
+		};
+
+	private:
+		Variant<SSOStorage, GeneralStorage>	m_storage	{ SSOStorage() };
+		size_type							m_size		{ 0 };
 	};
 
 	template<typename... Args>