[SandboxIR][Utils] Implement getMemoryLocation() (#109724)

This patch implements sandboxir::Utils::memoryLocationGetOrNone() that calls
MemoryLocation::getOrNone() internally.
Ideally this would require a sandboxir::MemoryLocation, but this should
be good enough for now.

5 files changed, 69 insertions, 0 deletions

diff --git a/llvm/include/llvm/SandboxIR/SandboxIR.h b/llvm/include/llvm/SandboxIR/SandboxIR.h
index d4c907c..d99d564 100644
--- a/llvm/include/llvm/SandboxIR/SandboxIR.h
+++ b/llvm/include/llvm/SandboxIR/SandboxIR.h
@@ -346,6 +346,7 @@ protected:
   friend class NoCFIValue;            // For `Val`.
   friend class ConstantPtrAuth;       // For `Val`.
   friend class ConstantExpr;          // For `Val`.
+  friend class Utils;                 // For `Val`.
 
   // Region needs to manipulate metadata in the underlying LLVM Value, we don't
   // expose metadata in sandboxir.
diff --git a/llvm/include/llvm/SandboxIR/Utils.h b/llvm/include/llvm/SandboxIR/Utils.h
index ccc0030..4e8a175 100644
--- a/llvm/include/llvm/SandboxIR/Utils.h
+++ b/llvm/include/llvm/SandboxIR/Utils.h
@@ -12,6 +12,9 @@
 #ifndef LLVM_SANDBOXIR_UTILS_H
 #define LLVM_SANDBOXIR_UTILS_H
 
+#include "llvm/Analysis/MemoryLocation.h"
+#include "llvm/SandboxIR/SandboxIR.h"
+
 namespace llvm::sandboxir {
 
 class Utils {
@@ -48,6 +51,12 @@ public:
     Type *Ty = getExpectedType(V);
     return DL.getTypeSizeInBits(Ty->LLVMTy);
   }
+
+  /// Equivalent to MemoryLocation::getOrNone(I).
+  static std::optional<llvm::MemoryLocation>
+  memoryLocationGetOrNone(const Instruction *I) {
+    return llvm::MemoryLocation::getOrNone(cast<llvm::Instruction>(I->Val));
+  }
 };
 } // namespace llvm::sandboxir
 
diff --git a/llvm/lib/SandboxIR/CMakeLists.txt b/llvm/lib/SandboxIR/CMakeLists.txt
index 03474be..b2e6f62 100644
--- a/llvm/lib/SandboxIR/CMakeLists.txt
+++ b/llvm/lib/SandboxIR/CMakeLists.txt
@@ -11,5 +11,6 @@ add_llvm_component_library(LLVMSandboxIR
   LINK_COMPONENTS
   Core
   Support
+  Analysis
   )
 
diff --git a/llvm/unittests/SandboxIR/CMakeLists.txt b/llvm/unittests/SandboxIR/CMakeLists.txt
index a228637..2ab284a 100644
--- a/llvm/unittests/SandboxIR/CMakeLists.txt
+++ b/llvm/unittests/SandboxIR/CMakeLists.txt
@@ -2,6 +2,7 @@ set(LLVM_LINK_COMPONENTS
   AsmParser
   SandboxIR
   Core
+  Analysis
   )
 
 add_llvm_unittest(SandboxIRTests
@@ -9,4 +10,5 @@ add_llvm_unittest(SandboxIRTests
   SandboxIRTest.cpp
   TrackerTest.cpp
   TypesTest.cpp
+  UtilsTest.cpp
   )
diff --git a/llvm/unittests/SandboxIR/UtilsTest.cpp b/llvm/unittests/SandboxIR/UtilsTest.cpp
new file mode 100644
index 0000000..ded3edf
--- /dev/null
+++ b/llvm/unittests/SandboxIR/UtilsTest.cpp
@@ -0,0 +1,56 @@
+//===- UtilsTest.cpp ------------------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/SandboxIR/Utils.h"
+#include "llvm/AsmParser/Parser.h"
+#include "llvm/IR/BasicBlock.h"
+#include "llvm/IR/DataLayout.h"
+#include "llvm/IR/Function.h"
+#include "llvm/IR/Instruction.h"
+#include "llvm/IR/Module.h"
+#include "llvm/SandboxIR/SandboxIR.h"
+#include "llvm/Support/SourceMgr.h"
+#include "gtest/gtest.h"
+
+using namespace llvm;
+
+struct UtilsTest : public testing::Test {
+  LLVMContext C;
+  std::unique_ptr<Module> M;
+
+  void parseIR(LLVMContext &C, const char *IR) {
+    SMDiagnostic Err;
+    M = parseAssemblyString(IR, Err, C);
+    if (!M)
+      Err.print("UtilsTest", errs());
+  }
+  BasicBlock *getBasicBlockByName(Function &F, StringRef Name) {
+    for (BasicBlock &BB : F)
+      if (BB.getName() == Name)
+        return &BB;
+    llvm_unreachable("Expected to find basic block!");
+  }
+};
+
+TEST_F(UtilsTest, getMemoryLocation) {
+  parseIR(C, R"IR(
+define void @foo(ptr %arg0) {
+  %ld = load i8, ptr %arg0
+  ret void
+}
+)IR");
+  llvm::Function *LLVMF = &*M->getFunction("foo");
+  auto *LLVMBB = &*LLVMF->begin();
+  auto *LLVMLd = cast<llvm::LoadInst>(&*LLVMBB->begin());
+  sandboxir::Context Ctx(C);
+  sandboxir::Function *F = Ctx.createFunction(LLVMF);
+  auto *BB = &*F->begin();
+  auto *Ld = cast<sandboxir::LoadInst>(&*BB->begin());
+  EXPECT_EQ(sandboxir::Utils::memoryLocationGetOrNone(Ld),
+            MemoryLocation::getOrNone(LLVMLd));
+}