Advent of Code 2020: Day 14

aoc_2020_day_14.exs

defmodule Advent2020.Day14 do
  defp parse_addr(<<"] = ", value::binary>>, acc),
    do: {:mem, String.to_integer(acc), String.to_integer(value)}

  defp parse_addr(<<d::binary-size(1), rest::binary>>, acc), do: parse_addr(rest, acc <> d)

  def parse("mask = " <> mask), do: {:mask, mask}
  def parse(<<"mem[", d::binary-size(1), rest::binary>>), do: parse_addr(rest, d)

  def input do
    Path.expand("day14.txt", :code.priv_dir(:advent_2020))
    |> File.stream!()
    |> Stream.map(fn line -> parse(String.trim_trailing(line)) end)
  end

  defmodule PartOne do
    use Bitwise

    defp or_mask(mask), do: String.replace(mask, "X", "0") |> String.to_integer(2)
    defp and_mask(mask), do: String.replace(mask, "X", "1") |> String.to_integer(2)

    defp process({:mask, mask}, state), do: %{state | mask: {or_mask(mask), and_mask(mask)}}

    defp process({:mem, addr, value}, %{mask: {or_mask, and_mask}, mem: mem} = state) do
      %{state | mem: Map.put(mem, addr, (value ||| or_mask) &&& and_mask)}
    end

    def run(input) do
      %{mem: mem} = Enum.reduce(input, %{mask: nil, mem: %{}}, &process/2)
      Map.values(mem) |> Enum.sum()
    end
  end

  defmodule PartTwo do
    use Bitwise

    defp or_mask(mask), do: String.replace(mask, "X", "0") |> String.to_integer(2)

    defp and_mask(mask) do
      String.replace(mask, "0", "1")
      |> String.replace("X", "0")
      |> String.to_integer(2)
    end

    defp bitmasks(mask) do
      to_charlist(mask)
      |> Enum.reverse()
      |> Enum.reduce({1, []}, fn
        ?X, {bitmask, masks} -> {bitmask <<< 1, [bitmask | masks]}
        _, {bitmask, masks} -> {bitmask <<< 1, masks}
      end)
      |> elem(1)
    end

    defp recur_bitmasks(addr, [], mem, value), do: Map.put(mem, addr, value)

    defp recur_bitmasks(addr, [bitmask | bitmasks], mem, value) do
      mem = recur_bitmasks(addr, bitmasks, mem, value)
      recur_bitmasks(addr ||| bitmask, bitmasks, mem, value)
    end

    defp process({:mask, mask}, state),
      do: %{state | mask: {or_mask(mask), and_mask(mask), bitmasks(mask)}}

    defp process({:mem, addr, value}, %{mask: mask, mem: mem} = state) do
      {or_mask, and_mask, bitmasks} = mask
      addr = (addr ||| or_mask) &&& and_mask
      %{state | mem: recur_bitmasks(addr, bitmasks, mem, value)}
    end

    def run(input) do
      %{mem: mem} = Enum.reduce(input, %{mask: nil, mem: %{}}, &process/2)
      Map.values(mem) |> Enum.sum()
    end
  end
end