MimIR 0.1
MimIR is my Intermediate Representation
Loading...
Searching...
No Matches
clos.cpp
Go to the documentation of this file.
2
3#include <mim/config.h>
4#include <mim/pass.h>
5
6#include <mim/pass/eta_exp.h>
7#include <mim/pass/eta_red.h>
9
16
17using namespace mim;
18using namespace mim::plug;
19
32
33extern "C" MIM_EXPORT Plugin mim_get_plugin() { return {"clos", clos::register_normalizers, reg_stages, nullptr}; }
34
35namespace mim::plug::clos {
36
37/*
38 * ClosLit
39 */
40
41const Def* ClosLit::env() {
42 assert(def_);
43 return std::get<2_u64>(clos_unpack(def_));
44}
45
46const Def* ClosLit::fnc() {
47 assert(def_);
48 return std::get<1_u64>(clos_unpack(def_));
49}
50
52 auto f = fnc();
53 if (auto a = Axm::isa<attr>(f)) f = a->arg();
54 return f->isa_mut<Lam>();
55}
56
58
59ClosLit isa_clos_lit(const Def* def, bool lambda_or_branch) {
60 auto tpl = def->isa<Tuple>();
61 if (tpl && isa_clos_type(def->type())) {
62 auto a = attr::bottom;
63 auto fnc = std::get<1_u64>(clos_unpack(tpl));
64 if (auto fa = Axm::isa<attr>(fnc)) {
65 fnc = fa->arg();
66 a = fa.id();
67 }
68 if (!lambda_or_branch || fnc->isa<Lam>()) return ClosLit(tpl, a);
69 }
70 return ClosLit(nullptr, attr::bottom);
71}
72
73const Def* clos_pack(const Def* env, const Def* lam, const Def* ct) {
74 assert(env && lam);
75 assert(!ct || isa_clos_type(ct));
76 auto& w = env->world();
77 auto pi = lam->type()->as<Pi>();
78 assert(env->type() == pi->dom(Clos_Env_Param));
79 ct = (ct) ? ct : clos_type(w.cn(clos_remove_env(pi->dom())));
80 return w.tuple(ct, {env->type(), lam, env})->isa<Tuple>();
81}
82
83std::tuple<const Def*, const Def*, const Def*> clos_unpack(const Def* c) {
84 assert(c && isa_clos_type(c->type()));
85 // auto& w = c->world();
86 // auto env_type = c->proj(0_u64);
87 // // auto pi = clos_type_to_pi(c->type(), env_type);
88 // auto fn = w.extract(c, w.lit_idx(3, 1));
89 // auto env = w.extract(c, w.lit_idx(3, 2));
90 // return {env_type, fn, env};
91 auto [ty, pi, env] = c->projs<3>();
92 return {ty, pi, env};
93}
94
95const Def* clos_apply(const Def* closure, const Def* args) {
96 auto& w = closure->world();
97 auto [_, fn, env] = clos_unpack(closure);
98 auto pi = fn->type()->as<Pi>();
99 return w.app(fn, DefVec(pi->num_doms(), [&](auto i) { return clos_insert_env(i, env, args); }));
100}
101
102/*
103 * closure types
104 */
105
106const Sigma* isa_clos_type(const Def* def) {
107 auto& w = def->world();
108 auto sig = def->isa_mut<Sigma>();
109 if (!sig || sig->num_ops() < 3 || sig->op(0_u64) != w.type()) return nullptr;
110 auto var = sig->var(0_u64);
111 if (sig->op(2_u64) != var) return nullptr;
112 auto pi = sig->op(1_u64)->isa<Pi>();
113 return (pi && Pi::isa_cn(pi) && pi->num_ops() > 1_u64 && pi->dom(Clos_Env_Param) == var) ? sig : nullptr;
114}
115
116Sigma* clos_type(const Pi* pi) {
117 auto& w = pi->world();
118 auto doms = pi->doms();
119 return ctype(w, doms, nullptr)->as_mut<Sigma>();
120}
121
122const Pi* clos_type_to_pi(const Def* ct, const Def* new_env_type) {
123 assert(isa_clos_type(ct));
124 auto& w = ct->world();
125 auto pi = ct->op(1_u64)->as<Pi>();
126 auto new_dom = new_env_type ? clos_sub_env(pi->dom(), new_env_type) : clos_remove_env(pi->dom());
127 return w.cn(new_dom);
128}
129
130/*
131 * closure environments
132 */
133
134const Def* clos_insert_env(size_t i, const Def* env, std::function<const Def*(size_t)> f) {
135 return (i == Clos_Env_Param) ? env : f(shift_env(i));
136}
137
138const Def* clos_remove_env(size_t i, std::function<const Def*(size_t)> f) { return f(skip_env(i)); }
139
140const Def* ctype(World& w, Defs doms, const Def* env_type) {
141 if (!env_type) {
142 auto sigma = w.mut_sigma(w.type(), 3_u64)->set("Clos");
143 sigma->set(0_u64, w.type());
144 sigma->set(1_u64, ctype(w, doms, sigma->var(0_u64)));
145 sigma->set(2_u64, sigma->var(0_u64));
146 return sigma;
147 }
148 return w.cn(
149 DefVec(doms.size() + 1, [&](auto i) { return clos_insert_env(i, env_type, [&](auto j) { return doms[j]; }); }));
150}
151
152} // namespace mim::plug::clos
void reg_stages(Flags2Stages &stages)
Definition affine.cpp:11
static auto isa(const Def *def)
Definition axm.h:107
Base class for all Defs.
Definition def.h:251
T * as_mut() const
Asserts that this is a mutable, casts constness away and performs a static_cast to T.
Definition def.h:494
World & world() const noexcept
Definition def.cpp:439
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 * op(size_t i) const noexcept
Definition def.h:308
const Def * var(nat_t a, nat_t i) noexcept
Definition def.h:429
const Def * type() const noexcept
Yields the "raw" type of this Def (maybe nullptr).
Definition def.cpp:447
A function.
Definition lam.h:111
A dependent function type.
Definition lam.h:15
static const Pi * isa_cn(const Def *d)
Is this a continuation - i.e. is the Pi::codom mim::Bottom?
Definition lam.h:48
A dependent tuple type.
Definition tuple.h:20
static void hook(Flags2Stages &stages)
Definition pass.h:57
Data constructor for a Sigma.
Definition tuple.h:68
The World represents the whole program and manages creation of MimIR nodes (Defs).
Definition world.h:32
const Def * fnc()
Definition clos.cpp:46
const Def * env_var()
Definition clos.cpp:57
const Def * env()
Definition clos.cpp:41
void reg_stages(Flags2Stages &stages)
Definition clos.cpp:20
#define MIM_EXPORT
Definition config.h:16
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:59
void register_normalizers(Normalizers &normalizers)
Sigma * clos_type(const Pi *pi)
Creates a typed closure type from pi.
Definition clos.cpp:116
size_t shift_env(size_t i)
Definition clos.h:110
const Def * clos_remove_env(size_t i, std::function< const Def *(size_t)> f)
Definition clos.cpp:138
const Def * clos_insert_env(size_t i, const Def *env, std::function< const Def *(size_t)> f)
Definition clos.cpp:134
const Def * ctype(World &w, Defs doms, const Def *env_type=nullptr)
Definition clos.cpp:140
static constexpr size_t Clos_Env_Param
Describes where the environment is placed in the argument list.
Definition clos.h:107
const Pi * clos_type_to_pi(const Def *ct, const Def *new_env_type=nullptr)
Convert a closure type to a Pi, where the environment type has been removed or replaced by new_env_ty...
Definition clos.cpp:122
std::tuple< const Def *, const Def *, const Def * > clos_unpack(const Def *c)
Deconstruct a closure into (env_type, function, env).
Definition clos.cpp:83
const Def * clos_sub_env(const Def *tup_or_sig, const Def *new_env)
Definition clos.h:139
size_t skip_env(size_t i)
Definition clos.h:113
const Def * clos_pack(const Def *env, const Def *fn, const Def *ct=nullptr)
Pack a typed closure. This assumes that fn expects the environment as its Clos_Env_Paramth argument.
Definition clos.cpp:73
const Def * clos_apply(const Def *closure, const Def *args)
Apply a closure to arguments.
Definition clos.cpp:95
const Sigma * isa_clos_type(const Def *def)
Definition clos.cpp:106
Definition ast.h:14
View< const Def * > Defs
Definition def.h:76
Vector< const Def * > DefVec
Definition def.h:77
absl::flat_hash_map< flags_t, std::function< std::unique_ptr< Stage >(World &)> > Flags2Stages
Maps an an axiom of a Stage to a function that creates one.
Definition plugin.h:22
mim::Plugin mim_get_plugin()
Basic info and registration function pointer to be returned from a specific plugin.
Definition plugin.h:30