[CodeGen] Fix a crash that occurs when attribute "naked" is attached to a c++ member function

[ahatanak]

Summary

clang crashes when it compiles a member function or a lambda is marked as "naked":

struct S {

__attribute__((naked)) S() {}

}

To fix this bug, this patch defines a function that cleans up the naked function removing all allocas and their users and all basic blocks except the entry block. Also, it makes changes so that, if a non-void function is marked as naked, llvm.trap isn't inserted into the function even if it doesn't have a return statement.

[ahatanak]

ping

[hans]

Just out of curiosity, where does this come up in practice?

It seems a little backward that we're first emitting a bunch of instructions, only to remove them later. It would be nice if for naked function we wouldn't emit them in the first place. But maybe that's not practical.

Anyway, this seems OK to me but I'd like to hear what Reid thinks too.

lib/CodeGen/CodeGenFunction.cpp

1973

would declaring I as llvm::Instruction* allow you to get rid off the "&*" stuff here and below? If so, same for "auto BB" further down.

1991

i don't think there's any need for curlies around the else branch here, especially since there are none for the then-branch

[ahatanak]

In D15599#325113, @hans wrote:

Just out of curiosity, where does this come up in practice?

I only have a short test case a user provided so t's not really clear whether it was necessary to mark a member function naked or there were other ways to get the same behavior. But unless we want to disallow marking member functions as naked, we shouldn't let clang crash when it compiles such functions.

It seems a little backward that we're first emitting a bunch of instructions, only to remove them later. It would be nice if for naked function we wouldn't emit them in the first place. But maybe that's not practical.

I initially tried to block the instructions from being emitted but I ended up checking hasAttr<NakedAttr>() in a lot of places. I agree approach taken in this patch might look a bit backward, but it seemed better than the alternative.

Anyway, this seems OK to me but I'd like to hear what Reid thinks too.

[ahatanak]

ping

[manmanren]

Thanks for working on this! A few comments inline.

Manman

lib/CodeGen/CodeGenFunction.cpp

1981

Do we need to check if the user has been erased? I am not sure if the same instruction can be pushed multiple times into InstrsToRemove.

1998

Is it possible that T has been erased in the above?

[ahatanak]

lib/CodeGen/CodeGenFunction.cpp

1981

I think you are right: it's incorrect to assume an instruction cannot be pushed twice. I'll change the type of InstrsToRemove to a set to prevent that from happenning.

1998

I'm not sure I understand. getTerminator will return null if there is no terminator or the basic block is empty, so I think this code is safe?

[ahatanak]

lib/CodeGen/CodeGenFunction.cpp

1973

We can git rid of the "&*" in the code here if we declare I as Instruction*, but we will need to add "&*" to Entry->begin() to convert it to an Instruction*. In addition, we will have to use getNextNode instead of ++II or I++, so it seems to me that using an iterator is better in this case.

[ahatanak]

Address review comments from Hans and Manman.

[manmanren]

lib/CodeGen/CodeGenFunction.cpp

1998

Yes, you are right.

[ahatanak]

ping and rebase.

[rnk]

I guess I'm OK with the approach.

lib/CodeGen/CodeGenFunction.cpp

1968

This isn't the right approach. Look at Function::dropAllReferences and try modelling this code on that. This does a lot of unnecessary RAUW when we can just null out all operands and delete all instructions except for the inline asm CallInst.

test/CodeGen/attr-naked.c

20

unrelated change?

test/CodeGenCXX/attr-naked.cpp

8

Can you add a test involving vtable thunks? I'm worried that we will clean up the thunk for a naked virtual method implementation. The code looks correct today, but I'd like to defend against that possibility in the future.

31

Won't this require asserts? We don't name BBs in NDEBUG builds.

[ahatanak]

Address review comments.

• Rewrite cleanupNakedFunction.
• Fix test case attr-naked.cpp. Check that the thunk function doesn't get removed when the virtual function is marked "naked".

[ahatanak]

test/CodeGen/attr-naked.c

20

This is a side effect of not exiting early in CodeGenFunction::EmitFunctionProlog.

test/CodeGenCXX/attr-naked.cpp

9

I've also added a check to make sure the thunk isn't marked naked.

lib/CodeGen/CGCall.cpp

if (NNAttr->isNonNull(ArgNo))
       return NNAttr;
   }
   return nullptr;
 }
 
 void CodeGenFunction::EmitFunctionProlog(const CGFunctionInfo &FI,
                                          llvm::Function *Fn,
                                          const FunctionArgList &Args) {
-  if (CurCodeDecl && CurCodeDecl->hasAttr<NakedAttr>())
-    // Naked functions don't have prologues.
-    return;
-
   // If this is an implicit-return-zero function, go ahead and
   // initialize the return value.  TODO: it might be nice to have
   // a more general mechanism for this that didn't require synthesized
   // return statements.
   if (const FunctionDecl *FD = dyn_cast_or_null<FunctionDecl>(CurCodeDecl)) {
     if (FD->hasImplicitReturnZero()) {
       QualType RetTy = FD->getReturnType().getUnqualifiedType();
       llvm::Type* LLVMTy = CGM.getTypes().ConvertType(RetTy);
static llvm::StoreInst *findDominatingStoreToReturnValue(CodeGenFunction &CGF) {
   // Okay, the store's basic block dominates the insertion point; we
   // can do our thing.
   return store;
 }
 
 void CodeGenFunction::EmitFunctionEpilog(const CGFunctionInfo &FI,
                                          bool EmitRetDbgLoc,
                                          SourceLocation EndLoc) {
-  if (CurCodeDecl && CurCodeDecl->hasAttr<NakedAttr>()) {
-    // Naked functions don't have epilogues.
-    Builder.CreateUnreachable();
-    return;
-  }
-
   // Functions with no result always return void.
   if (!ReturnValue.isValid()) {
     Builder.CreateRetVoid();
     return;
   }
 
   llvm::DebugLoc RetDbgLoc;
   llvm::Value *RV = nullptr;

lib/CodeGen/CGStmt.cpp

if (Info.isReadWrite()) {
       else
         InOutConstraints += OutputConstraint;
 
       InOutArgTypes.push_back(Arg->getType());
       InOutArgs.push_back(Arg);
     }
   }
 
-  // If this is a Microsoft-style asm blob, store the return registers (EAX:EDX)
-  // to the return value slot. Only do this when returning in registers.
-  if (isa<MSAsmStmt>(&S)) {
+  // If this is a Microsoft-style asm blob that is not naked, store the return
+  // registers (EAX:EDX) to the return value slot. Only do this when returning
+  // in registers.
+  if (isa<MSAsmStmt>(&S) && !isNakedFunction()) {
     const ABIArgInfo &RetAI = CurFnInfo->getReturnInfo();
     if (RetAI.isDirect() || RetAI.isExtend()) {
       // Make a fake lvalue for the return value slot.
       LValue ReturnSlot = MakeAddrLValue(ReturnValue, FnRetTy);
       CGM.getTargetCodeGenInfo().addReturnRegisterOutputs(
           *this, ReturnSlot, Constraints, ResultRegTypes, ResultTruncRegTypes,
           ResultRegDests, AsmString, S.getNumOutputs());
       SawAsmBlock = true;

lib/CodeGen/CodeGenFunction.h

llvm::Value *evaluateOrEmitBuiltinObjectSize(const Expr *E, unsigned Type,
                                                llvm::IntegerType *ResType);
 
   /// \brief Emits the size of E, as required by __builtin_object_size. This
   /// function is aware of pass_object_size parameters, and will act accordingly
   /// if E is a parameter with the pass_object_size attribute.
   llvm::Value *emitBuiltinObjectSize(const Expr *E, unsigned Type,
                                      llvm::IntegerType *ResType);
 
+  /// Return true if CurCodeDecl has attribute Naked.
+  bool isNakedFunction() const;
+
+  /// Clean up naked functions removing allocas and their users and all blocks
+  /// except the entry.
+  void cleanupNakedFunction();
+
 public:
 #ifndef NDEBUG
   // Determine whether the given argument is an Objective-C method
   // that may have type parameters in its signature.
   static bool isObjCMethodWithTypeParams(const ObjCMethodDecl *method) {
     const DeclContext *dc = method->getDeclContext();
     if (const ObjCInterfaceDecl *classDecl= dyn_cast<ObjCInterfaceDecl>(dc)) {
       return classDecl->getTypeParamListAsWritten();

lib/CodeGen/CodeGenFunction.cpp

if (!EscapedLocals.empty()) {
     EscapeArgs.resize(EscapedLocals.size());
     for (auto &Pair : EscapedLocals)
       EscapeArgs[Pair.second] = Pair.first;
     llvm::Function *FrameEscapeFn = llvm::Intrinsic::getDeclaration(
         &CGM.getModule(), llvm::Intrinsic::localescape);
     CGBuilderTy(*this, AllocaInsertPt).CreateCall(FrameEscapeFn, EscapeArgs);
   }
 
+  if (isNakedFunction())
+    cleanupNakedFunction();
+
   // Remove the AllocaInsertPt instruction, which is just a convenience for us.
   llvm::Instruction *Ptr = AllocaInsertPt;
   AllocaInsertPt = nullptr;
   Ptr->eraseFromParent();
 
   // If someone took the address of a label but never did an indirect goto, we
   // made a zero entry PHI node, which is illegal, zap it now.
   if (IndirectBranch) {
void CodeGenFunction::GenerateCode(GlobalDecl GD, llvm::Function *Fn,
 
   // C++11 [stmt.return]p2:
   //   Flowing off the end of a function [...] results in undefined behavior in
   //   a value-returning function.
   // C11 6.9.1p12:
   //   If the '}' that terminates a function is reached, and the value of the
   //   function call is used by the caller, the behavior is undefined.
   if (getLangOpts().CPlusPlus && !FD->hasImplicitReturnZero() && !SawAsmBlock &&
-      !FD->getReturnType()->isVoidType() && Builder.GetInsertBlock()) {
+      !FD->getReturnType()->isVoidType() && Builder.GetInsertBlock() &&
+      !isNakedFunction()) {
     if (SanOpts.has(SanitizerKind::Return)) {
       SanitizerScope SanScope(this);
       llvm::Value *IsFalse = Builder.getFalse();
       EmitCheck(std::make_pair(IsFalse, SanitizerKind::Return),
                 "missing_return", EmitCheckSourceLocation(FD->getLocation()),
                 None);
     } else if (CGM.getCodeGenOpts().OptimizationLevel == 0) {
       EmitTrapCall(llvm::Intrinsic::trap);
for (const auto &F : CalleeFeatureMap) {
         ReqFeatures.push_back(F.getKey());
     }
     if (!hasRequiredFeatures(ReqFeatures, CGM, FD, MissingFeature))
       CGM.getDiags().Report(E->getLocStart(), diag::err_function_needs_feature)
           << FD->getDeclName() << TargetDecl->getDeclName() << MissingFeature;
   }
 }
 
+bool CodeGenFunction::isNakedFunction() const {
+  return CurCodeDecl && CurCodeDecl->hasAttr<NakedAttr>();
+}
+
+void CodeGenFunction::cleanupNakedFunction() {
+  llvm::SmallPtrSet<llvm::Instruction *, 8> InstrsToRemove;
+  llvm::BasicBlock *EntryBB = AllocaInsertPt->getParent();
+
+  // Erase all allocas and their users.
+  for (auto I = EntryBB->begin(); &*I != &*AllocaInsertPt; ++I)
+    if (auto Alloca = dyn_cast<llvm::AllocaInst>(&*I))
+      InstrsToRemove.insert(Alloca);
+
+  while (!InstrsToRemove.empty()) {
+    auto *I = *InstrsToRemove.begin();
+
+    for (auto *U : I->users())
+      InstrsToRemove.insert(cast<llvm::Instruction>(U));
+
+    I->replaceAllUsesWith(llvm::UndefValue::get(I->getType()));
+    I->eraseFromParent();
+    InstrsToRemove.erase(I);
+  }
+
+  // Erase all blocks except the entry.
+  for (auto BB = CurFn->begin(); BB != CurFn->end();) {
+    if (&*BB == EntryBB)
+      ++BB;
+    else
+      BB++->eraseFromParent();
+  }
+
+  // Erase terminator of the entry block.
+  if (auto *T = EntryBB->getTerminator())
+    T->eraseFromParent();
+
+  // Insert an unreachable at the end of the entry block.
+  Builder.SetInsertPoint(EntryBB, EntryBB->end());
+  Builder.CreateUnreachable();
+}
+
 void CodeGenFunction::EmitSanitizerStatReport(llvm::SanitizerStatKind SSK) {
   if (!CGM.getCodeGenOpts().SanitizeStats)
     return;
 
   llvm::IRBuilder<> IRB(Builder.GetInsertBlock(), Builder.GetInsertPoint());
   IRB.SetCurrentDebugLocation(Builder.getCurrentDebugLocation());
   CGM.getSanStats().create(IRB, SSK);
 }

test/CodeGen/attr-naked.c

 // Make sure this doesn't explode in the verifier.
 // (It doesn't really make sense, but it isn't invalid.)
 // CHECK: define void @t2() [[NAKED]] {
 __attribute((naked, always_inline)) void t2() {
 }
 
 // Make sure not to generate prolog or epilog for naked functions.
 __attribute((naked)) void t3(int x) {
-// CHECK: define void @t3(i32)
+// CHECK: define void @t3(i32 %x)
 // CHECK-NOT: alloca
 // CHECK-NOT: store
 // CHECK: unreachable
 }
 
 // CHECK: attributes [[NAKED]] = { naked noinline nounwind{{.*}} }

test/CodeGenCXX/attr-naked.cpp

+// RUN: %clang_cc1 -triple=x86_64-apple-darwin -std=c++11 -emit-llvm -o - %s | FileCheck %s
+
+class Class1 {
+public:
+  void __attribute__((naked)) m1() { __asm__ volatile("retq"); }
+  int __attribute__((naked)) m2() { __asm__ volatile("retq"); }
+};
+
+int foo1(int a, Class1 *c1) {
+  auto Fn = []() __attribute__((naked))->int {
+    __asm__ volatile("retq");
+  };
+
+  if (a < 0)
+    return Fn();
+
+  if (a > 0)
+    return c1->m2();
+  
+  c1->m1();
+  return 1;
+}
+
+// Check for the absence of llvm.trap.
+int __attribute__((naked)) foo2(int a) {
+  __asm__ volatile("retq");
+}
+
+// CHECK: define internal i32 @"_ZZ4foo1iP6Class1ENK3$_0clEv"(%class.anon* %this) [[ATTR1:#[0-9]]]
+// CHECK-NEXT: entry:
+// CHECK-NEXT: call void asm sideeffect "retq
+// CHECK-NEXT: unreachable
+// CHECK-NEXT: }
+
+
+// CHECK: define linkonce_odr i32 @_ZN6Class12m2Ev(%class.Class1* %this) [[ATTR2:#[0-9]]]
+// CHECK-NEXT: entry:
+// CHECK-NEXT: call void asm sideeffect "retq
+// CHECK-NEXT: unreachable
+// CHECK-NEXT: }
+
+// CHECK: define linkonce_odr void @_ZN6Class12m1Ev(%class.Class1* %this) [[ATTR2]]
+// CHECK-NEXT: entry:
+// CHECK-NEXT: call void asm sideeffect "retq
+// CHECK-NEXT: unreachable
+// CHECK-NEXT: }
+
+// CHECK: define i32 @_Z4foo2i(i32 %a) [[ATTR2]]
+// CHECK-NEXT: entry:
+// CHECK-NEXT: call void asm sideeffect "retq
+// CHECK-NEXT: unreachable
+// CHECK-NEXT: }
+
+// CHECK: attributes [[ATTR1]] = { {{.*}}naked{{.*}} }
+// CHECK: attributes [[ATTR2]] = { {{.*}}naked{{.*}} }

test/CodeGenCXX/ms-inline-asm-return.cpp

 // REQUIRES: x86-registered-target
-// RUN: %clang_cc1 %s -triple i686-pc-windows-msvc -emit-llvm -o - -fasm-blocks | FileCheck %s
+// RUN: %clang_cc1 %s -triple i686-pc-windows-msvc -emit-llvm -o - -fasm-blocks -fms-compatibility | FileCheck %s
 
 // Check that we take EAX or EAX:EDX and return it from these functions for MSVC
 // compatibility.
 
 extern "C" {
 
 long long f_i64() {
   __asm {

 } // extern "C"
 
 int main() {
   __asm xor eax, eax
 }
 // CHECK-LABEL: define i32 @main()
 // CHECK: %[[r:[^ ]*]] = call i32 asm sideeffect inteldialect "xor eax, eax", "={eax},{{.*}}"
 // CHECK: ret i32 %[[r]]
+
+// Don't set the return value if the function is marked as naked. 
+__declspec(naked) int nakedFunc(int a, int b)
+{
+  __asm {
+    ret
+  }
+}
+
+// CHECK: define i32 @{{.*}}nakedFunc
+// CHECK-NEXT: entry:
+// CHECK-NEXT: call void asm sideeffect inteldialect "ret
+// CHECK-NEXT: unreachable
+// CHECK-NEXT: }