.

include/flatmemory/details/byte_buffer2.hpp

@@ -33,6 +33,7 @@ namespace flatmemory {
 class ByteBuffer2 {
 private:
     std::vector<uint8_t> m_data;
+    size_t m_size;
 
     void resize_to_fit(size_t pos, size_t amount) {
         size_t final_size = pos + amount;
@@ -72,10 +73,14 @@ class ByteBuffer2 {
         return amount;
     }
 
+    /// @brief Set the final size of the buffer.
+    /// @param count 
+    void set_size(size_t count) { m_size = count; }
+
     [[nodiscard]] uint8_t* data() { return m_data.data(); }
     [[nodiscard]] const uint8_t* data() const { return m_data.data(); }
 
-    [[nodiscard]] size_t size() const { return m_data.size(); }
+    [[nodiscard]] size_t size() const { return m_size; }
     [[nodiscard]] size_t capacity() const { return m_data.capacity(); }
 };
 

include/flatmemory/details/types/bitset.hpp

@@ -286,12 +286,14 @@ namespace flatmemory
             buffer_size_type buffer_size = Layout<Bitset<Block>>::blocks_position;
             // Write blocks
             m_blocks.finish();
-            buffer_size += m_buffer.write(Layout<Bitset<Block>>::blocks_position, m_blocks.buffer().data(), read_value<buffer_size_type>(m_blocks.buffer().data()));
+            buffer_size_type blocks_buffer_size = read_value<buffer_size_type>(m_blocks.buffer().data());
+            buffer_size += m_buffer.write(Layout<Bitset<Block>>::blocks_position, m_blocks.buffer().data(), blocks_buffer_size);
             // Write final padding
             buffer_size += m_buffer.write_padding(buffer_size, calculate_amoung_padding(buffer_size, Layout<Bitset<Block>>::final_alignment));
 
             /* Write buffer size */
             m_buffer.write(Layout<Bitset<Block>>::buffer_size_position, buffer_size);
+            m_buffer.set_size(buffer_size);
         }
 
         void clear_impl()

include/flatmemory/details/types/trivial.hpp

@@ -19,6 +19,7 @@
 #define FLATMEMORY_TYPES_STRUCT_HPP_
 
 #include "../byte_buffer.hpp"
+#include "../byte_buffer2.hpp"
 #include "../byte_buffer_utils.hpp"
 #include "../layout_utils.hpp"
 #include "../layout.hpp"
@@ -68,24 +69,19 @@ namespace flatmemory
         private:
             T m_trivial;
 
-            ByteBuffer m_buffer;
+            ByteBuffer2 m_buffer;
 
             /* Implement IBuilder interface. */
             template<typename>
             friend class IBuilder;
 
             void finish_impl() {
-                m_buffer.write(m_trivial);
+                m_buffer.write(0, m_trivial);
+                m_buffer.set_size(sizeof(T));
             }
 
-            void clear_impl() {
-                // Clear this builder.
-                m_buffer.clear();
-            }
-
-
-            [[nodiscard]] ByteBuffer& get_buffer_impl() { return m_buffer; }
-            [[nodiscard]] const ByteBuffer& get_buffer_impl() const { return m_buffer; }
+            [[nodiscard]] auto& get_buffer_impl() { return m_buffer; }
+            [[nodiscard]] const auto& get_buffer_impl() const { return m_buffer; }
 
         public:
             [[nodiscard]] T& operator*() { return m_trivial; }
@@ -130,6 +126,11 @@ namespace flatmemory
             assert(test_correct_alignment<T>(m_buf));
             return &read_value<T>(m_buf); 
         }
+
+        [[nodiscard]] size_t buffer_size() const { 
+            assert(m_buf);
+            return sizeof(T);
+        }
     };
 
 
@@ -168,6 +169,11 @@ namespace flatmemory
             assert(test_correct_alignment<T>(m_buf));
             return &read_value<T>(m_buf); 
         }
+
+        [[nodiscard]] size_t buffer_size() const { 
+            assert(m_buf);
+            return sizeof(T);
+        }
     };
 }
 

include/flatmemory/details/types/vector.hpp

@@ -122,22 +122,28 @@ namespace flatmemory
                     size_t offset_end = offset_pos + m_data.size() * sizeof(offset_type);
                     size_t offset_padding = calculate_amoung_padding(offset_end, Layout<T>::final_alignment);
                     // position of data
-                    size_t data_pos = offset_end + offset_padding; 
-                    // buffer size points to data_pos
-                    buffer_size = data_pos;
+                    offset_type data_offset = offset_end + offset_padding; 
+                    // buffer size points to data_offset
+                    buffer_size = data_offset;
                     for (size_t i = 0; i < m_data.size(); ++i) {
+                        // write offset
+                        m_buffer.write(offset_pos, data_offset);
+                        offset_pos += sizeof(offset_type);
+
+                        // write data
                         auto& nested_builder = m_data[i];
                         nested_builder.finish();
-                        m_buffer.write(offset_pos, static_cast<offset_type>(data_pos));
-                        offset_pos += sizeof(offset_type);
-                        buffer_size += m_buffer.write(data_pos, nested_builder.buffer().data(), read_value<buffer_size_type>(nested_builder.buffer().data()));
-                        data_pos += nested_builder.buffer().size();
+                        buffer_size_type nested_buffer_size = read_value<buffer_size_type>(nested_builder.buffer().data());
+                        m_buffer.write(data_offset, nested_builder.buffer().data(), nested_buffer_size);
+                        data_offset += nested_buffer_size;
+                        buffer_size += nested_buffer_size;
                     }
                 }
                 buffer_size += m_buffer.write_padding(buffer_size, calculate_amoung_padding(buffer_size, Layout<Vector<T>>::final_alignment));
 
                 /* Write buffer size */
                 m_buffer.write(Layout<Vector<T>>::buffer_size_position, buffer_size);
+                m_buffer.set_size(buffer_size);
             }
 
             /* clear stl */