AK: Give Vector the ability to have an inline capacity.

This makes Vector malloc-free as long as you stay within the templated
inline capacity. :^)

1 files changed, 141 insertions, 115 deletions

diff --git a/AK/Vector.h b/AK/Vector.h
index dc078fb886..1ddd9f032c 100644
--- a/AK/Vector.h
+++ b/AK/Vector.h
@@ -1,71 +1,40 @@
 #pragma once
 
-#include "Assertions.h"
-#include "OwnPtr.h"
-#include "kmalloc.h"
+#include <AK/Assertions.h>
+#include <AK/StdLibExtras.h>
+#include <AK/kmalloc.h>
 
 namespace AK {
 
-template<typename T> class Vector;
-
-template<typename T>
-class VectorImpl {
+template<typename T, int inline_capacity = 0>
+class Vector {
 public:
-    ~VectorImpl()
+    Vector()
+        : m_capacity(inline_capacity)
     {
-        for (int i = 0; i < m_size; ++i)
-            at(i).~T();
     }
-    static OwnPtr<VectorImpl> create(int capacity)
-    {
-        int size = sizeof(VectorImpl) + sizeof(T) * capacity;
-        void* slot = kmalloc(size);
-        new (slot) VectorImpl(capacity);
-        return OwnPtr<VectorImpl>((VectorImpl*)slot);
-    }
-
-    int size() const { return m_size; }
-    int capacity() const { return m_capacity; }
 
-    T& at(int i) { ASSERT(i < m_size); return *slot(i); }
-    const T& at(int i) const { ASSERT(i < m_size); return *slot(i); }
-
-    void remove(int index)
+    ~Vector()
     {
-        ASSERT(index < m_size);
-        at(index).~T();
-        for (int i = index + 1; i < m_size; ++i) {
-            new (slot(i - 1)) T(move(at(i)));
-            at(i).~T();
-        }
-
-        --m_size;
+        clear();
     }
 
-private:
-    friend class Vector<T>;
-
-    VectorImpl(int capacity) : m_capacity(capacity) { }
-
-    T* tail() { return reinterpret_cast<T*>(this + 1); }
-    T* slot(int i) { return &tail()[i]; }
-
-    const T* tail() const { return reinterpret_cast<const T*>(this + 1); }
-    const T* slot(int i) const { return &tail()[i]; }
-
-    int m_size { 0 };
-    int m_capacity;
-};
-
-template<typename T>
-class Vector {
-public:
-    Vector() { }
-    ~Vector() { clear(); }
-
     Vector(Vector&& other)
-        : m_impl(move(other.m_impl))
+        : m_size(other.m_size)
+        , m_capacity(other.m_capacity)
+        , m_outline_buffer(other.m_outline_buffer)
     {
+        if constexpr (inline_capacity > 0) {
+            if (!m_outline_buffer) {
+                for (int i = 0; i < m_size; ++i) {
+                    new (&inline_buffer()[i]) T(move(other.inline_buffer()[i]));
+                    other.inline_buffer()[i].~T();
+                }
+            }
+        }
+        other.m_outline_buffer = nullptr;
+        other.m_size = 0;
+        other.reset_capacity();
     }
 
     Vector(const Vector& other)
@@ -75,25 +44,44 @@ public:
             unchecked_append(other[i]);
     }
 
+    // FIXME: What about assigning from a vector with lower inline capacity?
     Vector& operator=(Vector&& other)
     {
-        if (this != &other)
-            m_impl = move(other.m_impl);
+        if (this != &other) {
+            clear();
+            m_size = other.m_size;
+            m_capacity = other.m_capacity;
+            m_outline_buffer = other.m_outline_buffer;
+            if constexpr (inline_capacity > 0) {
+                if (!m_outline_buffer) {
+                    for (int i = 0; i < m_size; ++i) {
+                        new (&inline_buffer()[i]) T(move(other.inline_buffer()[i]));
+                        other.inline_buffer()[i].~T();
+                    }
+                }
+            }
+            other.m_outline_buffer = nullptr;
+            other.m_size = 0;
+            other.reset_capacity();
+        }
         return *this;
     }
 
     void clear()
     {
-        m_impl = nullptr;
+        clear_with_capacity();
+        if (m_outline_buffer) {
+            kfree(m_outline_buffer);
+            m_outline_buffer = nullptr;
+        }
+        reset_capacity();
     }
 
     void clear_with_capacity()
     {
-        if (!m_impl)
-            return;
-        for (int i = 0; i < size(); ++i)
-            at(i).~T();
-        m_impl->m_size = 0;
+        for (int i = 0; i < m_size; ++i)
+            data()[i].~T();
+        m_size = 0;
     }
 
     bool contains_slow(const T& value) const
@@ -106,14 +94,24 @@ public:
     }
 
     bool is_empty() const { return size() == 0; }
-    int size() const { return m_impl ? m_impl->size() : 0; }
-    int capacity() const { return m_impl ? m_impl->capacity() : 0; }
+    int size() const { return m_size; }
+    int capacity() const { return m_capacity; }
 
-    T* data() { return m_impl ? m_impl->slot(0) : nullptr; }
-    const T* data() const { return m_impl ? m_impl->slot(0) : nullptr; }
+    T* data()
+    {
+        if constexpr (inline_capacity > 0)
+            return m_outline_buffer ? m_outline_buffer : inline_buffer();
+        return m_outline_buffer;
+    }
+    const T* data() const
+    {
+        if constexpr (inline_capacity > 0)
+            return m_outline_buffer ? m_outline_buffer : inline_buffer();
+        return m_outline_buffer;
+    }
 
-    const T& at(int i) const { return m_impl->at(i); }
-    T& at(int i) { return m_impl->at(i); }
+    const T& at(int i) const { ASSERT(i >= 0 && i < m_size); return data()[i]; }
+    T& at(int i) { ASSERT(i >= 0 && i < m_size); return data()[i]; }
 
     const T& operator[](int i) const { return at(i); }
     T& operator[](int i) { return at(i); }
@@ -129,7 +127,7 @@ public:
         ASSERT(!is_empty());
         T value = move(last());
         last().~T();
-        --m_impl->m_size;
+        --m_size;
         return value;
     }
 
@@ -143,7 +141,14 @@ public:
 
     void remove(int index)
     {
-        m_impl->remove(index);
+        ASSERT(index < m_size);
+        at(index).~T();
+        for (int i = index + 1; i < m_size; ++i) {
+            new (slot(i - 1)) T(move(at(i)));
+            at(i).~T();
+        }
+
+        --m_size;
     }
 
     void insert(int index, T&& value)
@@ -151,16 +156,16 @@ public:
         ASSERT(index <= size());
         if (index == size())
             return append(move(value));
-        ensure_capacity(size() + 1);
-        ++m_impl->m_size;
+        grow_capacity(size() + 1);
+        ++m_size;
         for (int i = size() - 1; i > index; --i) {
-            new (m_impl->slot(i)) T(move(m_impl->at(i - 1)));
-            m_impl->at(i - 1).~T();
+            new (slot(i)) T(move(at(i - 1)));
+            at(i - 1).~T();
         }
-        new (m_impl->slot(index)) T(move(value));
+        new (slot(index)) T(move(value));
     }
 
-    Vector& operator=(const Vector<T>& other)
+    Vector& operator=(const Vector& other)
     {
         if (this != &other) {
             clear();
@@ -171,17 +176,16 @@ public:
         return *this;
     }
 
-    void append(Vector<T>&& other)
+    void append(Vector&& other)
     {
-        if (!m_impl) {
-            m_impl = move(other.m_impl);
+        if (is_empty()) {
+            *this = move(other);
             return;
         }
-        Vector<T> tmp = move(other);
-        ensure_capacity(size() + tmp.size());
-        for (auto&& v : tmp) {
+        Vector tmp = move(other);
+        grow_capacity(size() + tmp.size());
+        for (auto&& v : tmp)
             unchecked_append(move(v));
-        }
     }
 
     template<typename Callback>
@@ -198,65 +202,72 @@ public:
     void unchecked_append(T&& value)
     {
         ASSERT((size() + 1) <= capacity());
-        new (m_impl->slot(m_impl->m_size)) T(move(value));
-        ++m_impl->m_size;
+        new (slot(m_size)) T(move(value));
+        ++m_size;
     }
 
     void unchecked_append(const T& value)
     {
-        new (m_impl->slot(m_impl->m_size)) T(value);
-        ++m_impl->m_size;
+        new (slot(m_size)) T(value);
+        ++m_size;
     }
 
     void append(T&& value)
     {
-        ensure_capacity(size() + 1);
-        new (m_impl->slot(m_impl->m_size)) T(move(value));
-        ++m_impl->m_size;
+        grow_capacity(size() + 1);
+        new (slot(m_size)) T(move(value));
+        ++m_size;
     }
 
     void append(const T& value)
     {
-        ensure_capacity(size() + 1);
-        new (m_impl->slot(m_impl->m_size)) T(value);
-        ++m_impl->m_size;
+        grow_capacity(size() + 1);
+        new (slot(m_size)) T(value);
+        ++m_size;
     }
 
     void prepend(const T& value)
     {
-        ensure_capacity(size() + 1);
+        grow_capacity(size() + 1);
         for (int i = size(); i > 0; --i) {
-            new (m_impl->slot(i)) T(move(at(i - 1)));
+            new (slot(i)) T(move(at(i - 1)));
             at(i - 1).~T();
         }
-        new (m_impl->slot(0)) T(value);
-        ++m_impl->m_size;
+        new (slot(0)) T(value);
+        ++m_size;
     }
 
     void append(const T* values, int count)
     {
         if (!count)
             return;
-        ensure_capacity(size() + count);
+        grow_capacity(size() + count);
         for (int i = 0; i < count; ++i)
-            new (m_impl->slot(m_impl->m_size + i)) T(values[i]);
-        m_impl->m_size += count;
+            new (slot(m_size + i)) T(values[i]);
+        m_size += count;
     }
 
-    void ensure_capacity(int neededCapacity)
+    void grow_capacity(int needed_capacity)
     {
-        if (capacity() >= neededCapacity)
+        if (m_capacity >= needed_capacity)
             return;
-        int new_capacity = padded_capacity(neededCapacity);
-        auto new_impl = VectorImpl<T>::create(new_capacity);
-        if (m_impl) {
-            new_impl->m_size = m_impl->m_size;
-            for (int i = 0; i < size(); ++i) {
-                new (new_impl->slot(i)) T(move(m_impl->at(i)));
-                m_impl->at(i).~T();
-            }
+        ensure_capacity(padded_capacity(needed_capacity));
+    }
+
+    void ensure_capacity(int needed_capacity)
+    {
+        if (m_capacity >= needed_capacity)
+            return;
+        int new_capacity = needed_capacity;
+        auto* new_buffer = (T*)kmalloc(new_capacity * sizeof(T));
+        for (int i = 0; i < m_size; ++i) {
+            new (&new_buffer[i]) T(move(at(i)));
+            at(i).~T();
         }
-        m_impl = move(new_impl);
+        if (m_outline_buffer)
+            kfree(m_outline_buffer);
+        m_outline_buffer = new_buffer;
+        m_capacity = new_capacity;
     }
 
     void resize(int new_size)
@@ -272,12 +283,12 @@ public:
         if (new_size > size()) {
             ensure_capacity(new_size);
             for (int i = size(); i < new_size; ++i)
-                new (m_impl->slot(i)) T;
+                new (slot(i)) T;
         } else {
             for (int i = new_size; i < size(); ++i)
-                m_impl->at(i).~T();
+                at(i).~T();
         }
-        m_impl->m_size = new_size;
+        m_size = new_size;
     }
 
     class Iterator {
@@ -319,12 +330,27 @@ public:
     ConstIterator end() const { return ConstIterator(*this, size()); }
 
 private:
+    void reset_capacity()
+    {
+        m_capacity = inline_capacity;
+    }
+
     static int padded_capacity(int capacity)
     {
         return max(int(4), capacity + (capacity / 4) + 4);
     }
 
-    OwnPtr<VectorImpl<T>> m_impl;
+    T* slot(int i) { return &data()[i]; }
+    const T* slot(int i) const { return &data()[i]; }
+
+    T* inline_buffer() { static_assert(inline_capacity > 0); return reinterpret_cast<T*>(m_inline_buffer_storage); }
+    const T* inline_buffer() const { static_assert(inline_capacity > 0); return reinterpret_cast<const T*>(m_inline_buffer_storage); }
+
+    int m_size { 0 };
+    int m_capacity { 0 };
+
+    alignas(T) byte m_inline_buffer_storage[sizeof(T) * inline_capacity];
+    T* m_outline_buffer { nullptr };
 };
 
 }