nixir

public

Import-resolving Nix IR plugin

clone

ssh://git@git.notashelf.dev:33/notashelf/nixir.git

#include "ir_gen.h"
#include <stack>
#include <unordered_map>
#include <algorithm>

namespace nix_irc {

struct NameResolver::Impl {
    std::vector<std::unordered_map<std::string, uint32_t>> scopes;
    std::vector<std::vector<std::string>> scope_names;

    Impl() {
        scopes.push_back({});
        scope_names.push_back({});
    }
};

NameResolver::NameResolver() : pImpl(std::make_unique<Impl>()) {}
NameResolver::~NameResolver() = default;

void NameResolver::enter_scope() {
    pImpl->scopes.push_back({});
    pImpl->scope_names.push_back({});
}

void NameResolver::exit_scope() {
    if (!pImpl->scopes.empty()) {
        pImpl->scopes.pop_back();
        pImpl->scope_names.pop_back();
    }
}

void NameResolver::bind(const std::string& name) {
    if (pImpl->scopes.empty()) return;
    uint32_t idx = pImpl->scope_names.back().size();
    pImpl->scopes.back()[name] = idx;
    pImpl->scope_names.back().push_back(name);
}

uint32_t NameResolver::resolve(const std::string& name) {
    for (int i = (int)pImpl->scopes.size() - 1; i >= 0; --i) {
        auto it = pImpl->scopes[i].find(name);
        if (it != pImpl->scopes[i].end()) {
            uint32_t depth = pImpl->scopes.size() - 1 - i;
            uint32_t offset = it->second;
            return depth << 16 | offset;
        }
    }
    return 0xFFFFFFFF;
}

bool NameResolver::is_bound(const std::string& name) const {
    for (auto it = pImpl->scopes.rbegin(); it != pImpl->scopes.rend(); ++it) {
        if (it->count(name)) return true;
    }
    return false;
}

struct IRGenerator::Impl {
    std::unordered_map<std::string, uint32_t> string_table;
    uint32_t next_string_id = 0;
    NameResolver name_resolver;

    Impl() {}

    uint32_t add_string(const std::string& str) {
        auto it = string_table.find(str);
        if (it != string_table.end()) {
            return it->second;
        }
        uint32_t id = next_string_id++;
        string_table[str] = id;
        return id;
    }

    std::shared_ptr<Node> convert(const std::shared_ptr<Node>& node_ptr) {
        if (!node_ptr) return std::make_shared<Node>(ConstNullNode{});

        const Node& node = *node_ptr;

        if (auto* n = node.get_if<ConstIntNode>()) {
            return std::make_shared<Node>(*n);
        }
        if (auto* n = node.get_if<ConstStringNode>()) {
            return std::make_shared<Node>(*n);
        }
        if (auto* n = node.get_if<ConstPathNode>()) {
            return std::make_shared<Node>(*n);
        }
        if (auto* n = node.get_if<ConstBoolNode>()) {
            return std::make_shared<Node>(*n);
        }
        if (auto* n = node.get_if<ConstNullNode>()) {
            return std::make_shared<Node>(*n);
        }
        if (auto* n = node.get_if<VarNode>()) {
            uint32_t idx = name_resolver.resolve(n->name.value_or(""));
            VarNode converted(idx);
            converted.name = n->name;
            converted.line = n->line;
            return std::make_shared<Node>(converted);
        }
        if (auto* n = node.get_if<LambdaNode>()) {
            name_resolver.enter_scope();
            if (n->param_name) {
                name_resolver.bind(*n->param_name);
            }
            auto body = convert(n->body);
            name_resolver.exit_scope();
            LambdaNode lambda(n->arity, body, n->line);
            lambda.param_name = n->param_name;
            return std::make_shared<Node>(lambda);
        }
        if (auto* n = node.get_if<AppNode>()) {
            auto func = convert(n->func);
            auto arg = convert(n->arg);
            return std::make_shared<Node>(AppNode(func, arg, n->line));
        }
        if (auto* n = node.get_if<AttrsetNode>()) {
            AttrsetNode attrs(n->recursive, n->line);
            name_resolver.enter_scope();
            for (const auto& [key, val] : n->attrs) {
                name_resolver.bind(key);
            }
            for (const auto& [key, val] : n->attrs) {
                attrs.attrs.push_back({key, convert(val)});
            }
            name_resolver.exit_scope();
            return std::make_shared<Node>(attrs);
        }
        if (auto* n = node.get_if<SelectNode>()) {
            auto expr = convert(n->expr);
            auto attr = convert(n->attr);
            SelectNode select(expr, attr, n->line);
            if (n->default_expr) {
                select.default_expr = convert(*n->default_expr);
            }
            return std::make_shared<Node>(select);
        }
        if (auto* n = node.get_if<WithNode>()) {
            auto attrs = convert(n->attrs);
            auto body = convert(n->body);
            return std::make_shared<Node>(WithNode(attrs, body, n->line));
        }
        if (auto* n = node.get_if<IfNode>()) {
            auto cond = convert(n->cond);
            auto then_b = convert(n->then_branch);
            auto else_b = convert(n->else_branch);
            return std::make_shared<Node>(IfNode(cond, then_b, else_b, n->line));
        }
        if (auto* n = node.get_if<LetNode>()) {
            name_resolver.enter_scope();
            for (const auto& [key, val] : n->bindings) {
                name_resolver.bind(key);
            }
            std::vector<std::pair<std::string, std::shared_ptr<Node>>> new_bindings;
            for (const auto& [key, val] : n->bindings) {
                new_bindings.push_back({key, convert(val)});
            }
            auto body = convert(n->body);
            name_resolver.exit_scope();
            LetNode let(body, n->line);
            let.bindings = std::move(new_bindings);
            return std::make_shared<Node>(let);
        }
        if (auto* n = node.get_if<LetRecNode>()) {
            name_resolver.enter_scope();
            for (const auto& [key, val] : n->bindings) {
                name_resolver.bind(key);
            }
            std::vector<std::pair<std::string, std::shared_ptr<Node>>> new_bindings;
            for (const auto& [key, val] : n->bindings) {
                new_bindings.push_back({key, convert(val)});
            }
            auto body = convert(n->body);
            name_resolver.exit_scope();
            LetRecNode letrec(body, n->line);
            letrec.bindings = std::move(new_bindings);
            return std::make_shared<Node>(letrec);
        }
        if (auto* n = node.get_if<AssertNode>()) {
            auto cond = convert(n->cond);
            auto body = convert(n->body);
            return std::make_shared<Node>(AssertNode(cond, body, n->line));
        }
        if (auto* n = node.get_if<BinaryOpNode>()) {
            auto left = convert(n->left);
            auto right = convert(n->right);
            return std::make_shared<Node>(BinaryOpNode(n->op, left, right, n->line));
        }
        if (auto* n = node.get_if<UnaryOpNode>()) {
            auto operand = convert(n->operand);
            return std::make_shared<Node>(UnaryOpNode(n->op, operand, n->line));
        }
        return std::make_shared<Node>(ConstNullNode{});
    }
};

IRGenerator::IRGenerator() : pImpl(std::make_unique<Impl>()) {}
IRGenerator::~IRGenerator() = default;

void IRGenerator::set_string_table(const std::unordered_map<std::string, uint32_t>& table) {
    pImpl->string_table = table;
}

uint32_t IRGenerator::add_string(const std::string& str) {
    return pImpl->add_string(str);
}

std::shared_ptr<Node> IRGenerator::generate(const std::shared_ptr<Node>& ast) {
    return pImpl->convert(ast);
}

}

1	#include "ir_gen.h"
2	#include <stack>
3	#include <unordered_map>
4	#include <algorithm>
5
6	namespace nix_irc {
7
8	struct NameResolver::Impl {
9	std::vector<std::unordered_map<std::string, uint32_t>> scopes;
10	std::vector<std::vector<std::string>> scope_names;
11
12	Impl() {
13	scopes.push_back({});
14	scope_names.push_back({});
15	}
16	};
17
18	NameResolver::NameResolver() : pImpl(std::make_unique<Impl>()) {}
19	NameResolver::~NameResolver() = default;
20
21	void NameResolver::enter_scope() {
22	pImpl->scopes.push_back({});
23	pImpl->scope_names.push_back({});
24	}
25
26	void NameResolver::exit_scope() {
27	if (!pImpl->scopes.empty()) {
28	pImpl->scopes.pop_back();
29	pImpl->scope_names.pop_back();
30	}
31	}
32
33	void NameResolver::bind(const std::string& name) {
34	if (pImpl->scopes.empty()) return;
35	uint32_t idx = pImpl->scope_names.back().size();
36	pImpl->scopes.back()[name] = idx;
37	pImpl->scope_names.back().push_back(name);
38	}
39
40	uint32_t NameResolver::resolve(const std::string& name) {
41	for (int i = (int)pImpl->scopes.size() - 1; i >= 0; --i) {
42	auto it = pImpl->scopes[i].find(name);
43	if (it != pImpl->scopes[i].end()) {
44	uint32_t depth = pImpl->scopes.size() - 1 - i;
45	uint32_t offset = it->second;
46	return depth << 16 \| offset;
47	}
48	}
49	return 0xFFFFFFFF;
50	}
51
52	bool NameResolver::is_bound(const std::string& name) const {
53	for (auto it = pImpl->scopes.rbegin(); it != pImpl->scopes.rend(); ++it) {
54	if (it->count(name)) return true;
55	}
56	return false;
57	}
58
59	struct IRGenerator::Impl {
60	std::unordered_map<std::string, uint32_t> string_table;
61	uint32_t next_string_id = 0;
62	NameResolver name_resolver;
63
64	Impl() {}
65
66	uint32_t add_string(const std::string& str) {
67	auto it = string_table.find(str);
68	if (it != string_table.end()) {
69	return it->second;
70	}
71	uint32_t id = next_string_id++;
72	string_table[str] = id;
73	return id;
74	}
75
76	std::shared_ptr<Node> convert(const std::shared_ptr<Node>& node_ptr) {
77	if (!node_ptr) return std::make_shared<Node>(ConstNullNode{});
78
79	const Node& node = *node_ptr;
80
81	if (auto* n = node.get_if<ConstIntNode>()) {
82	return std::make_shared<Node>(*n);
83	}
84	if (auto* n = node.get_if<ConstStringNode>()) {
85	return std::make_shared<Node>(*n);
86	}
87	if (auto* n = node.get_if<ConstPathNode>()) {
88	return std::make_shared<Node>(*n);
89	}
90	if (auto* n = node.get_if<ConstBoolNode>()) {
91	return std::make_shared<Node>(*n);
92	}
93	if (auto* n = node.get_if<ConstNullNode>()) {
94	return std::make_shared<Node>(*n);
95	}
96	if (auto* n = node.get_if<VarNode>()) {
97	uint32_t idx = name_resolver.resolve(n->name.value_or(""));
98	VarNode converted(idx);
99	converted.name = n->name;
100	converted.line = n->line;
101	return std::make_shared<Node>(converted);
102	}
103	if (auto* n = node.get_if<LambdaNode>()) {
104	name_resolver.enter_scope();
105	if (n->param_name) {
106	name_resolver.bind(*n->param_name);
107	}
108	auto body = convert(n->body);
109	name_resolver.exit_scope();
110	LambdaNode lambda(n->arity, body, n->line);
111	lambda.param_name = n->param_name;
112	return std::make_shared<Node>(lambda);
113	}
114	if (auto* n = node.get_if<AppNode>()) {
115	auto func = convert(n->func);
116	auto arg = convert(n->arg);
117	return std::make_shared<Node>(AppNode(func, arg, n->line));
118	}
119	if (auto* n = node.get_if<AttrsetNode>()) {
120	AttrsetNode attrs(n->recursive, n->line);
121	name_resolver.enter_scope();
122	for (const auto& [key, val] : n->attrs) {
123	name_resolver.bind(key);
124	}
125	for (const auto& [key, val] : n->attrs) {
126	attrs.attrs.push_back({key, convert(val)});
127	}
128	name_resolver.exit_scope();
129	return std::make_shared<Node>(attrs);
130	}
131	if (auto* n = node.get_if<SelectNode>()) {
132	auto expr = convert(n->expr);
133	auto attr = convert(n->attr);
134	SelectNode select(expr, attr, n->line);
135	if (n->default_expr) {
136	select.default_expr = convert(*n->default_expr);
137	}
138	return std::make_shared<Node>(select);
139	}
140	if (auto* n = node.get_if<WithNode>()) {
141	auto attrs = convert(n->attrs);
142	auto body = convert(n->body);
143	return std::make_shared<Node>(WithNode(attrs, body, n->line));
144	}
145	if (auto* n = node.get_if<IfNode>()) {
146	auto cond = convert(n->cond);
147	auto then_b = convert(n->then_branch);
148	auto else_b = convert(n->else_branch);
149	return std::make_shared<Node>(IfNode(cond, then_b, else_b, n->line));
150	}
151	if (auto* n = node.get_if<LetNode>()) {
152	name_resolver.enter_scope();
153	for (const auto& [key, val] : n->bindings) {
154	name_resolver.bind(key);
155	}
156	std::vector<std::pair<std::string, std::shared_ptr<Node>>> new_bindings;
157	for (const auto& [key, val] : n->bindings) {
158	new_bindings.push_back({key, convert(val)});
159	}
160	auto body = convert(n->body);
161	name_resolver.exit_scope();
162	LetNode let(body, n->line);
163	let.bindings = std::move(new_bindings);
164	return std::make_shared<Node>(let);
165	}
166	if (auto* n = node.get_if<LetRecNode>()) {
167	name_resolver.enter_scope();
168	for (const auto& [key, val] : n->bindings) {
169	name_resolver.bind(key);
170	}
171	std::vector<std::pair<std::string, std::shared_ptr<Node>>> new_bindings;
172	for (const auto& [key, val] : n->bindings) {
173	new_bindings.push_back({key, convert(val)});
174	}
175	auto body = convert(n->body);
176	name_resolver.exit_scope();
177	LetRecNode letrec(body, n->line);
178	letrec.bindings = std::move(new_bindings);
179	return std::make_shared<Node>(letrec);
180	}
181	if (auto* n = node.get_if<AssertNode>()) {
182	auto cond = convert(n->cond);
183	auto body = convert(n->body);
184	return std::make_shared<Node>(AssertNode(cond, body, n->line));
185	}
186	if (auto* n = node.get_if<BinaryOpNode>()) {
187	auto left = convert(n->left);
188	auto right = convert(n->right);
189	return std::make_shared<Node>(BinaryOpNode(n->op, left, right, n->line));
190	}
191	if (auto* n = node.get_if<UnaryOpNode>()) {
192	auto operand = convert(n->operand);
193	return std::make_shared<Node>(UnaryOpNode(n->op, operand, n->line));
194	}
195	return std::make_shared<Node>(ConstNullNode{});
196	}
197	};
198
199	IRGenerator::IRGenerator() : pImpl(std::make_unique<Impl>()) {}
200	IRGenerator::~IRGenerator() = default;
201
202	void IRGenerator::set_string_table(const std::unordered_map<std::string, uint32_t>& table) {
203	pImpl->string_table = table;
204	}
205
206	uint32_t IRGenerator::add_string(const std::string& str) {
207	return pImpl->add_string(str);
208	}
209
210	std::shared_ptr<Node> IRGenerator::generate(const std::shared_ptr<Node>& ast) {
211	return pImpl->convert(ast);
212	}
213
214	}