artichoke-lang/lib/src/Tokenizer/Tokenizer.cpp

#include <artichoke/Tokenizer/Tokenizer.hpp>

#include <utility>

#include <artichoke/Tokenizer/TokenizerRange.hpp>
#include <artichoke/Util/Strings.hpp>
#include <artichoke/Util/Demangle.hpp>
#include <artichoke/Util/TrieMap.hpp>

namespace arti::lang {

  Tokenizer::Tokenizer(std::string source) noexcept
    : line(0)
    , column(0)
    , tokensGenerator{ tokenize() }
    , tokensBuffer{}
    , source(std::move(source)) {
    iter = this->source.begin();
  }

  Tokenizer::Tokenizer(Tokenizer &&rhs) noexcept
    : line(rhs.line)
    , column(rhs.column)
    , iter(rhs.iter)
    , tokensGenerator(std::exchange(rhs.tokensGenerator, {}))
    , tokensBuffer(std::exchange(rhs.tokensBuffer, {}))
    , source(std::exchange(rhs.source, "")) { }

  Tokenizer &Tokenizer::operator=(Tokenizer &&rhs) noexcept {
    line = rhs.line;
    column = rhs.column;
    iter = std::move(rhs.iter);
    source = std::move(rhs.source);
    tokensBuffer = std::move(rhs.tokensBuffer);
    tokensGenerator = std::move(rhs.tokensGenerator);
    return *this;
  }

  TokenizerRange Tokenizer::range() noexcept {
    return TokenizerRange{ this };
  }

  Expected<void> Tokenizer::consume(std::size_t n) noexcept {
    while (n > 0) {
      if (tokensBuffer.empty()) {
        if (auto ok = peek(); ! ok) {
          return Unexpected<>{ ok.error() };
        }

        if (finished()) {
          break;
        }
      }

      tokensBuffer.pop_front();
      n -= 1;
    }

    if (n != 0) {
      return langException<ExceptCode::ecEOB>(line, column);
    }

    return {};
  }

  Expected<Token> Tokenizer::peek(std::size_t n) noexcept {
    while (not finished() && tokensBuffer.size() < (n + 1)) {
      try {
        tokensGenerator.next();
      }
      catch (std::exception &e) {
        const auto exceptionName = demangle(typeid(e).name());

        if (not exceptionName) {
          return langException<ExceptCode::ecStdException>(
            line,
            column,
            "abi::__cxa_demangle(...)",
            demangleErrorString(exceptionName.error())
          );
        }

        return langException<ExceptCode::ecStdException>(
          line,
          column,
          exceptionName.value(),
          e.what()
        );
      }

      auto tok = tokensGenerator.get();

      if (not tok) {
        return tok;
      }

      tokensBuffer.push_back(std::move(tok).value());
    }

    if (finished()) {
      if (tokensBuffer.size() < (n + 1)) {
        return langException<ExceptCode::ecEOF>(line, column);
      }
    }

    return tokensBuffer.at(n);
  }

  bool Tokenizer::finished() const noexcept {
    if (tokensGenerator.finished()) {
      if (!tokensBuffer.empty()) {
        return tokensBuffer.front().value == TokenV::tkEOF;
      }

      return true;
    }

    return false;
  }

  void Tokenizer::swap(Tokenizer &other) noexcept {
    std::swap(line, other.line);
    std::swap(column, other.column);
    std::swap(iter, other.iter);
    tokensGenerator.swap(other.tokensGenerator);
    tokensBuffer.swap(other.tokensBuffer);
    source.swap(other.source);
  }

  Generator<Expected<Token>> Tokenizer::tokenize() {
    while (iter != source.end()) {
      skip_whitespace();

      if (iter == source.end()) {
        break;
      }

      if (isDigit(*iter)) {
        yield readNumber();
      }
      else if (*iter == '"') {
        yield readString();
      }
      else if (*iter == '\'') {
        yield readCharacter();
      }
      else if (isFirstIdentChar(*iter)) {
        yield readIdentifier();
      }
      else if (*iter == '/') {
        if ((iter + 1) != source.end() && *(iter + 1) == '*') {
          if (auto ok = skip_comment(); !ok) {
            auto err = ok.error();
            yield Unexpected<>{err};
          }
        }
        else {
          yield readOperator();
        }
      }
      else {
        yield readOperator();
      }
    }

    coret Token{ .value = TokenV::tkEOF,
                 .line = line,
                 .column = column,
                 .strValue = "" };
  }

  void Tokenizer::skip_whitespace() {
    while (iter != source.end()) {
      switch (*iter) {
        case ' ':  column += 1; break;
        case '\t': column += CE_TAB_SIZE; break;
        case '\n':
          column = 0;
          line += 1;
          break;
        default: return;
      }

      ++iter;
    }
  }

  Expected<void> Tokenizer::skip_comment() {
    iter += 2;
    column += 2;

    bool isEnd = false;

    while (iter != source.end()) {
      if (*iter == '\n') {
        column = 0;
        line += 1;
      }
      else {
        column += 1;

        if (*iter == '*') {
          if ((iter + 1) == source.end()) {
            return langException<ExceptCode::ecInvalidComment>(line, column);
          }
          else if (*(iter + 1) == '/') {
            iter += 2;
            column += 2;
            return {};
          }
        }
      }

      ++iter;
    }
    return langException<ExceptCode::ecInvalidComment>(line, column);
  }

  Expected<Token> Tokenizer::readNumber() {
    auto stIter = iter;

    auto cLine = line;
    auto cColumn = column;

    auto forward = [&] {
      ++iter;
      ++column;
    };

    if (*iter == '0') {
      bool continueMatching = false;

      forward();

      if (iter == source.end()) {
        return Token{
          TokenV::tkInteger,
          cLine,
          cColumn,
          { stIter, iter }
        };
      }

      if (toLower(*iter) == 'x') {
        forward();

        if (iter == source.end()) {
          return langException<ExceptCode::ecInvalidLiteral>(
            line,
            column,
            "hex value",
            "EOF"
          );
        }

        if (not isHexChar(*iter)) {
          return langException<ExceptCode::ecInvalidLiteral>(
            line,
            column,
            "hex value",
            *iter
          );
        }

        while (iter != source.end()) {
          if (not isHexChar(*iter)) {
            break;
          }

          forward();
        }
      }
      else if (toLower(*iter) == 'b') {
        forward();

        if (iter == source.end()) {
          return langException<ExceptCode::ecInvalidLiteral>(
            line,
            column,
            "binary value",
            "EOF"
          );
        }

        if (not isBinaryChar(*iter)) {
          return langException<ExceptCode::ecInvalidLiteral>(
            line,
            column,
            "binary value",
            *iter
          );
        }

        while (iter != source.end()) {
          if (not isBinaryChar(*iter)) {
            if (isDigit(*iter)) {
              return langException<ExceptCode::ecInvalidLiteral>(
                line,
                column,
                "binary value",
                *iter
              );
            }

            break;
          }

          forward();
        }
      }
      else if (*iter == '.') {
        continueMatching = true;
      }
      else {
        while (iter != source.end()) {
          if (not isOctalChar(*iter)) {
            if (isDigit(*iter)) {
              return langException<ExceptCode::ecInvalidLiteral>(
                line,
                column,
                "octal value",
                *iter
              );
            }

            break;
          }

          forward();
        }
      }

      if (not continueMatching) {
        return Token{
          TokenV::tkInteger,
          cLine,
          cColumn,
          { stIter, iter }
        };
      }
    }

    while (iter != source.end()) {
      if (not isDigit(*iter)) {
        break;
      }

      forward();
    }

    if (iter == source.end()) {
      return Token{
        TokenV::tkInteger,
        cLine,
        cColumn,
        { stIter, iter }
      };
    }

    if (*iter == '.') {
      auto dotIter = iter;
      forward();

      while (iter != source.end()) {
        if (not isDigit(*iter)) {
          break;
        }

        forward();
      }

      if ((iter - dotIter) == 1) {
        /* Revert to dot */
        --iter;
        --column;

        return langException<ExceptCode::ecInvalidLiteral>(
          line,
          column,
          "digit",
          iter == source.end()
            ? "EOF"
            : std::string{ *(iter + 1) }
        );
      }

      return Token{
        TokenV::tkDecimal,
        cLine,
        cColumn,
        { stIter, iter }
      };
    }

    return Token{
      TokenV::tkInteger,
      cLine,
      cColumn,
      { stIter, iter }
    };
  }

  Expected<Token> Tokenizer::readString() {
    auto stIter = iter;

    auto cLine = line;
    auto cColumn = column;

    auto forward = [&] {
      ++iter;
      ++column;
    };

    forward();

    while (iter != source.end()) {
      if (*iter == '\n') {
        return langException<ExceptCode::ecInvalidLiteral>(
          line,
          column,
          "end of string (\")",
          "\\n"
        );
      }

      if (*iter == '"') {
        break;
      }

      if (*iter == '\\') {
        forward();

        if (iter == source.end()) {
          return langException<ExceptCode::ecInvalidLiteral>(
            line,
            column,
            "escaped char",
            "EOF"
          );
        }

        /* TODO: Validate char escape codes */
      }

      forward();
    }

    if (*iter == '"') {
      forward();

      return Token{
        TokenV::tkString,
        cLine,
        cColumn,
        { stIter, iter }
      };
    }

    return langException<ExceptCode::ecInvalidLiteral>(
      line,
      column,
      "end of string (\")",
      "EOF"
    );
  }

  Expected<Token> Tokenizer::readCharacter() {
    auto stIter = iter;

    auto cLine = line;
    auto cColumn = column;

    auto forward = [&] {
      ++iter;
      ++column;
    };

    forward();

    if (iter == source.end()) {
      return langException<ExceptCode::ecInvalidLiteral>(
        line,
        column,
        "char literal",
        "EOF"
      );
    }

    if (*iter == '\n') {
      return langException<ExceptCode::ecInvalidLiteral>(
        line,
        column,
        "char literal",
        "\\n"
      );
    }

    if (*iter == '\\') {
      forward();

      if (iter == source.end()) {
        return langException<ExceptCode::ecInvalidLiteral>(
          line,
          column,
          "char literal",
          "EOF"
        );
      }

      /* TODO: Validate char escape codes */
    }

    forward();

    if (iter == source.end()) {
      return langException<ExceptCode::ecInvalidLiteral>(
        line,
        column,
        "end of char literal (')",
        "EOF"
      );
    }

    if (*iter != '\'') {
      return langException<ExceptCode::ecInvalidLiteral>(
        line,
        column,
        "end char literal (')",
        *iter
      );
    }

    return Token{
      TokenV::tkCharacter,
      cLine,
      cColumn,
      { stIter, ++iter }
    };
  }

  Expected<Token> Tokenizer::readIdentifier() {
    auto stIter = iter;

    auto cLine = line;
    auto cColumn = column;

    auto forward = [&] {
      ++iter;
      ++column;
    };

    while (iter != source.end()) {
      if (not isIdentChar(*iter)) {
        break;
      }

      forward();
    }

    auto tok = Token{
      TokenV::tkIdentifier,
      cLine,
      cColumn,
      { stIter, iter }
    };

    if (tok.strValue.compare("_") == 0) {
      tok.value = TokenV::kwUnderscore;
    }
    if (tok.strValue.compare("or") == 0) {
      tok.value = TokenV::kwOr;
    }
    else if (tok.strValue.compare("not") == 0) {
      tok.value = TokenV::kwNot;
    }
    else if (tok.strValue.compare("and") == 0) {
      tok.value = TokenV::kwAnd;
    }
    else if (tok.strValue.compare("if") == 0) {
      tok.value = TokenV::kwIf;
    }
    else if (tok.strValue.compare("else") == 0) {
      tok.value = TokenV::kwElse;
    }
    else if (tok.strValue.compare("fn") == 0) {
      tok.value = TokenV::kwFn;
    }
    else if (tok.strValue.compare("enum") == 0) {
      tok.value = TokenV::kwEnum;
    }
    else if (tok.strValue.compare("struct") == 0) {
      tok.value = TokenV::kwStruct;
    }
    else if (tok.strValue.compare("def") == 0) {
      tok.value = TokenV::kwDef;
    }
    else if (tok.strValue.compare("let") == 0) {
      tok.value = TokenV::kwLet;
    }
    else if (tok.strValue.compare("for") == 0) {
      tok.value = TokenV::kwFor;
    }
    else if (tok.strValue.compare("loop") == 0) {
      tok.value = TokenV::kwLoop;
    }
    else if (tok.strValue.compare("break") == 0) {
      tok.value = TokenV::kwBreak;
    }
    else if (tok.strValue.compare("continue") == 0) {
      tok.value = TokenV::kwContinue;
    }
    else if (tok.strValue.compare("while") == 0) {
      tok.value = TokenV::kwWhile;
    }
    else if (tok.strValue.compare("match") == 0) {
      tok.value = TokenV::kwMatch;
    }
    else if (tok.strValue.compare("switch") == 0) {
      tok.value = TokenV::kwSwitch;
    }
    else if (tok.strValue.compare("return") == 0) {
      tok.value = TokenV::kwReturn;
    }
    else if (tok.strValue.compare("unreachable") == 0) {
      tok.value = TokenV::kwUnreachable;
    }
    else if (tok.strValue.compare("defer") == 0) {
      tok.value = TokenV::kwDefer;
    }
    else if (tok.strValue.compare("errdefer") == 0) {
      tok.value = TokenV::kwErrDefer;
    }
    else if (tok.strValue.compare("true") == 0) {
      tok.value = TokenV::kwTrue;
    }
    else if (tok.strValue.compare("false") == 0) {
      tok.value = TokenV::kwFalse;
    }
    else if (tok.strValue.compare("null") == 0) {
      tok.value = TokenV::kwNull;
    }
    else if (tok.strValue.compare("this") == 0) {
      tok.value = TokenV::kwThis;
    }
    else if (tok.strValue.compare("import") == 0) {
      tok.value = TokenV::kwImport;
    }
    else if (tok.strValue.compare("export") == 0) {
      tok.value = TokenV::kwExport;
    }
    else if (tok.strValue.compare("module") == 0) {
      tok.value = TokenV::kwModule;
    }
    else if (tok.strValue.compare("using") == 0) {
      tok.value = TokenV::kwUsing;
    }

    return tok;
  }

  TrieMap<TokenV> buildOperatorsTrieMap();

  Expected<Token> Tokenizer::readOperator() {
    auto stIter = iter;

    auto cLine = line;
    auto cColumn = column;

    auto forward = [&] {
      ++iter;
      ++column;
    };

    auto tm = buildOperatorsTrieMap();

    if (not tm.root.childs.contains(*iter)) {
      return langException<ExceptCode::ecInvalidCharacter>(line, column, *iter);
    }

    auto node = &tm.root.childs.at(*iter);
    auto lvNode = node->isLeaf ? node : nullptr;

    long lvNodeDiff = 0;

    forward();

    while (iter != source.end() and not (isWhiteSpace(*iter))) {
      if (not node->childs.contains(*iter)) {
        break;
      }

      node = &node->childs.at(*iter);
      lvNode = node->isLeaf ? node : lvNode;
      lvNodeDiff = node->isLeaf ? 0 : (lvNodeDiff + 1);

      forward();
    }

    Token tok{
      TokenV::opStar,
      line = cLine,
      column = cColumn,
      { stIter, iter }
    };

    if (node->isLeaf) {
      tok.value = *node->value;
    }
    else if (lvNode != nullptr) {
      iter -= lvNodeDiff;
      tok.value = *lvNode->value;
      tok.strValue = std::string_view{ stIter, iter };
    }
    else {
      return langException<ExceptCode::ecInvalidToken>(
        line,
        column,
        tok.strValue
      );
    }

    return tok;
  }

  TrieMap<TokenV> buildOperatorsTrieMap() {
    TrieMap<TokenV> tm;

    tm.insert(".", TokenV::opDot);
    tm.insert("%", TokenV::opMod);
    tm.insert("+", TokenV::opPlus);
    tm.insert("-", TokenV::opHyphen);
    tm.insert("/", TokenV::opSlash);
    tm.insert("!", TokenV::opBang);
    tm.insert("*", TokenV::opStar);
    tm.insert(":", TokenV::opColon);
    tm.insert(",", TokenV::opComma);
    tm.insert("=", TokenV::opAssign);
    tm.insert("::", TokenV::opAccess);
    tm.insert(";", TokenV::opSemicolon);
    tm.insert("^", TokenV::opCaret);
    tm.insert("~", TokenV::opTilde);
    tm.insert("==", TokenV::opEq);
    tm.insert("!=", TokenV::opNeq);
    tm.insert("<", TokenV::opLt);
    tm.insert(">", TokenV::opGt);
    tm.insert("<=", TokenV::opLtEq);
    tm.insert(">=", TokenV::opGtEq);
    tm.insert("<<", TokenV::opLShift);
    tm.insert(">>", TokenV::opRShift);
    tm.insert("&&", TokenV::opBoolAnd);
    tm.insert("||", TokenV::opBoolOr);
    tm.insert("&", TokenV::opAnd);
    tm.insert("|", TokenV::opOr);
    tm.insert("(", TokenV::opLParen);
    tm.insert(")", TokenV::opRParen);
    tm.insert("[", TokenV::opLBracket);
    tm.insert("]", TokenV::opRBracket);
    tm.insert("{", TokenV::opLSquirly);
    tm.insert("}", TokenV::opRSquirly);
    tm.insert("->", TokenV::opArrow);
    tm.insert("+=", TokenV::opPlusAssign);
    tm.insert("-=", TokenV::opHyphenAssign);
    tm.insert("*=", TokenV::opStarAssign);
    tm.insert("/=", TokenV::opSlashAssign);
    tm.insert("%=", TokenV::opModAssign);
    tm.insert("&=", TokenV::opAndAssign);
    tm.insert("|=", TokenV::opOrAssign);
    tm.insert("<<=", TokenV::opLShiftAssign);
    tm.insert(">>=", TokenV::opRShiftAssign);
    tm.insert("&&=", TokenV::opBoolAndAssign);
    tm.insert("||=", TokenV::opBoolORAssign);
    tm.insert("$", TokenV::opMut);
    tm.insert("?", TokenV::opOpt);
    tm.insert(".#", TokenV::opSliceSize);
    tm.insert(".[", TokenV::opPtrSlice);
    tm.insert(".*", TokenV::opSlicePtr);
    tm.insert(".@", TokenV::opReflect);
    tm.insert(":=", TokenV::opLabel);

    return tm;
  }

} // namespace arti::lang