//
// Created by erki on 02.03.21.
//

#include "OperationFinderAstVisitor.hpp"

#include "OperationFinder.hpp"
#include "ASTHelpers.hpp"

OperationFinderAstVisitor::OperationFinderAstVisitor(OperationFinder* op_finder)
  : _context(nullptr)
  , _op_finder(op_finder)
{ }

void OperationFinderAstVisitor::NewContext(clang::ASTContext* context)
{
  _context = context;
}

bool OperationFinderAstVisitor::VisitForStmt(clang::ForStmt* stmt)
{
  assert(_context);

  return true;
}

bool OperationFinderAstVisitor::VisitBinaryOperator(clang::BinaryOperator* op)
{
  assert(_context);

  if (!op->isCompoundAssignmentOp())
    _op_finder->processArithmetic(op, _context->getSourceManager());

  return true;
}

bool OperationFinderAstVisitor::VisitUnaryOperator(clang::UnaryOperator* op)
{
  assert(_context);

  _op_finder->processUnaryArithmetic(op, _context->getSourceManager());

  return true;
}

bool OperationFinderAstVisitor::VisitCallExpr(clang::CallExpr* call)
{
  assert(_context);

  _op_finder->processFunctionCall(call, _context->getSourceManager());

  return true;
}

bool OperationFinderAstVisitor::VisitArraySubscriptExpr(clang::ArraySubscriptExpr* subscript)
{
  assert(_context);

  _op_finder->processArraySubscript(subscript, _context->getSourceManager());

  return true;
}

bool OperationFinderAstVisitor::dataTraverseStmtPre(clang::Stmt* stmt)
{
  assert(_context);

  if (clang::Stmt* branch_entry = _isBranchEntry(stmt))
  {
    _branch_stack.push_back(branch_entry);
    _op_finder->branchEntered();
  }

  return true;
}

bool OperationFinderAstVisitor::dataTraverseStmtPost(clang::Stmt* stmt)
{
  assert(_context);

  if (_loop_header.in_loop_header && (_loop_header.init == stmt || _loop_header.header_end == stmt))
  {
    if (_loop_header.header_end == stmt)
    {
      assert(_branch_stack.back() == stmt);
      _branch_stack.pop_back();
      _op_finder->branchExited();

      if (_loop_header.init)
      {
        assert(_branch_stack.back() == _loop_header.init);
        _branch_stack.pop_back();
        _op_finder->branchExited();
      }

      _loop_header = {};
    }
    else if (_loop_header.init == stmt && !_loop_header.header_end)
    {
      assert(_branch_stack.back() == stmt);
      _branch_stack.pop_back();
      _op_finder->branchExited();

      _loop_header = {};
    }
  }
  else if (!_branch_stack.empty() && _branch_stack.back() == stmt)
  {
    _branch_stack.pop_back();
    _op_finder->branchExited();
  }
  else
  {
    assert(std::find(_branch_stack.cbegin(), _branch_stack.cend(), stmt)
           == _branch_stack.cend());
  }

  return true;
}

clang::Stmt* OperationFinderAstVisitor::_isBranchEntry(clang::Stmt* stmt)
{
  if (auto* loop = clang::dyn_cast<clang::ForStmt>(stmt);
      loop && loop->getInit() && (loop->getInc() || loop->getCond()))
  {
    _loop_header.init = loop->getInit();

    if (loop->getCond())
    {
      _loop_header.header_start = loop->getCond();
      _loop_header.header_end = loop->getCond();
    }

    if (loop->getInc())
    {
      if (!_loop_header.header_start)
        _loop_header.header_start = loop->getInc();

      _loop_header.header_end = loop->getInc();
    }

    _loop_header.in_loop_header = _loop_header.init || _loop_header.header_end;

    if (_loop_header.init)
      return _loop_header.init;
    else if (_loop_header.header_start)
      return _loop_header.header_end;
  }
  else if (_loop_header.in_loop_header && _loop_header.header_start == stmt)
  {
    return _loop_header.header_end;
  }

  return nullptr;
}