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_term() ? 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 }
89
90 auto& w = world();
91
92 if (auto i = old2new_.find(def); i != old2new_.end()) return i->second;
93 if (auto var = def->isa<Var>();
94 var && var->mut()->isa_mut<Lam>()) // TODO put this conditions inside the assert below
95 assert(false && "Lam vars should appear in a map!");
96
97 auto new_type = rewrite(def->type());
98
99 if (auto ct = isa_clos_type(def)) {
100 auto pi = rewrite(ct->op(1))->as<Pi>();
101 if (Pi::isa_basicblock(pi)) pi = w.cn(insert_ret(pi->dom(), dummy_ret_->type()));
102 auto env_type = rewrite(ct->op(2));
103 return map(def, w.sigma({pi, env_type}));
104 } else if (auto proj = def->isa<Extract>()) {
105 auto tuple = proj->tuple();
106 auto idx = Lit::isa(proj->index());
107 if (isa_clos_type(tuple->type())) {
108 assert(idx && idx <= 2 && "unknown proj from closure tuple");
109 if (*idx == 0)
110 return map(def, env_type());
111 else
112 return map(def, rewrite(tuple)->proj(*idx - 1));
113 } else if (auto var = tuple->isa<Var>(); var && isa_clos_type(var->mut())) {
114 assert(false && "proj fst type form closure type");
115 }
116 }
117
118 if (auto c = isa_clos_lit(def)) {
119 auto env = rewrite(c.env());
120 auto mode = (env->type()->isa<Idx>() || match<mem::Ptr>(env->type())) ? Unbox : Box;
121 const Def* fn = make_stub(c.fnc_as_lam(), mode, true);
122 if (env->type() == w.sigma()) {
123 // Optimize empty env
124 env = w.bot(env_type());
125 } else if (!mode) {
126 auto mem_ptr = (c.get() == attr::esc) ? mem::op_alloc(env->type(), lcm_) : mem::op_slot(env->type(), lcm_);
127 auto mem = w.extract(mem_ptr, 0_u64);
128 auto env_ptr = mem_ptr->proj(1_u64); //, w.dbg(fn->sym() + "_env"));
129 lcm_ = w.call<mem::store>(Defs{mem, env_ptr, env});
130 map(lvm_, lcm_);
131 env = env_ptr;
132 }
133 fn = w.call<core::bitcast>(new_type->op(0), fn);
134 env = w.call<core::bitcast>(new_type->op(1), env);
135 return map(def, w.tuple({fn, env}));
136 } else if (auto lam = def->isa_mut<Lam>()) {
137 return make_stub(lam, No_Env, false);
138 } else if (auto mut = def->isa_mut()) {
139 assert(!isa_clos_type(mut));
140 auto new_mut = mut->stub(new_type);
141 map(mut, new_mut);
142 for (size_t i = 0; i < mut->num_ops(); i++)
143 if (mut->op(i)) new_mut->set(i, rewrite(mut->op(i)));
144 if (!def->isa_mut<Global>() && Checker::alpha<Checker::Check>(mut, new_mut)) return map(mut, mut);
145 if (auto imm = new_mut->immutabilize()) return map(mut, imm);
146 return new_mut;
147 } else if (def->isa<Axiom>()) {
148 return def;
149 } else {
150 auto new_ops = DefVec(def->num_ops(), [&](auto i) { return rewrite(def->op(i)); });
151
152 if (auto app = def->isa<App>()) {
153 // Add dummy retcont to first-class BB
154 if (auto p = app->callee()->isa<Extract>();
155 p && isa_clos_type(p->tuple()->type()) && Pi::isa_basicblock(app->callee_type()))
156 new_ops[1] = insert_ret(new_ops[1], dummy_ret_);
157 }
158
159 auto new_def = def->rebuild(new_type, new_ops);
160
161 // We may need to update the mem token after all ops have been rewritten:
162 // F (m, a1, ..., (env, f):pct)
163 // ~>
164 // let [m', env_ptr] = :alloc T m'
165 // let m'' = :store env_ptr env
166 // F (m, a1', ..., (env_ptr, f'))
167 // ~>
168 // let ...
169 // F (m'', a1', ..., (env_ptr, f'))
170 for (size_t i = 0; i < new_def->num_ops(); i++)
171 if (new_def->op(i)->type() == w.annex<mem::M>()) new_def = new_def->refine(i, lcm_);
172
173 if (new_type == w.annex<mem::M>()) { // :store
174 lcm_ = new_def;
175 lvm_ = def;
176 } else if (new_type->isa<Sigma>()) { // :alloc, :slot, ...
177 for (size_t i = 0; i < new_type->num_ops(); i++) {
178 if (new_type->op(i) == w.annex<mem::M>()) {
179 lcm_ = w.extract(new_def, i);
180 lvm_ = w.extract(def, i);
181 break;
182 }
183 }
184 }
185
186 return map(def, new_def);
187 }
188}
189
190} // namespace mim::plug::clos
static bool alpha(Ref d1, Ref d2)
Definition check.h:65
Base class for all Defs.
Definition def.h:212
bool is_set() const
Yields true if empty or the last op is set.
Definition def.cpp:281
void make_internal()
Definition def.cpp:478
bool is_external() const
Definition def.h:416
Ref type() const noexcept
Yields the raw type of this Def, i.e. maybe nullptr.
Definition def.h:241
Ref var(nat_t a, nat_t i) noexcept
Definition def.h:381
A function.
Definition lam.h:105
Ref filter() const
Definition lam.h:117
Lam * stub(Ref type)
Definition lam.h:177
Ref body() const
Definition lam.h:118
Ref dom() const
Definition lam.h:126
static const Lam * isa_basicblock(Ref d)
Definition lam.h:137
static std::optional< T > isa(Ref def)
Definition def.h:731
World & world()
Definition phase.h:27
static const Pi * isa_basicblock(Ref d)
Is this a continuation (Pi::isa_cn) that is not Pi::isa_returning?
Definition lam.h:48
This is a thin wrapper for std::span<T, N> with the following additional features:
Definition span.h:28
Ref bot(Ref type)
Definition world.h:440
void start() override
Actual entry.
@ Type
Definition def.h:40
@ Univ
Definition def.h:40
@ Nat
Definition def.h:40
@ Bot
Definition def.h:40
@ Top
Definition def.h:40
@ Pi
Definition def.h:40
@ Lam
Definition def.h:40
The clos Plugin
Definition clos.h:7
ClosLit isa_clos_lit(Ref def, bool fn_isa_lam=true)
Tries to match a closure literal.
Definition clos.cpp:64
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(Ref def)
Definition clos.cpp:111
The mem Plugin
Definition mem.h:11
Ref op_alloc(Ref type, Ref mem)
Definition mem.h:136
Ref mem_var(Lam *lam)
Returns the memory argument of a function if it has one.
Definition mem.h:38
Ref op_slot(Ref type, Ref mem)
Definition mem.h:144
View< const Def * > Defs
Definition def.h:61
Vector< const Def * > DefVec
Definition def.h:62
auto match(Ref def)
Definition axiom.h:112