#include "common.h"

#define MEMZ 32
#define OUTZ 32
#define BUFZ 64

enum { ADV, BXL, BST, JNZ, BXC, OUT, BDV, CDV };

struct vm {
	int64_t mem[MEMZ]; int nm, pc;
	int64_t out[OUTZ]; int no;
	int64_t a,b,c;
} vm;

/* returns 0 if halted, 1 otherwise */
static int
step(void)
{
	int64_t op, ar, ac;	/* operator, lit. arg, combo arg */

	assert(vm.pc >= 0);
	assert(vm.pc+2 <= MEMZ);

	if (vm.pc >= vm.nm)
		return 0;

	op = vm.mem[vm.pc++];
	ar = vm.mem[vm.pc++];
	ac = ar==4 ? vm.a : ar==5 ? vm.b : ar==6 ? vm.c : ar;

	switch (op) {
	case ADV: vm.a = vm.a >> ac; break;
	case BDV: vm.b = vm.a >> ac; break;
	case CDV: vm.c = vm.a >> ac; break;
	case BXL: vm.b = vm.b ^ ar; break;
	case BXC: vm.b = vm.b ^ vm.c; break;
	case BST: vm.b = ac % 8; break;
	case JNZ: if (vm.a) vm.pc = ar; break;
	case OUT: assert(vm.no < OUTZ); vm.out[vm.no++] = ac%8; break;
	default: assert(!"invalid opcode"); return 0;
	}

	return 1;
}

static int64_t
recur_p2(int64_t a0, int pos)
{
	int64_t a, i;

	/*
	 * The code in the input uses up to 7 low bits of the A register
	 * to produce a single number, then chops off the low 3 bits and
	 * continues.
	 *
	 * That means bits above the current triplet influence what
	 * number it generates, but bits below don't. To generate a
	 * desired sequence then, we append, not prepend,  candidate
	 * triplets.
	 *
	 * We may end up in a situation where, given the prefix found
	 * for the numbers so far, no triplet yields the desired next
	 * number. Then we have to backtrack and find another prefix,
	 * hence the recursion.
	 */

	if (pos >= vm.nm)
		return a0 >> 3;

	for (i=0; i<8; i++) {
		vm.a=a= a0+i;
		vm.b=vm.c=vm.pc=vm.no= 0;

		while (step() && !vm.no)
			;

		if (vm.no && vm.out[0] == vm.mem[vm.nm-pos-1])
			if ((a = recur_p2(a << 3, pos+1)))
				return a;
	}

	return 0;
}

int
main(int argc, char **argv)
{
	char b[BUFZ], *tok, *rest;
	int i;

	if (argc > 1)
		DISCARD(freopen(argv[1], "r", stdin));

	fgets(b, BUFZ, stdin); sscanf(b, "Register A: %lld", &vm.a);
	fgets(b, BUFZ, stdin); sscanf(b, "Register B: %lld", &vm.b);
	fgets(b, BUFZ, stdin); sscanf(b, "Register C: %lld", &vm.c);
	fgets(b, BUFZ, stdin);

	assert(vm.b == 0);	/* assumption for part 2 */
	assert(vm.c == 0);

	rest = fgets(b, sizeof(b), stdin);
	strsep(&rest, ":");

	while ((tok = strsep(&rest, ","))) {
		assert(vm.nm < MEMZ);
		vm.mem[vm.nm++] = atoll(tok);
	}

	while (step())
		;
	for (i=0; i<vm.no; i++)
		printf(i ? ",%lld" : "17: %lld", vm.out[i]);

	printf(" %lld\n", recur_p2(0, 0));
	return 0;
}

import Control.Monad.State.Lazy
import Data.Bits (xor)
import Data.List (isSuffixOf)
import qualified Data.Vector as V

data Instr =
        ADV Int | BXL Int | BST Int | JNZ Int | BXC | OUT Int | BDV Int | CDV Int
type Machine = (Int, Int, Int, Int, V.Vector Int)

parse :: String -> Machine
parse s =
    let (la : lb : lc : _ : lp : _) = lines s
        [a, b, c] = map (read . drop 12) [la, lb, lc]
        p = V.fromList $ read $ ('[' :) $ (++ "]") $ drop 9 lp
    in  (a, b, c, 0, p)

getA, getB, getC, getIP :: State Machine Int
getA  = gets $ \(a, _, _, _ , _) -> a
getB  = gets $ \(_, b, _, _ , _) -> b
getC  = gets $ \(_, _, c, _ , _) -> c
getIP = gets $ \(_, _, _, ip, _) -> ip

setA, setB, setC, setIP :: Int -> State Machine ()
setA  a  = modify $ \(_, b, c, ip, p) -> (a, b, c, ip, p)
setB  b  = modify $ \(a, _, c, ip, p) -> (a, b, c, ip, p)
setC  c  = modify $ \(a, b, _, ip, p) -> (a, b, c, ip, p)
setIP ip = modify $ \(a, b, c, _ , p) -> (a, b, c, ip, p)

incIP :: State Machine ()
incIP = getIP >>= (setIP . succ)

getMem :: State Machine (Maybe Int)
getMem = gets (\(_, _, _, ip, p) -> p V.!? ip) <* incIP

getCombo :: State Machine (Maybe Int)
getCombo = do
    n <- getMem
    case n of
        Just 4          -> Just <$> getA
        Just 5          -> Just <$> getB
        Just 6          -> Just <$> getC
        Just n | n <= 3 -> return $ Just n
        _               -> return Nothing

getInstr :: State Machine (Maybe Instr)
getInstr = do
    opcode <- getMem
    case opcode of
        Just 0 -> fmap        ADV  <$> getCombo
        Just 1 -> fmap        BXL  <$> getMem
        Just 2 -> fmap        BST  <$> getCombo
        Just 3 -> fmap        JNZ  <$> getMem
        Just 4 -> fmap (const BXC) <$> getMem
        Just 5 -> fmap        OUT  <$> getCombo
        Just 6 -> fmap        BDV  <$> getCombo
        Just 7 -> fmap        CDV  <$> getCombo
        _      -> return Nothing

execInstr :: Instr -> State Machine (Maybe Int)
execInstr (ADV n) = (getA >>= (setA . (`div` (2^n)))) *> return Nothing
execInstr (BDV n) = (getA >>= (setB . (`div` (2^n)))) *> return Nothing
execInstr (CDV n) = (getA >>= (setC . (`div` (2^n)))) *> return Nothing
execInstr (BXL n) = (getB >>= (setB . xor n)) *> return Nothing
execInstr (BST n) = setB (n `mod` 8) *> return Nothing
execInstr (JNZ n) = do
    a <- getA
    case a of
        0 -> return ()
        _ -> setIP n
    return Nothing
execInstr  BXC    = ((xor <$> getB <*> getC) >>= setB) *> return Nothing
execInstr (OUT n) = return $ Just $ n `mod` 8

run :: State Machine [Int]
run = do
    mInstr <- getInstr
    case mInstr of
        Nothing    -> return []
        Just instr -> do
            mOut <- execInstr instr
            case mOut of
                Nothing ->           run
                Just n  -> (n :) <$> run

solve2 :: Machine -> Int
solve2 machine@(_, _, _, _, p') = head [a | x <- [1 .. 7], a <- go [x]]
    where
        p = V.toList p'
        go as =
            let a = foldl ((+) . (* 8)) 0 as
            in  case evalState (setA a *> run) machine of
                    ns  | ns == p           -> [a]
                        | ns `isSuffixOf` p ->
                            concatMap go [as ++ [a] | a <- [0 .. 7]]
                        | otherwise         -> []

main :: IO ()
main = do
    machine@(_, _, _, _, p) <- parse <$> getContents
    putStrLn $ init $ tail $ show $ evalState run machine
    print $ solve2 machine

use rustc_hash::FxHashMap;

fn parse(input: &str) -> Option<Program> {
    let mut lines = input.lines();
    let a = lines.next()?.split_once(": ")?.1.parse().ok()?;
    let b = lines.next()?.split_once(": ")?.1.parse().ok()?;
    let c = lines.next()?.split_once(": ")?.1.parse().ok()?;
    lines.next()?;
    let program = lines
        .next()?
        .split_once(": ")?
        .1
        .split(',')
        .map(|s| s.parse())
        .collect::<Result<Vec<u8>, _>>()
        .ok()?;
    Some(Program {
        a,
        b,
        c,
        out: vec![],
        program,
        ip: 0,
    })
}

#[derive(Debug, Clone, Default)]
struct Program {
    a: u64,
    b: u64,
    c: u64,
    out: Vec<u8>,
    program: Vec<u8>,
    ip: usize,
}

impl Program {
    fn run(&mut self) {
        while self.step() {}
    }

    // Returns true if a step was taken, false if it halted
    fn step(&mut self) -> bool {
        let Some(&[opcode, operand]) = &self.program.get(self.ip..self.ip + 2) else {
            return false;
        };
        self.ip += 2;
        match opcode {
            0 => self.adv(self.combo(operand)),
            1 => self.bxl(operand),
            2 => self.bst(self.combo(operand)),
            3 => self.jnz(operand),
            4 => self.bxc(),
            5 => self.out(self.combo(operand)),
            6 => self.bdv(self.combo(operand)),
            7 => self.cdv(self.combo(operand)),
            _ => panic!(),
        }
        true
    }

    fn combo(&self, operand: u8) -> u64 {
        match operand {
            0..=3 => operand as u64,
            4 => self.a,
            5 => self.b,
            6 => self.c,
            _ => unreachable!(),
        }
    }

    fn adv(&mut self, x: u64) {
        self.a >>= x
    }

    fn bxl(&mut self, x: u8) {
        self.b ^= x as u64
    }

    fn bst(&mut self, x: u64) {
        self.b = x % 8
    }

    fn jnz(&mut self, x: u8) {
        if self.a != 0 {
            self.ip = x as usize
        }
    }

    fn bxc(&mut self) {
        self.b ^= self.c
    }

    fn out(&mut self, x: u64) {
        self.out.push((x % 8) as u8)
    }

    fn bdv(&mut self, x: u64) {
        self.b = self.a >> x
    }

    fn cdv(&mut self, x: u64) {
        self.c = self.a >> x
    }
}

fn part1(input: String) {
    let mut program = parse(&input).unwrap();
    program.run();
    if let Some(e) = program.out.first() {
        print!("{e}")
    }
    for e in program.out.iter().skip(1) {
        print!(",{e}")
    }
    println!()
}

// Some assumptions on the input:
// * There is exactly one jump instruction at the end of the program, jumping to 0
// * Right before that, an output is generated
// * Right before that, register a is shifted right by 3: adv(3)
//
// Each output depends on at most 10 bits of a (it is used with a shift of at most 7).
// Therefore we look at all 10-bit a's and group them by the first number that is output.
// Then we just need to combine these generators into a chain that fits together.
fn number_generators(mut program: Program) -> [Vec<u16>; 8] {
    let mut out = [const { vec![] }; 8];
    for a in 1..(1 << 10) {
        program.a = a as u64;
        program.out.clear();
        program.ip = 0;
        program.run();
        let &output = program.out.first().unwrap();
        out[output as usize].push(a);
    }
    out
}

fn part2(input: String) {
    let mut program = parse(&input).unwrap();
    let generators = number_generators(program.clone());

    let output = program.program.clone();
    // a_candidates maps from 7-bit required prefixes to the lower bits of a that
    // generate the required numbers so far.
    let mut a_candidates: FxHashMap<u8, u64> = generators[output[0] as usize]
        .iter()
        .rev() // Collects the values for each prefix
        .map(|&a| ((a >> 3) as u8, a as u64 % 8))
        .collect();
    let len = output.len();
    for (i, x) in output.iter().enumerate().skip(1) {
        let mut next_candidates = FxHashMap::default();
        for (prefix, val) in generators[*x as usize]
            .iter()
            .filter(|&a| {
                // Take only short candidates in the end to ensure that not too many numbers are generated
                let max_bits = (len - i) * 3;
                (*a as u64) < (1u64 << max_bits)
            })
            // Only use generators that match any required prefix
            .filter(|&a| a_candidates.contains_key(&((a % (1 << 7)) as u8)))
            .map(|&a| {
                let prefix = (a >> 3) as u8;
                let val = a as u64 % 8;
                let prev = a_candidates[&((a % (1 << 7)) as u8)];
                (prefix, (val << (i * 3)) | prev)
            })
        {
            // Only insert first (smallest) encountered value
            next_candidates.entry(prefix).or_insert(val);
        }
        a_candidates = next_candidates;
    }
    println!("{}", a_candidates[&0]);

    // Verify result
    program.a = a_candidates[&0];
    program.run();
    assert_eq!(program.out, program.program);
}

util::aoc_main!();

import Control.Arrow
import Data.Bits
import qualified Data.Char as Char
import qualified Data.List as List

replace c r c' = if c' == c then r else c'

parse :: String -> ([Integer], [Int])
parse = map (replace ',' ' ')
        >>> filter ((Char.isDigit &&& Char.isSpace) >>> uncurry (||))
        >>> words
        >>> splitAt 3
        >>> (map read *** map read)

type InstructionPointer = Int

adv = 0
bxl = 1
bst = 2
jnz = 3
bxc = 4
out = 5
bdv = 6
cdv = 7

lookupCombo _    0 = 0
lookupCombo _    1 = 1
lookupCombo _    2 = 2
lookupCombo _    3 = 3
lookupCombo regs 4 = regs !! 0
lookupCombo regs 5 = regs !! 1
lookupCombo regs 6 = regs !! 2
lookupCombo regs 7 = error "Invalid operand"

execute :: InstructionPointer -> [Integer] -> [Int] -> [Int]
execute ip regs@(regA:regB:regC:[]) ops
        | ip >= length ops = []
        | instruction == adv = execute (ip + 2) [regA `div` (2 ^ comboValue), regB, regC] ops
        | instruction == bxl = execute (ip + 2) [regA, xor regB (toInteger operand), regC] ops
        | instruction == bst = execute (ip + 2) [regA, comboValue `mod` 8, regC] ops
        | instruction == jnz && regA == 0 = execute (ip + 2) regs ops
        | instruction == jnz && regA /= 0 = execute operand regs ops
        | instruction == bxc = execute (ip + 2) [regA, xor regB regC, regC] ops
        | instruction == out = (fromIntegral comboValue) `mod` 8 : execute (ip + 2) regs ops
        | instruction == bdv = execute (ip + 2) [regA, regA `div` (2 ^ comboValue), regC] ops
        | instruction == cdv = execute (ip + 2) [regA, regB, regA `div` (2 ^ comboValue)] ops
        where
                (instruction, operand) = (ops !! ip, ops !! (succ ip))
                comboValue             = lookupCombo regs operand

part1 = uncurry (execute 0)
        >>> List.map show
        >>> List.intercalate ","

valid i t n = ((n `div` (8^i)) `mod` 8) `xor` 7 `xor` (n `div` (4*(8^i))) == t

part2 = const 247839653009594

main = getContents
        >>= print
        . (part1 &&& part2)
        . parse
```haskell

import Control.Monad
import Control.Monad.RWS
import Data.Bits
import Data.Foldable
import Data.List
import Data.List.Split
import Data.Vector (Vector)
import Data.Vector qualified as Vector

type Machine = RWS (Vector Int) [Int] (Int, Int, Int)

readInput :: String -> ((Int, Int, Int), Vector Int)
readInput s =
  let (regs, _ : [prog]) = break null $ lines s
   in ( let [a, b, c] = map (read . last . words) regs in (a, b, c),
        let [_, s] = words prog in Vector.fromList $ map read $ splitOn "," s
      )

stepMachine :: Int -> Machine Int
stepMachine ip = do
  opcode <- asks (Vector.! ip)
  operand <- asks (Vector.! (ip + 1))
  (a, b, c) <- get
  let combo = [0, 1, 2, 3, a, b, c, undefined] !! operand
      ip' = ip + 2
      adv = a `div` (2 ^ combo)
      store 'A' v = modify (\(_, b, c) -> (v, b, c))
      store 'B' v = modify (\(a, _, c) -> (a, v, c))
      store 'C' v = modify (\(a, b, _) -> (a, b, v))
  case opcode of
    0 -> store 'A' adv >> return ip'
    1 -> store 'B' (b `xor` operand) >> return ip'
    2 -> store 'B' (combo .&. 7) >> return ip'
    3 -> return $ if a == 0 then ip' else operand
    4 -> store 'B' (b `xor` c) >> return ip'
    5 -> tell [combo .&. 7] >> return ip'
    6 -> store 'B' adv >> return ip'
    7 -> store 'C' adv >> return ip'

part1 (regs, prog) =
  let (a, s, w) = runRWS (go 0) prog regs
   in intercalate "," $ map show w
  where
    go ip = when (ip < Vector.length prog) $ stepMachine ip >>= go

part2 (_, prog) = minimum $ foldM go 0 $ reverse $ toList prog
  where
    go a d = do
      b <- [0 .. 7]
      let a' = (a `shiftL` 3) .|. b
          b1 = b `xor` 5
          b2 = b1 `xor` (a' `shiftR` b1)
          b3 = b2 `xor` 6
      guard $ b3 .&. 7 == d
      return a'

main = do
  input <- readInput <$> readFile "input17"
  putStrLn $ part1 input
  print $ part2 input