From 1b9102e8d19c3cc74ba4a78a4bebcb1d039220c3 Mon Sep 17 00:00:00 2001
From: peefy <xpf6677@163.com>
Date: Sat, 24 Aug 2024 23:37:23 +0800
Subject: [PATCH] feat: add location method for the operation structure

Signed-off-by: peefy <xpf6677@163.com>
---
 melior/src/ir/operation.rs | 33 ++++++++++++++++++++++++++-------
 1 file changed, 26 insertions(+), 7 deletions(-)

diff --git a/melior/src/ir/operation.rs b/melior/src/ir/operation.rs
index 6fa3d65a13..2cc4892fd2 100644
--- a/melior/src/ir/operation.rs
+++ b/melior/src/ir/operation.rs
@@ -7,7 +7,7 @@ mod result;
 pub use self::{
     builder::OperationBuilder, printing_flags::OperationPrintingFlags, result::OperationResult,
 };
-use super::{Attribute, AttributeLike, BlockRef, Identifier, RegionRef, Value};
+use super::{Attribute, AttributeLike, BlockRef, Identifier, Location, RegionRef, Value};
 use crate::{
     context::{Context, ContextRef},
     utility::{print_callback, print_string_callback},
@@ -20,12 +20,12 @@ use core::{
 use mlir_sys::{
     mlirOperationClone, mlirOperationDestroy, mlirOperationDump, mlirOperationEqual,
     mlirOperationGetAttribute, mlirOperationGetAttributeByName, mlirOperationGetBlock,
-    mlirOperationGetContext, mlirOperationGetName, mlirOperationGetNextInBlock,
-    mlirOperationGetNumAttributes, mlirOperationGetNumOperands, mlirOperationGetNumRegions,
-    mlirOperationGetNumResults, mlirOperationGetNumSuccessors, mlirOperationGetOperand,
-    mlirOperationGetParentOperation, mlirOperationGetRegion, mlirOperationGetResult,
-    mlirOperationGetSuccessor, mlirOperationPrint, mlirOperationPrintWithFlags,
-    mlirOperationRemoveAttributeByName, mlirOperationRemoveFromParent,
+    mlirOperationGetContext, mlirOperationGetLocation, mlirOperationGetName,
+    mlirOperationGetNextInBlock, mlirOperationGetNumAttributes, mlirOperationGetNumOperands,
+    mlirOperationGetNumRegions, mlirOperationGetNumResults, mlirOperationGetNumSuccessors,
+    mlirOperationGetOperand, mlirOperationGetParentOperation, mlirOperationGetRegion,
+    mlirOperationGetResult, mlirOperationGetSuccessor, mlirOperationPrint,
+    mlirOperationPrintWithFlags, mlirOperationRemoveAttributeByName, mlirOperationRemoveFromParent,
     mlirOperationSetAttributeByName, mlirOperationVerify, MlirOperation,
 };
 use std::{
@@ -130,6 +130,11 @@ impl<'c> Operation<'c> {
         (0..self.region_count()).map(|index| self.region(index).expect("valid result index"))
     }
 
+    /// Gets the location of the operation.
+    pub fn location(&self) -> Location<'c> {
+        unsafe { Location::from_raw(mlirOperationGetLocation(self.raw)) }
+    }
+
     /// Returns the number of successors.
     pub fn successor_count(&self) -> usize {
         unsafe { mlirOperationGetNumSuccessors(self.raw) as usize }
@@ -644,6 +649,20 @@ mod tests {
         );
     }
 
+    #[test]
+    fn location() {
+        let context = create_test_context();
+        context.set_allow_unregistered_dialects(true);
+        let location = Location::new(&context, "test", 1, 1);
+
+        let operation = OperationBuilder::new("foo", location)
+            .add_regions([Region::new()])
+            .build()
+            .unwrap();
+
+        assert_eq!(operation.location(), location);
+    }
+
     #[test]
     fn attribute() {
         let context = create_test_context();