17.rhm

#lang rhombus/static

import:
  "util/advent_of_code.rhm" as aoc
  "util/misc.rhm" open

@example_input(test_input1){
Register A: 729
Register B: 0
Register C: 0

Program: 0,1,5,4,3,0

}

class EMachine(a_reg :: Box.later_of(NonnegInt),
               b_reg :: Box.later_of(NonnegInt),
               c_reg :: Box.later_of(NonnegInt),
               program :: List):
  nonfinal
  field ip :: Int = 0
  field outp :: List = []
  import .bits open

  constructor (a_val :: NonnegInt,
               b_val :: NonnegInt,
               c_val :: NonnegInt,
               program :: List.of(satisfying((_ < 8)))):
    super(Box(a_val),
          Box(b_val),
          Box(c_val),
          program)

  method dv(dest :: Box, src0 :: NonnegInt, src1 :: NonnegInt):
    dest.value := src0 >> src1

  method xl(dest :: Box, src0 :: NonnegInt, src1 :: NonnegInt):
    dest.value := src0 xor src1

  method st(dest :: Box, src0 :: NonnegInt, _):
    dest.value := src0 and 0b111

  method jnz(_, arg0, arg1):
    unless arg0 == 0
    | ip := arg1 - 2

  method out(_, arg0, _):
    outp := outp.add(arg0 and 0b111)

  method
  | decode_operand(v :: satisfying((_ < 4))): v
  | decode_operand(4): a_reg.value
  | decode_operand(5): b_reg.value
  | decode_operand(6): c_reg.value

  method decode_op():
    fun operand():
      decode_operand(program[ip+1])
    fun lit_operand():
      program[ip+1]
    match program[ip]
    | 0: values(EMachine.dv, a_reg, a_reg.value, operand())
    | 1: values(EMachine.xl, b_reg, b_reg.value, lit_operand())
    | 2: values(EMachine.st, b_reg, operand(), #false)
    | 3: values(EMachine.jnz, #false, a_reg.value, lit_operand())
    | 4: values(EMachine.xl, b_reg, b_reg.value, c_reg.value)
    | 5: values(EMachine.out, #false, operand(), #false)
    | 6: values(EMachine.dv, b_reg, a_reg.value, operand())
    | 7: values(EMachine.dv, c_reg, a_reg.value, operand())

  method step():
    def values(op, dest, arg0, arg1) = decode_op()
    op(this, dest, arg0, arg1)
    ip := ip + 2

  method run() :~ EMachine:
    if ip < program.length()
    | step(); run()
    | this


fun get_puzzle_input():
  aoc.fetch_input(aoc.find_session(), 2024, 17)

fun run1(s): #void

module test:
  //check run1(test_input1) ~is #false
  check EMachine(0, 0, 9, [2,6]).run().b_reg.value ~is 1
  check EMachine(10, 0, 0, [5,0,5,1,5,4]).run().outp ~is [0,1,2]
  block:
    def m = EMachine(2024, 0, 0, [0,1,5,4,3,0]).run()
    check m.outp ~is [4,2,5,6,7,7,7,7,3,1,0]
    check m.a_reg.value ~is 0
  check EMachine(0, 29, 0, [1,7]).run().b_reg.value ~is 26
  check EMachine(0, 2024, 43690, [4,0]).run().b_reg.value ~is 44354
  check EMachine(729, 0, 0, [0,1,5,4,3,0]).run().outp ~is:
    [4,6,3,5,6,3,5,2,1,0]

module part1:
  def input = get_puzzle_input()
  run1(input)

class Matcher(pat :~ List, mutable state :~ Set):
  constructor (pat :: List):
    super(pat, for Set (i: 0..pat.length()): i)

  method match(v :: Int):
    def new_state:
      for Set (i: state):
        keep_when i < pat.length()
        keep_when pat[i] == v
        i + 1
    state := new_state
    state.length() > 0

  method is_complete():
    for any (i: state):
      i == pat.length()

class EVMachine(matcher :: Matcher, remember :: MutableMap):
  extends EMachine
  field trace :~ List = []
  field remembered = #false

  constructor (a_val :: NonnegInt,
               b_val :: NonnegInt,
               c_val :: NonnegInt,
               program :: List.of(satisfying((_ < 8))),
               remember :: MutableMap):
    def constr: super(a_val, b_val, c_val, program)
    constr(Matcher(program), remember)

  override method out(_, arg0, _):
    def v = arg0 bits.and 0b111
    if matcher.match(v)
    | super.out(#false, arg0, #false)
    | // jump to the end of the program to exit early
      ip := program.length()

  override method step():
    def state = Array(a_reg.value, b_reg.value, c_reg.value, ip).snapshot()
    match remember.get(state, #false)
    | #false:
        trace := trace.add(state)
        super.step()
    | r:
        remembered := r
        ip := program.length()

  override method run() :~ EVMachine:
    if ip < program.length()
    | step(); run()
    | def r:
        cond
        | remembered:
            match remembered
            | remembered :: List:
                for (r: remembered):
                  matcher.match(r)
                cond
                | matcher.is_complete():
                    outp := outp ++ remembered
                    outp
                | ~else:
                    #'fail
            | #'fail:
                #'fail
        | matcher.is_complete():
            outp
        | ~else:
            #'fail
      for (s: trace):
        remember[s] := r
      this

fun find_a_value(remember :: MutableMap, program :: List):
  recur rec(a = 0):
    def m = EVMachine(a, 0, 0, program, remember)
    m.run()
    cond
    | m.outp == program: a
    | ~else: rec(a + 1)

fun run2(s): #void

module test:
  check find_a_value(MutableMap(), [0,3,5,4,3,0]) ~is 117440


module part2:
  def input = get_puzzle_input()
  run2(input)