MimIR 0.1
MimIR is my Intermediate Representation
Loading...
Searching...
No Matches
lower_typed_clos.cpp
Go to the documentation of this file.
2
3#include <functional>
4
5#include "mim/check.h"
6
7namespace mim::plug::clos {
8
9namespace {
10const Def* insert_ret(const Def* def, const Def* ret) {
11 auto new_ops = DefVec(def->num_projs() + 1, [&](auto i) { return (i == def->num_projs()) ? ret : def->proj(i); });
12 auto& w = def->world();
13 return def->is_intro() ? w.tuple(new_ops) : w.sigma(new_ops);
14}
15} // namespace
16
18 // TODO put into c'tor - doesn't work right now, because world becomes invalid
19 dummy_ret_ = world().bot(world().cn(world().annex<mem::M>()));
20
21 for (auto mut : world().copy_externals()) rewrite(mut);
22 while (!worklist_.empty()) {
23 auto [lvm, lcm, lam] = worklist_.front();
24 worklist_.pop();
25 lcm_ = lcm;
26 lvm_ = lvm;
27 world().DLOG("in {} (lvm={}, lcm={})", lam, lvm_, lcm_);
28 if (lam->is_set()) {
29 auto new_f = rewrite(lam->filter());
30 auto new_b = rewrite(lam->body());
31 lam->reset({new_f, new_b});
32 }
33 }
34
35 for (auto lam : new_externals_) lam->make_external();
36}
37
38Lam* LowerTypedClos::make_stub(Lam* lam, enum Mode mode, bool adjust_bb_type) {
39 assert(lam && "make_stub: not a lam");
40 if (auto i = old2new_.find(lam); i != old2new_.end() && i->second->isa_mut<Lam>()) return i->second->as_mut<Lam>();
41 auto& w = world();
42 auto new_dom = w.sigma(DefVec(lam->num_doms(), [&](auto i) -> const Def* {
43 auto new_dom = rewrite(lam->dom(i));
44 if (i == Clos_Env_Param) {
45 if (mode == Unbox)
46 return env_type();
47 else if (mode == Box)
48 return world().call<mem::Ptr0>(new_dom);
49 }
50 return new_dom;
51 }));
52 if (Lam::isa_basicblock(lam) && adjust_bb_type) new_dom = insert_ret(new_dom, dummy_ret_->type());
53 auto new_type = w.cn(new_dom);
54 auto new_lam = lam->stub(new_type);
55 w.DLOG("stub {} ~> {}", lam, new_lam);
56 if (lam->is_set()) new_lam->set(lam->filter(), lam->body());
57 if (lam->is_external()) {
58 lam->make_internal();
59 new_externals_.emplace_back(new_lam);
60 }
61 const Def* lcm = mem::mem_var(new_lam);
62 const Def* env
63 = new_lam->var(Clos_Env_Param); //, (mode != No_Env) ? w.dbg("closure_env") : lam->var(Clos_Env_Param)->dbg());
64 if (mode == Box) {
65 auto env_mem = w.call<mem::load>(Defs{lcm, env});
66 lcm = w.extract(env_mem, 0_u64)->set("mem");
67 env = w.extract(env_mem, 1_u64)->set("closure_env");
68 } else if (mode == Unbox) {
69 env = w.call<core::bitcast>(lam->dom(Clos_Env_Param), env)->set("unboxed_env");
70 }
71 auto new_args = w.tuple(DefVec(lam->num_doms(), [&](auto i) {
72 return (i == Clos_Env_Param) ? env : (lam->var(i) == mem::mem_var(lam)) ? lcm : new_lam->var(i);
73 }));
74 assert(new_args->num_projs() == lam->num_doms());
75 assert(lam->num_doms() <= new_lam->num_doms());
76 map(lam->var(), new_args);
77 worklist_.emplace(mem::mem_var(lam), lcm, new_lam);
78 return map(lam, new_lam)->as<Lam>();
79}
80
81const Def* LowerTypedClos::rewrite(const Def* def) {
82 switch (def->node()) {
83 case Node::Bot:
84 case Node::Top:
85 case Node::Type:
86 case Node::Univ:
87 case Node::Nat: return def;
88 default: break;
89 }
90
91 auto& w = world();
92
93 if (auto i = old2new_.find(def); i != old2new_.end()) return i->second;
94 if (auto var = def->isa<Var>();
95 var && var->mut()->isa_mut<Lam>()) // TODO put this conditions inside the assert below
96 assert(false && "Lam vars should appear in a map!");
97
98 auto new_type = rewrite(def->type());
99
100 if (auto ct = isa_clos_type(def)) {
101 auto pi = rewrite(ct->op(1))->as<Pi>();
102 if (Pi::isa_basicblock(pi)) pi = w.cn(insert_ret(pi->dom(), dummy_ret_->type()));
103 auto env_type = rewrite(ct->op(2));
104 return map(def, w.sigma({pi, env_type}));
105 } else if (auto proj = def->isa<Extract>()) {
106 auto tuple = proj->tuple();
107 auto idx = Lit::isa(proj->index());
108 if (isa_clos_type(tuple->type())) {
109 assert(idx && idx <= 2 && "unknown proj from closure tuple");
110 if (*idx == 0)
111 return map(def, env_type());
112 else
113 return map(def, rewrite(tuple)->proj(*idx - 1));
114 } else if (auto var = tuple->isa<Var>(); var && isa_clos_type(var->mut())) {
115 assert(false && "proj fst type form closure type");
116 }
117 }
118
119 if (auto c = isa_clos_lit(def)) {
120 auto env = rewrite(c.env());
121 auto mode = (env->type()->isa<Idx>() || match<mem::Ptr>(env->type())) ? Unbox : Box;
122 const Def* fn = make_stub(c.fnc_as_lam(), mode, true);
123 if (env->type() == w.sigma()) {
124 // Optimize empty env
125 env = w.bot(env_type());
126 } else if (!mode) {
127 auto mem_ptr = (c.get() == attr::esc) ? mem::op_alloc(env->type(), lcm_) : mem::op_slot(env->type(), lcm_);
128 auto mem = w.extract(mem_ptr, 0_u64);
129 auto env_ptr = mem_ptr->proj(1_u64); //, w.dbg(fn->sym() + "_env"));
130 lcm_ = w.call<mem::store>(Defs{mem, env_ptr, env});
131 map(lvm_, lcm_);
132 env = env_ptr;
133 }
134 fn = w.call<core::bitcast>(new_type->op(0), fn);
135 env = w.call<core::bitcast>(new_type->op(1), env);
136 return map(def, w.tuple({fn, env}));
137 } else if (auto lam = def->isa_mut<Lam>()) {
138 return make_stub(lam, No_Env, false);
139 } else if (auto mut = def->isa_mut()) {
140 assert(!isa_clos_type(mut));
141 auto new_mut = mut->stub(new_type);
142 map(mut, new_mut);
143 for (size_t i = 0; i < mut->num_ops(); i++)
144 if (mut->op(i)) new_mut->set(i, rewrite(mut->op(i)));
145 if (!def->isa_mut<Global>() && Checker::alpha<Checker::Check>(mut, new_mut)) return map(mut, mut);
146 if (auto imm = new_mut->immutabilize()) return map(mut, imm);
147 return new_mut;
148 } else if (def->isa<Axiom>()) {
149 return def;
150 } else {
151 auto new_ops = DefVec(def->num_ops(), [&](auto i) { return rewrite(def->op(i)); });
152
153 if (auto app = def->isa<App>()) {
154 // Add dummy retcont to first-class BB
155 if (auto p = app->callee()->isa<Extract>();
156 p && isa_clos_type(p->tuple()->type()) && Pi::isa_basicblock(app->callee_type()))
157 new_ops[1] = insert_ret(new_ops[1], dummy_ret_);
158 }
159
160 auto new_def = def->rebuild(new_type, new_ops);
161
162 // We may need to update the mem token after all ops have been rewritten:
163 // F (m, a1, ..., (env, f):pct)
164 // ~>
165 // let [m', env_ptr] = :alloc T m'
166 // let m'' = :store env_ptr env
167 // F (m, a1', ..., (env_ptr, f'))
168 // ~>
169 // let ...
170 // F (m'', a1', ..., (env_ptr, f'))
171 for (size_t i = 0; i < new_def->num_ops(); i++)
172 if (new_def->op(i)->type() == w.annex<mem::M>()) new_def = new_def->refine(i, lcm_);
173
174 if (new_type == w.annex<mem::M>()) { // :store
175 lcm_ = new_def;
176 lvm_ = def;
177 } else if (new_type->isa<Sigma>()) { // :alloc, :slot, ...
178 for (size_t i = 0; i < new_type->num_ops(); i++) {
179 if (new_type->op(i) == w.annex<mem::M>()) {
180 lcm_ = w.extract(new_def, i);
181 lvm_ = w.extract(def, i);
182 break;
183 }
184 }
185 }
186
187 return map(def, new_def);
188 }
189}
190
191} // namespace mim::plug::clos
static bool alpha(const Def *d1, const Def *d2)
Definition check.h:71
Base class for all Defs.
Definition def.h:198
bool is_set() const
Yields true if empty or the last op is set.
Definition def.cpp:304
void make_internal()
Definition def.cpp:519
const Def * var(nat_t a, nat_t i) noexcept
Definition def.h:379
bool is_external() const
Definition def.h:414
const Def * type() const noexcept
Yields the "raw" type of this Def (maybe nullptr).
Definition def.h:242
A function.
Definition lam.h:105
const Def * filter() const
Definition lam.h:117
const Def * dom() const
Definition lam.h:126
static const Lam * isa_basicblock(const Def *d)
Definition lam.h:137
Lam * stub(const Def *type)
Definition lam.h:178
const Def * body() const
Definition lam.h:118
static std::optional< T > isa(const Def *def)
Definition def.h:730
World & world()
Definition phase.h:27
static const Pi * isa_basicblock(const Def *d)
Is this a continuation (Pi::isa_cn) that is not Pi::isa_returning?
Definition lam.h:48
const Def * bot(const Def *type)
Definition world.h:431
void start() override
Actual entry.
The clos Plugin
Definition clos.h:7
ClosLit isa_clos_lit(const Def *def, bool fn_isa_lam=true)
Tries to match a closure literal.
Definition clos.cpp:65
static constexpr size_t Clos_Env_Param
Describes where the environment is placed in the argument list.
Definition clos.h:107
const Sigma * isa_clos_type(const Def *def)
Definition clos.cpp:112
The mem Plugin
Definition mem.h:11
const Def * op_slot(const Def *type, const Def *mem)
Definition mem.h:144
const Def * mem_var(Lam *lam)
Returns the memory argument of a function if it has one.
Definition mem.h:38
const Def * op_alloc(const Def *type, const Def *mem)
Definition mem.h:136
View< const Def * > Defs
Definition def.h:49
Vector< const Def * > DefVec
Definition def.h:50
auto match(const Def *def)
Definition axiom.h:112
@ Nat
Definition def.h:85
@ Pi
Definition def.h:85
@ Univ
Definition def.h:85
@ Bot
Definition def.h:85
@ Axiom
Definition def.h:85
@ Lam
Definition def.h:85
@ Global
Definition def.h:85
@ Var
Definition def.h:85
@ Sigma
Definition def.h:85
@ Extract
Definition def.h:85
@ Type
Definition def.h:85
@ Top
Definition def.h:85
@ App
Definition def.h:85