MimIR 0.1
MimIR is my Intermediate Representation
Loading...
Searching...
No Matches
clos_conv_prep.cpp
Go to the documentation of this file.
2
3#include <mim/nest.h>
4
5#include <mim/pass/eta_exp.h>
6
8
9namespace mim::plug::clos {
10
11namespace {
12// FIXME: these guys do not work if another pass rewrites curr_mut()'s body
13bool isa_cnt(const App* body, const Def* def, size_t i) {
14 return Pi::isa_returning(body->callee_type()) && body->arg() == def && i == def->num_ops() - 1;
15}
16
17const Def* isa_br(const App* body, const Def* def) {
18 if (!Pi::isa_cn(body->callee_type())) return nullptr;
19 auto proj = body->callee()->isa<Extract>();
20 return (proj && proj->tuple() == def && proj->tuple()->isa<Tuple>()) ? proj->tuple() : nullptr;
21}
22
23bool isa_callee_br(const App* body, const Def* def, size_t i) {
24 if (!Pi::isa_cn(body->callee_type())) return false;
25 return isa_callee(def, i) || isa_br(body, def);
26}
27
28Lam* isa_retvar(const Def* def) {
29 if (auto [var, lam] = ca_isa_var<Lam>(def); var && lam && var == lam->ret_var()) return lam;
30 return nullptr;
31}
32
33} // namespace
34
37 eta_exp_ = man->find<EtaExp>();
38}
39
41 if (eta_exp_) lam = eta_exp_->new2old(lam);
42 return lam2fscope_[lam];
43}
44
47 lam2fscope_[curr_mut()] = curr_mut();
48 DLOG("scope {} -> {}", curr_mut(), curr_mut());
49 auto nest = Nest(curr_mut());
50 for (auto mut : nest.muts()) {
51 if (auto bb_lam = Lam::isa_mut_basicblock(mut)) {
52 DLOG("scope {} -> {}", bb_lam, curr_mut());
53 lam2fscope_[bb_lam] = curr_mut();
54 }
55 }
56 }
57
58 auto body = curr_mut()->body()->isa<App>();
59 // Skip if the mutable is already wrapped or the body is undefined/no continuation.
60 ignore_ = !(body && Pi::isa_cn(body->callee_type())) || wrapper_.contains(curr_mut());
61}
62
63const App* ClosConvPrep::rewrite_arg(const App* app) {
64 auto& w = world();
65 auto arg = app->arg();
66
67 if (arg->isa<Var>()) return app;
68
69 for (auto i = 0u; i < arg->num_projs(); i++) {
70 auto op = arg->proj(i);
71
72 auto refine = [&](const Def* new_op) {
73 const Def* args;
74 if (arg == op)
75 args = new_op;
76 else
77 args = arg->refine(i, new_op);
78 return app->refine(1, args)->as<App>();
79 };
80
81 if (auto lam = isa_retvar(op); lam && from_outer_scope(lam)) {
82 DLOG("found return var from enclosing scope: {}", op);
83 return refine(eta_wrap(op, attr::freeBB)->set("free_ret"));
84 }
85 if (auto bb_lam = Lam::isa_mut_basicblock(op); bb_lam && from_outer_scope(bb_lam)) {
86 DLOG("found BB from enclosing scope {}", op);
87 return refine(w.call(attr::freeBB, op));
88 }
89 if (isa_cnt(app, arg, i)) {
90 if (Axm::isa<attr>(attr::returning, op) || isa_retvar(op)) {
91 return app;
92 } else if (auto contlam = op->isa_mut<Lam>()) {
93 return refine(w.call(attr::returning, contlam));
94 } else {
95 auto wrapper = eta_wrap(op, attr::returning)->set("eta_cont");
96 DLOG("eta expanded return cont: {} -> {}", op, wrapper);
97 return refine(wrapper);
98 }
99 }
100
101 if (!isa_callee_br(app, arg, i)) {
102 if (auto bb_lam = Lam::isa_mut_basicblock(op)) {
103 DLOG("found firstclass use of BB: {}", bb_lam);
104 return refine(w.call(attr::fstclassBB, bb_lam));
105 }
106 // TODO: If EtaRed eta-reduces branches, we have to wrap them again!
107 if (isa_retvar(op)) {
108 DLOG("found firstclass use of return var: {}", op);
109 return refine(eta_wrap(op, attr::fstclassBB)->set("fstclass_ret"));
110 }
111 }
112 }
113
114 return app;
115}
116
118 if (Pi::isa_cn(app->callee_type())) {
119 if (auto br = app->callee()->isa<Extract>()) {
120 auto branches = br->tuple();
121 // Eta-Expand branches
122 if (branches->isa<Tuple>() && branches->type()->isa<Arr>()) {
123 for (size_t i = 0, e = branches->num_ops(); i != e; ++i) {
124 if (!branches->op(i)->isa_mut<Lam>()) {
125 auto wrapper = eta_wrap(branches->op(i), attr::bottom)->set("eta_br");
126 DLOG("eta wrap branch: {} -> {}", branches->op(i), wrapper);
127 branches = branches->refine(i, wrapper);
128 }
129 }
130 return app->refine(0, app->callee()->refine(0, branches))->as<App>();
131 }
132 }
133 }
134
135 return app;
136}
137
138const Def* ClosConvPrep::rewrite(const Def* def) {
139 if (ignore_ || Axm::isa<attr>(def)) return def;
140
141 if (auto app = def->isa<App>()) {
142 app = rewrite_arg(app);
143 app = rewrite_callee(app);
144 return app;
145 }
146
147 return def;
148}
149
150} // namespace mim::plug::clos
const Pi * callee_type() const
Definition lam.h:279
const Def * callee() const
Definition lam.h:277
const Def * arg() const
Definition lam.h:286
A (possibly paramterized) Array.
Definition tuple.h:117
static auto isa(const Def *def)
Definition axm.h:107
Base class for all Defs.
Definition def.h:251
const Def * proj(nat_t a, nat_t i) const
Similar to World::extract while assuming an arity of a, but also works on Sigmas and Arrays.
Definition def.cpp:588
Def * set(size_t i, const Def *)
Successively set from left to right.
Definition def.cpp:263
const Def * refine(size_t i, const Def *new_op) const
Definition def.cpp:229
T * isa_mut() const
If this is mutable, it will cast constness away and perform a dynamic_cast to T.
Definition def.h:485
const Def * type() const noexcept
Yields the "raw" type of this Def (maybe nullptr).
Definition def.cpp:447
Performs η-expansion: f -> λx.f x, if f is a Lam with more than one user and does not appear in calle...
Definition eta_exp.h:13
Extracts from a Sigma or Array-typed Extract::tuple the element at position Extract::index.
Definition tuple.h:206
A function.
Definition lam.h:111
static Lam * isa_mut_basicblock(const Def *d)
Only for mutables.
Definition lam.h:146
const Def * body() const
Definition lam.h:124
Builds a nesting tree of all immutables‍/binders.
Definition nest.h:11
virtual void init(PassMan *)
Definition pass.cpp:30
PassMan & man()
Definition pass.h:97
friend class PassMan
Definition pass.h:166
static const Pi * isa_cn(const Def *d)
Is this a continuation - i.e. is the Pi::codom mim::Bottom?
Definition lam.h:48
static const Pi * isa_returning(const Def *d)
Is this a continuation (Pi::isa_cn) which has a Pi::ret_pi?
Definition lam.h:50
World & world()
Definition pass.h:64
Data constructor for a Sigma.
Definition tuple.h:68
A variable introduced by a binder (mutable).
Definition def.h:700
const Def * rewrite(const Def *) override
const App * rewrite_callee(const App *app)
void enter() override
Invoked just before Pass::rewriteing PassMan::curr_mut's body.
const App * rewrite_arg(const App *app)
const Def * eta_wrap(const Def *def, attr a)
#define DLOG(...)
Vaporizes to nothingness in Debug build.
Definition log.h:95
The clos Plugin
Definition clos.h:7
std::tuple< const Extract *, N * > ca_isa_var(const Def *def)
Checks is def is the variable of a mut of type N.
Definition clos.h:79
const App * isa_callee(const Def *def, size_t i)
Definition lam.h:346
@ Lam
Definition def.h:114
@ Extract
Definition def.h:114
@ App
Definition def.h:114
@ Tuple
Definition def.h:114