| 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<HasAttrNode>()) { |
| 141 | auto expr = convert(n->expr); |
| 142 | auto attr = convert(n->attr); |
| 143 | return std::make_shared<Node>(HasAttrNode(expr, attr, n->line)); |
| 144 | } |
| 145 | if (auto* n = node.get_if<WithNode>()) { |
| 146 | auto attrs = convert(n->attrs); |
| 147 | auto body = convert(n->body); |
| 148 | return std::make_shared<Node>(WithNode(attrs, body, n->line)); |
| 149 | } |
| 150 | if (auto* n = node.get_if<IfNode>()) { |
| 151 | auto cond = convert(n->cond); |
| 152 | auto then_b = convert(n->then_branch); |
| 153 | auto else_b = convert(n->else_branch); |
| 154 | return std::make_shared<Node>(IfNode(cond, then_b, else_b, n->line)); |
| 155 | } |
| 156 | if (auto* n = node.get_if<LetNode>()) { |
| 157 | name_resolver.enter_scope(); |
| 158 | for (const auto& [key, val] : n->bindings) { |
| 159 | name_resolver.bind(key); |
| 160 | } |
| 161 | std::vector<std::pair<std::string, std::shared_ptr<Node>>> new_bindings; |
| 162 | for (const auto& [key, val] : n->bindings) { |
| 163 | new_bindings.push_back({key, convert(val)}); |
| 164 | } |
| 165 | auto body = convert(n->body); |
| 166 | name_resolver.exit_scope(); |
| 167 | LetNode let(body, n->line); |
| 168 | let.bindings = std::move(new_bindings); |
| 169 | return std::make_shared<Node>(let); |
| 170 | } |
| 171 | if (auto* n = node.get_if<LetRecNode>()) { |
| 172 | name_resolver.enter_scope(); |
| 173 | for (const auto& [key, val] : n->bindings) { |
| 174 | name_resolver.bind(key); |
| 175 | } |
| 176 | std::vector<std::pair<std::string, std::shared_ptr<Node>>> new_bindings; |
| 177 | for (const auto& [key, val] : n->bindings) { |
| 178 | new_bindings.push_back({key, convert(val)}); |
| 179 | } |
| 180 | auto body = convert(n->body); |
| 181 | name_resolver.exit_scope(); |
| 182 | LetRecNode letrec(body, n->line); |
| 183 | letrec.bindings = std::move(new_bindings); |
| 184 | return std::make_shared<Node>(letrec); |
| 185 | } |
| 186 | if (auto* n = node.get_if<AssertNode>()) { |
| 187 | auto cond = convert(n->cond); |
| 188 | auto body = convert(n->body); |
| 189 | return std::make_shared<Node>(AssertNode(cond, body, n->line)); |
| 190 | } |
| 191 | if (auto* n = node.get_if<BinaryOpNode>()) { |
| 192 | auto left = convert(n->left); |
| 193 | auto right = convert(n->right); |
| 194 | return std::make_shared<Node>(BinaryOpNode(n->op, left, right, n->line)); |
| 195 | } |
| 196 | if (auto* n = node.get_if<UnaryOpNode>()) { |
| 197 | auto operand = convert(n->operand); |
| 198 | return std::make_shared<Node>(UnaryOpNode(n->op, operand, n->line)); |
| 199 | } |
| 200 | return std::make_shared<Node>(ConstNullNode{}); |
| 201 | } |
| 202 | }; |
| 203 | |
| 204 | IRGenerator::IRGenerator() : pImpl(std::make_unique<Impl>()) {} |
| 205 | IRGenerator::~IRGenerator() = default; |
| 206 | |
| 207 | void IRGenerator::set_string_table(const std::unordered_map<std::string, uint32_t>& table) { |
| 208 | pImpl->string_table = table; |
| 209 | } |
| 210 | |
| 211 | uint32_t IRGenerator::add_string(const std::string& str) { |
| 212 | return pImpl->add_string(str); |
| 213 | } |
| 214 | |
| 215 | std::shared_ptr<Node> IRGenerator::generate(const std::shared_ptr<Node>& ast) { |
| 216 | return pImpl->convert(ast); |
| 217 | } |
| 218 | |
| 219 | } |