09.rhm

#lang rhombus/static/and_meta

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

@example_input(test_input1){
2333133121414131402

}

class Block(label :: maybe(NonnegInt),
            mutable size :: NonnegInt,
            mutable pos :: NonnegInt,
            mutable prev :~ maybe(Block),
            mutable next :~ maybe(Block)):
  implements Printable
  export: Free File Head

  annot.macro 'Free': 'Block && satisfying(Block.is_free_space(_))'
  annot.macro 'File': 'Block && satisfying(Block.is_file(_))'
  annot.macro 'Head': 'Block && satisfying(fun(b): Block.prev(b) == #false)'

  override describe(mode, rec):
    fun d4(n):
      str.d(n, ~align: #'right, ~width: 4)
    fun d3(n):
      str.d(n, ~align: #'right, ~width: 3)
    def [children, ...]:
      recur each_child(child :~ maybe(Block) = this.first(),
                       children :~ List = []):
        if child
        | def desc:
            PrintDesc.concat(PrintDesc.newline(),
                             "<< ",
                             match child!!.label
                             | #false: "free"
                             | lbl: d4(lbl),
                             " :: ",
                             d3(child!!.size),
                             " @ ",
                             d4(child!!.pos),
                             " >>",
                             if this == child | " <--" | "")
          each_child(child!!.next, children.add(desc))
        | children
    PrintDesc.concat("{[",
                     PrintDesc.nest(1, children), ...,
                     PrintDesc.newline(),
                     "]}")

  method is_file():
    !is_free_space()

  method is_free_space():
    !label

  method first() :~ Block:
    cond
    | prev: prev!!.first()
    | ~else: this

  method last() :~ Block:
    cond
    | next: next!!.last()
    | ~else: this

  method find(pred,
              ~dir: dir = Block.next,
              ~missing: missing = fun(who):
                          error(~who: who,
                                "cannot find block satisfying condition"))
    :~ Block:
      ~who: who
      recur do_find(b :~ maybe(Block) = this):
        cond
        | !b: match missing
              | missing :: Function: missing(who)
              | ~else: missing
        | pred(b): b
        | ~else: do_find(dir(b))

  method count(dir):
    recur do_count(b :~ maybe(Block) = this, c = 0):
      cond
      | !b: c
      | ~else: do_count(dir(b), c + 1)

  method relink():
    when prev | prev!!.next := this
    when next | next!!.prev := this

  method split(split_size :: PosInt, is_frontend :: Boolean) :~ Block:
    ~who: who
    unless split_size < size
    | error(~who: who, "split size too big",
            error.text(~label: "block size", size),
            error.text(~label: "split size", split_size))

    def values(this_size, other_size):
      if is_frontend:
      | values(split_size, size - split_size)
      | values(size - split_size, split_size)
    def new_block:
      Block(label,
            other_size,
            pos + this_size,
            this,
            next)

    size := this_size
    next := new_block

    this.relink()
    new_block.relink()

    if is_frontend | this | new_block

  method coalesce() :~ Block.Free:
    ~who: who
    unless this is_a Block.Free
    | error(~who: who, "cannot coalesce non-free block",
            error.text(~label: "block label", label),
            error.text(~label: "block position", pos))
    fun boundary(scan_dir) :~ Block:
      recur find_boundary(block :~ Block = this):
        match scan_dir(block)
        | free :: Block.Free:
            find_boundary(free)
        | ~else: block

    def first_free = boundary(Block.prev)
    def last_free = boundary(Block.next)
    when first_free != last_free
    | def new_size:
        recur calculate_size(block :~ Block = first_free, size = 0):
          cond
          | block == last_free: size + block.size
          | ~else: calculate_size(block.next, size + block.size)
      this.size := new_size
      this.pos := first_free.pos
      this.prev := first_free.prev
      this.next := last_free.next
      relink()
    this


  method copy_to(dest :: Block.Free):
    ~who: who
    unless this is_a Block.File
    | error(~who: who,
            "source is not a file",
            error.text(~label: "source position", pos),
            error.text(~label: "dest position", dest.pos))
    unless size == dest.size
    | error(~who: who,
            "cannot swap blocks of different sizes",
            error.text(~label: "block size", size),
            error.text(~label: "other block size", dest.size))

    def adjacent: dest.next == this

    def values(src_pos, src_prev, src_next):
      values(pos, if adjacent | this | prev, next)
    def values(dest_pos, dest_prev, dest_next):
      values(dest.pos, dest.prev, if adjacent | dest | dest.next)

    pos := dest_pos
    prev := dest_prev
    next := dest_next

    dest.pos := src_pos
    dest.prev := src_prev
    dest.next := src_next

    relink()
    dest.relink()
    values(this, dest)


def ZERO = Char"0".to_int()

fun get_puzzle_input() :~ String:
  aoc.fetch_input(aoc.find_session(), 2024, 9)

fun read_input(input :: String) :: Block:
  fun read_blocks(i :: NonnegInt,
                  pos :: NonnegInt,
                  label :: NonnegInt,
                  last_block :: maybe(Block))
    :: Block:
      fun done(): last_block!!.first()
      cond
      | i < input.length():
          def ch = input[i]
          cond
          | ch == Char"\n": done()
          | ~else:
              def size = input[i].to_int() - ZERO
              cond
              | size == 0: read_blocks(i+1, pos, label, last_block)
              // is_file
              | i mod 2 == 0:
                  def new_block:
                    Block(label, size, pos, last_block, #false)
                  when last_block
                  | last_block!!.next := new_block
                  read_blocks(i + 1, pos + size, label + 1, new_block)
              // is_free_space
              | ~else:
                  def new_block:
                    Block(#false, size, pos, last_block, #false)
                  when last_block
                  | last_block!!.next := new_block
                  read_blocks(i + 1, pos + size, label, new_block)
      | ~else: done()
  read_blocks(0, 0, 0, #false)

fun defrag_part1(head :: Block.Head):
  recur defrag1(free :: Block.Free = head.find(~dir: Block.next,
                                               Block.is_free_space),
                file :: Block.File = head.last()):
    when free.pos < file.pos
    | cond
      | free.size == file.size:
          file.copy_to(free)
          defrag1(file.find(~dir: Block.next, Block.is_free_space),
                  free.coalesce().find(~dir: Block.prev, Block.is_file))
      | free.size < file.size:
          defrag1(free, file.split(free.size, #false))
      | ~else:
          defrag1(free.split(file.size, #true), file)

fun checksum(head :: Block.Head):
  recur do_checksum(block :: maybe(Block) = head, sum = 0):
    cond
    | !block: sum
    | block!!.label:
        // XXX: if I was to go through this again I would calculate the
        // triangular number part instead of looping
        def size:
          for values(size = 0) (pos: block!!.pos .. block!!.pos + block!!.size):
            size + block!!.label * pos
        do_checksum(block!!.next, sum + size)
    | ~else: do_checksum(block!!.next, sum)

fun defrag_part2(head :: Block.Head):
  fun find_free_block(pos, size) :~ maybe(Block):
    def free :~ maybe(Block):
      head.find(~dir: Block.next,
                fun(b :~ Block):
                  b is_a Block.Free
                    && b.pos < pos
                    && b.size >= size,
                ~missing: #false)
    cond
    | !free: free
    | size < free!!.size: free.split(size, #true)
    | ~else: free

  recur defrag1(file :: maybe(Block.File) = head.last()):
    when file && file!!.label > 0
    | let file :~ Block = file
      fun is_next_file(b :~ Block):
        b is_a Block.File && (b.label :~ NonnegInt) < file.label
      match find_free_block(file.pos, file.size)
      | free :: Block.Free:
          let values(_, free :~ Block) = file.copy_to(free)
          defrag1(free.coalesce().find(~dir: Block.prev, is_next_file))
      | #false: defrag1(file.prev?.find(~dir: Block.prev, is_next_file))


fun run(input :: String, defrag):
  def head = read_input(input)
  defrag(head)
  checksum(head)

check run(test_input1, defrag_part1) ~is 1928

module part1:
  def input = get_puzzle_input()
  run(input, defrag_part1)

check run(test_input1, defrag_part2) ~is 2858

module part2:
  def input = get_puzzle_input()
  run(input, defrag_part2)