1#include "duckdb/storage/arena_allocator.hpp"
2
3#include "duckdb/common/assert.hpp"
4
5namespace duckdb {
6
7//===--------------------------------------------------------------------===//
8// Arena Chunk
9//===--------------------------------------------------------------------===//
10ArenaChunk::ArenaChunk(Allocator &allocator, idx_t size) : current_position(0), maximum_size(size), prev(nullptr) {
11 D_ASSERT(size > 0);
12 data = allocator.Allocate(size);
13}
14ArenaChunk::~ArenaChunk() {
15 if (next) {
16 auto current_next = std::move(next);
17 while (current_next) {
18 current_next = std::move(current_next->next);
19 }
20 }
21}
22
23//===--------------------------------------------------------------------===//
24// Allocator Wrapper
25//===--------------------------------------------------------------------===//
26struct ArenaAllocatorData : public PrivateAllocatorData {
27 explicit ArenaAllocatorData(ArenaAllocator &allocator) : allocator(allocator) {
28 }
29
30 ArenaAllocator &allocator;
31};
32
33static data_ptr_t ArenaAllocatorAllocate(PrivateAllocatorData *private_data, idx_t size) {
34 auto &allocator_data = private_data->Cast<ArenaAllocatorData>();
35 return allocator_data.allocator.Allocate(size);
36}
37
38static void ArenaAllocatorFree(PrivateAllocatorData *, data_ptr_t, idx_t) {
39 // nop
40}
41
42static data_ptr_t ArenaAllocateReallocate(PrivateAllocatorData *private_data, data_ptr_t pointer, idx_t old_size,
43 idx_t size) {
44 auto &allocator_data = private_data->Cast<ArenaAllocatorData>();
45 return allocator_data.allocator.Reallocate(pointer, old_size, size);
46}
47//===--------------------------------------------------------------------===//
48// Arena Allocator
49//===--------------------------------------------------------------------===//
50ArenaAllocator::ArenaAllocator(Allocator &allocator, idx_t initial_capacity)
51 : allocator(allocator), arena_allocator(ArenaAllocatorAllocate, ArenaAllocatorFree, ArenaAllocateReallocate,
52 make_uniq<ArenaAllocatorData>(args&: *this)) {
53 head = nullptr;
54 tail = nullptr;
55 current_capacity = initial_capacity;
56}
57
58ArenaAllocator::~ArenaAllocator() {
59}
60
61data_ptr_t ArenaAllocator::Allocate(idx_t len) {
62 D_ASSERT(!head || head->current_position <= head->maximum_size);
63 if (!head || head->current_position + len > head->maximum_size) {
64 do {
65 current_capacity *= 2;
66 } while (current_capacity < len);
67 auto new_chunk = make_unsafe_uniq<ArenaChunk>(args&: allocator, args&: current_capacity);
68 if (head) {
69 head->prev = new_chunk.get();
70 new_chunk->next = std::move(head);
71 } else {
72 tail = new_chunk.get();
73 }
74 head = std::move(new_chunk);
75 }
76 D_ASSERT(head->current_position + len <= head->maximum_size);
77 auto result = head->data.get() + head->current_position;
78 head->current_position += len;
79 return result;
80}
81
82data_ptr_t ArenaAllocator::Reallocate(data_ptr_t pointer, idx_t old_size, idx_t size) {
83 D_ASSERT(head);
84 if (old_size == size) {
85 // nothing to do
86 return pointer;
87 }
88
89 auto head_ptr = head->data.get() + head->current_position;
90 int64_t diff = size - old_size;
91 if (pointer == head_ptr && (size < old_size || head->current_position + diff <= head->maximum_size)) {
92 // passed pointer is the head pointer, and the diff fits on the current chunk
93 head->current_position += diff;
94 return pointer;
95 } else {
96 // allocate new memory
97 auto result = Allocate(len: size);
98 memcpy(dest: result, src: pointer, n: old_size);
99 return result;
100 }
101}
102
103data_ptr_t ArenaAllocator::AllocateAligned(idx_t size) {
104 return Allocate(len: AlignValue<idx_t>(n: size));
105}
106
107data_ptr_t ArenaAllocator::ReallocateAligned(data_ptr_t pointer, idx_t old_size, idx_t size) {
108 return Reallocate(pointer, old_size, size: AlignValue<idx_t>(n: size));
109}
110
111void ArenaAllocator::Reset() {
112
113 if (head) {
114 // destroy all chunks except the current one
115 if (head->next) {
116 auto current_next = std::move(head->next);
117 while (current_next) {
118 current_next = std::move(current_next->next);
119 }
120 }
121 tail = head.get();
122
123 // reset the head
124 head->current_position = 0;
125 head->prev = nullptr;
126 }
127}
128
129void ArenaAllocator::Destroy() {
130 head = nullptr;
131 tail = nullptr;
132 current_capacity = ARENA_ALLOCATOR_INITIAL_CAPACITY;
133}
134
135void ArenaAllocator::Move(ArenaAllocator &other) {
136 D_ASSERT(!other.head);
137 other.tail = tail;
138 other.head = std::move(head);
139 other.current_capacity = current_capacity;
140 Destroy();
141}
142
143ArenaChunk *ArenaAllocator::GetHead() {
144 return head.get();
145}
146
147ArenaChunk *ArenaAllocator::GetTail() {
148 return tail;
149}
150
151bool ArenaAllocator::IsEmpty() {
152 return head == nullptr;
153}
154
155} // namespace duckdb
156