54 lines · 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-exception6//7//===----------------------------------------------------------------------===//8 9#include "mlir/IR/Location.h"10#include "mlir/IR/Builders.h"11#include "gtest/gtest.h"12 13using namespace mlir;14 15// Check that we only walk *locations* and not non-location attributes.16TEST(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.35TEST(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}54