[clang][Interp] Implement lambda static invokers

[tbaeder]

Summary

Emit custom byte code for the static invoker function.

Unfortunately, there are a few (RVO-related) cases we can't handle yet. There's a bigger refactoring for the initialization stuff coming that will hopefully take care of that.

[tbaeder]

Ping

[aaron.ballman]

clang/lib/AST/Interp/ByteCodeEmitter.cpp

100–103

I don't like the isa followed by a cast code smell, but rewriting it to use dyn_cast is perhaps kind of ugly:

bool IsEligibleForCompilation = false;
if (const auto *MD = dyn_cast<CXXMethodDecl>())
  IsEligibleForCompilation = MD->isLambdaStaticInvoker();
if (!IsEligibleForCompilation)
  IsEligibleForCompilation = FuncDecl->isConstexpr();

WDYT?

clang/lib/AST/Interp/ByteCodeStmtGen.cpp

125–126

Why do we need to do lvalue-to-rvalue conversion at all? The caller of the lambda function call operator would have already performed that conversion, right?

[tbaeder]

clang/lib/AST/Interp/ByteCodeEmitter.cpp

100–103

I'm not a fan of that pattern either. I've seen it elsewhere in the code base, didn't like it, but couldn't come up with an objective downside either. If you're also not a fan, I have no problem with changing it.

[aaron.ballman]

LGTM!

clang/lib/AST/Interp/ByteCodeEmitter.cpp

Func = P.createFunction(FuncDecl, ParamOffset, std::move(ParamTypes),
                             std::move(ParamOffsets), HasThisPointer, HasRVO);
 
   assert(Func);
   // For not-yet-defined functions, we only create a Function instance and
   // compile their body later.
   if (!FuncDecl->isDefined())
     return Func;
 
+  // Lambda static invokers are a special case that we emit custom code for.
+  bool IsEligibleForCompilation = false;
+  if (const auto *MD = dyn_cast<CXXMethodDecl>(FuncDecl))
+    IsEligibleForCompilation = MD->isLambdaStaticInvoker();
+  if (!IsEligibleForCompilation)
+    IsEligibleForCompilation = FuncDecl->isConstexpr();
+
   // Compile the function body.
-  if (!FuncDecl->isConstexpr() || !visitFunc(FuncDecl)) {
+  if (!IsEligibleForCompilation || !visitFunc(FuncDecl)) {
     // Return a dummy function if compilation failed.
     if (BailLocation)
       return llvm::make_error<ByteCodeGenError>(*BailLocation);
     else {
       Func->setIsFullyCompiled(true);
       return Func;
     }
   } else {

clang/lib/AST/Interp/ByteCodeStmtGen.h

private:
   bool visitForStmt(const ForStmt *S);
   bool visitCXXForRangeStmt(const CXXForRangeStmt *S);
   bool visitBreakStmt(const BreakStmt *S);
   bool visitContinueStmt(const ContinueStmt *S);
   bool visitSwitchStmt(const SwitchStmt *S);
   bool visitCaseStmt(const CaseStmt *S);
   bool visitDefaultStmt(const DefaultStmt *S);
 
+  bool emitLambdaStaticInvokerBody(const CXXMethodDecl *MD);
+
   /// Type of the expression returned by the function.
   std::optional<PrimType> ReturnType;
 
   /// Switch case mapping.
   CaseMap CaseLabels;
 
   /// Point to break to.
   OptLabelTy BreakLabel;

clang/lib/AST/Interp/ByteCodeStmtGen.cpp

private:
   OptLabelTy OldDefaultLabel;
   CaseMap OldCaseLabels;
 };
 
 } // namespace interp
 } // namespace clang
 
 template <class Emitter>
+bool ByteCodeStmtGen<Emitter>::emitLambdaStaticInvokerBody(
+    const CXXMethodDecl *MD) {
+  assert(MD->isLambdaStaticInvoker());
+  assert(MD->hasBody());
+  assert(cast<CompoundStmt>(MD->getBody())->body_empty());
+
+  const CXXRecordDecl *ClosureClass = MD->getParent();
+  const CXXMethodDecl *LambdaCallOp = ClosureClass->getLambdaCallOperator();
+  assert(ClosureClass->captures_begin() == ClosureClass->captures_end());
+  const Function *Func = this->getFunction(LambdaCallOp);
+  if (!Func)
+    return false;
+  assert(Func->hasThisPointer());
+  assert(Func->getNumParams() == (MD->getNumParams() + 1 + Func->hasRVO()));
+
+  if (Func->hasRVO()) {
+    if (!this->emitRVOPtr(MD))
+      return false;
+  }
+
+  // The lambda call operator needs an instance pointer, but we don't have
+  // one here, and we don't need one either because the lambda cannot have
+  // any captures, as verified above. Emit a null pointer. This is then
+  // special-cased when interpreting to not emit any misleading diagnostics.
+  if (!this->emitNullPtr(MD))
+    return false;
+
+  // Forward all arguments from the static invoker to the lambda call operator.
+  for (const ParmVarDecl *PVD : MD->parameters()) {
+    auto It = this->Params.find(PVD);
+    assert(It != this->Params.end());
+
+    // We do the lvalue-to-rvalue conversion manually here, so no need
+    // to care about references.
+    PrimType ParamType = this->classify(PVD->getType()).value_or(PT_Ptr);
+    if (!this->emitGetParam(ParamType, It->second, MD))
+      return false;
+  }
+
+  if (!this->emitCall(Func, LambdaCallOp))
+    return false;
+
+  this->emitCleanup();
+  if (ReturnType)
+    return this->emitRet(*ReturnType, MD);
+
+  // Nothing to do, since we emitted the RVO pointer above.
+  return this->emitRetVoid(MD);
+}
+
+template <class Emitter>
 bool ByteCodeStmtGen<Emitter>::visitFunc(const FunctionDecl *F) {
   // Classify the return type.
   ReturnType = this->classify(F->getReturnType());
 
+  // Emit custom code if this is a lambda static invoker.
+  if (const auto *MD = dyn_cast<CXXMethodDecl>(F);
+      MD && MD->isLambdaStaticInvoker())
+    return this->emitLambdaStaticInvokerBody(MD);
+
   // Constructor. Set up field initializers.
   if (const auto Ctor = dyn_cast<CXXConstructorDecl>(F)) {
     const RecordDecl *RD = Ctor->getParent();
     const Record *R = this->getRecord(RD);
     if (!R)
       return false;
 
     for (const auto *Init : Ctor->inits()) {

clang/lib/AST/Interp/Function.h

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_CLANG_AST_INTERP_FUNCTION_H
 #define LLVM_CLANG_AST_INTERP_FUNCTION_H
 
 #include "Pointer.h"
 #include "Source.h"
+#include "clang/AST/ASTLambda.h"
 #include "clang/AST/Decl.h"
 #include "llvm/Support/raw_ostream.h"
 
 namespace clang {
 namespace interp {
 class Program;
 class ByteCodeEmitter;
 enum PrimType : uint32_t;

 /// # Calling Convention
 ///
 /// When calling a function, all argument values must be on the stack.
 ///
 /// If the function has a This pointer (i.e. hasThisPointer() returns true,
 /// the argument values need to be preceeded by a Pointer for the This object.
 ///
 /// If the function uses Return Value Optimization, the arguments (and
-/// potentially the This pointer) need to be proceeded by a Pointer pointing
+/// potentially the This pointer) need to be preceeded by a Pointer pointing
 /// to the location to construct the returned value.
 ///
 /// After the function has been called, it will remove all arguments,
 /// including RVO and This pointer, from the stack.
 ///
 class Function final {
 public:
   using ParamDescriptor = std::pair<PrimType, Descriptor *>;
public:
   /// Returns a specific scope.
   Scope &getScope(unsigned Idx) { return Scopes[Idx]; }
   const Scope &getScope(unsigned Idx) const { return Scopes[Idx]; }
 
   /// Returns the source information at a given PC.
   SourceInfo getSource(CodePtr PC) const;
 
   /// Checks if the function is valid to call in constexpr.
-  bool isConstexpr() const { return IsValid; }
+  bool isConstexpr() const { return IsValid || isLambdaStaticInvoker(); }
 
   /// Checks if the function is virtual.
   bool isVirtual() const;
 
   /// Checks if the function is a constructor.
   bool isConstructor() const { return isa<CXXConstructorDecl>(F); }
   /// Checks if the function is a destructor.
   bool isDestructor() const { return isa<CXXDestructorDecl>(F); }
 
   /// Returns the parent record decl, if any.
   const CXXRecordDecl *getParentDecl() const {
     if (const auto *MD = dyn_cast<CXXMethodDecl>(F))
       return MD->getParent();
     return nullptr;
   }
 
+  /// Returns whether this function is a lambda static invoker,
+  /// which we generate custom byte code for.
+  bool isLambdaStaticInvoker() const {
+    if (const auto *MD = dyn_cast<CXXMethodDecl>(F))
+      return MD->isLambdaStaticInvoker();
+    return false;
+  }
+
+  /// Returns whether this function is the call operator
+  /// of a lambda record decl.
+  bool isLambdaCallOperator() const {
+    if (const auto *MD = dyn_cast<CXXMethodDecl>(F))
+      return clang::isLambdaCallOperator(MD);
+    return false;
+  }
+
   /// Checks if the function is fully done compiling.
   bool isFullyCompiled() const { return IsFullyCompiled; }
 
   bool hasThisPointer() const { return HasThisPointer; }
 
   // Checks if the funtion already has a body attached.
   bool hasBody() const { return HasBody; }
 

clang/lib/AST/Interp/Interp.h

 
 inline bool Call(InterpState &S, CodePtr OpPC, const Function *Func) {
   if (Func->hasThisPointer()) {
     size_t ThisOffset =
         Func->getArgSize() + (Func->hasRVO() ? primSize(PT_Ptr) : 0);
 
     const Pointer &ThisPtr = S.Stk.peek<Pointer>(ThisOffset);
 
+    // If the current function is a lambda static invoker and
+    // the function we're about to call is a lambda call operator,
+    // skip the CheckInvoke, since the ThisPtr is a null pointer
+    // anyway.
+    if (!(S.Current->getFunction() &&
+          S.Current->getFunction()->isLambdaStaticInvoker() &&
+          Func->isLambdaCallOperator())) {
       if (!CheckInvoke(S, OpPC, ThisPtr))
         return false;
+    }
 
     if (S.checkingPotentialConstantExpression())
       return false;
   }
 
   if (!CheckCallable(S, OpPC, Func))
     return false;
 

clang/test/AST/Interp/lambda.cpp

constexpr int foo() {
     callThis(f);
 
     return a;
   }
   /// FIXME: This should work in the new interpreter.
   static_assert(foo() == 1); // expected-error {{not an integral constant expression}}
 }
 
+namespace StaticInvoker {
+  constexpr int sv1(int i) {
+    auto l = []() { return 12; };
+    int (*fp)() = l;
+    return fp();
+  }
+  static_assert(sv1(12) == 12);
+
+  constexpr int sv2(int i) {
+    auto l = [](int m, float f, void *A) { return m; };
+    int (*fp)(int, float, void*) = l;
+    return fp(i, 4.0f, nullptr);
+  }
+  static_assert(sv2(12) == 12);
+
+  constexpr int sv3(int i) {
+    auto l = [](int m, const int &n) { return m; };
+    int (*fp)(int, const int &) = l;
+    return fp(i, 3);
+  }
+  static_assert(sv3(12) == 12);
+
+  constexpr int sv4(int i) {
+    auto l = [](int &m) { return m; };
+    int (*fp)(int&) = l;
+    return fp(i);
+  }
+  static_assert(sv4(12) == 12);
+
+
+
+  /// FIXME: This is broken for lambda-unrelated reasons.
+#if 0
+  constexpr int sv5(int i) {
+    struct F { int a; float f; };
+    auto l = [](int m, F f) { return m; };
+    int (*fp)(int, F) = l;
+    return fp(i, F{12, 14.0});
+  }
+  static_assert(sv5(12) == 12);
+#endif
+
+  constexpr int sv6(int i) {
+    struct F { int a;
+      constexpr F(int a) : a(a) {}
+    };
+
+    auto l = [](int m) { return F(12); };
+    F (*fp)(int) = l;
+    F f = fp(i);
+
+    return fp(i).a;
+  }
+  static_assert(sv6(12) == 12);
+}