1 | //===- ADTExtras.h - Extra ADTs for use in MLIR -----------------*- C++ -*-===// |
2 | // |
3 | // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. |
4 | // See https://llvm.org/LICENSE.txt for license information. |
5 | // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception |
6 | // |
7 | //===----------------------------------------------------------------------===// |
8 | |
9 | #ifndef MLIR_SUPPORT_ADTEXTRAS_H |
10 | #define |
11 | |
12 | #include "mlir/Support/LLVM.h" |
13 | #include "llvm/ADT/ArrayRef.h" |
14 | #include "llvm/ADT/SmallVector.h" |
15 | |
16 | namespace mlir { |
17 | |
18 | //===----------------------------------------------------------------------===// |
19 | // CopyOnWriteArrayRef<T> |
20 | //===----------------------------------------------------------------------===// |
21 | |
22 | // A wrapper around an ArrayRef<T> that copies to a SmallVector<T> on |
23 | // modification. This is for use in the mlir::<Type>::Builders. |
24 | template <typename T> |
25 | class CopyOnWriteArrayRef { |
26 | public: |
27 | CopyOnWriteArrayRef(ArrayRef<T> array) : nonOwning(array){}; |
28 | |
29 | CopyOnWriteArrayRef &operator=(ArrayRef<T> array) { |
30 | nonOwning = array; |
31 | owningStorage = {}; |
32 | return *this; |
33 | } |
34 | |
35 | void insert(size_t index, T value) { |
36 | SmallVector<T> &vector = ensureCopy(); |
37 | vector.insert(vector.begin() + index, value); |
38 | } |
39 | |
40 | void erase(size_t index) { |
41 | // Note: A copy can be avoided when just dropping the front/back dims. |
42 | if (isNonOwning() && index == 0) { |
43 | nonOwning = nonOwning.drop_front(); |
44 | } else if (isNonOwning() && index == size() - 1) { |
45 | nonOwning = nonOwning.drop_back(); |
46 | } else { |
47 | SmallVector<T> &vector = ensureCopy(); |
48 | vector.erase(vector.begin() + index); |
49 | } |
50 | } |
51 | |
52 | void set(size_t index, T value) { ensureCopy()[index] = value; } |
53 | |
54 | size_t size() const { return ArrayRef<T>(*this).size(); } |
55 | |
56 | bool empty() const { return ArrayRef<T>(*this).empty(); } |
57 | |
58 | operator ArrayRef<T>() const { |
59 | return nonOwning.empty() ? ArrayRef<T>(owningStorage) : nonOwning; |
60 | } |
61 | |
62 | private: |
63 | bool isNonOwning() const { return !nonOwning.empty(); } |
64 | |
65 | SmallVector<T> &ensureCopy() { |
66 | // Empty non-owning storage signals the array has been copied to the owning |
67 | // storage (or both are empty). Note: `nonOwning` should never reference |
68 | // `owningStorage`. This can lead to dangling references if the |
69 | // CopyOnWriteArrayRef<T> is copied. |
70 | if (isNonOwning()) { |
71 | owningStorage = SmallVector<T>(nonOwning); |
72 | nonOwning = {}; |
73 | } |
74 | return owningStorage; |
75 | } |
76 | |
77 | ArrayRef<T> nonOwning; |
78 | SmallVector<T> owningStorage; |
79 | }; |
80 | |
81 | } // namespace mlir |
82 | |
83 | #endif |
84 | |