1 //===- TypeSwitch.h - Switch functionality for RTTI casting -*- 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 /// \file 10 /// This file implements the TypeSwitch template, which mimics a switch() 11 /// statement whose cases are type names. 12 /// 13 //===-----------------------------------------------------------------------===/ 14 15 #ifndef LLVM_ADT_TYPESWITCH_H 16 #define LLVM_ADT_TYPESWITCH_H 17 18 #include "llvm/ADT/STLExtras.h" 19 #include "llvm/Support/Casting.h" 20 #include <optional> 21 22 namespace llvm { 23 namespace detail { 24 25 template <typename DerivedT, typename T> class TypeSwitchBase { 26 public: TypeSwitchBase(const T & value)27 TypeSwitchBase(const T &value) : value(value) {} TypeSwitchBase(TypeSwitchBase && other)28 TypeSwitchBase(TypeSwitchBase &&other) : value(other.value) {} 29 ~TypeSwitchBase() = default; 30 31 /// TypeSwitchBase is not copyable. 32 TypeSwitchBase(const TypeSwitchBase &) = delete; 33 void operator=(const TypeSwitchBase &) = delete; 34 void operator=(TypeSwitchBase &&other) = delete; 35 36 /// Invoke a case on the derived class with multiple case types. 37 template <typename CaseT, typename CaseT2, typename... CaseTs, 38 typename CallableT> 39 // This is marked always_inline and nodebug so it doesn't show up in stack 40 // traces at -O0 (or other optimization levels). Large TypeSwitch's are 41 // common, are equivalent to a switch, and don't add any value to stack 42 // traces. 43 LLVM_ATTRIBUTE_ALWAYS_INLINE LLVM_ATTRIBUTE_NODEBUG DerivedT & Case(CallableT && caseFn)44 Case(CallableT &&caseFn) { 45 DerivedT &derived = static_cast<DerivedT &>(*this); 46 return derived.template Case<CaseT>(caseFn) 47 .template Case<CaseT2, CaseTs...>(caseFn); 48 } 49 50 /// Invoke a case on the derived class, inferring the type of the Case from 51 /// the first input of the given callable. 52 /// Note: This inference rules for this overload are very simple: strip 53 /// pointers and references. Case(CallableT && caseFn)54 template <typename CallableT> DerivedT &Case(CallableT &&caseFn) { 55 using Traits = function_traits<std::decay_t<CallableT>>; 56 using CaseT = std::remove_cv_t<std::remove_pointer_t< 57 std::remove_reference_t<typename Traits::template arg_t<0>>>>; 58 59 DerivedT &derived = static_cast<DerivedT &>(*this); 60 return derived.template Case<CaseT>(std::forward<CallableT>(caseFn)); 61 } 62 63 protected: 64 /// Attempt to dyn_cast the given `value` to `CastT`. 65 template <typename CastT, typename ValueT> decltype(auto)66 static decltype(auto) castValue(ValueT &&value) { 67 return dyn_cast<CastT>(value); 68 } 69 70 /// The root value we are switching on. 71 const T value; 72 }; 73 } // end namespace detail 74 75 /// This class implements a switch-like dispatch statement for a value of 'T' 76 /// using dyn_cast functionality. Each `Case<T>` takes a callable to be invoked 77 /// if the root value isa<T>, the callable is invoked with the result of 78 /// dyn_cast<T>() as a parameter. 79 /// 80 /// Example: 81 /// Operation *op = ...; 82 /// LogicalResult result = TypeSwitch<Operation *, LogicalResult>(op) 83 /// .Case<ConstantOp>([](ConstantOp op) { ... }) 84 /// .Default([](Operation *op) { ... }); 85 /// 86 template <typename T, typename ResultT = void> 87 class TypeSwitch : public detail::TypeSwitchBase<TypeSwitch<T, ResultT>, T> { 88 public: 89 using BaseT = detail::TypeSwitchBase<TypeSwitch<T, ResultT>, T>; 90 using BaseT::BaseT; 91 using BaseT::Case; 92 TypeSwitch(TypeSwitch &&other) = default; 93 94 /// Add a case on the given type. 95 template <typename CaseT, typename CallableT> Case(CallableT && caseFn)96 TypeSwitch<T, ResultT> &Case(CallableT &&caseFn) { 97 if (result) 98 return *this; 99 100 // Check to see if CaseT applies to 'value'. 101 if (auto caseValue = BaseT::template castValue<CaseT>(this->value)) 102 result.emplace(caseFn(caseValue)); 103 return *this; 104 } 105 106 /// As a default, invoke the given callable within the root value. 107 template <typename CallableT> Default(CallableT && defaultFn)108 [[nodiscard]] ResultT Default(CallableT &&defaultFn) { 109 if (result) 110 return std::move(*result); 111 return defaultFn(this->value); 112 } 113 /// As a default, return the given value. Default(ResultT defaultResult)114 [[nodiscard]] ResultT Default(ResultT defaultResult) { 115 if (result) 116 return std::move(*result); 117 return defaultResult; 118 } 119 ResultT()120 [[nodiscard]] operator ResultT() { 121 assert(result && "Fell off the end of a type-switch"); 122 return std::move(*result); 123 } 124 125 private: 126 /// The pointer to the result of this switch statement, once known, 127 /// null before that. 128 std::optional<ResultT> result; 129 }; 130 131 /// Specialization of TypeSwitch for void returning callables. 132 template <typename T> 133 class TypeSwitch<T, void> 134 : public detail::TypeSwitchBase<TypeSwitch<T, void>, T> { 135 public: 136 using BaseT = detail::TypeSwitchBase<TypeSwitch<T, void>, T>; 137 using BaseT::BaseT; 138 using BaseT::Case; 139 TypeSwitch(TypeSwitch &&other) = default; 140 141 /// Add a case on the given type. 142 template <typename CaseT, typename CallableT> Case(CallableT && caseFn)143 TypeSwitch<T, void> &Case(CallableT &&caseFn) { 144 if (foundMatch) 145 return *this; 146 147 // Check to see if any of the types apply to 'value'. 148 if (auto caseValue = BaseT::template castValue<CaseT>(this->value)) { 149 caseFn(caseValue); 150 foundMatch = true; 151 } 152 return *this; 153 } 154 155 /// As a default, invoke the given callable within the root value. Default(CallableT && defaultFn)156 template <typename CallableT> void Default(CallableT &&defaultFn) { 157 if (!foundMatch) 158 defaultFn(this->value); 159 } 160 161 private: 162 /// A flag detailing if we have already found a match. 163 bool foundMatch = false; 164 }; 165 } // end namespace llvm 166 167 #endif // LLVM_ADT_TYPESWITCH_H 168