Day14.scala

package sungkmi.aoc2020.day14

case class State(mask: Mask, memory: Memory)

object State:
  val initial: State = State(
    mask = Mask(Set.empty, Set.empty),
    memory = Map.empty
  )

case class Mask(zero: Set[Int], one: Set[Int])

object Mask:
  def fromString(s: String): Mask =
    val grouped: Map[Char, Seq[(Char, Int)]] = s.reverse.zipWithIndex.groupBy(_._1)
    def mkSet(key: Char): Set[Int] = grouped(key).map(_._2).toSet
    Mask(zero = mkSet('0'), one = mkSet('1'))

extension (m: Mask)
  def apply(n: BigInt): BigInt =
    val cleared = m.zero.foldLeft(n)(_ `clearBit` _)
    m.one.foldLeft(cleared)(_ `setBit` _)

type Memory = Map[Int, BigInt]

extension (m: Memory)
  def apply(mask: Mask): Memory = m.view.mapValues(mask(_)).toMap
  def sumValues: BigInt = m.values.sum

enum Instruction:
  case SetMask(maskString: String)
  case SetMemory(index: Int, value: BigInt)

extension (i: Instruction)
  def apply(s: State): State = i match
    case Instruction.SetMask(maskString) =>
      s.copy(mask = Mask.fromString(maskString))
    case Instruction.SetMemory(index, value) =>
      s.copy(memory = s.memory + (index -> s.mask(value)))
    
object Instruction:
  def parse(s: String): Seq[Instruction] = s.split("\n").toSeq.map:
    case line if line.startsWith("mask = ") =>
      SetMask(line.drop(7))
    case line =>
      val Array(k, v) = line `split` " = "
      SetMemory(k.drop(4).dropRight(1).toInt, BigInt(v))

case class State2(mask: Mask2, memory: Memory2)
object State2:
  val initial: State2 = State2(
    mask = Mask2(Set.empty, Set.empty),
    memory = Map.empty
  )

case class Mask2(sets: Set[Int], floats: Set[Int])

object Mask2:
  def fromString(s: String): Mask2 =
    val grouped: Map[Char, Seq[(Char, Int)]] =
      s.reverse.zipWithIndex.groupBy(_._1).withDefaultValue(Seq.empty)
    def mkSet(key: Char): Set[Int] = grouped(key).map(_._2).toSet
    Mask2(sets = mkSet('1'), floats = mkSet('X'))

extension (m: Mask2)
  def apply(n: BigInt): Seq[BigInt] =
    val afterSet = m.sets.foldLeft(n)((n1, index) => n1.setBit(index))
    m.floats.foldLeft(Seq(afterSet)):
      case (ints, float) => ints.map(_.setBit(float)) ++ ints.map(_.clearBit(float))

type Memory2 = Map[BigInt, BigInt]

enum Instruction2:
  case SetMask(maskString: String)
  case SetMemory(index: BigInt, value: BigInt)
  
object Instruction2:
  def parse(s: String): Seq[Instruction2] = s.split("\n").toSeq.map:
    case line if line.startsWith("mask = ") =>
      SetMask(line.drop(7))
    case line =>
      val Array(k, v) = line `split` " = "
      SetMemory(BigInt(k.drop(4).dropRight(1)), BigInt(v))

extension (i: Instruction2)
  def apply(s: State2): State2 = i match
    case Instruction2.SetMask(maskString) =>
      s.copy(mask = Mask2.fromString(maskString))
    case Instruction2.SetMemory(index, value) =>
      val newMemory = s.mask(index).foldLeft(s.memory):
        (memory, i) => memory + (i -> value)
      s.copy(memory = newMemory)

def solve1(input: String): BigInt =
  val instructions = Instruction.parse(input)
  instructions.foldLeft(State.initial):
    case (state, instruction) => instruction(state)
  ._2.sumValues

def solve2(input: String): BigInt =
  val instructions = Instruction2.parse(input)
  instructions.foldLeft(State2.initial):
    case (state, instruction) => instruction(state)
  .memory.values.sum

@main def part1: Unit =
  val ans = solve1(input)
  println(ans)

@main def part2: Unit = 
  val ans = solve2(input)
  println(ans)

//lazy val input: String = """mask = 0XX1XXX1101X101100101001010X1X110000

day14Test.scala

package sungkmi.aoc2020.day14

class Day14Test extends munit.FunSuite {

  val s = """mask = XXXXXXXXXXXXXXXXXXXXXXXXXXXXX1XXXX0X
mem[8] = 11
mem[7] = 101
mem[8] = 0"""

  test("Mask#fromString") {
    val mask = Mask.fromString("XXXXXXXXXXXXXXXXXXXXXXXXXXXXX1XXXX0X")
    val expected = Mask(zero = Set(1), one = Set(6))
    assertEquals(mask, expected)
  }

  test("Mask#apply") {
    val mask = Mask(zero = Set(1), one = Set(6))
    assertEquals(mask(BigInt(11)), BigInt(73))
    assertEquals(mask(BigInt(101)), BigInt(101))
    assertEquals(mask(BigInt(0)), BigInt(64))
  }

  test("parse") {
    import Instruction._
    val expected = Seq(
      SetMask("XXXXXXXXXXXXXXXXXXXXXXXXXXXXX1XXXX0X"),
      SetMemory(8, BigInt(11)),
      SetMemory(7, BigInt(101)),
      SetMemory(8, BigInt(0)),
    )
    assertEquals(Instruction.parse(s), expected)
  }

  test("solve1") {
    assertEquals(solve1(s), BigInt(165))
  }

  val s2 = """mask = 000000000000000000000000000000X1001X
mem[42] = 100
mask = 00000000000000000000000000000000X0XX
mem[26] = 1"""

  test("solve2") {
    assertEquals(solve2(s2), BigInt(208))
  }
}