Add numeric label support to assembler

Source/Core/Common/Assembler/GekkoIRGen.cpp

@@ -46,6 +46,7 @@ public:
   void OnLoaddr(std::string_view id) override;
   void OnCloseParen(ParenType type) override;
   void OnLabelDecl(std::string_view name) override;
+  void OnNumericLabelDecl(std::string_view name, u32 num) override;
   void OnVarDecl(std::string_view name) override;
   void PostParseAction() override;
 
@@ -75,6 +76,7 @@ public:
   void AddAbsoluteAddressConv();
   void AddLiteral(u32 lit);
   void AddSymbolResolve(std::string_view sym, bool absolute);
+  void AddNumLabelSymResolve(std::string_view sym, u32 num);
 
   void RunFixups();
 
@@ -97,6 +99,8 @@ private:
   u64* m_active_var;
   size_t m_operand_scan_begin;
 
+  // Ordered top-to-bottom, stores (label number, address)
+  std::vector<std::pair<u32, u32>> m_numlabs;
   std::map<std::string, u32, std::less<>> m_labels;
   std::map<std::string, u64, std::less<>> m_constants;
   std::set<std::string> m_symset;
@@ -362,6 +366,11 @@ void GekkoIRPlugin::OnLabelDecl(std::string_view name)
   m_symset.insert(name_str);
 }
 
+void GekkoIRPlugin::OnNumericLabelDecl(std::string_view, u32 num)
+{
+  m_numlabs.emplace_back(num, m_active_block->BlockEndAddress());
+}
+
 void GekkoIRPlugin::OnVarDecl(std::string_view name)
 {
   const std::string name_str(name);
@@ -555,6 +564,35 @@ void GekkoIRPlugin::AddSymbolResolve(std::string_view sym, bool absolute)
       });
 }
 
+void GekkoIRPlugin::AddNumLabelSymResolve(std::string_view sym, u32 num)
+{
+  const u32 source_address = m_active_block->BlockEndAddress();
+  AssemblerError err_on_fail = AssemblerError{
+      fmt::format("No numeric label '{}' found below here", num),
+      m_owner->lexer.CurrentLine(),
+      m_owner->lexer.LineNumber(),
+      // Lexer should currently point to the label, as it hasn't been eaten yet
+      m_owner->lexer.ColNumber(),
+      sym.size(),
+  };
+
+  // Searching forward only
+  size_t search_start_idx = static_cast<size_t>(m_numlabs.size());
+  m_fixup_stack.emplace(
+      [this, num, source_address, search_start_idx, err_on_fail = std::move(err_on_fail)]() {
+        for (size_t i = search_start_idx; i < m_numlabs.size(); i++)
+        {
+          if (num == m_numlabs[i].first)
+          {
+            return m_numlabs[i].second - source_address;
+          }
+        }
+
+        m_owner->error = std::move(err_on_fail);
+        return u32{0};
+      });
+}
+
 void GekkoIRPlugin::SaveOperandFixup(size_t str_left, size_t str_right)
 {
   m_operand_fixups.emplace_back(std::move(m_fixup_stack.top()));
@@ -715,6 +753,33 @@ void GekkoIRPlugin::EvalTerminalRel(Terminal type, const AssemblerToken& tok)
     break;
   }
 
+  case Terminal::NumLabFwd:
+  {
+    std::optional<u32> val = tok.EvalToken<u32>();
+    ASSERT(val.has_value());
+    AddNumLabelSymResolve(tok.token_val, *val);
+    break;
+  }
+
+  case Terminal::NumLabBwd:
+  {
+    std::optional<u32> mval = tok.EvalToken<u32>();
+    ASSERT(mval.has_value());
+    u32 val = *mval;
+    if (auto label_it = std::find_if(m_numlabs.rbegin(), m_numlabs.rend(),
+                                     [val](std::pair<u32, u32> p) { return p.first == val; });
+        label_it != m_numlabs.rend())
+    {
+      AddLiteral(label_it->second - CurrentAddress());
+    }
+    else
+    {
+      m_owner->EmitErrorHere(fmt::format("No numeric label '{}' found above here", val));
+      return;
+    }
+    break;
+  }
+
   // Parser should disallow this from happening
   default:
     ASSERT(false);
@@ -780,6 +845,31 @@ void GekkoIRPlugin::EvalTerminalAbs(Terminal type, const AssemblerToken& tok)
     break;
   }
 
+  case Terminal::NumLabFwd:
+    m_owner->EmitErrorHere(
+        fmt::format("Forward label references not supported in fully resolved expressons"));
+    break;
+
+  case Terminal::NumLabBwd:
+  {
+    std::optional<u32> mval = tok.EvalToken<u32>();
+    ASSERT(mval.has_value());
+    u32 val = *mval;
+    if (auto label_it = std::find_if(m_numlabs.rbegin(), m_numlabs.rend(),
+                                     [val](std::pair<u32, u32> p) { return p.first == val; });
+        label_it != m_numlabs.rend())
+    {
+      m_eval_stack.push_back(label_it->second);
+    }
+    else
+    {
+      m_owner->EmitErrorHere(fmt::format("No numeric label '{}' found above here", val));
+      return;
+    }
+
+    break;
+  }
+
   // Parser should disallow this from happening
   default:
     ASSERT(false);