mlir/unittests/IR/LocationTest.cpp

   1 //===- LocationTest.cpp - unit tests for affine map API -------------------===//
   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 #include "mlir/IR/Location.h"
  10 #include "mlir/IR/Builders.h"
  11 #include "gtest/gtest.h"
  12
  13 using namespace mlir;
  14
  15 // Check that we only walk *locations* and not non-location attributes.
  16 TEST(LocationTest, Walk) {
  17   MLIRContext ctx;
  18   Builder builder(&ctx);
  19   BoolAttr trueAttr = builder.getBoolAttr(true);
  20
  21   Location loc1 = FileLineColLoc::get(builder.getStringAttr("foo"), 1, 2);
  22   Location loc2 = FileLineColLoc::get(builder.getStringAttr("foo"), 3, 4);
  23   Location fused = builder.getFusedLoc({loc1, loc2}, trueAttr);
  24
  25   SmallVector<Attribute> visited;
  26   fused->walk([&](Location l) {
  27     visited.push_back(LocationAttr(l));
  28     return WalkResult::advance();
  29   });
  30
  31   EXPECT_EQ(llvm::ArrayRef(visited), ArrayRef<Attribute>({fused, loc1, loc2}));
  32 }
  33
  34 // Check that we skip location attrs nested under a non-location attr.
  35 TEST(LocationTest, SkipNested) {
  36   MLIRContext ctx;
  37   Builder builder(&ctx);
  38
  39   Location loc1 = FileLineColLoc::get(builder.getStringAttr("foo"), 1, 2);
  40   Location loc2 = FileLineColLoc::get(builder.getStringAttr("foo"), 3, 4);
  41   Location loc3 = FileLineColLoc::get(builder.getStringAttr("bar"), 1, 2);
  42   Location loc4 = FileLineColLoc::get(builder.getStringAttr("bar"), 3, 4);
  43   ArrayAttr arr = builder.getArrayAttr({loc3, loc4});
  44   Location fused = builder.getFusedLoc({loc1, loc2}, arr);
  45
  46   SmallVector<Attribute> visited;
  47   fused->walk([&](Location l) {
  48     visited.push_back(LocationAttr(l));
  49     return WalkResult::advance();
  50   });
  51
  52   EXPECT_EQ(llvm::ArrayRef(visited), ArrayRef<Attribute>({fused, loc1, loc2}));
  53 }