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 | |
34 | OwningOpRef(std::nullptr_t = nullptr) : op(nullptr) {} |
35 | OwningOpRef(OpTy op) : op(op) {} |
36 | OwningOpRef(OwningOpRef &&other) : op(other.release()) {} |
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. |
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. |
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 | |