recognize dotimes (#310)

decompiler/IR2/AtomicOp.h

@@ -148,6 +148,7 @@ class SimpleAtom {
     return m_int;
   }
   bool is_int() const { return m_kind == Kind::INTEGER_CONSTANT; };
+  bool is_int(s64 integer) const { return is_int() && get_int() == integer; }
   bool is_sym_ptr() const { return m_kind == Kind::SYMBOL_PTR; };
   bool is_sym_val() const { return m_kind == Kind::SYMBOL_VAL; };
   bool is_empty_list() const { return m_kind == Kind::EMPTY_LIST; };

decompiler/IR2/Form.cpp

@@ -1969,4 +1969,74 @@ void LetElement::set_body(Form* new_body) {
   m_body->parent_element = this;
 }
 
+/////////////////////////////
+// DoTimesElement
+/////////////////////////////
+
+DoTimesElement::DoTimesElement(RegisterAccess var_init,
+                               RegisterAccess var_check,
+                               RegisterAccess var_inc,
+                               Form* check_value,
+                               Form* body)
+    : m_var_init(var_init),
+      m_var_check(var_check),
+      m_var_inc(var_inc),
+      m_check_value(check_value),
+      m_body(body) {
+  m_body->parent_element = this;
+  m_check_value->parent_element = this;
+  assert(m_var_inc.reg() == m_var_check.reg());
+  assert(m_var_init.reg() == m_var_inc.reg());
+}
+
+goos::Object DoTimesElement::to_form_internal(const Env& env) const {
+  std::vector<goos::Object> outer = {
+      pretty_print::to_symbol("dotimes"),
+      pretty_print::build_list(m_var_init.to_form(env), m_check_value->to_form(env))};
+  m_body->inline_forms(outer, env);
+  return pretty_print::build_list(outer);
+}
+
+void DoTimesElement::apply(const std::function<void(FormElement*)>& f) {
+  f(this);
+  m_check_value->apply(f);
+  m_body->apply(f);
+}
+
+void DoTimesElement::apply_form(const std::function<void(Form*)>& f) {
+  m_check_value->apply_form(f);
+  m_body->apply_form(f);
+}
+
+void DoTimesElement::collect_vars(RegAccessSet& vars, bool recursive) const {
+  vars.insert(m_var_init);
+  vars.insert(m_var_check);
+  vars.insert(m_var_inc);
+  if (recursive) {
+    m_body->collect_vars(vars, recursive);
+    m_check_value->collect_vars(vars, recursive);
+  }
+}
+
+void DoTimesElement::get_modified_regs(RegSet& regs) const {
+  regs.insert(m_var_inc.reg());
+  m_body->get_modified_regs(regs);
+  m_check_value->get_modified_regs(regs);
+}
+
+std::optional<SimpleAtom> form_as_atom(const Form* f) {
+  auto as_single = f->try_as_single_element();
+  auto as_atom = dynamic_cast<SimpleAtomElement*>(as_single);
+  if (as_atom) {
+    return as_atom->atom();
+  }
+
+  auto as_se = dynamic_cast<SimpleExpressionElement*>(as_single);
+  if (as_se && as_se->expr().is_identity()) {
+    return as_se->expr().get_arg(0);
+  }
+
+  return {};
+}
+
 }  // namespace decompiler

decompiler/IR2/Form.h

@@ -1162,6 +1162,25 @@ class LetElement : public FormElement {
   bool m_star = false;
 };
 
+class DoTimesElement : public FormElement {
+ public:
+  DoTimesElement(RegisterAccess var_init,
+                 RegisterAccess var_check,
+                 RegisterAccess var_inc,
+                 Form* check_value,
+                 Form* body);
+  goos::Object to_form_internal(const Env& env) const override;
+  void apply(const std::function<void(FormElement*)>& f) override;
+  void apply_form(const std::function<void(Form*)>& f) override;
+  void collect_vars(RegAccessSet& vars, bool recursive) const override;
+  void get_modified_regs(RegSet& regs) const override;
+
+ private:
+  RegisterAccess m_var_init, m_var_check, m_var_inc;
+  Form* m_check_value = nullptr;
+  Form* m_body = nullptr;
+};
+
 /*!
  * A Form is a wrapper around one or more FormElements.
  * This is done for two reasons:
@@ -1306,4 +1325,6 @@ class FormPool {
   std::vector<Form*> m_forms;
   std::vector<FormElement*> m_elements;
 };
+
+std::optional<SimpleAtom> form_as_atom(const Form* f);
 }  // namespace decompiler

decompiler/IR2/GenericElementMatcher.cpp

@@ -123,6 +123,13 @@ Matcher Matcher::set(const Matcher& dst, const Matcher& src) {
   return m;
 }
 
+Matcher Matcher::while_loop(const Matcher& condition, const Matcher& body) {
+  Matcher m;
+  m.m_kind = Kind::WHILE_LOOP;
+  m.m_sub_matchers = {condition, body};
+  return m;
+}
+
 bool Matcher::do_match(Form* input, MatchResult::Maps* maps_out) const {
   switch (m_kind) {
     case Kind::ANY:
@@ -432,6 +439,22 @@ bool Matcher::do_match(Form* input, MatchResult::Maps* maps_out) const {
       return true;
     } break;
 
+    case Kind::WHILE_LOOP: {
+      auto as_while = dynamic_cast<WhileElement*>(input->try_as_single_element());
+      if (!as_while) {
+        return false;
+      }
+
+      if (!m_sub_matchers.at(0).do_match(as_while->condition, maps_out)) {
+        return false;
+      }
+
+      if (!m_sub_matchers.at(1).do_match(as_while->body, maps_out)) {
+        return false;
+      }
+      return true;
+    } break;
+
     default:
       assert(false);
       return false;

decompiler/IR2/GenericElementMatcher.h

@@ -45,6 +45,7 @@ class Matcher {
   static Matcher if_with_else(const Matcher& condition,
                               const Matcher& true_case,
                               const Matcher& false_case);
+  static Matcher while_loop(const Matcher& condition, const Matcher& body);
 
   enum class Kind {
     ANY_REG,     // matching any register
@@ -62,6 +63,7 @@ class Matcher {
     ANY_LABEL,
     SYMBOL,
     IF_WITH_ELSE,
+    WHILE_LOOP,
     INVALID
   };
 

decompiler/analysis/insert_lets.cpp

@@ -1,6 +1,7 @@
 #include <algorithm>
 
 #include "insert_lets.h"
+#include "decompiler/IR2/GenericElementMatcher.h"
 
 namespace decompiler {
 
@@ -79,6 +80,101 @@ Form* lca_form(Form* a, Form* b, const Env& env) {
   // fmt::print("{}\n\n", result->to_string(env));
   return result;
 }
+
+bool is_constant_int(const Form* f, int val) {
+  auto as_atom = form_as_atom(f);
+  return as_atom && as_atom->is_int(val);
+}
+
+FormElement* rewrite_as_dotimes(LetElement* in, const Env& env, FormPool& pool) {
+  // dotimes OpenGOAL:
+  /*
+     (defmacro dotimes (var &rest body)
+       "Loop like for (int i = 0; i < end; i++)"
+       `(let ((,(first var) 0))
+         (while (< ,(first var) ,(second var))
+                 ,@body
+                 (+1! ,(first var))
+                 )
+          ,@(cddr var)
+          )
+       )
+   */
+
+  // should have this anyway, but double check so we don't throw this away.
+  if (in->entries().size() != 1) {
+    return nullptr;
+  }
+
+  // look for setting a var to zero.
+  auto ra = in->entries().at(0).dest;
+  auto var = env.get_variable_name(ra);
+  if (!is_constant_int(in->entries().at(0).src, 0)) {
+    return nullptr;
+  }
+
+  // still have to check body for the increment and have to check that the lt operates on the right
+  // thing.
+  Matcher while_matcher =
+      Matcher::while_loop(Matcher::op(GenericOpMatcher::fixed(FixedOperatorKind::LT),
+                                      {Matcher::any_reg(0), Matcher::any(1)}),
+                          Matcher::any(2));
+
+  auto mr = match(while_matcher, in->body());
+  if (!mr.matched) {
+    return nullptr;
+  }
+
+  // check the lt operation:
+  auto lt_var = mr.maps.regs.at(0);
+  assert(lt_var);
+  if (env.get_variable_name(*lt_var) != var) {
+    return nullptr;  // wrong variable checked
+  }
+
+  // check the body
+  auto body = mr.maps.forms.at(2);
+  auto last_in_body = body->elts().back();
+
+  // kind hacky
+  Form fake_form;
+  fake_form.elts().push_back(last_in_body);
+  Matcher increment_matcher =
+      Matcher::op(GenericOpMatcher::fixed(FixedOperatorKind::ADDITION_IN_PLACE),
+                  {Matcher::any_reg(0), Matcher::integer(1)});
+
+  auto int_mr = match(increment_matcher, &fake_form);
+  if (!int_mr.matched) {
+    return nullptr;
+  }
+
+  auto inc_var = int_mr.maps.regs.at(0);
+  assert(inc_var);
+  if (env.get_variable_name(*inc_var) != var) {
+    return nullptr;  // wrong variable incremented
+  }
+
+  // success! here we commit to modifying this:
+
+  // first, remove the increment
+  body->pop_back();
+
+  return pool.alloc_element<DoTimesElement>(in->entries().at(0).dest, *lt_var, *inc_var,
+                                            mr.maps.forms.at(1), body);
+}
+
+/*!
+ * Attempt to rewrite a let as another form.  If it cannot be rewritten, this will return nullptr.
+ */
+FormElement* rewrite_let(LetElement* in, const Env& env, FormPool& pool) {
+  auto as_dotimes = rewrite_as_dotimes(in, env, pool);
+  if (as_dotimes) {
+    return as_dotimes;
+  }
+
+  // nothing matched.
+  return nullptr;
+}
 }  // namespace
 
 LetStats insert_lets(const Function& func, Env& env, FormPool& pool, Form* top_level_form) {
@@ -296,7 +392,19 @@ LetStats insert_lets(const Function& func, Env& env, FormPool& pool, Form* top_l
     group.first->claim_all_children();
   }
 
-  // Part 8: (todo) recognize loops and stuff.
+  // Part 8: recognize loop forms
+  top_level_form->apply_form([&](Form* f) {
+    for (auto& elt : f->elts()) {
+      auto as_let = dynamic_cast<LetElement*>(elt);
+      if (as_let) {
+        auto rewritten = rewrite_let(as_let, env, pool);
+        if (rewritten) {
+          rewritten->parent_form = f;
+          elt = rewritten;
+        }
+      }
+    }
+  });
 
   // Part 9: compact recursive lets:
   bool changed = true;