From c6e76eeb69bb23fe2d91e88945bb8c83b9ebf86f Mon Sep 17 00:00:00 2001
From: Vasileios Porpodas <vporpodas@google.com>
Date: Mon, 23 Sep 2024 13:59:51 -0700
Subject: [PATCH] [SandboxIR][Utils] Implement getMemoryLocation()

This patch implements sandboxir::Utils::getMemoryLocation() that calls
MemoryLocation::getOrNone() internally.
Ideally this would require a sandboxir::MemoryLocation, but this should
be good enough for now.
---
 llvm/include/llvm/SandboxIR/SandboxIR.h |  1 +
 llvm/include/llvm/SandboxIR/Utils.h     |  9 ++++
 llvm/lib/SandboxIR/CMakeLists.txt       |  1 +
 llvm/unittests/SandboxIR/CMakeLists.txt |  2 +
 llvm/unittests/SandboxIR/UtilsTest.cpp  | 56 +++++++++++++++++++++++++
 5 files changed, 69 insertions(+)
 create mode 100644 llvm/unittests/SandboxIR/UtilsTest.cpp

diff --git a/llvm/include/llvm/SandboxIR/SandboxIR.h b/llvm/include/llvm/SandboxIR/SandboxIR.h
index 84c32aeb0451f..d3c6132c6ceec 100644
--- a/llvm/include/llvm/SandboxIR/SandboxIR.h
+++ b/llvm/include/llvm/SandboxIR/SandboxIR.h
@@ -346,6 +346,7 @@ class Value {
   friend class NoCFIValue;            // For `Val`.
   friend class ConstantPtrAuth;       // For `Val`.
   friend class ConstantExpr;          // For `Val`.
+  friend class Utils;                 // For `Val`.
 
   /// All values point to the context.
   Context &Ctx;
diff --git a/llvm/include/llvm/SandboxIR/Utils.h b/llvm/include/llvm/SandboxIR/Utils.h
index ccc0030868a55..4e8a175f54705 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 @@ class Utils {
     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 03474be0c7b80..b2e6f6285fea5 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 a228637b062a4..2ab284a511fca 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 0000000000000..ded3edf1206a4
--- /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));
+}