[MSan] [AArch64] Fix vararg helper for >1 or non-int fixed arguments.

[• koriakin]

Summary

This fixes http://llvm.org/PR27646 on AArch64.

There are three issues here:

• The GR save area is 7 words in size, instead of 8. This is not enough if none of the fixed arguments is passed in GRs (they're all floats or aggregates).
• The first argument is ignored (which counteracts the above if it's passed in GR).
• Like x86_64, fixed arguments landing in the overflow area are wrongly counted towards the overflow offset.

[zatrazz]

LGTM, thanks for fixing it.

llvm/trunk/lib/Transforms/Instrumentation/MemorySanitizer.cpp

for (size_t i = 0, n = VAStartInstrumentationList.size(); i < n; i++) {
       IRB.CreateMemCpy(RegSaveAreaShadowPtr, VAArgTLSCopy, CopySize, 8);
     }
   }
 };
 
 
 /// \brief AArch64-specific implementation of VarArgHelper.
 struct VarArgAArch64Helper : public VarArgHelper {
-  static const unsigned kAArch64GrArgSize = 56;
+  static const unsigned kAArch64GrArgSize = 64;
   static const unsigned kAArch64VrArgSize = 128;
 
   static const unsigned AArch64GrBegOffset = 0;
   static const unsigned AArch64GrEndOffset = kAArch64GrArgSize;
   // Make VR space aligned to 16 bytes.
-  static const unsigned AArch64VrBegOffset = AArch64GrEndOffset + 8;
+  static const unsigned AArch64VrBegOffset = AArch64GrEndOffset;
   static const unsigned AArch64VrEndOffset = AArch64VrBegOffset
                                              + kAArch64VrArgSize;
   static const unsigned AArch64VAEndOffset = AArch64VrEndOffset;
 
   Function &F;
   MemorySanitizer &MS;
   MemorySanitizerVisitor &MSV;
   Value *VAArgTLSCopy;
struct VarArgAArch64Helper : public VarArgHelper {
   // Using constant offset within the va_arg TLS array allows fast copy
   // in the finalize instrumentation.
   void visitCallSite(CallSite &CS, IRBuilder<> &IRB) override {
     unsigned GrOffset = AArch64GrBegOffset;
     unsigned VrOffset = AArch64VrBegOffset;
     unsigned OverflowOffset = AArch64VAEndOffset;
 
     const DataLayout &DL = F.getParent()->getDataLayout();
-    for (CallSite::arg_iterator ArgIt = CS.arg_begin() + 1, End = CS.arg_end();
+    for (CallSite::arg_iterator ArgIt = CS.arg_begin(), End = CS.arg_end();
          ArgIt != End; ++ArgIt) {
       Value *A = *ArgIt;
+      unsigned ArgNo = CS.getArgumentNo(ArgIt);
+      bool IsFixed = ArgNo < CS.getFunctionType()->getNumParams();
       ArgKind AK = classifyArgument(A);
       if (AK == AK_GeneralPurpose && GrOffset >= AArch64GrEndOffset)
         AK = AK_Memory;
       if (AK == AK_FloatingPoint && VrOffset >= AArch64VrEndOffset)
         AK = AK_Memory;
       Value *Base;
       switch (AK) {
         case AK_GeneralPurpose:
           Base = getShadowPtrForVAArgument(A->getType(), IRB, GrOffset);
           GrOffset += 8;
           break;
         case AK_FloatingPoint:
           Base = getShadowPtrForVAArgument(A->getType(), IRB, VrOffset);
           VrOffset += 16;
           break;
         case AK_Memory:
+          // Don't count fixed arguments in the overflow area - va_start will
+          // skip right over them.
+          if (IsFixed)
+            continue;
           uint64_t ArgSize = DL.getTypeAllocSize(A->getType());
           Base = getShadowPtrForVAArgument(A->getType(), IRB, OverflowOffset);
           OverflowOffset += alignTo(ArgSize, 8);
           break;
       }
+      // Count Gp/Vr fixed arguments to their respective offsets, but don't
+      // bother to actually store a shadow.
+      if (IsFixed)
+        continue;
       IRB.CreateAlignedStore(MSV.getShadow(A), Base, kShadowTLSAlignment);
     }
     Constant *OverflowSize =
       ConstantInt::get(IRB.getInt64Ty(), OverflowOffset - AArch64VAEndOffset);
     IRB.CreateStore(OverflowSize, MS.VAArgOverflowSizeTLS);
   }
 
   /// Compute the shadow address for a given va_arg.

llvm/trunk/test/Instrumentation/MemorySanitizer/AArch64/vararg.ll

define i32 @foo(i32 %guard, ...) {
   call void @llvm.lifetime.start(i64 32, i8* %1)
   call void @llvm.va_start(i8* %1)
   call void @llvm.va_end(i8* %1)
   call void @llvm.lifetime.end(i64 32, i8* %1)
   ret i32 0
 }
 
 ; First check if the variadic shadow values are saved in stack with correct
-; size (192 is total of general purpose registers size, 56, rounded to 16
-; plus total of floating-point registers size, 128).
+; size (192 is total of general purpose registers size, 64, plus total of
+; floating-point registers size, 128).
 
 ; CHECK-LABEL: @foo
 ; CHECK: [[A:%.*]] = load {{.*}} @__msan_va_arg_overflow_size_tls
 ; CHECK: [[B:%.*]] = add i64 192, [[A]]
 ; CHECK: alloca {{.*}} [[B]]
 
 ; We expect three memcpy operations: one for the general purpose registers,
 ; one for floating-point/SIMD ones, and one for thre remaining arguments.
 
 ; Propagate the GR shadow values on for the va_list::__gp_top, adjust the 
 ; offset in the __msan_va_arg_tls based on va_list:__gp_off, and finally
 ; issue the memcpy.
 ; CHECK: [[GRP:%.*]] = getelementptr inbounds i8, i8* {{%.*}}, i64 {{%.*}}
-; CHECK: [[GRSIZE:%.*]] = sub i64 56, {{%.*}}
+; CHECK: [[GRSIZE:%.*]] = sub i64 64, {{%.*}}
 ; CHECK: call void @llvm.memcpy.p0i8.p0i8.i64(i8* {{%.*}}, i8* [[GRP]], i64 [[GRSIZE]], i32 8, i1 false)
 
 ; Propagate the VR shadow values on for the va_list::__vr_top, adjust the 
 ; offset in the __msan_va_arg_tls based on va_list:__vr_off, and finally
 ; issue the memcpy.
 ; CHECK: [[VRP:%.*]] = getelementptr inbounds i8, i8* {{%.*}}, i64 {{%.*}}
 ; CHECK: [[VRSIZE:%.*]] = sub i64 128, {{%.*}}
 ; CHECK: call void @llvm.memcpy.p0i8.p0i8.i64(i8* {{%.*}}, i8* [[VRP]], i64 [[VRSIZE]], i32 8, i1 false)

 define i32 @bar() {
   %1 = call i32 (i32, ...) @foo(i32 0, i32 1, i32 2, double 3.000000e+00, 
                                 double 4.000000e+00, i32 5, i32 6,
                                 double 7.000000e+00, i32 8, i32 9, i32 10, i32 11)
   ret i32 %1
 }
 
 ; Save the incoming shadow value from the arguments in the __msan_va_arg_tls
-; array.  General purpose registers are saved at positions from 0 to 56, Floating
+; array.  General purpose registers are saved at positions from 0 to 64, Floating
 ; point and SIMD are saved from 64 to 192, and the remaining from 192.
 ; CHECK-LABEL: @bar
 ; CHECK: store {{.*}} @__msan_va_arg_tls {{.*}} 8
+; CHECK: store {{.*}} @__msan_va_arg_tls {{.*}} 16
 ; CHECK: store {{.*}} @__msan_va_arg_tls {{.*}} 64
 ; CHECK: store {{.*}} @__msan_va_arg_tls {{.*}} 80
-; CHECK: store {{.*}} @__msan_va_arg_tls {{.*}} 16
 ; CHECK: store {{.*}} @__msan_va_arg_tls {{.*}} 24
-; CHECK: store {{.*}} @__msan_va_arg_tls {{.*}} 96
 ; CHECK: store {{.*}} @__msan_va_arg_tls {{.*}} 32
+; CHECK: store {{.*}} @__msan_va_arg_tls {{.*}} 96
 ; CHECK: store {{.*}} @__msan_va_arg_tls {{.*}} 40
 ; CHECK: store {{.*}} @__msan_va_arg_tls {{.*}} 48
+; CHECK: store {{.*}} @__msan_va_arg_tls {{.*}} 56
 ; CHECK: store {{.*}} @__msan_va_arg_tls {{.*}} 192
 ; CHECK: store {{.*}} 8, {{.*}} @__msan_va_arg_overflow_size_tls