Make WholeProgramDevirt understand ConstStruct vtables.

[LemonBoy]

Summary

I've added another test to the non-array-vtable file, I hope that's the correct place.

[pcc]

lib/Transforms/IPO/WholeProgramDevirt.cpp

427–433

All this code could be factored out into a separate function.

434

I think you also want to check that SL->getElementOffset(Op) == GlobalSlotOffset.

Please add negative tests for that case and the out-of-bounds case you have above.

test/Transforms/WholeProgramDevirt/non-array-vtable.ll

44

Please put your test in a separate file and rename this to something like non-aggregate-vtable.ll.

[LemonBoy]

lib/Transforms/IPO/WholeProgramDevirt.cpp

434

I haven't been able to come up with a way to access the table mid-element via a GEP but have added the check anyway.

[pcc]

lib/Transforms/IPO/WholeProgramDevirt.cpp

434

See e.g. the unaligned test in bad-read-from-vtable.ll.

[pcc]

lib/Transforms/IPO/WholeProgramDevirt.cpp

386

Could be simpler to return a const Constant * from this function.

428–431

These two could be folded into the only use (after applying my suggestion above).

test/Transforms/WholeProgramDevirt/non-array-vtable.ll

1

struct-vtable.ll would be a better name for this test.

[LemonBoy]

lib/Transforms/IPO/WholeProgramDevirt.cpp

386

I suppose you meant const Value * since that's what getOperand returns.

[pcc]

LGTM with nit

lib/Transforms/IPO/WholeProgramDevirt.cpp

386

getOperand on a constant will only ever return a constant; you could cast it below.

[pcc]

It doesn't look like you have commit access, so I committed in r289162. I had to remove some of the const qualification to make it compile.

lib/Transforms/IPO/WholeProgramDevirt.cpp

struct DevirtModule {
   bool areRemarksEnabled();
 
   void scanTypeTestUsers(Function *TypeTestFunc, Function *AssumeFunc);
   void scanTypeCheckedLoadUsers(Function *TypeCheckedLoadFunc);
 
   void buildTypeIdentifierMap(
       std::vector<VTableBits> &Bits,
       DenseMap<Metadata *, std::set<TypeMemberInfo>> &TypeIdMap);
+  Optional<unsigned> getOpAtOffset(const Constant *I, const uint64_t Offset);
   bool
   tryFindVirtualCallTargets(std::vector<VirtualCallTarget> &TargetsForSlot,
                             const std::set<TypeMemberInfo> &TypeMemberInfos,
                             uint64_t ByteOffset);
   bool trySingleImplDevirt(MutableArrayRef<VirtualCallTarget> TargetsForSlot,
                            MutableArrayRef<VirtualCallSite> CallSites);
   bool tryEvaluateFunctionsWithArgs(
       MutableArrayRef<VirtualCallTarget> TargetsForSlot,
for (MDNode *Type : Types) {
               cast<ConstantAsMetadata>(Type->getOperand(0))->getValue())
               ->getZExtValue();
 
       TypeIdMap[TypeID].insert({BitsPtr, Offset});
     }
   }
 }
 
+Optional<unsigned>
+DevirtModule::getOpAtOffset(const Constant *I, const uint64_t Offset) {
+  const DataLayout &DL = M.getDataLayout();
+  unsigned Op;
+
+  if (auto *C = dyn_cast<ConstantStruct>(I)) {
+    const StructLayout *SL = DL.getStructLayout(C->getType());
+
+    if (Offset >= SL->getSizeInBytes())
+      return None;
+
+    Op = SL->getElementContainingOffset(Offset);
+
+    if (Offset != SL->getElementOffset(Op))
+      return None;
+
+    return Op;
+  } else if (auto *C = dyn_cast<ConstantArray>(I)) {
+    ArrayType *VTableTy = C->getType();
+    const uint64_t ElemSize = DL.getTypeAllocSize(VTableTy->getElementType());
+
+    if (Offset % ElemSize != 0)
+      return None;
+
+    Op = Offset / ElemSize;
+
+    if (Op >= C->getNumOperands())
+      return None;
+
+    return Op;
+  }
+
+  return None;
+}
+
 bool DevirtModule::tryFindVirtualCallTargets(
     std::vector<VirtualCallTarget> &TargetsForSlot,
     const std::set<TypeMemberInfo> &TypeMemberInfos, uint64_t ByteOffset) {
   for (const TypeMemberInfo &TM : TypeMemberInfos) {
     if (!TM.Bits->GV->isConstant())
       return false;
 
-    auto Init = dyn_cast<ConstantArray>(TM.Bits->GV->getInitializer());
-    if (!Init)
-      return false;
-    ArrayType *VTableTy = Init->getType();
-
-    uint64_t ElemSize =
-        M.getDataLayout().getTypeAllocSize(VTableTy->getElementType());
-    uint64_t GlobalSlotOffset = TM.Offset + ByteOffset;
-    if (GlobalSlotOffset % ElemSize != 0)
-      return false;
+    const Constant *I = TM.Bits->GV->getInitializer();
+    const uint64_t GlobalSlotOffset = TM.Offset + ByteOffset;
+    Optional<unsigned> Op = getOpAtOffset(I, GlobalSlotOffset);
 
-    unsigned Op = GlobalSlotOffset / ElemSize;
-    if (Op >= Init->getNumOperands())
+    if (!Op.hasValue())
       return false;
 
-    auto Fn = dyn_cast<Function>(Init->getOperand(Op)->stripPointerCasts());
+    auto Fn = dyn_cast<Function>(I->getOperand(*Op)->stripPointerCasts());
     if (!Fn)
       return false;
 
     // We can disregard __cxa_pure_virtual as a possible call target, as
     // calls to pure virtuals are UB.
     if (Fn->getName() == "__cxa_pure_virtual")
       continue;
 

test/Transforms/WholeProgramDevirt/non-aggregate-vtable.ll

+; RUN: opt -S -wholeprogramdevirt %s | FileCheck %s
+
+target datalayout = "e-p:64:64"
+target triple = "x86_64-unknown-linux-gnu"
+
+@vt = constant i8* bitcast (void (i8*)* @vf to i8*), !type !0
+
+define void @vf(i8* %this) {
+  ret void
+}
+
+; CHECK: define void @call
+define void @call(i8* %obj) {
+  %vtableptr = bitcast i8* %obj to [1 x i8*]**
+  %vtable = load [1 x i8*]*, [1 x i8*]** %vtableptr
+  %vtablei8 = bitcast [1 x i8*]* %vtable to i8*
+  %p = call i1 @llvm.type.test(i8* %vtablei8, metadata !"typeid")
+  call void @llvm.assume(i1 %p)
+  %fptrptr = getelementptr [1 x i8*], [1 x i8*]* %vtable, i32 0, i32 0
+  %fptr = load i8*, i8** %fptrptr
+  %fptr_casted = bitcast i8* %fptr to void (i8*)*
+  ; CHECK: call void %
+  call void %fptr_casted(i8* %obj)
+  ret void
+}
+
+declare i1 @llvm.type.test(i8*, metadata)
+declare void @llvm.assume(i1)
+
+!0 = !{i32 0, !"typeid"}

test/Transforms/WholeProgramDevirt/non-array-vtable.ll

 ; RUN: opt -S -wholeprogramdevirt %s | FileCheck %s
 
 target datalayout = "e-p:64:64"
 target triple = "x86_64-unknown-linux-gnu"
 
-@vt = constant i8* bitcast (void (i8*)* @vf to i8*), !type !0
+%vtTy = type { void (i8*)* }
+
+@vt = constant %vtTy { void (i8*)* @vf }, !type !0
 
 define void @vf(i8* %this) {
   ret void
 }
 
 ; CHECK: define void @call
 define void @call(i8* %obj) {
   %vtableptr = bitcast i8* %obj to [1 x i8*]**
   %vtable = load [1 x i8*]*, [1 x i8*]** %vtableptr
   %vtablei8 = bitcast [1 x i8*]* %vtable to i8*
   %p = call i1 @llvm.type.test(i8* %vtablei8, metadata !"typeid")
   call void @llvm.assume(i1 %p)
   %fptrptr = getelementptr [1 x i8*], [1 x i8*]* %vtable, i32 0, i32 0
   %fptr = load i8*, i8** %fptrptr
   %fptr_casted = bitcast i8* %fptr to void (i8*)*
+  ; CHECK: call void @vf(
+  call void %fptr_casted(i8* %obj)
+  ret void
+}
+
+; CHECK: define void @call_oob
+define void @call_oob(i8* %obj) {
+  %vtableptr = bitcast i8* %obj to [1 x i8*]**
+  %vtable = load [1 x i8*]*, [1 x i8*]** %vtableptr
+  %vtablei8 = bitcast [1 x i8*]* %vtable to i8*
+  %p = call i1 @llvm.type.test(i8* %vtablei8, metadata !"typeid")
+  call void @llvm.assume(i1 %p)
+  %fptrptr = getelementptr [1 x i8*], [1 x i8*]* %vtable, i32 0, i32 4
+  %fptr = load i8*, i8** %fptrptr
+  %fptr_casted = bitcast i8* %fptr to void (i8*)*
   ; CHECK: call void %
   call void %fptr_casted(i8* %obj)
   ret void
 }
 
 declare i1 @llvm.type.test(i8*, metadata)
 declare void @llvm.assume(i1)
 
 !0 = !{i32 0, !"typeid"}