//===--- Forest.cpp - Parse forest  ------------------------------*- 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/Forest.h"
#include "clang-pseudo/Token.h"
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/Support/FormatVariadic.h"
#include <optional>

namespace clang {
namespace pseudo {

void ForestNode::RecursiveIterator::operator++() {
  auto C = Cur->children();
  // Try to find a child of the current node to descend into.
  for (unsigned I = 0; I < C.size(); ++I) {
    if (Seen.insert(C[I]).second) {
      Stack.push_back({Cur, I});
      Cur = C[I];
      return;
    }
  }
  // Try to find a sibling af an ancestor to advance to.
  for (; !Stack.empty(); Stack.pop_back()) {
    C = Stack.back().Parent->children();
    unsigned &Index = Stack.back().ChildIndex;
    while (++Index < C.size()) {
      if (Seen.insert(C[Index]).second) {
        Cur = C[Index];
        return;
      }
    }
  }
  Cur = nullptr;
}

llvm::iterator_range<ForestNode::RecursiveIterator>
ForestNode::descendants() const {
  return {RecursiveIterator(this), RecursiveIterator()};
}

std::string ForestNode::dump(const Grammar &G) const {
  switch (kind()) {
  case Ambiguous:
    return llvm::formatv("{0} := <ambiguous>", G.symbolName(symbol()));
  case Terminal:
    return llvm::formatv("{0} := tok[{1}]", G.symbolName(symbol()),
                         startTokenIndex());
  case Sequence:
    return G.dumpRule(rule());
  case Opaque:
    return llvm::formatv("{0} := <opaque>", G.symbolName(symbol()));
  }
  llvm_unreachable("Unhandled node kind!");
}

std::string ForestNode::dumpRecursive(const Grammar &G,
                                      bool Abbreviated) const {
  using llvm::formatv;
  Token::Index MaxToken = 0;
  // Count visits of nodes so we can mark those seen multiple times.
  llvm::DenseMap<const ForestNode *, /*VisitCount*/ unsigned> VisitCounts;
  std::function<void(const ForestNode *)> CountVisits =
      [&](const ForestNode *P) {
        MaxToken = std::max(MaxToken, P->startTokenIndex());
        if (VisitCounts[P]++ > 0)
          return; // Don't count children as multiply visited.
        if (P->kind() == Ambiguous)
          llvm::for_each(P->alternatives(), CountVisits);
        else if (P->kind() == Sequence)
          llvm::for_each(P->elements(), CountVisits);
      };
  CountVisits(this);

  unsigned IndexWidth = std::max(3, (int)std::to_string(MaxToken).size());
  // e.g. "[{0,4}, {1,4})" if MaxToken is 5742.
  std::string RangeFormat = formatv("[{{0,{0}}, {{1,{0}}) ", IndexWidth);

  // The box-drawing characters that should be added as a child is rendered.
  struct LineDecoration {
    std::string Prefix;         // Prepended to every line.
    llvm::StringRef First;      // added to the child's line.
    llvm::StringRef Subsequent; // added to descendants' lines.
  };

  // We print a "#<id>" for nonterminal forest nodes that are being dumped
  // multiple times.
  llvm::DenseMap<const ForestNode *, size_t> ReferenceIds;
  std::string Result;
  constexpr Token::Index KEnd = std::numeric_limits<Token::Index>::max();
  std::function<void(const ForestNode *, Token::Index, std::optional<SymbolID>,
                     LineDecoration &LineDec)>
      Dump = [&](const ForestNode *P, Token::Index End,
                 std::optional<SymbolID> ElidedParent, LineDecoration LineDec) {
        bool SharedNode = VisitCounts.find(P)->getSecond() > 1;
        llvm::ArrayRef<const ForestNode *> Children;
        auto EndOfElement = [&](size_t ChildIndex) {
          return ChildIndex + 1 == Children.size()
                     ? End
                     : Children[ChildIndex + 1]->startTokenIndex();
        };
        if (P->kind() == Ambiguous) {
          Children = P->alternatives();
        } else if (P->kind() == Sequence) {
          Children = P->elements();
          if (Abbreviated) {
            // Abbreviate chains of trivial sequence nodes.
            //    A := B, B := C, C := D, D := X Y Z
            // becomes
            //    A~D := X Y Z
            //
            // We can't hide nodes that appear multiple times in the tree,
            // because we need to call out their identity with IDs.
            if (Children.size() == 1 && !SharedNode) {
              assert(Children[0]->startTokenIndex() == P->startTokenIndex() &&
                     EndOfElement(0) == End);
              return Dump(Children[0], End,
                          /*ElidedParent=*/ElidedParent.value_or(P->symbol()),
                          LineDec);
            }
          }
        }

        if (End == KEnd)
          Result += formatv(RangeFormat.c_str(), P->startTokenIndex(), "end");
        else
          Result += formatv(RangeFormat.c_str(), P->startTokenIndex(), End);
        Result += LineDec.Prefix;
        Result += LineDec.First;
        if (ElidedParent) {
          Result += G.symbolName(*ElidedParent);
          Result += "~";
        }

        if (SharedNode && P->kind() != ForestNode::Terminal) {
          auto It = ReferenceIds.try_emplace(P, ReferenceIds.size() + 1);
          bool First = It.second;
          unsigned ID = It.first->second;

          // The first time, print as #1. Later, =#1.
          if (First) {
            Result += formatv("{0} #{1}", P->dump(G), ID);
          } else {
            Result += formatv("{0} =#{1}", G.symbolName(P->symbol()), ID);
            Children = {}; // Don't walk the children again.
          }
        } else {
          Result.append(P->dump(G));
        }
        Result.push_back('\n');

        auto OldPrefixSize = LineDec.Prefix.size();
        LineDec.Prefix += LineDec.Subsequent;
        for (size_t I = 0; I < Children.size(); ++I) {
          if (I == Children.size() - 1) {
            LineDec.First = "└─";
            LineDec.Subsequent = "  ";
          } else {
            LineDec.First = "├─";
            LineDec.Subsequent = "│ ";
          }
          Dump(Children[I], P->kind() == Sequence ? EndOfElement(I) : End,
               std::nullopt, LineDec);
        }
        LineDec.Prefix.resize(OldPrefixSize);
      };
  LineDecoration LineDec;
  Dump(this, KEnd, std::nullopt, LineDec);
  return Result;
}

llvm::ArrayRef<ForestNode>
ForestArena::createTerminals(const TokenStream &Code) {
  ForestNode *Terminals = Arena.Allocate<ForestNode>(Code.tokens().size() + 1);
  size_t Index = 0;
  for (const auto &T : Code.tokens()) {
    new (&Terminals[Index])
        ForestNode(ForestNode::Terminal, tokenSymbol(T.Kind),
                   /*Start=*/Index, /*TerminalData*/ 0);
    ++Index;
  }
  // Include an `eof` terminal.
  // This is important to drive the final shift/recover/reduce loop.
  new (&Terminals[Index])
      ForestNode(ForestNode::Terminal, tokenSymbol(tok::eof),
                 /*Start=*/Index, /*TerminalData*/ 0);
  ++Index;
  NodeCount = Index;
  return llvm::ArrayRef(Terminals, Index);
}

} // namespace pseudo
} // namespace clang