AOC 2021 Day 24 – No U - TrebledJ's Pages

This post is a writeup on solving the Advent of Code 2021 Day 24 challenge: Arithmetic Logic Unit.

The Problem

At first glance, I thought it was something similar to the intcode problems from two years ago. After a more thorough read, it turns out this is a completely different problem. To my surprise, the day 24 challenge actually resembles the reverse engineering found in some capture the flag competitions!

To summarise the challenge, we're given a list of assembly-like instructions like so:

inp w
add z w
mod z 2
div w 2

Our submarine's arithmetic logic unit (ALU) has four registers: w, x, y, and z.

There are six different types of instructions in total:

inp a reads the leftmost digit from input into register a (if the input is 123, inp w will store 1 in w),
add a b adds register a with register/value b,
mul a b multiplies register a by register/value b,
div a b divides register a by register/value b (integer division—i.e. divide + truncate—is performed, so div 10 3 is 3 and div -5 2 is -2),
mod a b takes the remainder of dividing register a by register/value b, and
eql a b evaluates to 1 if a equals b, and 0 otherwise.

You may notice that there are no jump instructions, so the ALU program is quite literally straightforward. We just need to run each instruction once until we hit the end.

Given the ALU program, we need to find the maximum (part 1) and minimum (part 2) ALU input which will cause the program to terminate with a 0 in register z. The ALU input is constrained to 14-digit integers, and all digits are non-zero.

Solving

In the midst of reading the challenge, I had already decided on the tool to use. Z3!

Z3 is a logic programming engine (theorem prover, actually) developed by Microsoft Research. It's pretty useful for working with symbols and constraints and more importantly, deriving possible solutions for said symbols. For example, say we want to know what are the solutions of $x^2 = x$. Z3 can solve this for us:

>>> import z3
>>> x = z3.Int('x')         # Create a symbol.
>>> z3.solve(x*x == x)      # Tell Z3 to solve for x.
[x = 1]                     # Z3 found a solution!
>>> z3.solve(x*x == x, x != 1) # Any other solutions?
[x = 0]                     # uwu

Initially, since I wanted to stick to my resolution of using Haskell or Rust, I tried looking for Z3 bindings for either language. They do exist. But they seem annoying to wrestle with.

So I decided to make an exception... just for this day... I'll use Python.

(I might return to write a Haskell/Rust Z3 solution later.)

Sokath, his eyes uncovered!

In my experience, whenever starting a reverse challenge and given some assembly, it's useful to start by analysing the control flow of the assembly and look for patterns. Well, the ALU program has ~~super simple~~ barely any control flow, every single instruction will run one after another.

Now if we were given actual binaries, we would be able to throw them into a decompiler and get a control flow graph. But since the ALU instructions here are contrived, I decided not to bother and go with the fallback method of analysis:

STARING AT THE ASSEMBLY.

Here are the first 40 lines (out of 252) of the given ALU program. Now stare with me.

inp w
mul x 0
add x z
mod x 26
div z 1
add x 13
eql x w
eql x 0
mul y 0
add y 25
mul y x
add y 1
mul z y
mul y 0
add y w
add y 8
mul y x
add z y
inp w
mul x 0
add x z
mod x 26
div z 1
add x 12
eql x w
eql x 0
mul y 0
add y 25
mul y x
add y 1
mul z y
mul y 0
add y w
add y 16
mul y x
add z y
inp w
mul x 0
add x z
mod x 26
[...]

^{(You can view the full ALU program I was given here.)}

Done staring? I find that it also helps to add annotations on the right of the assembly.

inp w           ; Read input.
mul x 0         ; x = 0
add x z         ; x = z
mod x 26        ; x = z % 26
div z 1         ; z = z???
add x 13        ; x += 13

After enough staring, your eyes will be uncovered. Or exhausted. Or perhaps both.

Here were my observations:

There are 14 repeated blocks. One block handling each digit from input.
Blocks are demarcated with an opening inp w.
Each block contains 18 instructions.
There are only three differences across each block:
- Instruction 5: div z D, where D is either 1 or 26.
- Instruction 6: add x A, where A can be any integer (negatives also appear!).
- Instruction 16: add y B, where B can be any integer (seems like only non-negatives here).
The other 15 instructions are the same in each block.
Between each block, only the value of z is propagated. w, x, and y are all local within a block.
We can simplify each iteration to the following expression:
- For each input digit I,
  - x = (I != (z % 26 + A)).
  - If x, then
    - z = (z / D) * 26 + I + B.
  - Else,
    - z = z / D.

Nifty!

This is really all the intelligence we need. The rest is just hacking together the solution with the Z3 API.

Snoring with ZZZ

Firstly, we read our file and gather our intelligence like a shepherd gathering their flock or a mother hen gathering her chicks.

with open(file) as f:
    instrs = [l.strip().split() for l in f.read().splitlines()]

n = 14  # Number of digits.
block = 18

# Gather the "magic numbers" from the ALU program.
addx = []
addy = []
divz = []
for i in range(n):
    divz.append(int(instrs[i*block + 4][-1]))
    addx.append(int(instrs[i*block + 5][-1]))
    addy.append(int(instrs[i*block + 15][-1]))

Since all the numbers are at the end of the line, it's sufficient to split each line by whitespace, index the last token, and convert it to an int.

Next, we'll start doing our Z3 things. We'll define one symbol for each digit, and constrain each symbol to non-zero digits.

# Make input ints.
inp = [z3.Int(f'inp_{x}') for x in range(n)]

# Create a Z3 solver.
s = z3.Optimize()

# Constrain input to non-zero digits.
s.add(*[1 <= i for i in inp], *[i <= 9 for i in inp])

We use a special Z3 Optimize solver. This provides .maximize and .minimize methods which will be useful for tackling parts 1 and 2 respectively.

We then implement the iteration logic and add constraints. To keep things clear (and since I'm not too sure how modifying Z3 symbols work), I used a separate z symbol (zs[i]) for each iteration.

# Chain constraints. Each iteration will have a separate z variable.
# The constraints added will connect z[i+1] to z[i].
zs = [z3.Int(f'z_{i}') for i in range(n + 1)]
s.add(zs[0] == 0)
for i in range(n):
    x = inp[i] != (zs[i] % 26 + addx[i])
    s.add(zs[i+1] == z3.If(x, (zs[i] / divz[i]) * 26 + inp[i] + addy[i],
                              zs[i] / divz[i]))

s.add(zs[-1] == 0) # Victory constraint.

Finally, we chain together the symbol we want to maximise/minimise (which is our ALU input), then we call the respective .maximise/.minimise functions.

# Construct full input to optimise (with place value).
full_inp = reduce(lambda acc, x: acc*10 + x, inp, 0)


def get_inp(model):
    return ''.join(str(model[i]) for i in inp)


# Part 1.
s.push()
s.maximize(full_inp)
s.check()
part1 = get_inp(s.model())
print('part1:', part1)
s.pop()

# Part 2.
s.push()
s.minimize(full_inp)
s.check()
part2 = get_inp(s.model())
print('part2:', part2)
s.pop()

Here, I used s.push() and s.pop() to create a sort of "scope". Without the "scopes", .maximise and .minimise work a bit weirdly in that once I call the first, the second will still return the same .maximised number.

And that's pretty much it. On my laptop, it spits out both solutions within half a minute. I would've hoped for better, but eh, at least it's much faster—both coding-wise and runtime-wise—than a brute force scan.