Revert "[mlir] Make single value ValueRanges memory safer (#121996)"

This reverts commit ab6e63a.

Author: cota

mlir/include/mlir/IR/TypeRange.h

@@ -29,12 +29,11 @@ namespace mlir {
 /// a SmallVector/std::vector. This class should be used in places that are not
 /// suitable for a more derived type (e.g. ArrayRef) or a template range
 /// parameter.
-class TypeRange
-    : public llvm::detail::indexed_accessor_range_base<
-          TypeRange,
-          llvm::PointerUnion<const Value *, const Type *, OpOperand *,
-                             detail::OpResultImpl *, Type>,
-          Type, Type, Type> {
+class TypeRange : public llvm::detail::indexed_accessor_range_base<
+                      TypeRange,
+                      llvm::PointerUnion<const Value *, const Type *,
+                                         OpOperand *, detail::OpResultImpl *>,
+                      Type, Type, Type> {
 public:
   using RangeBaseT::RangeBaseT;
   TypeRange(ArrayRef<Type> types = std::nullopt);
@@ -45,11 +44,8 @@ class TypeRange
   TypeRange(ValueTypeRange<ValueRangeT> values)
       : TypeRange(ValueRange(ValueRangeT(values.begin().getCurrent(),
                                          values.end().getCurrent()))) {}
-
-  TypeRange(Type type) : TypeRange(type, /*count=*/1) {}
-  template <typename Arg, typename = std::enable_if_t<
-                              std::is_constructible_v<ArrayRef<Type>, Arg> &&
-                              !std::is_constructible_v<Type, Arg>>>
+  template <typename Arg, typename = std::enable_if_t<std::is_constructible<
+                              ArrayRef<Type>, Arg>::value>>
   TypeRange(Arg &&arg) : TypeRange(ArrayRef<Type>(std::forward<Arg>(arg))) {}
   TypeRange(std::initializer_list<Type> types)
       : TypeRange(ArrayRef<Type>(types)) {}
@@ -60,9 +56,8 @@ class TypeRange
   /// * A pointer to the first element of an array of types.
   /// * A pointer to the first element of an array of operands.
   /// * A pointer to the first element of an array of results.
-  /// * A single 'Type' instance.
   using OwnerT = llvm::PointerUnion<const Value *, const Type *, OpOperand *,
-                                    detail::OpResultImpl *, Type>;
+                                    detail::OpResultImpl *>;
 
   /// See `llvm::detail::indexed_accessor_range_base` for details.
   static OwnerT offset_base(OwnerT object, ptrdiff_t index);

mlir/include/mlir/IR/ValueRange.h

@@ -374,16 +374,16 @@ class ResultRange::UseIterator final
 /// SmallVector/std::vector. This class should be used in places that are not
 /// suitable for a more derived type (e.g. ArrayRef) or a template range
 /// parameter.
-class ValueRange final : public llvm::detail::indexed_accessor_range_base<
-                             ValueRange,
-                             PointerUnion<const Value *, OpOperand *,
-                                          detail::OpResultImpl *, Value>,
-                             Value, Value, Value> {
+class ValueRange final
+    : public llvm::detail::indexed_accessor_range_base<
+          ValueRange,
+          PointerUnion<const Value *, OpOperand *, detail::OpResultImpl *>,
+          Value, Value, Value> {
 public:
   /// The type representing the owner of a ValueRange. This is either a list of
-  /// values, operands, or results or a single value.
+  /// values, operands, or results.
   using OwnerT =
-      PointerUnion<const Value *, OpOperand *, detail::OpResultImpl *, Value>;
+      PointerUnion<const Value *, OpOperand *, detail::OpResultImpl *>;
 
   using RangeBaseT::RangeBaseT;
 
@@ -392,7 +392,7 @@ class ValueRange final : public llvm::detail::indexed_accessor_range_base<
                 std::is_constructible<ArrayRef<Value>, Arg>::value &&
                 !std::is_convertible<Arg, Value>::value>>
   ValueRange(Arg &&arg) : ValueRange(ArrayRef<Value>(std::forward<Arg>(arg))) {}
-  ValueRange(Value value) : ValueRange(value, /*count=*/1) {}
+  ValueRange(const Value &value) : ValueRange(&value, /*count=*/1) {}
   ValueRange(const std::initializer_list<Value> &values)
       : ValueRange(ArrayRef<Value>(values)) {}
   ValueRange(iterator_range<OperandRange::iterator> values)

mlir/lib/IR/OperationSupport.cpp

@@ -653,15 +653,6 @@ ValueRange::ValueRange(ResultRange values)
 /// See `llvm::detail::indexed_accessor_range_base` for details.
 ValueRange::OwnerT ValueRange::offset_base(const OwnerT &owner,
                                            ptrdiff_t index) {
-  if (llvm::isa_and_nonnull<Value>(owner)) {
-    // Prevent out-of-bounds indexing for single values.
-    // Note that we do allow an index of 1 as is required by 'slice'ing that
-    // returns an empty range. This also matches the usual rules of C++ of being
-    // allowed to index past the last element of an array.
-    assert(index <= 1 && "out-of-bound offset into single-value 'ValueRange'");
-    // Return nullptr to quickly cause segmentation faults on misuse.
-    return index == 0 ? owner : nullptr;
-  }
   if (const auto *value = llvm::dyn_cast_if_present<const Value *>(owner))
     return {value + index};
   if (auto *operand = llvm::dyn_cast_if_present<OpOperand *>(owner))
@@ -670,10 +661,6 @@ ValueRange::OwnerT ValueRange::offset_base(const OwnerT &owner,
 }
 /// See `llvm::detail::indexed_accessor_range_base` for details.
 Value ValueRange::dereference_iterator(const OwnerT &owner, ptrdiff_t index) {
-  if (auto value = llvm::dyn_cast_if_present<Value>(owner)) {
-    assert(index == 0 && "cannot offset into single-value 'ValueRange'");
-    return value;
-  }
   if (const auto *value = llvm::dyn_cast_if_present<const Value *>(owner))
     return value[index];
   if (auto *operand = llvm::dyn_cast_if_present<OpOperand *>(owner))

mlir/lib/IR/TypeRange.cpp

@@ -31,23 +31,12 @@ TypeRange::TypeRange(ValueRange values) : TypeRange(OwnerT(), values.size()) {
     this->base = result;
   else if (auto *operand = llvm::dyn_cast_if_present<OpOperand *>(owner))
     this->base = operand;
-  else if (auto value = llvm::dyn_cast_if_present<Value>(owner))
-    this->base = value.getType();
   else
     this->base = cast<const Value *>(owner);
 }
 
 /// See `llvm::detail::indexed_accessor_range_base` for details.
 TypeRange::OwnerT TypeRange::offset_base(OwnerT object, ptrdiff_t index) {
-  if (llvm::isa_and_nonnull<Type>(object)) {
-    // Prevent out-of-bounds indexing for single values.
-    // Note that we do allow an index of 1 as is required by 'slice'ing that
-    // returns an empty range. This also matches the usual rules of C++ of being
-    // allowed to index past the last element of an array.
-    assert(index <= 1 && "out-of-bound offset into single-value 'ValueRange'");
-    // Return nullptr to quickly cause segmentation faults on misuse.
-    return index == 0 ? object : nullptr;
-  }
   if (const auto *value = llvm::dyn_cast_if_present<const Value *>(object))
     return {value + index};
   if (auto *operand = llvm::dyn_cast_if_present<OpOperand *>(object))
@@ -59,10 +48,6 @@ TypeRange::OwnerT TypeRange::offset_base(OwnerT object, ptrdiff_t index) {
 
 /// See `llvm::detail::indexed_accessor_range_base` for details.
 Type TypeRange::dereference_iterator(OwnerT object, ptrdiff_t index) {
-  if (auto type = llvm::dyn_cast_if_present<Type>(object)) {
-    assert(index == 0 && "cannot offset into single-value 'TypeRange'");
-    return type;
-  }
   if (const auto *value = llvm::dyn_cast_if_present<const Value *>(object))
     return (value + index)->getType();
   if (auto *operand = llvm::dyn_cast_if_present<OpOperand *>(object))

mlir/unittests/IR/OperationSupportTest.cpp

@@ -313,21 +313,4 @@ TEST(OperationEquivalenceTest, HashWorksWithFlags) {
   op2->destroy();
 }
 
-TEST(ValueRangeTest, ValueConstructable) {
-  MLIRContext context;
-  Builder builder(&context);
-
-  Operation *useOp =
-      createOp(&context, /*operands=*/std::nullopt, builder.getIntegerType(16));
-  // Valid construction despite a temporary 'OpResult'.
-  ValueRange operands = useOp->getResult(0);
-
-  useOp->setOperands(operands);
-  EXPECT_EQ(useOp->getNumOperands(), 1u);
-  EXPECT_EQ(useOp->getOperand(0), useOp->getResult(0));
-
-  useOp->dropAllUses();
-  useOp->destroy();
-}
-
 } // namespace