[Shape] Simplify getShapeVec a bit. NFCI.

diff --git a/mlir/lib/Dialect/Shape/IR/Shape.cpp b/mlir/lib/Dialect/Shape/IR/Shape.cpp
index 89a9766..5c851f5 100644 (file)
--- a/mlir/lib/Dialect/Shape/IR/Shape.cpp
+++ b/mlir/lib/Dialect/Shape/IR/Shape.cpp
@@ -47,19 +47,15 @@ bool shape::isExtentTensorType(Type type) {
 LogicalResult shape::getShapeVec(Value input,
                                  SmallVectorImpl<int64_t> &shapeValues) {
   if (auto inputOp = input.getDefiningOp<ShapeOfOp>()) {
-    auto type = inputOp.getArg().getType().dyn_cast<ShapedType>();
+    auto type = inputOp.getArg().getType().cast<ShapedType>();
     if (!type.hasRank())
       return failure();
-    shapeValues = llvm::to_vector<6>(type.getShape());
+    llvm::append_range(shapeValues, type.getShape());
     return success();
   }
-  if (auto inputOp = input.getDefiningOp<ConstShapeOp>()) {
-    shapeValues = llvm::to_vector<6>(inputOp.getShape().getValues<int64_t>());
-    return success();
-  }
-  if (auto inputOp = input.getDefiningOp<arith::ConstantOp>()) {
-    shapeValues = llvm::to_vector<6>(
-        inputOp.getValue().cast<DenseIntElementsAttr>().getValues<int64_t>());
+  DenseIntElementsAttr attr;
+  if (matchPattern(input, m_Constant(&attr))) {
+    llvm::append_range(shapeValues, attr.getValues<int64_t>());
     return success();
   }
   return failure();