{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import defaultdict, deque\n",
    "\n",
    "\n",
    "class Jump(Exception):\n",
    "    def __init__(self, offset):\n",
    "        self.offset = offset\n",
    "        super().__init__(offset)\n",
    "\n",
    "\n",
    "def opcode(operands):\n",
    "    def decorator(f):\n",
    "        class Opcode:\n",
    "            def __set_name__(self, owner, name):\n",
    "                self.opcode = name[3:]\n",
    "                owner.opcodes[self.opcode] = self\n",
    "\n",
    "            def __repr__(self):\n",
    "                return f\"<opcode {self.opcode} {operands!r}>\"\n",
    "\n",
    "            def value(self, operand, type_):\n",
    "                if type_ == \"r\":\n",
    "                    return operand\n",
    "                try:\n",
    "                    return int(operand)\n",
    "                except ValueError:\n",
    "                    return self.registers[operand]\n",
    "\n",
    "            def __call__(self, cpu, *ops):\n",
    "                self.registers = cpu.registers\n",
    "                try:\n",
    "                    result = f(cpu, *map(self.value, ops, operands))\n",
    "                    cpu.pos += 1\n",
    "                except Jump as j:\n",
    "                    cpu.pos += j.offset\n",
    "                    result = None\n",
    "                return self.opcode, result\n",
    "\n",
    "        return Opcode()\n",
    "\n",
    "    return decorator\n",
    "\n",
    "\n",
    "class Proc:\n",
    "    opcodes = {}\n",
    "\n",
    "    def __init__(self):\n",
    "        self.reset()\n",
    "\n",
    "    def reset(self):\n",
    "        self.registers = defaultdict(int)\n",
    "        self.sound_freq = 0\n",
    "        self.pos = 0\n",
    "\n",
    "    def run(self, instructions):\n",
    "        while True:\n",
    "            opcode, *ops = instructions[self.pos]\n",
    "            yield self.opcodes[opcode](self, *ops)\n",
    "\n",
    "    def run_until_rcv(self, instructions):\n",
    "        return next(\n",
    "            val for op, val in self.run(instructions) if op == \"rcv\" and val is not None\n",
    "        )\n",
    "\n",
    "    @opcode(\"v\")\n",
    "    def op_snd(self, x):\n",
    "        self.sound_freq = x\n",
    "\n",
    "    @opcode(\"rv\")\n",
    "    def op_set(self, x, y):\n",
    "        self.registers[x] = y\n",
    "\n",
    "    @opcode(\"rv\")\n",
    "    def op_add(self, x, y):\n",
    "        self.registers[x] += y\n",
    "\n",
    "    @opcode(\"rv\")\n",
    "    def op_mul(self, x, y):\n",
    "        self.registers[x] *= y\n",
    "\n",
    "    @opcode(\"rv\")\n",
    "    def op_mod(self, x, y):\n",
    "        self.registers[x] %= y\n",
    "\n",
    "    @opcode(\"r\")\n",
    "    def op_rcv(self, x):\n",
    "        if self.registers[x]:\n",
    "            return self.sound_freq\n",
    "\n",
    "    @opcode(\"vv\")\n",
    "    def op_jgz(self, x, y):\n",
    "        if x > 0:\n",
    "            raise Jump(y)\n",
    "\n",
    "\n",
    "class SendingProc(Proc):\n",
    "    opcodes = Proc.opcodes.copy()\n",
    "\n",
    "    def __init__(self, cpu_id):\n",
    "        self.cpu_id = cpu_id\n",
    "        super().__init__()\n",
    "\n",
    "    def set_pair(self, cpu):\n",
    "        self.paired = cpu\n",
    "\n",
    "    def reset(self):\n",
    "        super().reset()\n",
    "        self.message_queue = deque()\n",
    "        self.registers[\"p\"] = self.cpu_id\n",
    "\n",
    "    @opcode(\"v\")\n",
    "    def op_snd(self, x):\n",
    "        self.paired.message_queue.append(x)\n",
    "\n",
    "    @opcode(\"r\")\n",
    "    def op_rcv(self, x):\n",
    "        if not self.message_queue:\n",
    "            raise Jump(0)\n",
    "        value = self.message_queue.popleft()\n",
    "        self.registers[x] = value\n",
    "        return value\n",
    "\n",
    "\n",
    "def parallel_run(instructions):\n",
    "    cpu0 = SendingProc(0)\n",
    "    cpu1 = SendingProc(1)\n",
    "    cpu0.set_pair(cpu1)\n",
    "    cpu1.set_pair(cpu0)\n",
    "    sendcount = 0\n",
    "    for (op0, res0), (op1, res1) in zip(cpu0.run(instructions), cpu1.run(instructions)):\n",
    "        if op1 == \"snd\":\n",
    "            sendcount += 1\n",
    "        if (op0, res0, op1, res1) == (\"rcv\", None, \"rcv\", None):\n",
    "            # deadlock\n",
    "            return sendcount"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "proc = Proc()\n",
    "test_instr = [\n",
    "    instr.split()\n",
    "    for instr in \"\"\"\\\n",
    "set a 1\n",
    "add a 2\n",
    "mul a a\n",
    "mod a 5\n",
    "snd a\n",
    "set a 0\n",
    "rcv a\n",
    "jgz a -1\n",
    "set a 1\n",
    "jgz a -2\n",
    "\"\"\".splitlines()\n",
    "]\n",
    "assert proc.run_until_rcv(test_instr) == 4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_instr = [\n",
    "    instr.split()\n",
    "    for instr in \"\"\"\\\n",
    "snd 1\n",
    "snd 2\n",
    "snd p\n",
    "rcv a\n",
    "rcv b\n",
    "rcv c\n",
    "rcv d\n",
    "\"\"\".splitlines()\n",
    "]\n",
    "assert parallel_run(test_instr) == 3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import aocd\n",
    "\n",
    "data = aocd.get_data(day=18, year=2017)\n",
    "instructions = [line.split() for line in data.splitlines()]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Part 1: 4601\n"
     ]
    }
   ],
   "source": [
    "proc = Proc()\n",
    "print(\"Part 1:\", proc.run_until_rcv(instructions))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Part 2: 6858\n"
     ]
    }
   ],
   "source": [
    "print(\"Part 2:\", parallel_run(instructions))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}