[Clang][CodeGen][Sema] Fix crash when compiling naked lambdas

clang/lib/CodeGen/CodeGenFunction.cpp

@@ -1277,44 +1277,58 @@ void CodeGenFunction::StartFunction(GlobalDecl GD, QualType RetTy,
       MD && !MD->isStatic()) {
     bool IsInLambda =
         MD->getParent()->isLambda() && MD->getOverloadedOperator() == OO_Call;
-    if (MD->isImplicitObjectMemberFunction())
-      CGM.getCXXABI().EmitInstanceFunctionProlog(*this);
-    if (IsInLambda) {
-      // We're in a lambda; figure out the captures.
-      MD->getParent()->getCaptureFields(LambdaCaptureFields,
-                                        LambdaThisCaptureField);
-      if (LambdaThisCaptureField) {
-        // If the lambda captures the object referred to by '*this' - either by
-        // value or by reference, make sure CXXThisValue points to the correct
-        // object.
-
-        // Get the lvalue for the field (which is a copy of the enclosing object
-        // or contains the address of the enclosing object).
-        LValue ThisFieldLValue = EmitLValueForLambdaField(LambdaThisCaptureField);
-        if (!LambdaThisCaptureField->getType()->isPointerType()) {
-          // If the enclosing object was captured by value, just use its
-          // address. Sign this pointer.
-          CXXThisValue = ThisFieldLValue.getPointer(*this);
-        } else {
-          // Load the lvalue pointed to by the field, since '*this' was captured
-          // by reference.
-          CXXThisValue =
-              EmitLoadOfLValue(ThisFieldLValue, SourceLocation()).getScalarVal();
+
+    const FunctionDecl *FD = dyn_cast_if_present<FunctionDecl>(D);
+    bool IsNaked = FD && FD->hasAttr<NakedAttr>();
+    if (!IsNaked) {
+      if (MD->isImplicitObjectMemberFunction())
+        CGM.getCXXABI().EmitInstanceFunctionProlog(*this);
+
+      if (IsInLambda) {
+        // We're in a lambda; figure out the captures.
+        MD->getParent()->getCaptureFields(LambdaCaptureFields,
+                                          LambdaThisCaptureField);
+        if (LambdaThisCaptureField) {
+          // If the lambda captures the object referred to by '*this' - either
+          // by value or by reference, make sure CXXThisValue points to the
+          // correct object.
+
+          // Get the lvalue for the field (which is a copy of the enclosing
+          // object or contains the address of the enclosing object).
+          LValue ThisFieldLValue =
+              EmitLValueForLambdaField(LambdaThisCaptureField);
+          if (!LambdaThisCaptureField->getType()->isPointerType()) {
+            // If the enclosing object was captured by value, just use its
+            // address. Sign this pointer.
+            CXXThisValue = ThisFieldLValue.getPointer(*this);
+          } else {
+            // Load the lvalue pointed to by the field, since '*this' was
+            // captured by reference.
+            CXXThisValue = EmitLoadOfLValue(ThisFieldLValue, SourceLocation())
+                               .getScalarVal();
+          }
         }
-      }
-      for (auto *FD : MD->getParent()->fields()) {
-        if (FD->hasCapturedVLAType()) {
-          auto *ExprArg = EmitLoadOfLValue(EmitLValueForLambdaField(FD),
-                                           SourceLocation()).getScalarVal();
-          auto VAT = FD->getCapturedVLAType();
-          VLASizeMap[VAT->getSizeExpr()] = ExprArg;
+
+        for (auto *FD : MD->getParent()->fields()) {
+          if (FD->hasCapturedVLAType()) {
+            auto *ExprArg =
+                EmitLoadOfLValue(EmitLValueForLambdaField(FD), SourceLocation())
+                    .getScalarVal();
+            auto VAT = FD->getCapturedVLAType();
+            VLASizeMap[VAT->getSizeExpr()] = ExprArg;
+          }
         }
+      } else if (MD->isImplicitObjectMemberFunction()) {
+        // Not in a lambda; just use 'this' from the method.
+        // FIXME: Should we generate a new load for each use of 'this'?  The
+        // fast register allocator would be happier...
+        CXXThisValue = CXXABIThisValue;
       }
-    } else if (MD->isImplicitObjectMemberFunction()) {
-      // Not in a lambda; just use 'this' from the method.
-      // FIXME: Should we generate a new load for each use of 'this'?  The
-      // fast register allocator would be happier...
-      CXXThisValue = CXXABIThisValue;
+    } else if (IsInLambda && MD->isImplicitObjectMemberFunction()) {
+      // Populate capture fields metadata for analysis. We skip
+      // EmitInstanceProlog to avoid emitting prologue code.
+      MD->getParent()->getCaptureFields(LambdaCaptureFields,
+                                        LambdaThisCaptureField);
     }
 
     // Check the 'this' pointer once per function, if it's available.

clang/lib/Sema/SemaLambda.cpp

@@ -2332,6 +2332,20 @@ ExprResult Sema::BuildLambdaExpr(SourceLocation StartLoc,
     maybeAddDeclWithEffects(LSI->CallOperator);
   }
 
+  // This is for GCC compatibility. If any lambda captures are actually used in the
+  // function body. GCC silently removes the naked attribute when captures are
+  // ODR-used, as naked functions cannot have prologues to set up the closure.
+  if (CallOperator->hasAttr<NakedAttr>() && !Captures.empty()) {
+    // If any captures are ODR-used by examining the capture list
+    // that was already analyzed during semantic analysis, drop it.
+    for (const Capture &Cap : LSI->Captures) {
+      if (Cap.isODRUsed()) {
+        CallOperator->dropAttr<NakedAttr>();
+        break;
+      }
+    }
+  }
+
   return MaybeBindToTemporary(Lambda);
 }
 

clang/test/CodeGenCXX/naked-lambda-capture-multi.cpp

@@ -0,0 +1,18 @@
+// RUN: %clang_cc1 -triple x86_64-pc-linux-gnu -emit-llvm %s -o - | FileCheck %s
+
+void test_naked_lambda_capture_multi() {
+  int x = 42;
+  int y = 100;
+  auto l = [&x, y]() __attribute__((naked)) {
+    asm volatile("retq");
+  };
+  l();
+}
+
+// CHECK-LABEL: define {{.*}} @"_ZZ31test_naked_lambda_capture_multivENK3$_0clEv"
+// CHECK-NOT: load i32
+// CHECK-NOT: load ptr
+// CHECK-NOT: getelementptr
+// CHECK-NOT: alloca
+// CHECK: call void asm sideeffect "retq"
+// CHECK-NEXT: unreachable

clang/test/CodeGenCXX/naked-lambda-capture-this.cpp

@@ -0,0 +1,28 @@
+// RUN: %clang_cc1 -triple x86_64-pc-linux-gnu -emit-llvm %s -o - | FileCheck %s
+// RUN: %clang_cc1 -triple x86_64-pc-linux-gnu -S %s -o - | FileCheck %s --check-prefix=ASM
+
+struct S {
+  int member;
+  void test_naked_lambda_capture_this() {
+    auto l = [this]() __attribute__((naked)) {
+      asm volatile("retq");
+    };
+    l();
+  }
+};
+
+void test() {
+  S s;
+  s.test_naked_lambda_capture_this();
+}
+
+// CHECK-LABEL: define {{.*}} @_ZZN1S30test_naked_lambda_capture_thisEvENKUlvE_clEv
+// CHECK-NOT: load ptr
+// CHECK-NOT: getelementptr
+// CHECK-NOT: alloca
+// CHECK: call void asm sideeffect "retq"
+
+// ASM-LABEL: _ZZN1S30test_naked_lambda_capture_thisEvENKUlvE_clEv:
+// ASM-NOT: push
+// ASM-NOT: mov
+// ASM: retq

clang/test/CodeGenCXX/naked-lambda-capture-var.cpp

@@ -0,0 +1,21 @@
+// RUN: %clang_cc1 -triple x86_64-pc-linux-gnu -emit-llvm %s -o - | FileCheck %s
+// RUN: %clang_cc1 -triple x86_64-pc-linux-gnu -S %s -o - | FileCheck %s --check-prefix=ASM
+
+void test_naked_lambda_capture_var() {
+  int x = 42;
+  auto l = [x]() __attribute__((naked)) {
+    asm volatile("retq");
+  };
+  l();
+}
+
+// CHECK-LABEL: define {{.*}} @"_ZZ29test_naked_lambda_capture_varvENK3$_0clEv"
+// CHECK-NOT: load i32
+// CHECK-NOT: alloca
+// CHECK-NOT: getelementptr
+// CHECK: call void asm sideeffect "retq"
+
+// ASM-LABEL: _ZZ29test_naked_lambda_capture_varvENK3$_0clEv:
+// ASM-NOT: push
+// ASM-NOT: mov
+// ASM: retq

clang/test/CodeGenCXX/naked-lambda-odr-used-captures.cpp

@@ -0,0 +1,17 @@
+// RUN: %clang_cc1 -triple x86_64-pc-linux-gnu -emit-llvm %s -o - | FileCheck %s
+
+// Test that naked attribute is removed when captures are ODR-used (GCC compat)
+void test_odr_used_captures() {
+  int x = 42;
+  int y = 6;
+  auto l = [x, &y]() __attribute__((naked)) {
+    asm volatile("movl %0, %%eax\n\tmovl %1, %%ebx\n\tretq" : : "r"(x), "r"(y));
+  };
+  l();
+}
+
+// CHECK-LABEL: define internal void @"_ZZ22test_odr_used_capturesvENK3$_0clEv"
+// CHECK-NOT: naked
+// CHECK: alloca
+// CHECK: store
+

clang/test/CodeGenCXX/naked-lambda.cpp

@@ -0,0 +1,18 @@
+// RUN: %clang_cc1 -triple x86_64-pc-linux-gnu -emit-llvm %s -o - | FileCheck %s
+// RUN: %clang_cc1 -triple x86_64-pc-linux-gnu -S %s -o - | FileCheck %s --check-prefix=ASM
+
+void test_naked_lambda() {
+  auto l = []() __attribute__((naked)) {
+    asm volatile("retq");
+  };
+  l();
+}
+
+// CHECK: define internal void @"_ZZ17test_naked_lambdavENK3$_0clEv"
+// CHECK-NOT: alloca
+// CHECK-NOT: store
+// CHECK-NOT: call void @_ZN
+// ASM-LABEL: _ZZ17test_naked_lambdavENK3$_0clEv:
+// ASM-NOT: push
+// ASM-NOT: pop
+// ASM: retq

clang/test/SemaCXX/naked-lambda-odr.cpp

@@ -0,0 +1,18 @@
+// RUN: %clang_cc1 -fsyntax-only -verify %s -triple x86_64-pc-linux-gnu
+
+void uses_capture() {
+  int x = 42;
+  int y = 6;
+  auto l = [x, &y]() __attribute__((naked)) { // expected-no-diagnostics
+    asm volatile("movl %0, %%eax\n\tmovl %1, %%ebx\n\tretq" : : "r"(x), "r"(y));
+  };
+  l();
+}
+
+void unused_captures() {
+  int x = 42;
+  auto l = [x]() __attribute__((naked)) { // expected-no-diagnostics
+    asm volatile("retq");
+  };
+  l();
+}