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

[Thurston Dang via llvm-commits]

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

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 make it easier to know when CreateShadowCast is being used with unsupported types.

>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] [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.