#include <ast/all.hh>
#include <ast/default-visitor.hh>
#include <ast/non-assert-visitor.hh>
#include <ast/non-object-visitor.hh>
#include <llvmtranslate/escapes-collector.hh>

namespace llvmtranslate
{
  /// LLVM IR doesn't support static link and nested functions.
  /// In order to translate those functions to LLVM IR, we use a technique
  /// called Lambda Lifting, which consists in passing a pointer to
  /// the escaped variables to the nested function using that variable.
  ///
  /// In order to do that, we need a visitor to collect these kind of
  /// variables and associate them to each function.

  class EscapesCollector
    : public ast::DefaultConstVisitor
    , public ast::NonObjectConstVisitor
    , public ast::NonAssertConstVisitor
  {
  public:
    /// Super class.
    using super_type = ast::DefaultConstVisitor;
    /// Import overloaded operator() methods.
    using super_type::operator();

    EscapesCollector()
      : did_modify_{false}
      , escaped_{}
    {}

    escaped_map_type& escaped_get() { return escaped_; }

    void operator()(const ast::FunctionChunk& e) override
    {
      const bool saved_did_modify = did_modify_;

      // Iterate on the chunk in order to iteratively collect all the callee
      // functions' escaped variables.
      did_modify_ = !e.empty();
      while (did_modify_)
        {
          did_modify_ = false;
          super_type::operator()(e);
        }

      did_modify_ = saved_did_modify;
    }

    void operator()(const ast::FunctionDec& e) override
    {
      // Keep track of the current function
      // FIXME DONE: Some code was deleted here.
      const auto function_type = \
        dynamic_cast<const type::Function*>(&e.type_get()->actual());

      precondition(function_type != nullptr);

      // Last visited function
      const type::Function* previous_scope = current_;

      if (previous_scope == nullptr)
        {
          // Then we are at the first depth of the program, hence _main
          main_function_ = function_type;
        }

      current_ = function_type;

      if (!escaped_.contains(current_))
      {
        escaped_[current_] = misc::set<const ast::VarDec*>{};
      }

      super_type::operator()(e);

      // Guess we only really need to go back one step before
      current_ = previous_scope;
    }

    void operator()(const ast::CallExp& e) override
    {
      super_type::operator()(e);

      // FIXME DONE: Some code was deleted here.
      const auto escaped_reference = escaped_.find(current_);
      precondition(escaped_reference != escaped_.end());
      const auto escaped_state = escaped_reference->second;

      const size_t before_size = escaped_state.size();

      super_type::operator()(e.def_get());

      // Check whether there are any newly collected escaped variables.
      // If there are, mark the iteration as modified.
      // FIXME DONE: Some code was deleted here.
      const size_t after_size = escaped_state.size();

      did_modify_ = before_size != after_size;
    }

    void operator()(const ast::SimpleVar& e) override
    {
      // Associate escaped variables declared in parent frames with their
      // functions
      // FIXME DONE: Some code was deleted here.
      if (current_ == main_function_ || !e.def_get()->escape_get())
      {
          super_type::operator()(e);
          return;
      }

      // Checking that the variable is escaped because not all variables are
      auto attached_to = escaped_.find(current_);

      // We have a function to attach this to
      if (attached_to != escaped_.end()
          && !attached_to->second.contains(e.def_get()))
      {
          attached_to->second.insert(e.def_get());
          did_modify_ = true;
      }

      super_type::operator()(e);
    }

  private:
    /// Whether any modification was done during the iteration.
    bool did_modify_ = false;

    /// Associate a set of variables with their function.
    escaped_map_type escaped_;

    /// Current visiting function.
    // FIXME DONE: Some code was deleted here.
    const type::Function* current_ = nullptr;

    /// The reference to the main function
    const type::Function* main_function_ = nullptr;
  };

  escaped_map_type collect_escapes(const ast::Ast& ast)
  {
    EscapesCollector collect;
    collect(ast);

    return std::move(collect.escaped_get());
  }

} // namespace llvmtranslate