//===--- Lex.cpp - extract token stream from source code ---------*- 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 "clang-pseudo/Token.h"

#include "clang/Basic/IdentifierTable.h"

#include "clang/Basic/SourceLocation.h"

#include "clang/Basic/TokenKinds.h"

#include "clang/Lex/Lexer.h"

#include "clang/Lex/LiteralSupport.h"

namespace clang {

namespace pseudo {

TokenStream lex(const std::string &Code, const clang::LangOptions &LangOpts) {

  clang::SourceLocation Start;

  // Tokenize using clang's lexer in raw mode.

  // std::string guarantees null-termination, which the lexer needs.

  clang::Lexer Lexer(Start, LangOpts, Code.data(), Code.data(),

                     Code.data() + Code.size());

  Lexer.SetCommentRetentionState(true);

  TokenStream Result;

  clang::Token CT;

  // Index into the token stream of original source code.

  Token::Index TokenIndex = 0;

  unsigned LastOffset = 0;

  unsigned Line = 0;

  unsigned Indent = 0;

  for (Lexer.LexFromRawLexer(CT); CT.getKind() != clang::tok::eof;

       Lexer.LexFromRawLexer(CT)) {

    unsigned Offset =

        CT.getLocation().getRawEncoding() - Start.getRawEncoding();

    Token Tok;

    Tok.Data = &Code[Offset];

    Tok.Length = CT.getLength();

    Tok.Kind = CT.getKind();

    // Update current line number and indentation from raw source code.

    unsigned NewLineStart = 0;

    for (unsigned I = LastOffset; I < Offset; ++I) {

      if (Code[I] == '\n') {

        NewLineStart = I + 1;

        ++Line;

      }

    }

    if (NewLineStart || !LastOffset) {

      Indent = 0;

      for (char C : StringRef(Code).slice(NewLineStart, Offset)) {

        if (C == ' ')

          ++Indent;

        else if (C == '\t')

          Indent += 8;

        else

          break;

      }

    }

    Tok.Indent = Indent;

    Tok.Line = Line;

    if (CT.isAtStartOfLine())

      Tok.setFlag(LexFlags::StartsPPLine);

    if (CT.needsCleaning() || CT.hasUCN())

      Tok.setFlag(LexFlags::NeedsCleaning);

    Tok.OriginalIndex = TokenIndex++;

    Result.push(Tok);

    LastOffset = Offset;

  }

  Result.finalize();

  return Result;

}

TokenStream cook(const TokenStream &Code, const LangOptions &LangOpts) {

  auto CleanedStorage = std::make_shared<llvm::BumpPtrAllocator>();

  clang::IdentifierTable Identifiers(LangOpts);

  TokenStream Result(CleanedStorage);

  Result.addPayload(Code.getPayload());

  for (auto Tok : Code.tokens()) {

    if (Tok.flag(LexFlags::NeedsCleaning)) {

      // Remove escaped newlines and trigraphs.

      llvm::SmallString<64> CleanBuffer;

      const char *Pos = Tok.text().begin();

      while (Pos < Tok.text().end()) {

        auto [Char, CharSize] =

            clang::Lexer::getCharAndSizeNoWarn(Pos, LangOpts);

        CleanBuffer.push_back(Char);

        assert(CharSize != 0 && "no progress!");

        Pos += CharSize;

      }

      llvm::StringRef Text = CleanBuffer;

      llvm::SmallString<64> UCNBuffer;

      // A surface reading of the standard suggests UCNs might appear anywhere.

      // But we need only decode them in raw_identifiers.

      //  - they cannot appear in punctuation/keyword tokens, because UCNs

      //    cannot encode basic characters outside of literals [lex.charset]

      //  - they can appear in literals, but we need not unescape them now.

      //    We treat them as escape sequences when evaluating the literal.

      //  - comments are handled similarly to literals

      // This is good fortune, because expandUCNs requires its input to be a

      // reasonably valid identifier (e.g. without stray backslashes).

      if (Tok.Kind == tok::raw_identifier) {

        clang::expandUCNs(UCNBuffer, CleanBuffer);

        Text = UCNBuffer;

      }

      Tok.Data = Text.copy(*CleanedStorage).data();

      Tok.Length = Text.size();

      Tok.Flags &= ~static_cast<decltype(Tok.Flags)>(LexFlags::NeedsCleaning);

    }

    if (Tok.Kind == tok::raw_identifier) {

      // Cook raw_identifiers into identifier, keyword, etc.

      Tok.Kind = Identifiers.get(Tok.text()).getTokenID();

    } else if (Tok.Kind == tok::greatergreater) {

      // Split the greatergreater token.

      // FIXME: split lessless token to support Cuda triple angle brackets <<<.

      assert(Tok.text() == ">>");

      Tok.Kind = tok::greater;

      Tok.Length = 1;

      Result.push(Tok);

      // Line is wrong if the first greater is followed by an escaped newline!

      Tok.Data = Tok.text().data() + 1;

    }

    Result.push(std::move(Tok));

  }

  Result.finalize();

  return Result;

}

} // namespace pseudo

} // namespace clang