1 | #include "duckdb/storage/arena_allocator.hpp" |
2 | |
3 | #include "duckdb/common/assert.hpp" |
4 | |
5 | namespace duckdb { |
6 | |
7 | //===--------------------------------------------------------------------===// |
8 | // Arena Chunk |
9 | //===--------------------------------------------------------------------===// |
10 | ArenaChunk::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 | } |
14 | ArenaChunk::~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 | //===--------------------------------------------------------------------===// |
26 | struct ArenaAllocatorData : public PrivateAllocatorData { |
27 | explicit ArenaAllocatorData(ArenaAllocator &allocator) : allocator(allocator) { |
28 | } |
29 | |
30 | ArenaAllocator &allocator; |
31 | }; |
32 | |
33 | static 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 | |
38 | static void ArenaAllocatorFree(PrivateAllocatorData *, data_ptr_t, idx_t) { |
39 | // nop |
40 | } |
41 | |
42 | static 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 | //===--------------------------------------------------------------------===// |
50 | ArenaAllocator::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 | |
58 | ArenaAllocator::~ArenaAllocator() { |
59 | } |
60 | |
61 | data_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 | |
82 | data_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 | |
103 | data_ptr_t ArenaAllocator::AllocateAligned(idx_t size) { |
104 | return Allocate(len: AlignValue<idx_t>(n: size)); |
105 | } |
106 | |
107 | data_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 | |
111 | void 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 | |
129 | void ArenaAllocator::Destroy() { |
130 | head = nullptr; |
131 | tail = nullptr; |
132 | current_capacity = ARENA_ALLOCATOR_INITIAL_CAPACITY; |
133 | } |
134 | |
135 | void 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 | |
143 | ArenaChunk *ArenaAllocator::GetHead() { |
144 | return head.get(); |
145 | } |
146 | |
147 | ArenaChunk *ArenaAllocator::GetTail() { |
148 | return tail; |
149 | } |
150 | |
151 | bool ArenaAllocator::IsEmpty() { |
152 | return head == nullptr; |
153 | } |
154 | |
155 | } // namespace duckdb |
156 | |