xref: /llvm-project/mlir/include/mlir/IR/OwningOpRef.h (revision 89cfae41ecc043f8c47be4dea4b7c740d4f950b3)
1 //===- OwningOpRef.h - MLIR OwningOpRef -------------------------*- 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 // This file provides a base class for owning op refs.
10 //
11 //===----------------------------------------------------------------------===//
12 
13 #ifndef MLIR_IR_OWNINGOPREF_H
14 #define MLIR_IR_OWNINGOPREF_H
15 
16 #include <type_traits>
17 #include <utility>
18 
19 namespace mlir {
20 class Operation;
21 
22 /// This class acts as an owning reference to an op, and will automatically
23 /// destroy the held op on destruction if the held op is valid.
24 ///
25 /// Note that OpBuilder and related functionality should be highly preferred
26 /// instead, and this should only be used in situations where existing solutions
27 /// are not viable.
28 template <typename OpTy>
29 class OwningOpRef {
30 public:
31   /// The underlying operation type stored in this reference.
32   using OperationT = OpTy;
33 
op(nullptr)34   OwningOpRef(std::nullptr_t = nullptr) : op(nullptr) {}
OwningOpRef(OpTy op)35   OwningOpRef(OpTy op) : op(op) {}
OwningOpRef(OwningOpRef && other)36   OwningOpRef(OwningOpRef &&other) : op(other.release()) {}
~OwningOpRef()37   ~OwningOpRef() {
38     if (op)
39       op->erase();
40   }
41 
42   /// Assign from another op reference.
43   OwningOpRef &operator=(OwningOpRef &&other) {
44     if (op)
45       op->erase();
46     op = other.release();
47     return *this;
48   }
49 
50   /// Allow accessing the internal op.
get()51   OpTy get() const { return op; }
52   OpTy operator*() const { return op; }
53   auto operator->() {
54     // Specialize for the case where OpTy is a pointer, to allow using
55     // OwningOpRef<Operation*>.
56     if constexpr (std::is_pointer<OpTy>::value)
57       return op;
58     else
59       return &op;
60   }
61   explicit operator bool() const { return op; }
62 
63   /// Downcast to generic operation.
64   operator OwningOpRef<Operation *>() && { return release().getOperation(); }
65 
66   /// Release the referenced op.
release()67   OpTy release() {
68     OpTy released(nullptr);
69     std::swap(released, op);
70     return released;
71   }
72 
73 private:
74   OpTy op;
75 };
76 
77 } // namespace mlir
78 
79 #endif // MLIR_IR_OWNINGOPREF_H
80