1// Copyright 2019 The Marl Authors.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7// https://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15#ifndef marl_containers_h
16#define marl_containers_h
17
18#include "debug.h"
19#include "memory.h"
20
21#include <algorithm> // std::max
22#include <cstddef> // size_t
23#include <utility> // std::move
24
25#include <deque>
26#include <map>
27#include <set>
28#include <unordered_map>
29#include <unordered_set>
30
31namespace marl {
32namespace containers {
33
34////////////////////////////////////////////////////////////////////////////////
35// STL wrappers
36// STL containers that use a marl::StlAllocator backed by a marl::Allocator.
37// Note: These may be re-implemented to optimize for marl's usage cases.
38// See: https://github.com/google/marl/issues/129
39////////////////////////////////////////////////////////////////////////////////
40template <typename T>
41using deque = std::deque<T, StlAllocator<T>>;
42
43template <typename K, typename V, typename C = std::less<K>>
44using map = std::map<K, V, C, StlAllocator<std::pair<const K, V>>>;
45
46template <typename K, typename C = std::less<K>>
47using set = std::set<K, C, StlAllocator<K>>;
48
49template <typename K,
50 typename V,
51 typename H = std::hash<K>,
52 typename E = std::equal_to<K>>
53using unordered_map =
54 std::unordered_map<K, V, H, E, StlAllocator<std::pair<const K, V>>>;
55
56template <typename K, typename H = std::hash<K>, typename E = std::equal_to<K>>
57using unordered_set = std::unordered_set<K, H, E, StlAllocator<K>>;
58
59// take() takes and returns the front value from the deque.
60template <typename T>
61MARL_NO_EXPORT inline T take(deque<T>& queue) {
62 auto out = std::move(queue.front());
63 queue.pop_front();
64 return out;
65}
66
67// take() takes and returns the first value from the unordered_set.
68template <typename T, typename H, typename E>
69MARL_NO_EXPORT inline T take(unordered_set<T, H, E>& set) {
70 auto it = set.begin();
71 auto out = std::move(*it);
72 set.erase(it);
73 return out;
74}
75
76////////////////////////////////////////////////////////////////////////////////
77// vector<T, BASE_CAPACITY>
78////////////////////////////////////////////////////////////////////////////////
79
80// vector is a container of contiguously stored elements.
81// Unlike std::vector, marl::containers::vector keeps the first
82// BASE_CAPACITY elements internally, which will avoid dynamic heap
83// allocations. Once the vector exceeds BASE_CAPACITY elements, vector will
84// allocate storage from the heap.
85template <typename T, int BASE_CAPACITY>
86class vector {
87 public:
88 MARL_NO_EXPORT inline vector(Allocator* allocator = Allocator::Default);
89
90 template <int BASE_CAPACITY_2>
91 MARL_NO_EXPORT inline vector(const vector<T, BASE_CAPACITY_2>& other,
92 Allocator* allocator = Allocator::Default);
93
94 template <int BASE_CAPACITY_2>
95 MARL_NO_EXPORT inline vector(vector<T, BASE_CAPACITY_2>&& other,
96 Allocator* allocator = Allocator::Default);
97
98 MARL_NO_EXPORT inline ~vector();
99
100 MARL_NO_EXPORT inline vector& operator=(const vector&);
101
102 template <int BASE_CAPACITY_2>
103 MARL_NO_EXPORT inline vector<T, BASE_CAPACITY>& operator=(
104 const vector<T, BASE_CAPACITY_2>&);
105
106 template <int BASE_CAPACITY_2>
107 MARL_NO_EXPORT inline vector<T, BASE_CAPACITY>& operator=(
108 vector<T, BASE_CAPACITY_2>&&);
109
110 MARL_NO_EXPORT inline void push_back(const T& el);
111 MARL_NO_EXPORT inline void emplace_back(T&& el);
112 MARL_NO_EXPORT inline void pop_back();
113 MARL_NO_EXPORT inline T& front();
114 MARL_NO_EXPORT inline T& back();
115 MARL_NO_EXPORT inline const T& front() const;
116 MARL_NO_EXPORT inline const T& back() const;
117 MARL_NO_EXPORT inline T* begin();
118 MARL_NO_EXPORT inline T* end();
119 MARL_NO_EXPORT inline const T* begin() const;
120 MARL_NO_EXPORT inline const T* end() const;
121 MARL_NO_EXPORT inline T& operator[](size_t i);
122 MARL_NO_EXPORT inline const T& operator[](size_t i) const;
123 MARL_NO_EXPORT inline size_t size() const;
124 MARL_NO_EXPORT inline size_t cap() const;
125 MARL_NO_EXPORT inline void resize(size_t n);
126 MARL_NO_EXPORT inline void reserve(size_t n);
127 MARL_NO_EXPORT inline T* data();
128 MARL_NO_EXPORT inline const T* data() const;
129
130 Allocator* const allocator;
131
132 private:
133 using TStorage = typename marl::aligned_storage<sizeof(T), alignof(T)>::type;
134
135 vector(const vector&) = delete;
136
137 MARL_NO_EXPORT inline void free();
138
139 size_t count = 0;
140 size_t capacity = BASE_CAPACITY;
141 TStorage buffer[BASE_CAPACITY];
142 TStorage* elements = buffer;
143 Allocation allocation;
144};
145
146template <typename T, int BASE_CAPACITY>
147vector<T, BASE_CAPACITY>::vector(
148 Allocator* allocator /* = Allocator::Default */)
149 : allocator(allocator) {}
150
151template <typename T, int BASE_CAPACITY>
152template <int BASE_CAPACITY_2>
153vector<T, BASE_CAPACITY>::vector(
154 const vector<T, BASE_CAPACITY_2>& other,
155 Allocator* allocator /* = Allocator::Default */)
156 : allocator(allocator) {
157 *this = other;
158}
159
160template <typename T, int BASE_CAPACITY>
161template <int BASE_CAPACITY_2>
162vector<T, BASE_CAPACITY>::vector(
163 vector<T, BASE_CAPACITY_2>&& other,
164 Allocator* allocator /* = Allocator::Default */)
165 : allocator(allocator) {
166 *this = std::move(other);
167}
168
169template <typename T, int BASE_CAPACITY>
170vector<T, BASE_CAPACITY>::~vector() {
171 free();
172}
173
174template <typename T, int BASE_CAPACITY>
175vector<T, BASE_CAPACITY>& vector<T, BASE_CAPACITY>::operator=(
176 const vector& other) {
177 free();
178 reserve(other.size());
179 count = other.size();
180 for (size_t i = 0; i < count; i++) {
181 new (&reinterpret_cast<T*>(elements)[i]) T(other[i]);
182 }
183 return *this;
184}
185
186template <typename T, int BASE_CAPACITY>
187template <int BASE_CAPACITY_2>
188vector<T, BASE_CAPACITY>& vector<T, BASE_CAPACITY>::operator=(
189 const vector<T, BASE_CAPACITY_2>& other) {
190 free();
191 reserve(other.size());
192 count = other.size();
193 for (size_t i = 0; i < count; i++) {
194 new (&reinterpret_cast<T*>(elements)[i]) T(other[i]);
195 }
196 return *this;
197}
198
199template <typename T, int BASE_CAPACITY>
200template <int BASE_CAPACITY_2>
201vector<T, BASE_CAPACITY>& vector<T, BASE_CAPACITY>::operator=(
202 vector<T, BASE_CAPACITY_2>&& other) {
203 free();
204 reserve(other.size());
205 count = other.size();
206 for (size_t i = 0; i < count; i++) {
207 new (&reinterpret_cast<T*>(elements)[i]) T(std::move(other[i]));
208 }
209 other.resize(0);
210 return *this;
211}
212
213template <typename T, int BASE_CAPACITY>
214void vector<T, BASE_CAPACITY>::push_back(const T& el) {
215 reserve(count + 1);
216 new (&reinterpret_cast<T*>(elements)[count]) T(el);
217 count++;
218}
219
220template <typename T, int BASE_CAPACITY>
221void vector<T, BASE_CAPACITY>::emplace_back(T&& el) {
222 reserve(count + 1);
223 new (&reinterpret_cast<T*>(elements)[count]) T(std::move(el));
224 count++;
225}
226
227template <typename T, int BASE_CAPACITY>
228void vector<T, BASE_CAPACITY>::pop_back() {
229 MARL_ASSERT(count > 0, "pop_back() called on empty vector");
230 count--;
231 reinterpret_cast<T*>(elements)[count].~T();
232}
233
234template <typename T, int BASE_CAPACITY>
235T& vector<T, BASE_CAPACITY>::front() {
236 MARL_ASSERT(count > 0, "front() called on empty vector");
237 return reinterpret_cast<T*>(elements)[0];
238}
239
240template <typename T, int BASE_CAPACITY>
241T& vector<T, BASE_CAPACITY>::back() {
242 MARL_ASSERT(count > 0, "back() called on empty vector");
243 return reinterpret_cast<T*>(elements)[count - 1];
244}
245
246template <typename T, int BASE_CAPACITY>
247const T& vector<T, BASE_CAPACITY>::front() const {
248 MARL_ASSERT(count > 0, "front() called on empty vector");
249 return reinterpret_cast<T*>(elements)[0];
250}
251
252template <typename T, int BASE_CAPACITY>
253const T& vector<T, BASE_CAPACITY>::back() const {
254 MARL_ASSERT(count > 0, "back() called on empty vector");
255 return reinterpret_cast<T*>(elements)[count - 1];
256}
257
258template <typename T, int BASE_CAPACITY>
259T* vector<T, BASE_CAPACITY>::begin() {
260 return reinterpret_cast<T*>(elements);
261}
262
263template <typename T, int BASE_CAPACITY>
264T* vector<T, BASE_CAPACITY>::end() {
265 return reinterpret_cast<T*>(elements) + count;
266}
267
268template <typename T, int BASE_CAPACITY>
269const T* vector<T, BASE_CAPACITY>::begin() const {
270 return reinterpret_cast<T*>(elements);
271}
272
273template <typename T, int BASE_CAPACITY>
274const T* vector<T, BASE_CAPACITY>::end() const {
275 return reinterpret_cast<T*>(elements) + count;
276}
277
278template <typename T, int BASE_CAPACITY>
279T& vector<T, BASE_CAPACITY>::operator[](size_t i) {
280 MARL_ASSERT(i < count, "index %d exceeds vector size %d", int(i), int(count));
281 return reinterpret_cast<T*>(elements)[i];
282}
283
284template <typename T, int BASE_CAPACITY>
285const T& vector<T, BASE_CAPACITY>::operator[](size_t i) const {
286 MARL_ASSERT(i < count, "index %d exceeds vector size %d", int(i), int(count));
287 return reinterpret_cast<T*>(elements)[i];
288}
289
290template <typename T, int BASE_CAPACITY>
291size_t vector<T, BASE_CAPACITY>::size() const {
292 return count;
293}
294
295template <typename T, int BASE_CAPACITY>
296void vector<T, BASE_CAPACITY>::resize(size_t n) {
297 reserve(n);
298 while (count < n) {
299 new (&reinterpret_cast<T*>(elements)[count++]) T();
300 }
301 while (n < count) {
302 reinterpret_cast<T*>(elements)[--count].~T();
303 }
304}
305
306template <typename T, int BASE_CAPACITY>
307void vector<T, BASE_CAPACITY>::reserve(size_t n) {
308 if (n > capacity) {
309 capacity = std::max<size_t>(n * 2, 8);
310
311 Allocation::Request request;
312 request.size = sizeof(T) * capacity;
313 request.alignment = alignof(T);
314 request.usage = Allocation::Usage::Vector;
315
316 auto alloc = allocator->allocate(request);
317 auto grown = reinterpret_cast<TStorage*>(alloc.ptr);
318 for (size_t i = 0; i < count; i++) {
319 new (&reinterpret_cast<T*>(grown)[i])
320 T(std::move(reinterpret_cast<T*>(elements)[i]));
321 }
322 free();
323 elements = grown;
324 allocation = alloc;
325 }
326}
327
328template <typename T, int BASE_CAPACITY>
329T* vector<T, BASE_CAPACITY>::data() {
330 return elements;
331}
332
333template <typename T, int BASE_CAPACITY>
334const T* vector<T, BASE_CAPACITY>::data() const {
335 return elements;
336}
337
338template <typename T, int BASE_CAPACITY>
339void vector<T, BASE_CAPACITY>::free() {
340 for (size_t i = 0; i < count; i++) {
341 reinterpret_cast<T*>(elements)[i].~T();
342 }
343
344 if (allocation.ptr != nullptr) {
345 allocator->free(allocation);
346 allocation = {};
347 elements = nullptr;
348 }
349}
350
351////////////////////////////////////////////////////////////////////////////////
352// list<T, BASE_CAPACITY>
353////////////////////////////////////////////////////////////////////////////////
354
355// list is a minimal std::list like container that supports constant time
356// insertion and removal of elements.
357// list keeps hold of allocations (it only releases allocations on destruction),
358// to avoid repeated heap allocations and frees when frequently inserting and
359// removing elements.
360template <typename T>
361class list {
362 struct Entry {
363 T data;
364 Entry* next;
365 Entry* prev;
366 };
367
368 public:
369 class iterator {
370 public:
371 MARL_NO_EXPORT inline iterator(Entry*);
372 MARL_NO_EXPORT inline T* operator->();
373 MARL_NO_EXPORT inline T& operator*();
374 MARL_NO_EXPORT inline iterator& operator++();
375 MARL_NO_EXPORT inline bool operator==(const iterator&) const;
376 MARL_NO_EXPORT inline bool operator!=(const iterator&) const;
377
378 private:
379 friend list;
380 Entry* entry;
381 };
382
383 MARL_NO_EXPORT inline list(Allocator* allocator = Allocator::Default);
384 MARL_NO_EXPORT inline ~list();
385
386 MARL_NO_EXPORT inline iterator begin();
387 MARL_NO_EXPORT inline iterator end();
388 MARL_NO_EXPORT inline size_t size() const;
389
390 template <typename... Args>
391 MARL_NO_EXPORT inline iterator emplace_front(Args&&... args);
392 MARL_NO_EXPORT inline void erase(iterator);
393
394 private:
395 // copy / move is currently unsupported.
396 list(const list&) = delete;
397 list(list&&) = delete;
398 list& operator=(const list&) = delete;
399 list& operator=(list&&) = delete;
400
401 struct AllocationChain {
402 Allocation allocation;
403 AllocationChain* next;
404 };
405
406 MARL_NO_EXPORT inline void grow(size_t count);
407
408 MARL_NO_EXPORT static inline void unlink(Entry* entry, Entry*& list);
409 MARL_NO_EXPORT static inline void link(Entry* entry, Entry*& list);
410
411 Allocator* const allocator;
412 size_t size_ = 0;
413 size_t capacity = 0;
414 AllocationChain* allocations = nullptr;
415 Entry* free = nullptr;
416 Entry* head = nullptr;
417};
418
419template <typename T>
420list<T>::iterator::iterator(Entry* entry) : entry(entry) {}
421
422template <typename T>
423T* list<T>::iterator::operator->() {
424 return &entry->data;
425}
426
427template <typename T>
428T& list<T>::iterator::operator*() {
429 return entry->data;
430}
431
432template <typename T>
433typename list<T>::iterator& list<T>::iterator::operator++() {
434 entry = entry->next;
435 return *this;
436}
437
438template <typename T>
439bool list<T>::iterator::operator==(const iterator& rhs) const {
440 return entry == rhs.entry;
441}
442
443template <typename T>
444bool list<T>::iterator::operator!=(const iterator& rhs) const {
445 return entry != rhs.entry;
446}
447
448template <typename T>
449list<T>::list(Allocator* allocator /* = Allocator::Default */)
450 : allocator(allocator) {}
451
452template <typename T>
453list<T>::~list() {
454 for (auto el = head; el != nullptr; el = el->next) {
455 el->data.~T();
456 }
457
458 auto curr = allocations;
459 while (curr != nullptr) {
460 auto next = curr->next;
461 allocator->free(curr->allocation);
462 curr = next;
463 }
464}
465
466template <typename T>
467typename list<T>::iterator list<T>::begin() {
468 return {head};
469}
470
471template <typename T>
472typename list<T>::iterator list<T>::end() {
473 return {nullptr};
474}
475
476template <typename T>
477size_t list<T>::size() const {
478 return size_;
479}
480
481template <typename T>
482template <typename... Args>
483typename list<T>::iterator list<T>::emplace_front(Args&&... args) {
484 if (free == nullptr) {
485 grow(std::max<size_t>(capacity, 8));
486 }
487
488 auto entry = free;
489
490 unlink(entry, free);
491 link(entry, head);
492
493 new (&entry->data) T(std::forward<T>(args)...);
494 size_++;
495
496 return entry;
497}
498
499template <typename T>
500void list<T>::erase(iterator it) {
501 auto entry = it.entry;
502 unlink(entry, head);
503 link(entry, free);
504
505 entry->data.~T();
506 size_--;
507}
508
509template <typename T>
510void list<T>::grow(size_t count) {
511 auto const entriesSize = sizeof(Entry) * count;
512 auto const allocChainOffset = alignUp(entriesSize, alignof(AllocationChain));
513 auto const allocSize = allocChainOffset + sizeof(AllocationChain);
514
515 Allocation::Request request;
516 request.size = allocSize;
517 request.alignment = std::max(alignof(Entry), alignof(AllocationChain));
518 request.usage = Allocation::Usage::List;
519 auto alloc = allocator->allocate(request);
520
521 auto entries = reinterpret_cast<Entry*>(alloc.ptr);
522 for (size_t i = 0; i < count; i++) {
523 auto entry = &entries[i];
524 entry->prev = nullptr;
525 entry->next = free;
526 if (free) {
527 free->prev = entry;
528 }
529 free = entry;
530 }
531
532 auto allocChain = reinterpret_cast<AllocationChain*>(
533 reinterpret_cast<uint8_t*>(alloc.ptr) + allocChainOffset);
534
535 allocChain->allocation = alloc;
536 allocChain->next = allocations;
537 allocations = allocChain;
538
539 capacity += count;
540}
541
542template <typename T>
543void list<T>::unlink(Entry* entry, Entry*& list) {
544 if (list == entry) {
545 list = list->next;
546 }
547 if (entry->prev) {
548 entry->prev->next = entry->next;
549 }
550 if (entry->next) {
551 entry->next->prev = entry->prev;
552 }
553 entry->prev = nullptr;
554 entry->next = nullptr;
555}
556
557template <typename T>
558void list<T>::link(Entry* entry, Entry*& list) {
559 MARL_ASSERT(entry->next == nullptr, "link() called on entry already linked");
560 MARL_ASSERT(entry->prev == nullptr, "link() called on entry already linked");
561 if (list) {
562 entry->next = list;
563 list->prev = entry;
564 }
565 list = entry;
566}
567
568} // namespace containers
569} // namespace marl
570
571#endif // marl_containers_h
572