From 745fc404ac76d8baf87306a4107af2f75dec61e3 Mon Sep 17 00:00:00 2001
From: Harald van Dijk <harald.vandijk@codeplay.com>
Date: Thu, 24 Oct 2024 14:52:05 +0100
Subject: [PATCH] [AlignModuleStructsPass] Update types in constants

Execution.Regression_35_Constant_Struct_Alignment reveals on LLVM 20
that if a constant expression references a type that needs to be
replaced, we were ignoring it unless the constant expression also
happened to contain references to globals that need to be replaced. Do
a pass over constant operands and update them to instructions as needed.
---
 .../source/align_module_structs_pass.cpp      | 37 +++++++++++++++++++
 1 file changed, 37 insertions(+)

diff --git a/modules/compiler/compiler_pipeline/source/align_module_structs_pass.cpp b/modules/compiler/compiler_pipeline/source/align_module_structs_pass.cpp
index 8355b0633..c79e2cc58 100644
--- a/modules/compiler/compiler_pipeline/source/align_module_structs_pass.cpp
+++ b/modules/compiler/compiler_pipeline/source/align_module_structs_pass.cpp
@@ -356,6 +356,30 @@ Type *getNewType(Value *v, const StructReplacementMap &typeMap) {
   return getNewType(v->getType(), typeMap);
 }
 
+bool replaceConstantsForRemapping(Constant *C,
+                                  const StructReplacementMap &typeMap,
+                                  const ValueToValueMapTy &valMap) {
+  if (valMap.find(C) != valMap.end()) {
+    return false;
+  }
+  if (getNewType(C->getType(), typeMap)) {
+    compiler::utils::replaceConstantExpressionWithInstruction(C);
+    return true;
+  }
+  if (auto *GEP = dyn_cast<GEPOperator>(C)) {
+    if (getNewType(GEP->getSourceElementType(), typeMap)) {
+      compiler::utils::replaceConstantExpressionWithInstruction(C);
+      return true;
+    }
+  }
+  for (auto *Op : C->operand_values()) {
+    if (replaceConstantsForRemapping(cast<Constant>(Op), typeMap, valMap)) {
+      return true;
+    }
+  }
+  return false;
+}
+
 /// @brief Performs the update on the LLVM module to replace all the Values
 ///        using the old struct type with Values using our padded variant.
 ///
@@ -378,6 +402,19 @@ void replaceModuleTypes(const StructReplacementMap &typeMap, Module &module) {
     replaceGlobalVariable(global, typeMap, valMap);
   }
 
+  // Replace constants if they are using types that need to be remapped.
+  for (auto &func : module.functions()) {
+    for (BasicBlock &block : func) {
+      for (Instruction &inst : block) {
+        for (auto *op : inst.operand_values()) {
+          if (auto *c = dyn_cast<Constant>(op)) {
+            replaceConstantsForRemapping(c, typeMap, valMap);
+          }
+        }
+      }
+    }
+  }
+
   // Identify all functions which use a struct type and need to be cloned,
   // this is to avoid unnecessary work in the clone.
   SmallVector<Function *, 4> funcs;