1*759a7b59SAman LaChapelle //===- LocationTest.cpp - unit tests for affine map API -------------------===// 2*759a7b59SAman LaChapelle // 3*759a7b59SAman LaChapelle // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. 4*759a7b59SAman LaChapelle // See https://llvm.org/LICENSE.txt for license information. 5*759a7b59SAman LaChapelle // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception 6*759a7b59SAman LaChapelle // 7*759a7b59SAman LaChapelle //===----------------------------------------------------------------------===// 8*759a7b59SAman LaChapelle 9*759a7b59SAman LaChapelle #include "mlir/IR/Location.h" 10*759a7b59SAman LaChapelle #include "mlir/IR/Builders.h" 11*759a7b59SAman LaChapelle #include "gtest/gtest.h" 12*759a7b59SAman LaChapelle 13*759a7b59SAman LaChapelle using namespace mlir; 14*759a7b59SAman LaChapelle 15*759a7b59SAman LaChapelle // Check that we only walk *locations* and not non-location attributes. 16*759a7b59SAman LaChapelle TEST(LocationTest, Walk) { 17*759a7b59SAman LaChapelle MLIRContext ctx; 18*759a7b59SAman LaChapelle Builder builder(&ctx); 19*759a7b59SAman LaChapelle BoolAttr trueAttr = builder.getBoolAttr(true); 20*759a7b59SAman LaChapelle 21*759a7b59SAman LaChapelle Location loc1 = FileLineColLoc::get(builder.getStringAttr("foo"), 1, 2); 22*759a7b59SAman LaChapelle Location loc2 = FileLineColLoc::get(builder.getStringAttr("foo"), 3, 4); 23*759a7b59SAman LaChapelle Location fused = builder.getFusedLoc({loc1, loc2}, trueAttr); 24*759a7b59SAman LaChapelle 25*759a7b59SAman LaChapelle SmallVector<Attribute> visited; 26*759a7b59SAman LaChapelle fused->walk([&](Location l) { 27*759a7b59SAman LaChapelle visited.push_back(LocationAttr(l)); 28*759a7b59SAman LaChapelle return WalkResult::advance(); 29*759a7b59SAman LaChapelle }); 30*759a7b59SAman LaChapelle 31*759a7b59SAman LaChapelle EXPECT_EQ(llvm::ArrayRef(visited), ArrayRef<Attribute>({fused, loc1, loc2})); 32*759a7b59SAman LaChapelle } 33*759a7b59SAman LaChapelle 34*759a7b59SAman LaChapelle // Check that we skip location attrs nested under a non-location attr. 35*759a7b59SAman LaChapelle TEST(LocationTest, SkipNested) { 36*759a7b59SAman LaChapelle MLIRContext ctx; 37*759a7b59SAman LaChapelle Builder builder(&ctx); 38*759a7b59SAman LaChapelle 39*759a7b59SAman LaChapelle Location loc1 = FileLineColLoc::get(builder.getStringAttr("foo"), 1, 2); 40*759a7b59SAman LaChapelle Location loc2 = FileLineColLoc::get(builder.getStringAttr("foo"), 3, 4); 41*759a7b59SAman LaChapelle Location loc3 = FileLineColLoc::get(builder.getStringAttr("bar"), 1, 2); 42*759a7b59SAman LaChapelle Location loc4 = FileLineColLoc::get(builder.getStringAttr("bar"), 3, 4); 43*759a7b59SAman LaChapelle ArrayAttr arr = builder.getArrayAttr({loc3, loc4}); 44*759a7b59SAman LaChapelle Location fused = builder.getFusedLoc({loc1, loc2}, arr); 45*759a7b59SAman LaChapelle 46*759a7b59SAman LaChapelle SmallVector<Attribute> visited; 47*759a7b59SAman LaChapelle fused->walk([&](Location l) { 48*759a7b59SAman LaChapelle visited.push_back(LocationAttr(l)); 49*759a7b59SAman LaChapelle return WalkResult::advance(); 50*759a7b59SAman LaChapelle }); 51*759a7b59SAman LaChapelle 52*759a7b59SAman LaChapelle EXPECT_EQ(llvm::ArrayRef(visited), ArrayRef<Attribute>({fused, loc1, loc2})); 53*759a7b59SAman LaChapelle } 54