advent.py

from __future__ import annotations
from collections import defaultdict

from os import path
from typing import DefaultDict, List, NamedTuple, Optional, Sequence, Set, Tuple, cast
import requests
import re
from copy import copy
from pyrsistent import m, s, v, PMap, PSet, PVector

COOKIE_PATH = path.expanduser("~/.advent_cookie")

def to_int_or_str(x: str):
	try:
		return int(x)
	except:
		return x

class Input:
	content: str

	def __init__(self, day: int, year: int = 2020, sample: bool = False, strip: Optional[str] = "\n"):
		if sample:
			input_path = path.join(path.dirname(__file__), "{:02d}/sample.txt".format(day))
			with open(input_path, "r") as f:
				self.content = f.read()
		else:
			input_path = path.join(path.dirname(__file__), "{:02d}/input.txt".format(day))

			if not path.exists(input_path):
				print("Downloading input file...")
				with open(COOKIE_PATH) as f:
					cookie = f.read()
				r = requests.get( # type: ignore
					"https://adventofcode.com/{}/day/{}/input".format(year, day),
					headers={ "Cookie": "session={}".format(cookie) }
				)
				with open(input_path, "w") as f:
					f.write(r.text)
					self.content = r.text
			else:
				with open(input_path, "r") as f:
					self.content = f.read()

		if strip is not None:
			self.content = self.content.strip(strip)

		print(f"\x1b[33mLoaded input: {len(self.content.splitlines())} lines\x1b[0m")

	def all(self) -> str:
		return self.content

	def ints(self) -> Sequence[int]:
		return [int(x) for x in self.content.split()]

	def lines(self) -> Sequence[str]:
		return self.content.splitlines()

	def tokens(self, sep: str = r"[\s\n]+") -> Sequence[str]:
		return re.compile(sep).split(self.content)

	def line_tokens(self, sep: str = r"[\s]+", line_sep: str = r"\n") -> Sequence[Sequence[str]]:
		return [
			re.compile(sep).split(line)
			for line in re.compile(line_sep).split(self.content)
		]

	def program(self) -> Program:
		return Program([
			Instruction(line[0], [int(x) for x in line[1:]]) # change this to to_int_or_str if need be (and cry)
			for line in self.line_tokens()
		])

class BadPc(Exception):
	execution: Execution

	def __init__(self, execution: Execution):
		super().__init__(f"Bad program counter at {execution}")
		self.execution = execution

class BadInstruction(Exception):
	execution: Execution

	def __init__(self, execution: Execution):
		super().__init__(f"Bad instruction at {execution}")
		self.execution = execution

class InfiniteLoop(Exception):
	execution: Execution

	def __init__(self, execution: Execution):
		super().__init__(f"Infinite loop at {execution}")
		self.execution = execution

class ExecutionComplete(Exception):
	execution: Execution

	def __init__(self, execution: Execution):
		super().__init__(f"Execution complete at {execution}")
		self.execution = execution

class Instruction(NamedTuple):
	operation: str
	arguments: List[int]

	def __repr__(self):
		arg_str = "".join([str(arg) for arg in self.arguments])
		return f"\x1b[36m{self.operation} \x1b[33m{arg_str}\x1b[39m"

class Program:
	instructions: List[Instruction]

	def __init__(self, instructions: List[Instruction]):
		self.instructions = instructions

	def replace(self, index: int, instruction: Instruction) -> Program:
		return Program([
			self.instructions[i] if i != index else instruction
			for i in range(len(self.instructions))
		])

	def create_execution(self):
		return Execution(self)

	def execute(self):
		return self.create_execution().execute()

	def show(self):
		for i, instr in enumerate(self.instructions):
			print(f"\x1b[90m{i:4d} " + repr(instr))

Position = int
default_acc = 0

class Execution:
	program: Program
	pc: int
	acc: PMap[int, int]
	visited: PSet[Position]
	history: PVector[Execution]

	def __init__(
		self,
		program: Program,
		pc: int = 0,
		acc: Optional[PMap[int, int]] = None,
		visited: Optional[PSet[Position]] = None,
		history: Optional[PVector[Execution]] = None
	):
		self.program = program
		self.pc = pc
		self.acc = acc or cast("PMap[int, int]", m()).set(default_acc, 0)
		self.visited = visited or s()
		self.history = history or v()

	def get_position(self) -> Position:
		return self.pc

	def step(self) -> Execution:
		position = self.get_position()

		if self.pc == len(self.program.instructions):
			raise ExecutionComplete(self)
		elif self.pc < 0 or self.pc > len(self.program.instructions):
			raise BadPc(self)
		elif position in self.visited:
			raise InfiniteLoop(self)

		visited = self.visited.add(position)
		history = self.history.append(self)
		op, args = self.program.instructions[self.pc]

		if op == "nop":
			return Execution(
				self.program,
				self.pc + 1,
				self.acc,
				visited,
				history
			)
		elif op == "acc":
			new_value = cast(int, self.acc.get(default_acc)) + args[0]
			acc = self.acc.set(0, new_value)
			return Execution(
				self.program,
				self.pc + 1,
				acc,
				visited,
				history
			)
		elif op == "jmp":
			return Execution(
				self.program,
				self.pc + args[0],
				self.acc,
				visited,
				history
			)
		else:
			raise BadInstruction(self)

	def execute(self) -> Execution:
		state = self
		while True:
			state = state.step()

	def show(self):
		min_instr = max(self.pc - 4, 0)
		max_instr = min(self.pc + 4, len(self.program.instructions) - 1)
		instructions_to_show = self.program.instructions[min_instr:max_instr+1]

		print()
		print(f" pc: \x1b[34m\x1b[1m{self.pc:5d}\x1b[0m")
		print(f"acc: \x1b[35m\x1b[1m{self.acc[0]:5d}\x1b[22m  {dict(self.acc)}\x1b[0m")
		print()

		for i, instruction in enumerate(instructions_to_show):
			pointer = "\x1b[48;5;236m\x1b[32m>>>" if min_instr + i == self.pc else "   "
			to_print = f"{pointer} \x1b[38;5;243m{min_instr+i:4d} {instruction}"
			print(to_print + " " * (80 - len(to_print)), end = "\x1b[0m\n")

		print()

	def __repr__(self):
		return f"<Execution pc={self.pc} acc={self.acc}>"