We have proof automation now

I've long had a soft spot for dependently-typed languages like Coq Rocq and Lean. They offer the possibility of a type system capable of encoding and enforcing arbitrarily subtle invariants. The sort of thing that, in regular languages, ends up (at best) as a comment, and which quickly gets lost as the size of the team grows. Then you get subtle misunderstandings and components that don't quite fit together. It's often the case that those components have grown to a sufficient size that, when the problem is noticed, aligning either of them is a wearying prospect. Perhaps, say dependent types seductively, you could write those invariants formally and have a machine check them.

(p.s. Coq changed its name! I remember many years ago at a Coq conference in Princeton, I tried suggesting that, in an English-speaking world, having a programming language called Coq was an impediment. I don't think the audience agreed at the time. I also joked that many of the talks there sounded like a speech by Tyrion Lannister, there being so many Coqs and Hoares. A joke that was hilarious and timely, even though it fell completely flat, coming as it did before the final season of that show and our collective memory-holing of it.)

The problem has always been that with great type-system power comes great proof effort. I can certainly attest to entire days spent proving really quite simple things. Doing proofs is actually quite fun: it's challenging, interactive, and there's a clear goal. But gosh, does it take a lot of time, especially if, like me, you don't know what you're doing. There's also the periodic, galling experience, at the end of many hours of effort, where you realise that the goal that you're trying to prove is, in fact, false. The classic result here is the retrospective from the seL4 effort that found that, even though the project was large enough for the engineers to develop considerable experience, they spent about 10 times as much time proving as they did designing and implementing. They ended up with more than 20 times as many lines of proof code as they did C code.

That overhead has made programming in dependently-typed languages extremely niche. It has also spurred people to try and automate it away. The attempt I'm passingly familiar with is F*, where the system tries to have an SMT solver automatically discharge the obligations. That certainly works for simple cases, but it's very easy to craft something that causes the SMT solver to go off into space and run for hours, leaving you wondering whether it's ever going to finish. I've seen that people who use these languages a lot have to develop a sixth sense for what is going to make the solver happy, and then craft everything around that. It can help, but to an extent it converts the problem into mysticism: you end up serving a complex and fickle god.

A critical fact is that, at least in theory, once the statement is correct, the contents of its proof are irrelevant: only its existence matters. This is not entirely true because of two complicating factors: first, what the seL4 group called “proof engineering”: the need to structure proofs so that the effort of realigning them after code changes is reduced. And, second, sufficiently complicated proofs can cause even type checkers to blow up and consume vast amounts of memory.

We now have LLMs which, combined with proof irrelevance, promise to be an extremely capable form of proof automation. With sufficient amounts of automation perhaps you don't need to worry about proof engineering nearly so much. You still need to avoid blowing up the type checker but, in my limited tests, LLMs can avoid that. Potentially, LLMs suddenly make dependent-type systems dramatically more practical. I wanted to play around with this so built a Zstandard decompressor in Lean, mostly because I was also curious about Zstandard.

Zstandard seems like it's winning the competition to replace gzip as the canonical compression utility. It's another LZ77-style compressor, but it offers better entropy coding and a careful design that allows it to achieve very impressive decompression speeds. It will never be as beautiful as bzip2, but the shining elegance of the Burrows–Wheeler transform doesn't count for too much in the face of significant practical advantages:

zstd bzip2 gzip lzma (XZ/LZMA2) 50 100 200 500 1000 2000 zstd bzip2 gzip lzma (XZ/LZMA2) 70 72 74 76 78 80 82 84 86 Compression tradeoff on 64 MiB of Lean/mathlib source Space saved (%) — farther right is more compression Decompression throughput (MiB/s, log scale) — higher is faster

(Measurements taken on the standard reference computer, i.e. whatever the author was using at the time. And note the log scale on the y-axis: gzip and Zstandard are in their own speed class. This is an Apple machine and Apple's gzip is especially optimised; expect gzip to be slower elsewhere.)

Zstandard (by Yann Collet, building on the seminal ANS work by Jarek Duda) has an RFC, but it is quite terse. It contains all the information you need to implement a decompressor, but unless you're already quite familiar with compression, I think you'll need to re-read it a few times to understand what's going on. I, at least, had to read section 4.1 half a dozen times before I felt that I had a decent grasp of it. Too late into this process, I discovered that my colleague, Nigel Tao, has written a better write-up of Zstandard than I was going to manage anyway. So, if you want to understand Zstandard, you should read that. I'm just going to give an explanation of the most interesting bit, the entropy encoder, and mix that in with some evangelism about Lean.

The job of an entropy encoder is, given a set of symbols with non-uniform probabilities, to encode a sequence of those symbols using the fewest number of bits. The classic entropy coder is a Huffman encoder. Huffman encoders build a binary tree with symbols at the leaf nodes, and Huffman showed that a very simple algorithm produces an optimal prefix-tree: you take the list of symbols, you find the two with the least probability, and you form a tree node with them as children. That tree node then has a probability that is the sum of its two children, and then you repeat the algorithm with two fewer symbols, but now with a tree node in the mix. Obviously each step of this algorithm reduces the size of the set of elements by one, so it terminates, and it also produces an optimal tree. Huffman trees are very fast because you can build a table indexed by the next n bits (where n is the length of the longest code). The table entry tells you what symbol you've decoded and how many bits to unread. The drawback of Huffman trees is that they can only use a whole number of bits for each symbol: if you have a symbol where -log₂(p) = 2.3 then ideally you want to use 2.3 bits to encode it. But Huffman forces you either to round up to 3 bits or to round down, which will force some other symbols to consume more bits.

Zstandard uses Huffman trees, but it also has a higher-compression entropy encoder called FSE. FSE is a state machine. There are more states than symbols, and each symbol gets a fraction of the states that mirrors its probability of occurrence in the stream. So if there's some symbol that is expected to appear 50% of the time, it gets ~50% of the states. Each state has three values: the symbol for that state, a number of bits to read from the bitstream when in that state, and a baseline state number that is added to those bits to get the next state. Now, if you recall, the problem with Huffman trees was that they could only use a whole number of bits, and these states also read a whole number of bits. But the trick is that if you are aiming to read one and a half bits for a given symbol, then half of its states will read one bit and half of them will read two bits. Then you hit your target on average. The table of states is never transmitted. The RFC prescribes an algorithm for building the table from a list of symbol probabilities, and so only the probabilities need to be transmitted.

Let's do an example. Let's say we have four symbols and we're going to use 16 states. So we have to approximate the symbol probabilities in terms of 16ths. (If you want a more accurate approximation of the probabilities, you can use a larger number of states; zstd actually never uses fewer than 32 states.)

state	0	1	2	3	4	5	6	7	8	9	10	11	12	13	14	15
Symbol	A	A	B	D	A	B	C	A	B	C	A	B	C	A	A	B
Num_Bits	2	1	2	4	1	2	3	1	2	2	1	1	2	1	1	1
Baseline	12	0	4	0	2	8	8	4	12	0	6	0	4	8	10	2

Any symbol may follow any other symbol, and a symbol might only have a single state. So every symbol must be able to reach every state. Take a look at state three, which is the only state for symbol D. Because it's the only one, it has to read four bits, which is sufficient to encode any other state. But if you look at a symbol like B, its states only demand that you read one or two bits. However, the set of 16 possible next states is exactly partitioned between those states for symbol B. So, for any particular state, there is exactly one state for symbol B that can reach it.

FSE decoding table cells tiling the state space A 16-cell decoding table with symbol B occupying states 2, 5, 8, 11 and 15. Arrows connect each of B's cells to its destination range below: state 11 covers next states 0 to 1, state 15 covers 2 to 3, state 2 covers 4 to 7, state 5 covers 8 to 11, and state 8 covers 12 to 15. Together the five ranges cover all sixteen states. 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 A A B D A B C A B C A B C A A B 0–11 bit 2–31 bit 4–72 bits 8–112 bits 12–152 bits

Again, consider symbol B, which we said had a probability of 5/16. The ideal number of bits to encode that symbol is -log₂(5/16) = 1.68. There are three symbol B states that read two bits and two that read one bit. The states aren't used equally often and, weighted by how often they're used, the average comes out to almost exactly the right value for the quantised probabilities. If you want to capture the true symbol probabilities with more accuracy, use a bigger table.

The central trick is that, by giving multiple states to more common symbols, the encoder doesn't just pick a symbol: it also picks which of that symbol's states to land in, and that choice carries information forward to the next symbol. That's where the fractional bits of information go. But this entropy encoder is still just table based, and so it runs very quickly.

The wrinkle is that you can't work forwards. Assume that you want to encode C, D. Which C state do you start in? Well, D only has one state so it has to be the C state which can reach that one. If D had multiple states then you would need to worry about what came after D to know which of those you needed. FSE forces you to start at the end of the sequence and work backwards. (That's not too bad because you usually need to know the whole sequence in order to calculate the symbol probabilities anyway.) Furthermore, a Zstandard compressor thus encodes symbols back-to-front, but writes output incrementally, so the decompressor has to seek to the end of a block and read the bits backwards in order to straighten it out! That's getting into broader details of the format that I'm not going to cover; see Nigel's piece.

Basic entropy encoders do not care about inter-symbol probabilities. I.e. they can't use the fact that the letter Q is disproportionately followed by the letter U (in English). There has to be some other encoding that is exploiting those redundancies. In Zstandard, that's a traditional Lempel–Ziv structure where it encodes either literal bytes or back references to previously decoded data. So FSE is primarily used for efficiently encoding these back reference offsets and lengths.

Lean

Let's talk about Lean! Above I said that it's a dependently-typed language, and that is a concept better articulated in examples than in a complicated definition. So here's the type of a function that reads n bytes from a stream and, if it doesn't throw, returns a byte array that the type system knows is n bytes long.

def IO.FS.Stream.readExact (st : Stream) (n : Nat) :
    IO {ba : ByteArray // ba.size = n} := …

Here is a function that returns two numbers and a byte array such that the first number is prime, the sum of the two numbers is divisible by six, and the byte array is at least as long as the smaller of those two numbers.

def getResult :
    IO (Σ a b : Nat, { bytes : ByteArray //
      Nat.Prime a ∧
      6 ∣ a + b ∧
      Nat.min a b ≤ bytes.size }) := …

That is not a type that anyone will ever need. It's just demonstrating that you can go as wild with this as you want. Dependently-typed languages are sufficient to encode even very complicated mathematical structures, and Lean's dominant use at the moment is as a formal language for stating and proving mathematics. The recent book, The Proof in the Code, is a short, well-written articulation of the story of how Lean came to be. The author does completely butcher constructive mathematics for a few paragraphs but, other than that, I enjoyed it!

Lean is a purely functional language like Haskell, although it has a few properties that make it potentially a lot more convenient as a programming language. Firstly, Lean is strict, while Haskell is lazy. Strictness means that arguments to functions are evaluated before the call happens, whereas in Haskell the evaluation of arguments is deferred until the value is actually required. So in Haskell it's free to write expensive expressions and pass them into functions, because they'll only actually be computed if they end up being used. But it also means that computation can happen in very surprising places in the program. This is a contentious topic but, while I appreciate the elegance of laziness, boy, it can make the performance of programs hard to reason about.

Next, Lean has some nice helpings of sugar. Its monadic do notation contains for loops and return statements and break statements. If you want to program in an imperative style, you can do so pretty reasonably!

Lastly, Lean has an optimisation where it will make mutating updates to objects as long as their reference count is equal to one. So you can mutate an array in place as efficiently as in an imperative language, as long as you are careful not to have a reference to it someplace else. Unfortunately, Lean does not have any aspects of a linear type system that I'm aware of, so it does not help you in ensuring that there is only a single reference to a value. It's a bit of a sharp edge that a seemingly minor tweak to the code can completely crater its performance by holding on to a reference to a large array somewhere inconspicuous. But it does mean that if you are trying to optimise the performance of something, you have a lot more tools at your disposal.

Here's an example of some of this, from the zstd decoder I sketched:

while true do
    let some blockHeaderBytes ← input.readExactOrEof 3 | break
    let some blockHeader := BlockHeader.fromBytes blockHeaderBytes frameHeader
      | throw (.userError "invalid block header")
    let blockBytes ← input.readExact blockHeader.contentSize

    match hty : blockHeader.type with
    | .rle =>
      let b := blockBytes.val[0]'(by
        rw [blockBytes.property, blockHeader.contentSize_rle hty]; omega)

Focus on line 9. There's an array index there, which is exactly the sort of place that implicit invariants live: blockBytes had better not be empty! C-like languages will give you undefined behaviour in that case. Modern languages will throw at run-time, or only give you an optional value to avoid that. Lean has another option: prove that it's not empty. That's what line 10 does. blockBytes.property is the fact that it's as long as the requested read, i.e. exactly blockHeader.contentSize bytes long. blockHeader.contentSize_rle is this:

theorem BlockHeader.contentSize_rle (h : BlockHeader) (hty : h.type = .rle) :
    h.contentSize = 1 := by
  simp [contentSize, hty]

That's a proof that, when the type is rle, the contentSize is always one. With those facts, Lean can figure out the rest.

It's a really short proof and probably I could have figured that out myself, but we can aim much higher:

I wrote an implementation of the FSE table construction algorithm from the RFC. The RFC contains “test vectors” for it: three sample outputs from given probabilities. Obviously those go into unit tests. But, in Lean, we can also prove universal properties of the function:

theorem ofDistribution_wellFormed (h : ofDistribution accuracyLog probs = some t) :
    t.entries.size = 2 ^ accuracyLog ∧
    (∀ s : Fin probs.size,
      t.entries.toList.countP (fun e => e.symbol == s.val) = probCells probs[s]) ∧
    (∀ (i : Nat) (hi : i < t.entries.size) (v : Nat), v < 2 ^ (t.entries[i]'hi).nbBits →
      (t.entries[i]'hi).baseline + v < 2 ^ accuracyLog) ∧
    (∀ (s : Fin probs.size), 0 < probCells probs[s] → ∀ x < 2 ^ accuracyLog,
      ∃! i : Nat, ∃ hi : i < t.entries.size,
        (t.entries[i]'hi).symbol = s.val ∧ (t.entries[i]'hi).baseline ≤ x ∧
          x < (t.entries[i]'hi).baseline + 2 ^ (t.entries[i]'hi).nbBits) := …

Repeating that, in words:

Assuming that the table construction function, when given the “accuracy” constant and a list of symbol probabilities, produces a value, then:

1. The table has the correct size for that accuracy.
2. The number of states for a given symbol is correct given its probability.
3. For all states, reading nbBits bits and adding the baseline value for that state produces a valid state number.
4. For all symbols with non-zero probability, and for all target states, there's exactly one state for that symbol which can reach the target state.

These are the subtle assumptions that an optimised decoding inner-loop requires, and things that can only ever be implicit or mere comments in weaker type systems. Proving strong statements like that is part of the 10× effort that the seL4 retrospective described, and a major barrier to the adoption of dependent types in regular software. Several LLMs can do it automatically now in about 20 minutes, and using only a fraction of a $20/month subscription quota. It'll probably be table-stakes next year. I must admit that they needed to change the table-generating code when doing so: I had used too much Id.run (i.e. dropping into imperative mode) and that's harder for the proof machinery to work with. (But Lean are working on it.) I confirmed that the proofs type-check and that there are no sorrys.

Combining dependent types and LLMs is not a new idea, but not much has been done on applying the combination to quotidian software engineering. Lots more experience would be needed. Very strong types can amplify the scope of changes as they have to be propagated out through all the derived types. Perhaps the proof effort scales poorly in larger systems, such that even modern LLMs can't keep up. Lean is a high-level language, and that's not suited to everything. (My toy Zstandard decoder is 10× slower than zstd on the command line.) Still, proof automation is here now and we, practically speaking, have a new type of programming language available to us. That's exciting!

(I'm not publishing the code because, frankly, for a small, well-defined case such as this, the LLMs can probably do a better job than I did. I did this to learn Lean a little and I don't hold my explorations up as an exemplar. This was inspired by lean-zip which does much more, includes a compressor, and proves round-tripping!)

Aside: verified assembly

AWS made LNSym: a semantics and simulator for AArch64. That's cool. Perhaps we could use it to show equivalence between an optimised assembly implementation of some functions, and their Lean counterparts, and then use the assembly code at run-time? Then we could let LLMs rip at optimisation and they couldn't introduce any functional bugs. Verified assembly is well-trodden in crypto implementations, but perhaps now it could be cheap?

I put some (mostly LLM) time into trying this. The small popcount example from the repo uses bv_decide, which is a certifying SAT solver, and that example requires more memory than my system has, which doesn't bode well. Tiny functions do work, and it is possible to get an equivalence proof to tiny Lean functions, and then to use extern to call them at run-time! But I, and a few LLMs, couldn't get it to scale any further.