make assembly more ergonomic;

NateSeymour · Dec 31, 2024 · c699f0b · c699f0b
1 parent 9c78aa8
commit c699f0b
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 11 deletions.
diff --git a/src/Assembly.cpp b/src/Assembly.cpp
@@ -2,6 +2,11 @@
 
 using namespace city;
 
+Symbol Assembly::operator[](std::string const &symbol) const
+{
+    return this->symbol_table_.at(symbol);
+}
+
 Symbol Assembly::Lookup(std::string const &symbol)
 {
     return this->symbol_table_[symbol];

diff --git a/src/Assembly.h b/src/Assembly.h
@@ -19,6 +19,7 @@ namespace city
         std::unordered_map<std::string, Symbol> symbol_table_;
 
     public:
+        [[nodiscard]] Symbol operator[](std::string const &symbol) const;
         [[nodiscard]] Symbol Lookup(std::string const &symbol);
 
         Assembly(NativeMemoryHandle native_memory);

diff --git a/test/amd64.test.cpp b/test/amd64.test.cpp
@@ -16,16 +16,13 @@ TEST_F(Amd64TestRunner, ReturnVoidFunction)
     city::IRModule module{"test"};
     auto builder = module.CreateBuilder();
 
-    auto function = builder.CreateFunction("return_void");
-    builder.SetInsertPoint(function);
-
+    (void)builder.CreateFunction("return_void");
     builder.InsertRetInst(nullptr);
 
     this->jit.InsertIRModule(std::move(module));
     auto assembly = this->jit.CompileAndLink();
 
-    auto test_symbol = assembly.Lookup("return_void");
-    auto test = test_symbol.ToPointer<void()>();
+    auto test = assembly["return_void"].ToPointer<void()>();
 
     test();
 
@@ -34,22 +31,23 @@ TEST_F(Amd64TestRunner, ReturnVoidFunction)
 
 TEST_F(Amd64TestRunner, ReturnConstantFunction)
 {
+    int const EXPECTED_RETURN_VALUE = 69;
+
     city::IRModule module{"test"};
     auto builder = module.CreateBuilder();
 
-    auto function = builder.CreateFunction("return_constant");
-    builder.SetInsertPoint(function);
+    auto return_type = builder.GetType<int>();
+    (void)builder.CreateFunction("return_constant", return_type);
 
-    auto return_value = builder.CreateConstant<int>(69);
+    auto return_value = builder.CreateConstant<int>(EXPECTED_RETURN_VALUE);
     builder.InsertRetInst(return_value);
 
     this->jit.InsertIRModule(std::move(module));
     auto assembly = this->jit.CompileAndLink();
 
-    auto test_symbol = assembly.Lookup("return_constant");
-    auto test = test_symbol.ToPointer<int()>();
+    auto test = assembly["return_constant"].ToPointer<int()>();
 
-    ASSERT_EQ(test(), 69);
+    ASSERT_EQ(test(), EXPECTED_RETURN_VALUE);
 
     this->jit.RemoveModule("test");
 }