nixir

public

Import-resolving Nix IR plugin

clone

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

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

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<HasAttrNode>()) {
      auto expr = convert(n->expr);
      auto attr = convert(n->attr);
      return std::make_shared<Node>(HasAttrNode(expr, attr, n->line));
    }
    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;
      new_bindings.reserve(n->bindings.size());
      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;
      new_bindings.reserve(n->bindings.size());
      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);
}

} // namespace nix_irc

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