[Mlir-commits] [mlir] [mlir] Add C and Python interface for file range (PR #123276)

[Jacques Pienaar]

https://github.com/jpienaar created https://github.com/llvm/llvm-project/pull/123276

Plumbs through creating file ranges to C and Python.

>From 5dd46b1aeb9c50c8811d65e90001da217cdfb1c7 Mon Sep 17 00:00:00 2001
From: Jacques Pienaar <jpienaar at google.com>
Date: Fri, 17 Jan 2025 02:58:48 +0000
Subject: [PATCH] [mlir] Add C and Python interface for file range

Plumbs through creating file ranges to C and Python.
---
 mlir/include/mlir-c/IR.h            |  5 +++++
 mlir/lib/Bindings/Python/IRCore.cpp | 15 +++++++++++++++
 mlir/lib/CAPI/IR/IR.cpp             |  9 +++++++++
 mlir/test/CAPI/ir.c                 |  7 +++++++
 mlir/test/python/ir/location.py     | 21 +++++++++++++--------
 5 files changed, 49 insertions(+), 8 deletions(-)

diff --git a/mlir/include/mlir-c/IR.h b/mlir/include/mlir-c/IR.h
index 0a515bbea3b504..7d2fd89e8560fc 100644
--- a/mlir/include/mlir-c/IR.h
+++ b/mlir/include/mlir-c/IR.h
@@ -256,6 +256,11 @@ mlirLocationFromAttribute(MlirAttribute attribute);
 MLIR_CAPI_EXPORTED MlirLocation mlirLocationFileLineColGet(
     MlirContext context, MlirStringRef filename, unsigned line, unsigned col);
 
+/// Creates an File/Line/Column range location owned by the given context.
+MLIR_CAPI_EXPORTED MlirLocation mlirLocationFileLineColRangeGet(
+    MlirContext context, MlirStringRef filename, unsigned start_line,
+    unsigned start_col, unsigned end_line, unsigned end_col);
+
 /// Creates a call site location with a callee and a caller.
 MLIR_CAPI_EXPORTED MlirLocation mlirLocationCallSiteGet(MlirLocation callee,
                                                         MlirLocation caller);
diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index 53806ca9f04a49..d6c2fdd5eac98b 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -50,6 +50,9 @@ static const char kContextGetCallSiteLocationDocstring[] =
 static const char kContextGetFileLocationDocstring[] =
     R"(Gets a Location representing a file, line and column)";
 
+static const char kContextGetFileRangeDocstring[] =
+    R"(Gets a Location representing a file, line and column range)";
+
 static const char kContextGetFusedLocationDocstring[] =
     R"(Gets a Location representing a fused location with optional metadata)";
 
@@ -2902,6 +2905,18 @@ void mlir::python::populateIRCore(nb::module_ &m) {
           nb::arg("filename"), nb::arg("line"), nb::arg("col"),
           nb::arg("context").none() = nb::none(),
           kContextGetFileLocationDocstring)
+      .def_static(
+          "file",
+          [](std::string filename, int startLine, int startCol, int endLine,
+             int endCol, DefaultingPyMlirContext context) {
+            return PyLocation(context->getRef(),
+                              mlirLocationFileLineColRangeGet(
+                                  context->get(), toMlirStringRef(filename),
+                                  startLine, startCol, endLine, endCol));
+          },
+          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_static(
           "fused",
           [](const std::vector<PyLocation> &pyLocations,
diff --git a/mlir/lib/CAPI/IR/IR.cpp b/mlir/lib/CAPI/IR/IR.cpp
index 24dc8854048532..f27af0ca9a2c78 100644
--- a/mlir/lib/CAPI/IR/IR.cpp
+++ b/mlir/lib/CAPI/IR/IR.cpp
@@ -264,6 +264,15 @@ MlirLocation mlirLocationFileLineColGet(MlirContext context,
       FileLineColLoc::get(unwrap(context), unwrap(filename), line, col)));
 }
 
+MlirLocation
+mlirLocationFileLineColRangeGet(MlirContext context, MlirStringRef filename,
+                                unsigned startLine, unsigned startCol,
+                                unsigned endLine, unsigned endCol) {
+  return wrap(
+      Location(FileLineColRange::get(unwrap(context), unwrap(filename),
+                                     startLine, startCol, endLine, endCol)));
+}
+
 MlirLocation mlirLocationCallSiteGet(MlirLocation callee, MlirLocation caller) {
   return wrap(Location(CallSiteLoc::get(unwrap(callee), unwrap(caller))));
 }
diff --git a/mlir/test/CAPI/ir.c b/mlir/test/CAPI/ir.c
index 15a3a1fb50dc9e..68da79f69cc0ad 100644
--- a/mlir/test/CAPI/ir.c
+++ b/mlir/test/CAPI/ir.c
@@ -2389,6 +2389,9 @@ void testDiagnostics(void) {
   MlirLocation fileLineColLoc = mlirLocationFileLineColGet(
       ctx, mlirStringRefCreateFromCString("file.c"), 1, 2);
   mlirEmitError(fileLineColLoc, "test diagnostics");
+  MlirLocation fileLineColRange = mlirLocationFileLineColRangeGet(
+      ctx, mlirStringRefCreateFromCString("other-file.c"), 1, 2, 3, 4);
+  mlirEmitError(fileLineColRange, "test diagnostics");
   MlirLocation callSiteLoc = mlirLocationCallSiteGet(
       mlirLocationFileLineColGet(
           ctx, mlirStringRefCreateFromCString("other-file.c"), 2, 3),
@@ -2418,6 +2421,10 @@ void testDiagnostics(void) {
   // CHECK: >> end of diagnostic (userData: 42)
   // CHECK: processing diagnostic (userData: 42) <<
   // CHECK:   test diagnostics
+  // CHECK:   loc("other-file.c":1:2 to 3:4)
+  // CHECK: >> end of diagnostic (userData: 42)
+  // CHECK: processing diagnostic (userData: 42) <<
+  // CHECK:   test diagnostics
   // CHECK:   loc(callsite("other-file.c":2:3 at "file.c":1:2))
   // CHECK: >> end of diagnostic (userData: 42)
   // CHECK: processing diagnostic (userData: 42) <<
diff --git a/mlir/test/python/ir/location.py b/mlir/test/python/ir/location.py
index f66d6c501dcf5c..fc74fbb6c39eb8 100644
--- a/mlir/test/python/ir/location.py
+++ b/mlir/test/python/ir/location.py
@@ -45,14 +45,19 @@ def testLocationAttr():
 
 # CHECK-LABEL: TEST: testFileLineCol
 def testFileLineCol():
-    with Context() as ctx:
-        loc = Location.file("foo.txt", 123, 56)
-    ctx = None
-    gc.collect()
-    # CHECK: file str: loc("foo.txt":123:56)
-    print("file str:", str(loc))
-    # CHECK: file repr: loc("foo.txt":123:56)
-    print("file repr:", repr(loc))
+  with Context() as ctx:
+    loc = Location.file("foo.txt", 123, 56)
+    range = Location.file("foo.txt", 123, 56, 123, 100)
+  ctx = None
+  gc.collect()
+  # CHECK: file str: loc("foo.txt":123:56)
+  print("file str:", str(loc))
+  # CHECK: file repr: loc("foo.txt":123:56)
+  print("file repr:", repr(loc))
+  # CHECK: file range str: loc("foo.txt":123:56 to :100)
+  print("file range str:", str(range))
+  # CHECK: file range repr: loc("foo.txt":123:56 to :100)
+  print("file range repr:", repr(range))
 
 
 run(testFileLineCol)