[Mlir-commits] [mlir] [mlir][CAPI][python] bind CallSiteLoc, FileLineColRange, FusedLoc, NameLoc, OpaqueLoc (PR #129351)

Maksim Levental llvmlistbot at llvm.org
Sat Mar 8 19:47:30 PST 2025


https://github.com/makslevental updated https://github.com/llvm/llvm-project/pull/129351

>From f595ba3544ce9307756010a4729da83e224ed48a Mon Sep 17 00:00:00 2001
From: Maksim Levental <maksim.levental at gmail.com>
Date: Fri, 28 Feb 2025 22:05:45 -0500
Subject: [PATCH] [mlir][CAPI][python] bind CallSiteLoc, FileLineColRange,
 FusedLoc, NameLoc, OpaqueLoc

---
 mlir/include/mlir-c/IR.h                      |  80 ++++++++
 .../mlir/Bindings/Python/NanobindAdaptors.h   |  12 +-
 mlir/lib/Bindings/Python/IRCore.cpp           |  51 ++++-
 mlir/lib/CAPI/IR/IR.cpp                       | 100 +++++++++-
 mlir/test/python/ir/location.py               | 174 +++++++++++++-----
 5 files changed, 364 insertions(+), 53 deletions(-)

diff --git a/mlir/include/mlir-c/IR.h b/mlir/include/mlir-c/IR.h
index 7d2fd89e8560f..b0b8e63fecec7 100644
--- a/mlir/include/mlir-c/IR.h
+++ b/mlir/include/mlir-c/IR.h
@@ -261,15 +261,75 @@ MLIR_CAPI_EXPORTED MlirLocation mlirLocationFileLineColRangeGet(
     MlirContext context, MlirStringRef filename, unsigned start_line,
     unsigned start_col, unsigned end_line, unsigned end_col);
 
+/// Getter for filename of FileLineColRange.
+MLIR_CAPI_EXPORTED MlirIdentifier
+mlirLocationFileLineColRangeGetFilename(MlirLocation location);
+
+/// Getter for start_line of FileLineColRange.
+MLIR_CAPI_EXPORTED int
+mlirLocationFileLineColRangeGetStartLine(MlirLocation location);
+
+/// Getter for start_column of FileLineColRange.
+MLIR_CAPI_EXPORTED int
+mlirLocationFileLineColRangeGetStartColumn(MlirLocation location);
+
+/// Getter for end_line of FileLineColRange.
+MLIR_CAPI_EXPORTED int
+mlirLocationFileLineColRangeGetEndLine(MlirLocation location);
+
+/// Getter for end_column of FileLineColRange.
+MLIR_CAPI_EXPORTED int
+mlirLocationFileLineColRangeGetEndColumn(MlirLocation location);
+
+/// TypeID Getter for FileLineColRange.
+MLIR_CAPI_EXPORTED MlirTypeID mlirLocationFileLineColRangeGetTypeID(void);
+
+/// Checks whether the given location is an FileLineColRange.
+MLIR_CAPI_EXPORTED bool mlirLocationIsAFileLineColRange(MlirLocation location);
+
 /// Creates a call site location with a callee and a caller.
 MLIR_CAPI_EXPORTED MlirLocation mlirLocationCallSiteGet(MlirLocation callee,
                                                         MlirLocation caller);
 
+/// Getter for callee of CallSite.
+MLIR_CAPI_EXPORTED MlirLocation
+mlirLocationCallSiteGetCallee(MlirLocation location);
+
+/// Getter for caller of CallSite.
+MLIR_CAPI_EXPORTED MlirLocation
+mlirLocationCallSiteGetCaller(MlirLocation location);
+
+/// TypeID Getter for CallSite.
+MLIR_CAPI_EXPORTED MlirTypeID mlirLocationCallSiteGetTypeID(void);
+
+/// Checks whether the given location is an CallSite.
+MLIR_CAPI_EXPORTED bool mlirLocationIsACallSite(MlirLocation location);
+
 /// Creates a fused location with an array of locations and metadata.
 MLIR_CAPI_EXPORTED MlirLocation
 mlirLocationFusedGet(MlirContext ctx, intptr_t nLocations,
                      MlirLocation const *locations, MlirAttribute metadata);
 
+/// Getter for number of locations fused together.
+MLIR_CAPI_EXPORTED unsigned
+mlirLocationFusedGetNumLocations(MlirLocation location);
+
+/// Getter for locations of Fused. Requires pre-allocated memory of
+/// #fusedLocations X sizeof(MlirLocation).
+MLIR_CAPI_EXPORTED void
+mlirLocationFusedGetLocations(MlirLocation location,
+                              MlirLocation *locationsCPtr);
+
+/// Getter for metadata of Fused.
+MLIR_CAPI_EXPORTED MlirAttribute
+mlirLocationFusedGetMetadata(MlirLocation location);
+
+/// TypeID Getter for Fused.
+MLIR_CAPI_EXPORTED MlirTypeID mlirLocationFusedGetTypeID(void);
+
+/// Checks whether the given location is an Fused.
+MLIR_CAPI_EXPORTED bool mlirLocationIsAFused(MlirLocation location);
+
 /// Creates a name location owned by the given context. Providing null location
 /// for childLoc is allowed and if childLoc is null location, then the behavior
 /// is the same as having unknown child location.
@@ -277,6 +337,20 @@ MLIR_CAPI_EXPORTED MlirLocation mlirLocationNameGet(MlirContext context,
                                                     MlirStringRef name,
                                                     MlirLocation childLoc);
 
+/// Getter for name of Name.
+MLIR_CAPI_EXPORTED MlirIdentifier
+mlirLocationNameGetName(MlirLocation location);
+
+/// Getter for childLoc of Name.
+MLIR_CAPI_EXPORTED MlirLocation
+mlirLocationNameGetChildLoc(MlirLocation location);
+
+/// TypeID Getter for Name.
+MLIR_CAPI_EXPORTED MlirTypeID mlirLocationNameGetTypeID(void);
+
+/// Checks whether the given location is an Name.
+MLIR_CAPI_EXPORTED bool mlirLocationIsAName(MlirLocation location);
+
 /// Creates a location with unknown position owned by the given context.
 MLIR_CAPI_EXPORTED MlirLocation mlirLocationUnknownGet(MlirContext context);
 
@@ -970,6 +1044,12 @@ mlirValueReplaceAllUsesExcept(MlirValue of, MlirValue with,
                               intptr_t numExceptions,
                               MlirOperation *exceptions);
 
+/// Gets the location of the value.
+MLIR_CAPI_EXPORTED MlirLocation mlirValueGetLocation(MlirValue v);
+
+/// Gets the context that a value was created with.
+MLIR_CAPI_EXPORTED MlirContext mlirValueGetContext(MlirValue v);
+
 //===----------------------------------------------------------------------===//
 // OpOperand API.
 //===----------------------------------------------------------------------===//
diff --git a/mlir/include/mlir/Bindings/Python/NanobindAdaptors.h b/mlir/include/mlir/Bindings/Python/NanobindAdaptors.h
index 517351cac6dbc..78d13b725bb7f 100644
--- a/mlir/include/mlir/Bindings/Python/NanobindAdaptors.h
+++ b/mlir/include/mlir/Bindings/Python/NanobindAdaptors.h
@@ -21,10 +21,10 @@
 
 #include <cstdint>
 
+#include "mlir-c/Bindings/Python/Interop.h" // This is expected after nanobind.
 #include "mlir-c/Diagnostics.h"
 #include "mlir-c/IR.h"
 #include "mlir/Bindings/Python/Nanobind.h"
-#include "mlir-c/Bindings/Python/Interop.h" // This is expected after nanobind.
 #include "llvm/ADT/Twine.h"
 
 // Raw CAPI type casters need to be declared before use, so always include them
@@ -319,6 +319,16 @@ struct type_caster<MlirType> {
   }
 };
 
+/// Casts MlirStringRef -> object.
+template <>
+struct type_caster<MlirStringRef> {
+  NB_TYPE_CASTER(MlirStringRef, const_name("MlirStringRef"))
+  static handle from_cpp(MlirStringRef s, rv_policy,
+                         cleanup_list *cleanup) noexcept {
+    return nanobind::str(s.data, s.length).release();
+  }
+};
+
 } // namespace detail
 } // namespace nanobind
 
diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index 8e351cb22eb94..a4c517e859a35 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -12,6 +12,7 @@
 #include "Globals.h"
 #include "IRModule.h"
 #include "NanobindUtils.h"
+#include "mlir-c/Bindings/Python/Interop.h" // This is expected after nanobind.
 #include "mlir-c/BuiltinAttributes.h"
 #include "mlir-c/Debug.h"
 #include "mlir-c/Diagnostics.h"
@@ -19,7 +20,6 @@
 #include "mlir-c/Support.h"
 #include "mlir/Bindings/Python/Nanobind.h"
 #include "mlir/Bindings/Python/NanobindAdaptors.h"
-#include "mlir-c/Bindings/Python/Interop.h" // This is expected after nanobind.
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/SmallVector.h"
 
@@ -299,7 +299,7 @@ struct PyAttrBuilderMap {
     return *builder;
   }
   static void dunderSetItemNamed(const std::string &attributeKind,
-                                nb::callable func, bool replace) {
+                                 nb::callable func, bool replace) {
     PyGlobals::get().registerAttributeBuilder(attributeKind, std::move(func),
                                               replace);
   }
@@ -2933,6 +2933,9 @@ void mlir::python::populateIRCore(nb::module_ &m) {
           nb::arg("callee"), nb::arg("frames"),
           nb::arg("context").none() = nb::none(),
           kContextGetCallSiteLocationDocstring)
+      .def("is_a_callsite", mlirLocationIsACallSite)
+      .def_prop_ro("callee", mlirLocationCallSiteGetCallee)
+      .def_prop_ro("caller", mlirLocationCallSiteGetCaller)
       .def_static(
           "file",
           [](std::string filename, int line, int col,
@@ -2957,6 +2960,16 @@ void mlir::python::populateIRCore(nb::module_ &m) {
           nb::arg("filename"), nb::arg("start_line"), nb::arg("start_col"),
           nb::arg("end_line"), nb::arg("end_col"),
           nb::arg("context").none() = nb::none(), kContextGetFileRangeDocstring)
+      .def("is_a_file", mlirLocationIsAFileLineColRange)
+      .def_prop_ro("filename",
+                   [](MlirLocation loc) {
+                     return mlirIdentifierStr(
+                         mlirLocationFileLineColRangeGetFilename(loc));
+                   })
+      .def_prop_ro("start_line", mlirLocationFileLineColRangeGetStartLine)
+      .def_prop_ro("start_col", mlirLocationFileLineColRangeGetStartColumn)
+      .def_prop_ro("end_line", mlirLocationFileLineColRangeGetEndLine)
+      .def_prop_ro("end_col", mlirLocationFileLineColRangeGetEndColumn)
       .def_static(
           "fused",
           [](const std::vector<PyLocation> &pyLocations,
@@ -2974,6 +2987,16 @@ void mlir::python::populateIRCore(nb::module_ &m) {
           nb::arg("locations"), nb::arg("metadata").none() = nb::none(),
           nb::arg("context").none() = nb::none(),
           kContextGetFusedLocationDocstring)
+      .def("is_a_fused", mlirLocationIsAFused)
+      .def_prop_ro("locations",
+                   [](MlirLocation loc) {
+                     unsigned numLocations =
+                         mlirLocationFusedGetNumLocations(loc);
+                     std::vector<MlirLocation> locations(numLocations);
+                     if (numLocations)
+                       mlirLocationFusedGetLocations(loc, locations.data());
+                     return locations;
+                   })
       .def_static(
           "name",
           [](std::string name, std::optional<PyLocation> childLoc,
@@ -2988,6 +3011,12 @@ void mlir::python::populateIRCore(nb::module_ &m) {
           nb::arg("name"), nb::arg("childLoc").none() = nb::none(),
           nb::arg("context").none() = nb::none(),
           kContextGetNameLocationDocString)
+      .def("is_a_name", mlirLocationIsAName)
+      .def_prop_ro("name_str",
+                   [](MlirLocation loc) {
+                     return mlirIdentifierStr(mlirLocationNameGetName(loc));
+                   })
+      .def_prop_ro("child_loc", mlirLocationNameGetChildLoc)
       .def_static(
           "from_attr",
           [](PyAttribute &attribute, DefaultingPyMlirContext context) {
@@ -3126,9 +3155,7 @@ void mlir::python::populateIRCore(nb::module_ &m) {
                      auto &concreteOperation = self.getOperation();
                      concreteOperation.checkValid();
                      MlirOperation operation = concreteOperation.get();
-                     MlirStringRef name =
-                         mlirIdentifierStr(mlirOperationGetName(operation));
-                     return nb::str(name.data, name.length);
+                     return mlirIdentifierStr(mlirOperationGetName(operation));
                    })
       .def_prop_ro("operands",
                    [](PyOperationBase &self) {
@@ -3713,8 +3740,7 @@ void mlir::python::populateIRCore(nb::module_ &m) {
       .def_prop_ro(
           "name",
           [](PyNamedAttribute &self) {
-            return nb::str(mlirIdentifierStr(self.namedAttr.name).data,
-                           mlirIdentifierStr(self.namedAttr.name).length);
+            return mlirIdentifierStr(self.namedAttr.name);
           },
           "The name of the NamedAttribute binding")
       .def_prop_ro(
@@ -3947,7 +3973,16 @@ void mlir::python::populateIRCore(nb::module_ &m) {
           nb::arg("with"), nb::arg("exceptions"),
           kValueReplaceAllUsesExceptDocstring)
       .def(MLIR_PYTHON_MAYBE_DOWNCAST_ATTR,
-           [](PyValue &self) { return self.maybeDownCast(); });
+           [](PyValue &self) { return self.maybeDownCast(); })
+      .def_prop_ro(
+          "location",
+          [](MlirValue self) {
+            return PyLocation(
+                PyMlirContext::forContext(mlirValueGetContext(self)),
+                mlirValueGetLocation(self));
+          },
+          "Returns the source location the value");
+
   PyBlockArgument::bind(m);
   PyOpResult::bind(m);
   PyOpOperand::bind(m);
diff --git a/mlir/lib/CAPI/IR/IR.cpp b/mlir/lib/CAPI/IR/IR.cpp
index f27af0ca9a2c7..b2a058892e8cd 100644
--- a/mlir/lib/CAPI/IR/IR.cpp
+++ b/mlir/lib/CAPI/IR/IR.cpp
@@ -254,7 +254,7 @@ MlirAttribute mlirLocationGetAttribute(MlirLocation location) {
 }
 
 MlirLocation mlirLocationFromAttribute(MlirAttribute attribute) {
-  return wrap(Location(llvm::cast<LocationAttr>(unwrap(attribute))));
+  return wrap(Location(llvm::dyn_cast<LocationAttr>(unwrap(attribute))));
 }
 
 MlirLocation mlirLocationFileLineColGet(MlirContext context,
@@ -273,10 +273,62 @@ mlirLocationFileLineColRangeGet(MlirContext context, MlirStringRef filename,
                                      startLine, startCol, endLine, endCol)));
 }
 
+MlirIdentifier mlirLocationFileLineColRangeGetFilename(MlirLocation location) {
+  return wrap(llvm::dyn_cast<FileLineColRange>(unwrap(location)).getFilename());
+}
+
+int mlirLocationFileLineColRangeGetStartLine(MlirLocation location) {
+  if (auto loc = llvm::dyn_cast<FileLineColRange>(unwrap(location)))
+    return loc.getStartLine();
+  return -1;
+}
+
+int mlirLocationFileLineColRangeGetStartColumn(MlirLocation location) {
+  if (auto loc = llvm::dyn_cast<FileLineColRange>(unwrap(location)))
+    return loc.getStartColumn();
+  return -1;
+}
+
+int mlirLocationFileLineColRangeGetEndLine(MlirLocation location) {
+  if (auto loc = llvm::dyn_cast<FileLineColRange>(unwrap(location)))
+    return loc.getEndLine();
+  return -1;
+}
+
+int mlirLocationFileLineColRangeGetEndColumn(MlirLocation location) {
+  if (auto loc = llvm::dyn_cast<FileLineColRange>(unwrap(location)))
+    return loc.getEndColumn();
+  return -1;
+}
+
+MlirTypeID mlirLocationFileLineColRangeGetTypeID() {
+  return wrap(FileLineColRange::getTypeID());
+}
+
+bool mlirLocationIsAFileLineColRange(MlirLocation location) {
+  return isa<FileLineColRange>(unwrap(location));
+}
+
 MlirLocation mlirLocationCallSiteGet(MlirLocation callee, MlirLocation caller) {
   return wrap(Location(CallSiteLoc::get(unwrap(callee), unwrap(caller))));
 }
 
+MlirLocation mlirLocationCallSiteGetCallee(MlirLocation location) {
+  return wrap(Location(llvm::cast<CallSiteLoc>(unwrap(location)).getCallee()));
+}
+
+MlirLocation mlirLocationCallSiteGetCaller(MlirLocation location) {
+  return wrap(Location(llvm::cast<CallSiteLoc>(unwrap(location)).getCaller()));
+}
+
+MlirTypeID mlirLocationCallSiteGetTypeID() {
+  return wrap(CallSiteLoc::getTypeID());
+}
+
+bool mlirLocationIsACallSite(MlirLocation location) {
+  return isa<CallSiteLoc>(unwrap(location));
+}
+
 MlirLocation mlirLocationFusedGet(MlirContext ctx, intptr_t nLocations,
                                   MlirLocation const *locations,
                                   MlirAttribute metadata) {
@@ -285,6 +337,30 @@ MlirLocation mlirLocationFusedGet(MlirContext ctx, intptr_t nLocations,
   return wrap(FusedLoc::get(unwrappedLocs, unwrap(metadata), unwrap(ctx)));
 }
 
+unsigned mlirLocationFusedGetNumLocations(MlirLocation location) {
+  if (auto locationsArrRef = llvm::dyn_cast<FusedLoc>(unwrap(location)))
+    return locationsArrRef.getLocations().size();
+  return 0;
+}
+
+void mlirLocationFusedGetLocations(MlirLocation location,
+                                   MlirLocation *locationsCPtr) {
+  if (auto locationsArrRef = llvm::dyn_cast<FusedLoc>(unwrap(location))) {
+    for (auto [i, location] : llvm::enumerate(locationsArrRef.getLocations()))
+      locationsCPtr[i] = wrap(location);
+  }
+}
+
+MlirAttribute mlirLocationFusedGetMetadata(MlirLocation location) {
+  return wrap(llvm::cast<FusedLoc>(unwrap(location)).getMetadata());
+}
+
+MlirTypeID mlirLocationFusedGetTypeID() { return wrap(FusedLoc::getTypeID()); }
+
+bool mlirLocationIsAFused(MlirLocation location) {
+  return isa<FusedLoc>(unwrap(location));
+}
+
 MlirLocation mlirLocationNameGet(MlirContext context, MlirStringRef name,
                                  MlirLocation childLoc) {
   if (mlirLocationIsNull(childLoc))
@@ -294,6 +370,20 @@ MlirLocation mlirLocationNameGet(MlirContext context, MlirStringRef name,
       StringAttr::get(unwrap(context), unwrap(name)), unwrap(childLoc))));
 }
 
+MlirIdentifier mlirLocationNameGetName(MlirLocation location) {
+  return wrap((llvm::cast<NameLoc>(unwrap(location)).getName()));
+}
+
+MlirLocation mlirLocationNameGetChildLoc(MlirLocation location) {
+  return wrap(Location(llvm::cast<NameLoc>(unwrap(location)).getChildLoc()));
+}
+
+MlirTypeID mlirLocationNameGetTypeID() { return wrap(NameLoc::getTypeID()); }
+
+bool mlirLocationIsAName(MlirLocation location) {
+  return isa<NameLoc>(unwrap(location));
+}
+
 MlirLocation mlirLocationUnknownGet(MlirContext context) {
   return wrap(Location(UnknownLoc::get(unwrap(context))));
 }
@@ -1033,6 +1123,14 @@ void mlirValueReplaceAllUsesExcept(MlirValue oldValue, MlirValue newValue,
   oldValueCpp.replaceAllUsesExcept(newValueCpp, exceptionSet);
 }
 
+MlirLocation mlirValueGetLocation(MlirValue v) {
+  return wrap(unwrap(v).getLoc());
+}
+
+MlirContext mlirValueGetContext(MlirValue v) {
+  return wrap(unwrap(v).getContext());
+}
+
 //===----------------------------------------------------------------------===//
 // OpOperand API.
 //===----------------------------------------------------------------------===//
diff --git a/mlir/test/python/ir/location.py b/mlir/test/python/ir/location.py
index 59d8a89e770dd..3e54dc922cd67 100644
--- a/mlir/test/python/ir/location.py
+++ b/mlir/test/python/ir/location.py
@@ -43,22 +43,64 @@ def testLocationAttr():
 
 run(testLocationAttr)
 
+
 # CHECK-LABEL: TEST: testFileLineCol
 def testFileLineCol():
     with Context() as ctx:
-        loc = Location.file("foo.txt", 123, 56)
-        range = Location.file("foo.txt", 123, 56, 123, 100)
+        loc = Location.file("foo1.txt", 123, 56)
+        range = Location.file("foo2.txt", 123, 56, 124, 100)
+
     ctx = None
     gc.collect()
-    # CHECK: file str: loc("foo.txt":123:56)
+
+    # CHECK: file str: loc("foo1.txt":123:56)
     print("file str:", str(loc))
-    # CHECK: file repr: loc("foo.txt":123:56)
+    # CHECK: file repr: loc("foo1.txt":123:56)
     print("file repr:", repr(loc))
-    # CHECK: file range str: loc("foo.txt":123:56 to :100)
+    # CHECK: file range str: loc("foo2.txt":123:56 to 124:100)
     print("file range str:", str(range))
-    # CHECK: file range repr: loc("foo.txt":123:56 to :100)
+    # CHECK: file range repr: loc("foo2.txt":123:56 to 124:100)
     print("file range repr:", repr(range))
 
+    assert loc.is_a_file()
+    assert not loc.is_a_name()
+    assert not loc.is_a_callsite()
+    assert not loc.is_a_fused()
+
+    # CHECK: file filename: foo1.txt
+    print("file filename:", loc.filename)
+    # CHECK: file start_line: 123
+    print("file start_line:", loc.start_line)
+    # CHECK: file start_col: 56
+    print("file start_col:", loc.start_col)
+    # CHECK: file end_line: 123
+    print("file end_line:", loc.end_line)
+    # CHECK: file end_col: 56
+    print("file end_col:", loc.end_col)
+
+    assert range.is_a_file()
+    # CHECK: file filename: foo2.txt
+    print("file filename:", range.filename)
+    # CHECK: file start_line: 123
+    print("file start_line:", range.start_line)
+    # CHECK: file start_col: 56
+    print("file start_col:", range.start_col)
+    # CHECK: file end_line: 124
+    print("file end_line:", range.end_line)
+    # CHECK: file end_col: 100
+    print("file end_col:", range.end_col)
+
+    with Context() as ctx:
+        ctx.allow_unregistered_dialects = True
+        loc = Location.file("foo3.txt", 127, 61)
+        with loc:
+            i32 = IntegerType.get_signless(32)
+            module = Module.create()
+            with InsertionPoint(module.body):
+                new_value = Operation.create("custom.op1", results=[i32]).result
+                # CHECK: new_value location: loc("foo3.txt":127:61)
+                print("new_value location: ", new_value.location)
+
 
 run(testFileLineCol)
 
@@ -67,17 +109,31 @@ def testFileLineCol():
 def testName():
     with Context() as ctx:
         loc = Location.name("nombre")
-        locWithChildLoc = Location.name("naam", loc)
+        loc_with_child_loc = Location.name("naam", loc)
+
     ctx = None
     gc.collect()
-    # CHECK: file str: loc("nombre")
-    print("file str:", str(loc))
-    # CHECK: file repr: loc("nombre")
-    print("file repr:", repr(loc))
-    # CHECK: file str: loc("naam"("nombre"))
-    print("file str:", str(locWithChildLoc))
-    # CHECK: file repr: loc("naam"("nombre"))
-    print("file repr:", repr(locWithChildLoc))
+
+    # CHECK: name str: loc("nombre")
+    print("name str:", str(loc))
+    # CHECK: name repr: loc("nombre")
+    print("name repr:", repr(loc))
+    # CHECK: name str: loc("naam"("nombre"))
+    print("name str:", str(loc_with_child_loc))
+    # CHECK: name repr: loc("naam"("nombre"))
+    print("name repr:", repr(loc_with_child_loc))
+
+    assert loc.is_a_name()
+    # CHECK: name name_str: nombre
+    print("name name_str:", loc.name_str)
+    # CHECK: name child_loc: loc(unknown)
+    print("name child_loc:", loc.child_loc)
+
+    assert loc_with_child_loc.is_a_name()
+    # CHECK: name name_str: naam
+    print("name name_str:", loc_with_child_loc.name_str)
+    # CHECK: name child_loc_with_child_loc: loc("nombre")
+    print("name child_loc_with_child_loc:", loc_with_child_loc.child_loc)
 
 
 run(testName)
@@ -91,10 +147,17 @@ def testCallSite():
             [Location.file("util.foo", 379, 21), Location.file("main.foo", 100, 63)],
         )
     ctx = None
-    # CHECK: file str: loc(callsite("foo.text":123:45 at callsite("util.foo":379:21 at "main.foo":100:63))
-    print("file str:", str(loc))
-    # CHECK: file repr: loc(callsite("foo.text":123:45 at callsite("util.foo":379:21 at "main.foo":100:63))
-    print("file repr:", repr(loc))
+    # CHECK: callsite str: loc(callsite("foo.text":123:45 at callsite("util.foo":379:21 at "main.foo":100:63))
+    print("callsite str:", str(loc))
+    # CHECK: callsite repr: loc(callsite("foo.text":123:45 at callsite("util.foo":379:21 at "main.foo":100:63))
+    print("callsite repr:", repr(loc))
+
+    assert loc.is_a_callsite()
+
+    # CHECK: callsite callee: loc("foo.text":123:45)
+    print("callsite callee:", loc.callee)
+    # CHECK: callsite caller: loc(callsite("util.foo":379:21 at "main.foo":100:63))
+    print("callsite caller:", loc.caller)
 
 
 run(testCallSite)
@@ -112,31 +175,56 @@ def testFused():
         loc_empty = Location.fused([])
         loc_empty_attr = Location.fused([], attr)
         loc_single_attr = Location.fused([Location.name("apple")], attr)
+
     ctx = None
-    # CHECK: file str: loc("apple")
-    print("file str:", str(loc_single))
-    # CHECK: file repr: loc("apple")
-    print("file repr:", repr(loc_single))
-    # CHECK: file str: loc(fused["apple", "banana"])
-    print("file str:", str(loc))
-    # CHECK: file repr: loc(fused["apple", "banana"])
-    print("file repr:", repr(loc))
-    # CHECK: file str: loc(fused<"sauteed">["carrot", "potatoes"])
-    print("file str:", str(loc_attr))
-    # CHECK: file repr: loc(fused<"sauteed">["carrot", "potatoes"])
-    print("file repr:", repr(loc_attr))
-    # CHECK: file str: loc(unknown)
-    print("file str:", str(loc_empty))
-    # CHECK: file repr: loc(unknown)
-    print("file repr:", repr(loc_empty))
-    # CHECK: file str: loc(fused<"sauteed">[unknown])
-    print("file str:", str(loc_empty_attr))
-    # CHECK: file repr: loc(fused<"sauteed">[unknown])
-    print("file repr:", repr(loc_empty_attr))
-    # CHECK: file str: loc(fused<"sauteed">["apple"])
-    print("file str:", str(loc_single_attr))
-    # CHECK: file repr: loc(fused<"sauteed">["apple"])
-    print("file repr:", repr(loc_single_attr))
+
+    assert not loc_single.is_a_fused()
+    # CHECK: fused str: loc("apple")
+    print("fused str:", str(loc_single))
+    # CHECK: fused repr: loc("apple")
+    print("fused repr:", repr(loc_single))
+    # # CHECK: fused locations: []
+    print("fused locations:", loc_single.locations)
+
+    assert loc.is_a_fused()
+    # CHECK: fused str: loc(fused["apple", "banana"])
+    print("fused str:", str(loc))
+    # CHECK: fused repr: loc(fused["apple", "banana"])
+    print("fused repr:", repr(loc))
+    # CHECK: fused locations: [loc("apple"), loc("banana")]
+    print("fused locations:", loc.locations)
+
+    assert loc_attr.is_a_fused()
+    # CHECK: fused str: loc(fused<"sauteed">["carrot", "potatoes"])
+    print("fused str:", str(loc_attr))
+    # CHECK: fused repr: loc(fused<"sauteed">["carrot", "potatoes"])
+    print("fused repr:", repr(loc_attr))
+    # CHECK: fused locations: [loc("carrot"), loc("potatoes")]
+    print("fused locations:", loc_attr.locations)
+
+    assert not loc_empty.is_a_fused()
+    # CHECK: fused str: loc(unknown)
+    print("fused str:", str(loc_empty))
+    # CHECK: fused repr: loc(unknown)
+    print("fused repr:", repr(loc_empty))
+    # CHECK: fused locations: []
+    print("fused locations:", loc_empty.locations)
+
+    assert loc_empty_attr.is_a_fused()
+    # CHECK: fused str: loc(fused<"sauteed">[unknown])
+    print("fused str:", str(loc_empty_attr))
+    # CHECK: fused repr: loc(fused<"sauteed">[unknown])
+    print("fused repr:", repr(loc_empty_attr))
+    # CHECK: fused locations: [loc(unknown)]
+    print("fused locations:", loc_empty_attr.locations)
+
+    assert loc_single_attr.is_a_fused()
+    # CHECK: fused str: loc(fused<"sauteed">["apple"])
+    print("fused str:", str(loc_single_attr))
+    # CHECK: fused repr: loc(fused<"sauteed">["apple"])
+    print("fused repr:", repr(loc_single_attr))
+    # CHECK: fused locations: [loc("apple")]
+    print("fused locations:", loc_single_attr.locations)
 
 
 run(testFused)



More information about the Mlir-commits mailing list