1 files changed, 151 insertions, 0 deletions
diff --git a/tiger-compiler/src/llvmtranslate/escapes-collector.cc b/tiger-compiler/src/llvmtranslate/escapes-collector.cc
new file mode 100644
index 0000000..4984805
--- /dev/null
+++ b/tiger-compiler/src/llvmtranslate/escapes-collector.cc
@@ -0,0 +1,151 @@
+#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