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 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: CopyOnWriteArrayRef(ArrayRef<T> array)27 CopyOnWriteArrayRef(ArrayRef<T> array) : nonOwning(array){}; 28 29 CopyOnWriteArrayRef &operator=(ArrayRef<T> array) { 30 nonOwning = array; 31 owningStorage = {}; 32 return *this; 33 } 34 insert(size_t index,T value)35 void insert(size_t index, T value) { 36 SmallVector<T> &vector = ensureCopy(); 37 vector.insert(vector.begin() + index, value); 38 } 39 erase(size_t index)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 set(size_t index,T value)52 void set(size_t index, T value) { ensureCopy()[index] = value; } 53 size()54 size_t size() const { return ArrayRef<T>(*this).size(); } 55 empty()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: isNonOwning()63 bool isNonOwning() const { return !nonOwning.empty(); } 64 ensureCopy()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