fix JitCall codegen for function with lambda as return value

Author: Vipul-Cariappa

lib/CppInterOp/CppInterOp.cpp

@@ -1935,9 +1935,27 @@ void collect_type_info(const FunctionDecl* FD, QualType& QT,
   PrintingPolicy Policy(C.getPrintingPolicy());
   Policy.SuppressElaboration = true;
   refType = kNotReference;
-  if (QT->isRecordType() && forArgument) {
-    get_type_as_string(QT, type_name, C, Policy);
-    return;
+  if (QT->isRecordType()) {
+    if (forArgument) {
+      get_type_as_string(QT, type_name, C, Policy);
+      return;
+    }
+    if (auto* CXXRD = QT->getAsCXXRecordDecl()) {
+      if (CXXRD->isLambda()) {
+        DeclarationName DFunction = &getASTContext().Idents.get("function");
+        auto result = getSema().getStdNamespace()->lookup(DFunction);
+        if (result.empty())
+          Cpp::Declare("#include <functional>");
+        std::string fn_name;
+        llvm::raw_string_ostream stream(fn_name);
+        Policy.FullyQualifiedName = true;
+        Policy.SuppressUnwrittenScope = true;
+        FD->getNameForDiagnostic(stream, Policy,
+                                 /*Qualified=*/false);
+        type_name = "std::function<decltype(" + fn_name + ")>::result_type";
+        return;
+      }
+    }
   }
   if (QT->isFunctionPointerType()) {
     std::string fp_typedef_name;

unittests/CppInterOp/FunctionReflectionTest.cpp

@@ -2028,6 +2028,19 @@ TEST(FunctionReflectionTest, GetFunctionCallWrapper) {
 
   auto consume_callable = Cpp::MakeFunctionCallable(consume);
   EXPECT_EQ(consume_callable.getKind(), Cpp::JitCall::kGenericCall);
+
+  Cpp::Declare(R"(
+    auto get_fn(int x) { return [x](int y){ return x + y; }; }
+  )");
+
+  Cpp::TCppScope_t get_fn = Cpp::GetNamed("get_fn");
+  EXPECT_TRUE(get_fn);
+
+  auto get_fn_callable = Cpp::MakeFunctionCallable(get_fn);
+  EXPECT_EQ(get_fn_callable.getKind(), Cpp::JitCall::kGenericCall);
+
+  EXPECT_TRUE(Cpp::IsLambdaClass(Cpp::GetFunctionReturnType(get_fn)));
+  EXPECT_FALSE(Cpp::IsLambdaClass(Cpp::GetFunctionReturnType(bar)));
 }
 
 TEST(FunctionReflectionTest, IsConstMethod) {