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 MLIR_SUPPORT_ADTEXTRAS_H
11
12#include "mlir/Support/LLVM.h"
13#include "llvm/ADT/ArrayRef.h"
14#include "llvm/ADT/SmallVector.h"
15
16namespace 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.
24template <typename T>
25class CopyOnWriteArrayRef {
26public:
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
62private:
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

source code of mlir/include/mlir/Support/ADTExtras.h