[llvm] [msan] Add struct support to CreateShadowCast (PR #130440)

[Thurston Dang via llvm-commits]

https://github.com/thurstond updated https://github.com/llvm/llvm-project/pull/130440

>From c901f03079ad5bb5a88fe4da624e036deaa7b3a8 Mon Sep 17 00:00:00 2001
From: Thurston Dang <thurston at google.com>
Date: Thu, 6 Mar 2025 18:47:23 +0000
Subject: [PATCH 1/2] [msan] Add struct support to CreateShadowCast

This fixes a TODO in CreateShadowCast, which could be useful for
applying handleIntrinsicByApplyingToShadow to intrinsics that return a
struct type. Note that getShadowTy already supports structs, hence
struct support in CreateShadowCast is the last missing piece.

Additionally, this adds assertions to prevent misuse of CreateShadowCast
with floating-point types (which is not currently used, and was never fully
supported).
---
 .../Instrumentation/MemorySanitizer.cpp       | 30 ++++++++++++++++++-
 1 file changed, 29 insertions(+), 1 deletion(-)

diff --git a/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp b/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp
index ba925275faba6..54e34bf9a0b85 100644
--- a/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp
+++ b/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp
@@ -2562,6 +2562,15 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
   Value *CreateShadowCast(IRBuilder<> &IRB, Value *V, Type *dstTy,
                           bool Signed = false) {
     Type *srcTy = V->getType();
+
+    // Shadow types should never be floating-point.
+    assert(!srcTy->isFPOrFPVectorTy());
+    assert(!dstTy->isFPOrFPVectorTy());
+
+    // TODO (though never needed in practice)
+    assert(!srcTy->isArrayTy());
+    assert(!dstTy->isArrayTy());
+
     if (srcTy == dstTy)
       return V;
     size_t srcSizeInBits = VectorOrPrimitiveTypeSizeInBits(srcTy);
@@ -2571,15 +2580,34 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
 
     if (dstTy->isIntegerTy() && srcTy->isIntegerTy())
       return IRB.CreateIntCast(V, dstTy, Signed);
+
     if (dstTy->isVectorTy() && srcTy->isVectorTy() &&
         cast<VectorType>(dstTy)->getElementCount() ==
             cast<VectorType>(srcTy)->getElementCount())
+      // We rely on the assertion at the start of the function that neither
+      // type is a floating-point vector.
       return IRB.CreateIntCast(V, dstTy, Signed);
+
+    if (dstTy->isStructTy() && srcTy->isStructTy() &&
+        cast<StructType>(dstTy)->getNumElements() ==
+            cast<StructType>(srcTy)->getNumElements()) {
+      Value *DStruct = llvm::UndefValue::get(dstTy);
+
+      for (unsigned i = 0, n = cast<StructType>(dstTy)->getNumElements(); i < n;
+           i++) {
+        Value *Elem = IRB.CreateExtractValue(V, i);
+        Type *ElemDstTy = cast<StructType>(dstTy)->getElementType(i);
+        Value *ElemDst = CreateShadowCast(IRB, Elem, ElemDstTy, Signed);
+        DStruct = IRB.CreateInsertValue(DStruct, ElemDst, i);
+      }
+
+      return DStruct;
+    }
+
     Value *V1 = IRB.CreateBitCast(V, Type::getIntNTy(*MS.C, srcSizeInBits));
     Value *V2 =
         IRB.CreateIntCast(V1, Type::getIntNTy(*MS.C, dstSizeInBits), Signed);
     return IRB.CreateBitCast(V2, dstTy);
-    // TODO: handle struct types.
   }
 
   /// Cast an application value to the type of its own shadow.

>From e343e0bc9069fb47335d59935c95fe6732ed8d30 Mon Sep 17 00:00:00 2001
From: Thurston Dang <thurston at google.com>
Date: Sat, 8 Mar 2025 22:50:36 +0000
Subject: [PATCH 2/2] Switch from undef to aspartame

---
 llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp b/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp
index 54e34bf9a0b85..097f561575a0c 100644
--- a/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp
+++ b/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp
@@ -2591,7 +2591,7 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
     if (dstTy->isStructTy() && srcTy->isStructTy() &&
         cast<StructType>(dstTy)->getNumElements() ==
             cast<StructType>(srcTy)->getNumElements()) {
-      Value *DStruct = llvm::UndefValue::get(dstTy);
+      Value *DStruct = PoisonValue::get(dstTy);
 
       for (unsigned i = 0, n = cast<StructType>(dstTy)->getNumElements(); i < n;
            i++) {