slip/opt.cpp at master · maxime-tournier/slip · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
#include "opt.hpp"

#include "ir.hpp"
#include "repr.hpp"

namespace ir {

  struct map_visitor {

    template<class Func>
    expr operator()(const expr& self, const Func& func) const {
      return func(self);
    }

    template<class Func>
    expr operator()(const block& self, const Func& func) const {
      vector<expr> items; items.reserve(self.items.size());

      for(const expr& e: self.items) {
        items.emplace_back(e.visit(*this, func));
      }

      return func(block{std::move(items)});
    }


    template<class Func>
    expr operator()(const ref<closure>& self, const Func& func) const {

      vector<expr> captures; captures.reserve(self->captures.size());
      for(const expr& c: self->captures) {
        captures.emplace_back(c.visit(*this, func));
      }

      expr body = operator()(self->body, func);

      vector<expr> items = body.match([&](const expr& self) {
          return vector<expr>(1, self);
        },
        [&](const block& self) {
          return self.items;
        });

      return func(make_ref<closure>(self->argc, std::move(captures),
                                    block{std::move(items)}));
    }


    template<class Func>
    expr operator()(const ref<branch>& self, const Func& func) const {
      expr then = self->then.visit(*this, func);
      expr alt = self->alt.visit(*this, func);

      return func(make_ref<branch>(std::move(then), std::move(alt)));
    }
  };


  // map an optimization pass recursively
  static expr map(const expr& self, expr (*pass)(const expr&)) {
    return self.visit(map_visitor(), pass);
  };


  // flatten nested blocks
  static expr flatten_blocks(const block& self) {
    vector<expr> items;

    for(const expr& e: self.items) {
      e.match([&](const expr& e) {
          items.emplace_back(e);
        },
        [&](const block& e) {
          for(const expr& sub: e.items) {
            items.emplace_back(sub);
          }
        });
    }

    return block{std::move(items)};
  }

  static expr flatten_blocks(const expr& self) {
    return self.match([&](const expr& self) {
        return self;
      }, [&](const block& self) {
        return flatten_blocks(self);
      });
  }


  // optimize expression
  expr opt(const expr& self) {
    return map(self, flatten_blocks);
  }

}