//===--- Context.cpp - Context for the constexpr VM -------------*- C++ -*-===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//

#include "Context.h"
#include "Boolean.h"
#include "ByteCodeEmitter.h"
#include "Char.h"
#include "Compiler.h"
#include "EvalEmitter.h"
#include "Integral.h"
#include "InterpFrame.h"
#include "InterpHelpers.h"
#include "InterpStack.h"
#include "Pointer.h"
#include "PrimType.h"
#include "Program.h"
#include "clang/AST/ASTLambda.h"
#include "clang/AST/Expr.h"
#include "clang/Basic/TargetInfo.h"

using namespace clang;
using namespace clang::interp;

Context::Context(ASTContext &Ctx) : Ctx(Ctx), P(new Program(*this)) {
  this->ShortWidth = Ctx.getTargetInfo().getShortWidth();
  this->IntWidth = Ctx.getTargetInfo().getIntWidth();
  this->LongWidth = Ctx.getTargetInfo().getLongWidth();
  this->LongLongWidth = Ctx.getTargetInfo().getLongLongWidth();
  assert(Ctx.getTargetInfo().getCharWidth() == 8 &&
         "We're assuming 8 bit chars");
}

Context::~Context() = default;

bool Context::isPotentialConstantExpr(State &Parent, const FunctionDecl *FD) {
  assert(Stk.empty());

  // Get a function handle.
  const Function *Func = getOrCreateFunction(FD);
  if (!Func)
    return false;

  // Compile the function.
  Compiler<ByteCodeEmitter>(*this, *P).compileFunc(
      FD, const_cast<Function *>(Func));

  if (!Func->isValid())
    return false;

  ++EvalID;
  // And run it.
  return Run(Parent, Func);
}

void Context::isPotentialConstantExprUnevaluated(State &Parent, const Expr *E,
                                                 const FunctionDecl *FD) {
  assert(Stk.empty());
  ++EvalID;
  size_t StackSizeBefore = Stk.size();
  Compiler<EvalEmitter> C(*this, *P, Parent, Stk);

  if (!C.interpretCall(FD, E)) {
    C.cleanup();
    Stk.clearTo(StackSizeBefore);
  }
}

bool Context::evaluateAsRValue(State &Parent, const Expr *E, APValue &Result) {
  ++EvalID;
  bool Recursing = !Stk.empty();
  size_t StackSizeBefore = Stk.size();
  Compiler<EvalEmitter> C(*this, *P, Parent, Stk);

  auto Res = C.interpretExpr(E, /*ConvertResultToRValue=*/E->isGLValue());

  if (Res.isInvalid()) {
    C.cleanup();
    Stk.clearTo(StackSizeBefore);
    return false;
  }

  if (!Recursing) {
    // We *can* actually get here with a non-empty stack, since
    // things like InterpState::noteSideEffect() exist.
    C.cleanup();
#ifndef NDEBUG
    // Make sure we don't rely on some value being still alive in
    // InterpStack memory.
    Stk.clearTo(StackSizeBefore);
#endif
  }

  Result = Res.stealAPValue();

  return true;
}

bool Context::evaluate(State &Parent, const Expr *E, APValue &Result,
                       ConstantExprKind Kind) {
  ++EvalID;
  bool Recursing = !Stk.empty();
  size_t StackSizeBefore = Stk.size();
  Compiler<EvalEmitter> C(*this, *P, Parent, Stk);

  auto Res = C.interpretExpr(E, /*ConvertResultToRValue=*/false,
                             /*DestroyToplevelScope=*/true);
  if (Res.isInvalid()) {
    C.cleanup();
    Stk.clearTo(StackSizeBefore);
    return false;
  }

  if (!Recursing) {
    assert(Stk.empty());
    C.cleanup();
#ifndef NDEBUG
    // Make sure we don't rely on some value being still alive in
    // InterpStack memory.
    Stk.clearTo(StackSizeBefore);
#endif
  }

  Result = Res.stealAPValue();
  return true;
}

bool Context::evaluateAsInitializer(State &Parent, const VarDecl *VD,
                                    const Expr *Init, APValue &Result) {
  ++EvalID;
  bool Recursing = !Stk.empty();
  size_t StackSizeBefore = Stk.size();
  Compiler<EvalEmitter> C(*this, *P, Parent, Stk);

  bool CheckGlobalInitialized =
      shouldBeGloballyIndexed(VD) &&
      (VD->getType()->isRecordType() || VD->getType()->isArrayType());
  auto Res = C.interpretDecl(VD, Init, CheckGlobalInitialized);
  if (Res.isInvalid()) {
    C.cleanup();
    Stk.clearTo(StackSizeBefore);

    return false;
  }

  if (!Recursing) {
    assert(Stk.empty());
    C.cleanup();
#ifndef NDEBUG
    // Make sure we don't rely on some value being still alive in
    // InterpStack memory.
    Stk.clearTo(StackSizeBefore);
#endif
  }

  Result = Res.stealAPValue();
  return true;
}

bool Context::evaluateDestruction(State &Parent, const VarDecl *VD,
                                  APValue Value) {
  assert(Stk.empty());
  Compiler<EvalEmitter> C(*this, *P, Parent, Stk);

  auto Res = C.interpretDestructor(VD, Value);

  if (Res.isInvalid()) {
    C.cleanup();
    Stk.clear();
    return false;
  }

  return true;
}

template <typename ResultT>
bool Context::evaluateStringRepr(State &Parent, const Expr *SizeExpr,
                                 const Expr *PtrExpr, ResultT &Result) {
  assert(Stk.empty());
  Compiler<EvalEmitter> C(*this, *P, Parent, Stk);

  // Evaluate size value.
  APValue SizeValue;
  if (!evaluateAsRValue(Parent, SizeExpr, SizeValue))
    return false;

  if (!SizeValue.isInt())
    return false;
  uint64_t Size = SizeValue.getInt().getZExtValue();

  auto PtrRes = C.interpretAsPointer(PtrExpr, [&](const Pointer &Ptr) {
    if (Size == 0) {
      if constexpr (std::is_same_v<ResultT, APValue>)
        Result = APValue(APValue::UninitArray{}, 0, 0);
      return true;
    }

    if (!Ptr.isLive() || !Ptr.getFieldDesc()->isPrimitiveArray())
      return false;

    // Must be char.
    if (Ptr.getFieldDesc()->getElemDataSize() != 1 /*bytes*/)
      return false;

    if (Size > Ptr.getNumElems()) {
      Parent.FFDiag(SizeExpr, diag::note_constexpr_access_past_end) << AK_Read;
      Size = Ptr.getNumElems();
    }

    if constexpr (std::is_same_v<ResultT, APValue>) {
      QualType CharTy = PtrExpr->getType()->getPointeeType();
      Result = APValue(APValue::UninitArray{}, Size, Size);
      for (uint64_t I = 0; I != Size; ++I) {
        if (std::optional<APValue> ElemVal =
                Ptr.atIndex(I).toRValue(*this, CharTy))
          Result.getArrayInitializedElt(I) = *ElemVal;
        else
          return false;
      }
    } else {
      assert((std::is_same_v<ResultT, std::string>));
      if (Size < Result.max_size())
        Result.resize(Size);
      Result.assign(reinterpret_cast<const char *>(Ptr.getRawAddress()), Size);
    }

    return true;
  });

  if (PtrRes.isInvalid()) {
    C.cleanup();
    Stk.clear();
    return false;
  }

  return true;
}

bool Context::evaluateCharRange(State &Parent, const Expr *SizeExpr,
                                const Expr *PtrExpr, APValue &Result) {
  assert(SizeExpr);
  assert(PtrExpr);

  return evaluateStringRepr(Parent, SizeExpr, PtrExpr, Result);
}

bool Context::evaluateCharRange(State &Parent, const Expr *SizeExpr,
                                const Expr *PtrExpr, std::string &Result) {
  assert(SizeExpr);
  assert(PtrExpr);

  return evaluateStringRepr(Parent, SizeExpr, PtrExpr, Result);
}

bool Context::evaluateString(State &Parent, const Expr *E,
                             std::string &Result) {
  assert(Stk.empty());
  Compiler<EvalEmitter> C(*this, *P, Parent, Stk);

  auto PtrRes = C.interpretAsPointer(E, [&](const Pointer &Ptr) {
    if (!Ptr.isBlockPointer())
      return false;

    const Descriptor *FieldDesc = Ptr.getFieldDesc();
    if (!FieldDesc->isPrimitiveArray())
      return false;

    if (!Ptr.isConst())
      return false;

    unsigned N = Ptr.getNumElems();

    if (Ptr.elemSize() == 1 /* bytes */) {
      const char *Chars = reinterpret_cast<const char *>(Ptr.getRawAddress());
      unsigned Length = strnlen(Chars, N);
      // Wasn't null terminated.
      if (N == Length)
        return false;
      Result.assign(Chars, Length);
      return true;
    }

    PrimType ElemT = FieldDesc->getPrimType();
    for (unsigned I = Ptr.getIndex(); I != N; ++I) {
      INT_TYPE_SWITCH(ElemT, {
        auto Elem = Ptr.elem<T>(I);
        if (Elem.isZero())
          return true;
        Result.push_back(static_cast<char>(Elem));
      });
    }
    // We didn't find a 0 byte.
    return false;
  });

  if (PtrRes.isInvalid()) {
    C.cleanup();
    Stk.clear();
    return false;
  }
  return true;
}

std::optional<uint64_t> Context::evaluateStrlen(State &Parent, const Expr *E) {
  assert(Stk.empty());
  Compiler<EvalEmitter> C(*this, *P, Parent, Stk);

  std::optional<uint64_t> Result;
  auto PtrRes = C.interpretAsPointer(E, [&](const Pointer &Ptr) {
    if (!Ptr.isBlockPointer())
      return false;

    const Descriptor *FieldDesc = Ptr.getFieldDesc();
    if (!FieldDesc->isPrimitiveArray())
      return false;

    if (Ptr.isDummy() || Ptr.isUnknownSizeArray() || Ptr.isPastEnd())
      return false;

    PrimType ElemT = FieldDesc->getPrimType();
    if (!isIntegerType(ElemT))
      return false;

    unsigned N = Ptr.getNumElems();
    if (Ptr.elemSize() == 1) {
      unsigned Size = N - Ptr.getIndex();
      Result =
          strnlen(reinterpret_cast<const char *>(Ptr.getRawAddress()), Size);
      return Result != Size;
    }

    Result = 0;
    for (unsigned I = Ptr.getIndex(); I != N; ++I) {
      INT_TYPE_SWITCH(ElemT, {
        auto Elem = Ptr.elem<T>(I);
        if (Elem.isZero())
          return true;
        ++(*Result);
      });
    }
    // We didn't find a 0 byte.
    return false;
  });

  if (PtrRes.isInvalid()) {
    C.cleanup();
    Stk.clear();
    return std::nullopt;
  }
  return Result;
}

std::optional<uint64_t>
Context::tryEvaluateObjectSize(State &Parent, const Expr *E, unsigned Kind) {
  assert(Stk.empty());
  Compiler<EvalEmitter> C(*this, *P, Parent, Stk);

  std::optional<uint64_t> Result;

  auto PtrRes = C.interpretAsLValuePointer(E, [&](const Pointer &Ptr) {
    const Descriptor *DeclDesc = Ptr.getDeclDesc();
    if (!DeclDesc)
      return false;

    QualType T = DeclDesc->getType().getNonReferenceType();
    if (T->isIncompleteType() || T->isFunctionType() ||
        !T->isConstantSizeType())
      return false;

    Pointer P = Ptr;
    if (auto ObjectSize = evaluateBuiltinObjectSize(getASTContext(), Kind, P)) {
      Result = *ObjectSize;
      return true;
    }
    return false;
  });

  if (PtrRes.isInvalid()) {
    C.cleanup();
    Stk.clear();
    return std::nullopt;
  }
  return Result;
}

const LangOptions &Context::getLangOpts() const { return Ctx.getLangOpts(); }

static PrimType integralTypeToPrimTypeS(unsigned BitWidth) {
  switch (BitWidth) {
  case 64:
    return PT_Sint64;
  case 32:
    return PT_Sint32;
  case 16:
    return PT_Sint16;
  case 8:
    return PT_Sint8;
  default:
    return PT_IntAPS;
  }
  llvm_unreachable("Unhandled BitWidth");
}

static PrimType integralTypeToPrimTypeU(unsigned BitWidth) {
  switch (BitWidth) {
  case 64:
    return PT_Uint64;
  case 32:
    return PT_Uint32;
  case 16:
    return PT_Uint16;
  case 8:
    return PT_Uint8;
  default:
    return PT_IntAP;
  }
  llvm_unreachable("Unhandled BitWidth");
}

OptPrimType Context::classify(QualType T) const {

  if (const auto *BT = dyn_cast<BuiltinType>(T.getCanonicalType())) {
    auto Kind = BT->getKind();
    if (Kind == BuiltinType::Bool)
      return PT_Bool;
    if (Kind == BuiltinType::NullPtr)
      return PT_Ptr;
    if (Kind == BuiltinType::BoundMember)
      return PT_MemberPtr;

    // Just trying to avoid the ASTContext::getIntWidth call below.
    if (Kind == BuiltinType::Short)
      return integralTypeToPrimTypeS(this->ShortWidth);
    if (Kind == BuiltinType::UShort)
      return integralTypeToPrimTypeU(this->ShortWidth);

    if (Kind == BuiltinType::Int)
      return integralTypeToPrimTypeS(this->IntWidth);
    if (Kind == BuiltinType::UInt)
      return integralTypeToPrimTypeU(this->IntWidth);
    if (Kind == BuiltinType::Long)
      return integralTypeToPrimTypeS(this->LongWidth);
    if (Kind == BuiltinType::ULong)
      return integralTypeToPrimTypeU(this->LongWidth);
    if (Kind == BuiltinType::LongLong)
      return integralTypeToPrimTypeS(this->LongLongWidth);
    if (Kind == BuiltinType::ULongLong)
      return integralTypeToPrimTypeU(this->LongLongWidth);

    if (Kind == BuiltinType::SChar || Kind == BuiltinType::Char_S)
      return integralTypeToPrimTypeS(8);
    if (Kind == BuiltinType::UChar || Kind == BuiltinType::Char_U ||
        Kind == BuiltinType::Char8)
      return integralTypeToPrimTypeU(8);

    if (BT->isSignedInteger())
      return integralTypeToPrimTypeS(Ctx.getIntWidth(T));
    if (BT->isUnsignedInteger())
      return integralTypeToPrimTypeU(Ctx.getIntWidth(T));

    if (BT->isFloatingPoint())
      return PT_Float;
  }

  if (T->isPointerOrReferenceType())
    return PT_Ptr;

  if (T->isMemberPointerType())
    return PT_MemberPtr;

  if (const auto *BT = T->getAs<BitIntType>()) {
    if (BT->isSigned())
      return integralTypeToPrimTypeS(BT->getNumBits());
    return integralTypeToPrimTypeU(BT->getNumBits());
  }

  if (const auto *D = T->getAsEnumDecl()) {
    if (!D->isComplete())
      return std::nullopt;
    return classify(D->getIntegerType());
  }

  if (const auto *AT = T->getAs<AtomicType>())
    return classify(AT->getValueType());

  if (const auto *DT = dyn_cast<DecltypeType>(T))
    return classify(DT->getUnderlyingType());

  if (const auto *OBT = T.getCanonicalType()->getAs<OverflowBehaviorType>())
    return classify(OBT->getUnderlyingType());

  if (T->isObjCObjectPointerType() || T->isBlockPointerType())
    return PT_Ptr;

  if (T->isFixedPointType())
    return PT_FixedPoint;

  // Vector and complex types get here.
  return std::nullopt;
}

unsigned Context::getCharBit() const {
  return Ctx.getTargetInfo().getCharWidth();
}

/// Simple wrapper around getFloatTypeSemantics() to make code a
/// little shorter.
const llvm::fltSemantics &Context::getFloatSemantics(QualType T) const {
  return Ctx.getFloatTypeSemantics(T);
}

bool Context::Run(State &Parent, const Function *Func) {
  InterpState State(Parent, *P, Stk, *this, Func);
  auto Memory = std::make_unique<char[]>(InterpFrame::allocSize(Func));
  InterpFrame *Frame = new (Memory.get()) InterpFrame(
      State, Func, /*Caller=*/nullptr, CodePtr(), Func->getArgSize());
  State.Current = Frame;

  if (Interpret(State)) {
    assert(Stk.empty());
    return true;
  }

  Stk.clear();
  Frame->~InterpFrame();
  State.Current = &State.BottomFrame;
  return false;
}

// TODO: Virtual bases?
const CXXMethodDecl *
Context::getOverridingFunction(const CXXRecordDecl *DynamicDecl,
                               const CXXRecordDecl *StaticDecl,
                               const CXXMethodDecl *InitialFunction) const {
  assert(DynamicDecl);
  assert(StaticDecl);
  assert(InitialFunction);

  const CXXRecordDecl *CurRecord = DynamicDecl;
  const CXXMethodDecl *FoundFunction = InitialFunction;
  for (;;) {
    const CXXMethodDecl *Overrider =
        FoundFunction->getCorrespondingMethodDeclaredInClass(CurRecord, false);
    if (Overrider)
      return Overrider;

    // Common case of only one base class.
    if (CurRecord->getNumBases() == 1) {
      CurRecord = CurRecord->bases_begin()->getType()->getAsCXXRecordDecl();
      continue;
    }

    // Otherwise, go to the base class that will lead to the StaticDecl.
    for (const CXXBaseSpecifier &Spec : CurRecord->bases()) {
      const CXXRecordDecl *Base = Spec.getType()->getAsCXXRecordDecl();
      if (Base == StaticDecl || Base->isDerivedFrom(StaticDecl)) {
        CurRecord = Base;
        break;
      }
    }
  }

  llvm_unreachable(
      "Couldn't find an overriding function in the class hierarchy?");
  return nullptr;
}

const Function *Context::getOrCreateFunction(const FunctionDecl *FuncDecl) {
  assert(FuncDecl);
  if (const Function *Func = P->getFunction(FuncDecl))
    return Func;

  // Manually created functions that haven't been assigned proper
  // parameters yet.
  if (!FuncDecl->param_empty() && !FuncDecl->param_begin())
    return nullptr;

  bool IsLambdaStaticInvoker = false;
  if (const auto *MD = dyn_cast<CXXMethodDecl>(FuncDecl);
      MD && MD->isLambdaStaticInvoker()) {
    // For a lambda static invoker, we might have to pick a specialized
    // version if the lambda is generic. In that case, the picked function
    // will *NOT* be a static invoker anymore. However, it will still
    // be a non-static member function, this (usually) requiring an
    // instance pointer. We suppress that later in this function.
    IsLambdaStaticInvoker = true;
  }
  // Set up argument indices.
  unsigned ParamOffset = 0;
  llvm::SmallVector<Function::ParamDescriptor> ParamDescriptors;

  // If the return is not a primitive, a pointer to the storage where the
  // value is initialized in is passed as the first argument. See 'RVO'
  // elsewhere in the code.
  QualType Ty = FuncDecl->getReturnType();
  bool HasRVO = false;
  if (!Ty->isVoidType() && !canClassify(Ty)) {
    HasRVO = true;
    ParamOffset += align(primSize(PT_Ptr));
  }

  // If the function decl is a member decl, the next parameter is
  // the 'this' pointer. This parameter is pop()ed from the
  // InterpStack when calling the function.
  bool HasThisPointer = false;
  if (const auto *MD = dyn_cast<CXXMethodDecl>(FuncDecl)) {
    if (!IsLambdaStaticInvoker) {
      HasThisPointer = MD->isInstance();
      if (MD->isImplicitObjectMemberFunction())
        ParamOffset += align(primSize(PT_Ptr));
    }

    if (isLambdaCallOperator(MD)) {
      // The parent record needs to be complete, we need to know about all
      // the lambda captures.
      if (!MD->getParent()->isCompleteDefinition())
        return nullptr;
      if (MD->isStatic()) {
        llvm::DenseMap<const ValueDecl *, FieldDecl *> LC;
        FieldDecl *LTC;

        MD->getParent()->getCaptureFields(LC, LTC);
        // Static lambdas cannot have any captures. If this one does,
        // it has already been diagnosed and we can only ignore it.
        if (!LC.empty())
          return nullptr;
      }
    }
  }

  // Assign descriptors to all parameters.
  // Composite objects are lowered to pointers.
  const auto *FuncProto = FuncDecl->getType()->getAs<FunctionProtoType>();
  unsigned BlockOffset = 0;
  for (auto [ParamIndex, PD] : llvm::enumerate(FuncDecl->parameters())) {
    bool IsConst = PD->getType().isConstQualified();
    bool IsVolatile = PD->getType().isVolatileQualified();

    if (!getASTContext().hasSameType(PD->getType(),
                                     FuncProto->getParamType(ParamIndex)))
      return nullptr;

    OptPrimType T = classify(PD->getType());
    PrimType PT = T.value_or(PT_Ptr);
    Descriptor *Desc = P->createDescriptor(PD, PT, nullptr, std::nullopt,
                                           IsConst, /*IsTemporary=*/false,
                                           /*IsMutable=*/false, IsVolatile);
    unsigned PrimTSize = align(primSize(PT));
    ParamDescriptors.emplace_back(Desc, ParamOffset, BlockOffset, PT);
    ParamOffset += PrimTSize;
    BlockOffset += sizeof(Block) + PrimTSize;
  }

  // Create a handle over the emitted code.
  assert(!P->getFunction(FuncDecl));
  const Function *Func =
      P->createFunction(FuncDecl, ParamOffset, std::move(ParamDescriptors),
                        HasThisPointer, HasRVO, IsLambdaStaticInvoker);
  return Func;
}

const Function *Context::getOrCreateObjCBlock(const BlockExpr *E) {
  const BlockDecl *BD = E->getBlockDecl();
  // Set up argument indices.
  unsigned ParamOffset = 0;
  llvm::SmallVector<Function::ParamDescriptor> ParamDescriptors;

  // Assign descriptors to all parameters.
  // Composite objects are lowered to pointers.
  for (const ParmVarDecl *PD : BD->parameters()) {
    bool IsConst = PD->getType().isConstQualified();
    bool IsVolatile = PD->getType().isVolatileQualified();

    OptPrimType T = classify(PD->getType());
    PrimType PT = T.value_or(PT_Ptr);
    Descriptor *Desc = P->createDescriptor(PD, PT, nullptr, std::nullopt,
                                           IsConst, /*IsTemporary=*/false,
                                           /*IsMutable=*/false, IsVolatile);
    ParamDescriptors.emplace_back(Desc, ParamOffset, ~0u, PT);
    ParamOffset += align(primSize(PT));
  }

  if (BD->hasCaptures())
    return nullptr;

  // Create a handle over the emitted code.
  Function *Func =
      P->createFunction(E, ParamOffset, std::move(ParamDescriptors),
                        /*HasThisPointer=*/false, /*HasRVO=*/false,
                        /*IsLambdaStaticInvoker=*/false);

  assert(Func);
  Func->setDefined(true);
  // We don't compile the BlockDecl code at all right now.
  Func->setIsFullyCompiled(true);

  return Func;
}

unsigned Context::collectBaseOffset(const RecordDecl *BaseDecl,
                                    const RecordDecl *DerivedDecl) const {
  assert(BaseDecl);
  assert(DerivedDecl);
  const auto *FinalDecl = cast<CXXRecordDecl>(BaseDecl);
  const RecordDecl *CurDecl = DerivedDecl;
  const Record *CurRecord = P->getOrCreateRecord(CurDecl);
  assert(CurDecl && FinalDecl);

  unsigned OffsetSum = 0;
  for (;;) {
    assert(CurRecord->getNumBases() > 0);
    // One level up
    for (const Record::Base &B : CurRecord->bases()) {
      const auto *BaseDecl = cast<CXXRecordDecl>(B.Decl);

      if (BaseDecl == FinalDecl || BaseDecl->isDerivedFrom(FinalDecl)) {
        OffsetSum += B.Offset;
        CurRecord = B.R;
        CurDecl = BaseDecl;
        break;
      }
    }
    if (CurDecl == FinalDecl)
      break;
  }

  assert(OffsetSum > 0);
  return OffsetSum;
}

const Record *Context::getRecord(const RecordDecl *D) const {
  return P->getOrCreateRecord(D);
}

bool Context::isUnevaluatedBuiltin(unsigned ID) {
  return ID == Builtin::BI__builtin_classify_type ||
         ID == Builtin::BI__builtin_os_log_format_buffer_size ||
         ID == Builtin::BI__builtin_constant_p || ID == Builtin::BI__noop;
}