AoC19/lib/intcode/assembler.ex

defmodule Assembler.Helpers do
  defmacro assemble_macro(match, asm) do
    quote do
      def assemble_insn(unquote(match), memory) do
        macro_mem = assemble_for_macro(unquote(asm), length(memory))
        macro_mem ++ memory
      end
    end
  end
end

defmodule Assembler do
  import NimbleParsec
  import Assembler.Helpers

  label =
    ascii_string([?a..?z, ?A..?Z, ?_], min: 1)
    |> ignore(ascii_char([?:]))
    |> ignore(repeat(ascii_char([?\s])))

  param =
    optional(ignore(ascii_string([?\s], min: 1)))
    |> ascii_string([{:not, ?,}], min: 1)

  data_label =
    label
    |> concat(param)
    |> tag(:data_label)

  instruction =
    ignore(ascii_string([?\s, ?\t], min: 0))
    |> ascii_string([?a..?z], min: 1)
    |> repeat(param |> ignore(ascii_char([?,])))
    |> optional(param)
    |> tag(:insn)

  defparsec(
    :parse_line,
    choice([
      data_label,
      label |> tag(:label),
      instruction,
      eos()
    ])
  )

  def test do
    File.read!("test/intcode/cmp8.asm")
    |> assemble()
    |> Day9.run_cli()
  end

  def assemble(asm) do
    {memory, labels} = do_assemble(asm)

    memory
    |> Enum.reverse()
    |> Enum.with_index()
    |> Enum.map(fn
      {{:expr, expr}, index} ->
        expr_tokens =
          if is_binary(expr) do
            {:ok, tokens, _, _, _, _} = ExpressionEvaluator.parse(expr)
            tokens
          else
            expr
          end

        ExpressionEvaluator.do_eval(expr_tokens, fn
          "_self" -> index
          name -> Map.fetch!(labels, name)
        end)

      {it, _} ->
        it
    end)
  end

  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")
    |> Enum.map(fn line ->
      case parse_line(line) do
        {:ok, res, _, _, _, _} ->
          res

        err ->
          raise "Unable to parse line: '#{line}':\n#{inspect(err)}"
      end
    end)
    |> Enum.reduce({[], %{}}, fn line, acc ->
      assemble_line(line, acc)
    end)
  end

  def assemble_line([], acc), do: acc

  def assemble_line([label: [name]], {memory, labels}) do
    {
      memory,
      Map.put(labels, name, length(memory))
    }
  end

  def assemble_line([data_label: [name, value]], {memory, labels}) do
    value = {:expr, value}

    {
      [value | memory],
      Map.put(labels, name, length(memory))
    }
  end

  def assemble_line([insn: insn], {memory, labels}) do
    {
      assemble_insn(insn, memory),
      labels
    }
  end

  def assemble_insn(["add" | params], memory) when length(params) == 3 do
    {[a, b, dest], modes} = parse_params(params, [:read, :read, :write])
    [dest, b, a, 1 + modes | memory]
  end

  def assemble_insn(["mul" | params], memory) when length(params) == 3 do
    {[a, b, dest], modes} = parse_params(params, [:read, :read, :write])
    [dest, b, a, 2 + modes | memory]
  end

  def assemble_insn(["in", dest], memory) do
    {[dest], modes} = parse_params([dest], [:write])
    [dest, 3 + modes | memory]
  end

  def assemble_insn(["out", src], memory) do
    {[src], modes} = parse_params([src], [:read])
    [src, 4 + modes | memory]
  end

  def assemble_insn(["jnz" | params], memory) when length(params) == 2 do
    {[target, dest], modes} = parse_params(params, [:read, :read])
    [dest, target, 5 + modes | memory]
  end

  def assemble_insn(["jez" | params], memory) when length(params) == 2 do
    {[target, dest], modes} = parse_params(params, [:read, :read])
    [dest, target, 6 + modes | memory]
  end

  def assemble_insn(["clt" | params], memory) when length(params) == 3 do
    {[a, b, dest], modes} = parse_params(params, [:read, :read, :write])
    [dest, b, a, 7 + modes | memory]
  end

  def assemble_insn(["ceq" | params], memory) when length(params) == 3 do
    {[a, b, dest], modes} = parse_params(params, [:read, :read, :write])
    [dest, b, a, 8 + modes | memory]
  end

  def assemble_insn(["srel", val], memory) do
    {[val], modes} = parse_params([val], [:read])
    [val, 9 + modes | memory]
  end

  def assemble_insn(["hlt"], memory) do
    [99 | memory]
  end

  def assemble_insn(["data" | data], memory) do
    data =
      data
      |> Enum.map(&String.to_integer/1)
      |> Enum.reverse()

    data ++ memory
  end

  # macros
  assemble_macro ["mov", src, dest],
                 """
                 add 0, #{src}, #{dest}
                 """

  assemble_macro ["jmp", dest],
                 """
                 jnz _self, #{dest}
                 """

  assemble_macro ["sub", a, b, dest],
                 """
                 mul #{b}, -1, #{dest}
                 add #{a}, $#{dest}, #{dest}
                 """

  assemble_macro ["not", src, dest],
                 """
                 sub 1, #{src}, #{dest}
                 """

  assemble_macro ["cle", a, b, dest],
                 """
                   clt #{a}, #{b}, #{dest}
                   jnz $#{dest}, end
                   ceq #{a}, #{b}, #{dest}

                 end:
                 """

  assemble_macro ["cgt", a, b, dest],
                 """
                 cle #{a}, #{b}, #{dest}
                 not $#{dest}, #{dest}
                 """

  assemble_macro ["cge", a, b, dest],
                 """
                 clt #{a}, #{b}, #{dest}
                 not $#{dest}, #{dest}
                 """

  @doc """
  Raises the base to to the given power.

  ## Examples
    iex> Assembler.pow(10, 3)
    1000
    iex> Assembler.pow(5, 4)
    625
  """
  def pow(base, 1), do: base
  def pow(base, exp), do: base * pow(base, exp - 1)

  @doc """
  Parses assembly instruction parameter values and modes.

  ## Examples
    iex> Assembler.parse_params(["1", "2", "3"], [:read, :read, :read])
    {[1, 2, 3], 11100}
    iex> Assembler.parse_params(["1", "2", "3"], [:read, :read, :write])
    {[1, 2, 3], 1100}
    iex> Assembler.parse_params(["1", "2", "$3"], [:read, :read, :write])
    {[1, 2, 3], 1100}
    iex> Assembler.parse_params(["$1", "2", "3"], [:read, :read, :write])
    {[1, 2, 3], 1000}
    iex> Assembler.parse_params(["$label", "2", "3"], [:read, :read, :write])
    {[{:expr, ["label"]}, 2, 3], 1000}
    iex> Assembler.parse_params(["1", "label", "3"], [:read, :read, :write])
    {[1, {:expr, ["label"]}, 3], 1100}
  """
  def parse_params(params, param_types) do
    params
    |> Enum.zip(param_types)
    |> Enum.with_index()
    |> Enum.map_reduce(0, fn {{param, type}, index}, modes ->
      val =
        case Regex.run(~r/^[\$#]?(\d+)$/, param) do
          [_, digits] ->
            String.to_integer(digits)

          _ ->
            case param do
              "$" <> param ->
                {:expr, parse_param_expr(param)}

              "#" <> param ->
                {:expr, parse_param_expr(param)}

              _ ->
                {:expr, parse_param_expr(param)}
            end
        end

      modes =
        case type do
          :read ->
            cond do
              String.starts_with?(param, "$") ->
                modes

              String.starts_with?(param, "#") ->
                modes + 2 * pow(10, index + 2)

              true ->
                modes + pow(10, index + 2)
            end

          :write ->
            cond do
              String.starts_with?(param, "#") ->
                modes + 2 * pow(10, index + 2)

              true ->
                modes
            end
        end

      {val, modes}
    end)
  end

  def parse_param_expr(param) do
    {:ok, expr, _, _, _, _} = ExpressionEvaluator.parse(param)
    expr
  end
end