Intcode assembler: Fix macro-scoped labels not working as intended

lib/intcode/assembler.ex

@@ -2,7 +2,7 @@ defmodule Assembler.Helpers do
   defmacro assemble_macro(match, asm) do
     quote do
       def assemble_insn(unquote(match), memory) do
-        {macro_mem, _} = do_assemble(unquote(asm), length(memory))
+        macro_mem = assemble_for_macro(unquote(asm), length(memory))
         macro_mem ++ memory
       end
     end
@@ -55,15 +55,17 @@ defmodule Assembler do
     |> Day5.run_cli()
   end
 
-  def assemble(asm, offset \\ 0) do
+  def assemble(asm) do
-    {memory, labels} = do_assemble(asm, offset)
+    {memory, labels} = do_assemble(asm)
 
     memory
     |> Enum.reverse()
     |> Enum.with_index()
     |> Enum.map(fn
       {{:expr, expr}, index} ->
-        ExpressionEvaluator.eval(expr, fn
+        expr_tokens = if is_binary(expr), do: ExpressionEvaluator.parse(expr), else: expr
+
+        ExpressionEvaluator.do_eval(expr_tokens, fn
           "_self" -> index
           name -> Map.fetch!(labels, name)
         end)
@@ -73,7 +75,38 @@ defmodule Assembler do
     end)
   end
 
-  def do_assemble(asm, offset) do
+  def assemble_for_macro(asm, macro_offset) do
+    {memory, labels} = do_assemble(asm)
+
+    memory
+    |> Enum.map(fn
+      {:expr, expr} -> {:expr, expand_macro_expression(expr, labels, macro_offset)}
+      it -> it
+    end)
+  end
+
+  def expand_macro_expression(expr, _, _) when is_integer(expr) do
+    expr
+  end
+
+  def expand_macro_expression(expr, macro_labels, macro_offset) when is_binary(expr) do
+    case Map.get(macro_labels, expr) do
+      nil -> expr
+      val -> macro_offset + val
+    end
+  end
+
+  def expand_macro_expression(expr, macro_labels, macro_offset) when is_list(expr) do
+    Enum.map(expr, fn expr ->
+      expand_macro_expression(expr, macro_labels, macro_offset)
+    end)
+  end
+
+  def expand_macro_expression({op, values}, macro_labels, macro_offset) do
+    {op, expand_macro_expression(values, macro_labels, macro_offset)}
+  end
+
+  def do_assemble(asm) do
     asm
     |> String.trim_trailing("\n")
     |> String.split("\n")
@@ -87,20 +120,20 @@ defmodule Assembler do
       end
     end)
     |> Enum.reduce({[], %{}}, fn line, acc ->
-      assemble_line(line, acc, offset)
+      assemble_line(line, acc)
     end)
   end
 
-  def assemble_line([], acc, _), do: acc
+  def assemble_line([], acc), do: acc
 
-  def assemble_line([label: [name]], {memory, labels}, offset) do
+  def assemble_line([label: [name]], {memory, labels}) do
     {
       memory,
-      Map.put(labels, name, length(memory) + offset)
+      Map.put(labels, name, length(memory))
     }
   end
 
-  def assemble_line([data_label: [name, value]], {memory, labels}, _) do
+  def assemble_line([data_label: [name, value]], {memory, labels}) do
     value =
       case value do
         {:neg, [val]} -> -val
@@ -113,7 +146,7 @@ defmodule Assembler do
     }
   end
 
-  def assemble_line([insn: insn], {memory, labels}, _) do
+  def assemble_line([insn: insn], {memory, labels}) do
     {
       assemble_insn(insn, memory),
       labels
@@ -246,10 +279,10 @@ defmodule Assembler do
           _ ->
             case param do
               "$" <> param ->
-                {:expr, param}
+                {:expr, parse_param_expr(param)}
 
               _ ->
-                {:expr, param}
+                {:expr, parse_param_expr(param)}
             end
         end
 
@@ -269,4 +302,9 @@ defmodule Assembler do
       {val, modes}
     end)
   end
+
+  def parse_param_expr(param) do
+    {:ok, expr, _, _, _, _} = ExpressionEvaluator.parse(param)
+    expr
+  end
 end