Leaning into the Coding Interview 3: completing our spec with tacticals and metaprogramming

Welcome back!

Last time, we got our dependently-typed fizzbuzz implementation to a point where we felt ready to prove our problem specifications:

import Mathlib.Data.Nat.Basic

-- implementation

inductive FB : Type where
  | Fizz
  | Buzz
  | FizzBuzz
  | Num (i : Nat)

instance : ToString FB where
  toString fb := match fb with
    | .Fizz => "Fizz"
    | .Buzz => "Buzz"
    | .FizzBuzz => "Fizzbuzz"
    | .Num i => toString i

def fb_one (i : Nat) :=
    if i % 15 = 0 then FB.FizzBuzz else
    if i % 5 = 0 then FB.Buzz else
    if i % 3 = 0 then FB.Fizz else
    FB.Num i

def fb_vec (n : Nat) : Vector FB n :=
  Vector.range' 1 n |> Vector.map fb_one

-- specification

theorem fb_one_to_fb_vec :
    ∀ (i n : Nat), (h : i < n) → (fb_vec n)[i]'h = fb_one (i + 1) := by
  intros i n h
  unfold fb_vec; simp
  rw [Nat.add_comm]

theorem thm1 : ∀ (i n : Nat) (H : i < n), 
    (i + 1) % 3 = 0 → (i + 1) % 5 ≠ 0 → 
    (fb_vec n)[i]'H = FB.Fizz := by sorry

theorem thm2 : ∀ (i n : Nat) (H : i < n), 
    (i + 1) % 3 ≠ 0 → (i + 1) % 5 = 0 → 
    (fb_vec n)[i]'H = FB.Buzz := by sorry

theorem thm3 : ∀ (i n : Nat) (H : i < n), 
    (i + 1) % 3 = 0 → (i + 1) % 5 = 0 → 
    (fb_vec n)[i]'H = FB.FizzBuzz := by sorry

theorem thm4 : ∀ (i n : Nat) (H : i < n), 
    (i + 1) % 3 ≠ 0 → (i + 1) % 5 ≠ 0 → 
    (fb_vec n)[i]'H = FB.Num (i + 1) := by sorry

At last, a real specification for Fizzbuzz

Let’s try writing one a proof of one of our actual Fizzbuzz specifications - following the design of our “program”, we’ll first prove that the call to the helper function fb_one i, when i is appropriately-chosen, is Fizzbuzz. We will then extrapolate that to vb_vec[i] under the same constraint for i using the fb_one_to_fb_vec theorem we proved last time.

theorem fb_one_i_mod_15_is_fizzbuzz:
  ∀ (i : Nat), i % 3 = 0 → i % 5 = 0 → fb_one i = FB.FizzBuzz := by

Intro and unfold, like usual

As usual, our first tactic involves introducing statements into our context and unfolding some basic definitions. You’ve seen this a bunch before so we’ll move quickly at the beginning.

 theorem mod_15_is_fizzbuzz : 
   ∀ (i : Nat), i % 3 = 0 → i % 5 = 0 → fb_one i = FB.FizzBuzz := by
+  intros i H3 H5
+  unfold fb_one

In which cases does the left-hand side of the goal equality match the right? There’s clearly one situation that makes sense here: the left-hand side only matches the right-hand side when i % 15 = 0.

More subtly, though, it also needs to be the case that if i % 15 ≠ 0 (that is, if we were to fall through the first if...then), it better not be the case that i % 3 = 0 and i % 5 = 0 also hold. (Given that those are the two hypotheses H3 and H5 that we’ve introduced, this is kind of like saying “any other case than i % 15 = 0 is a contradiction and impossible”. We’ll have a lot to say about contradictions shortly.)

When we have control flow like an if...then or match expression that we want to handle the different cases for, using the split tactic will decompose the goal into the different paths that the program could go

 theorem mod_15_is_fizzbuzz : 
   ∀ (i : Nat), i % 3 = 0 → i % 5 = 0 → fb_one i = FB.FizzBuzz := by
   intros i H3 H5
   unfold fb_one
+  split

You might expect to see four goals here because we have four branches in fb_one’s nested-if. It turns out that this tactic will automatically simplify down and eliminate branches that it knows are logically impossible given the assumptions H3 and H5. You can play around with handing the full if/else ladder manually by generalizing the proof: don’t introduce H3 and H5 right off the bat but wait until after you’ve started splitting each if; you’ll see goals with both, for instance, i % 5 = 0 and i % 5 ≠ 0; these are the ones that have been pruned away for us automatically here.

Focus in on subgoals with the · dot operator

This is, I think, our first example of a structured proof, where a tactic has broken our goal down into multiple subgoals for us to prove individually.

We can focus on the first goal with the · “bullet” operator - this will take our context, which currently has two goals to prove, and “zooms in” on the first one in the list. (This should remind you of proving a lemma that you’ve defined inline with have.)

 theorem mod_15_is_fizzbuzz : 
   ∀ (i : Nat), i % 3 = 0 → i % 5 = 0 → fb_one i = FB.FizzBuzz := by
   intros i H3 H5
   unfold fb_one
   split
+  ·

This is the “then” side of the if expression: when the conditional i % 15 = 0 is true, then fb_one produces FB.FizzBuzz, which is exactly the expression our theorem states it should. A rfl or simp is enough here.

(Note that we didn’t actually have to use anyhting in our context to prove this goal! The solution was right there in front of us, the whole time.)

Proving false means hunting for a contradiction

The second subgoal, the “else” side of the conditional, is more complicated to solve, and introduces a really interesting concept about what falsity means:

 theorem mod_15_is_fizzbuzz : 
   ∀ (i : Nat), i % 3 = 0 → i % 5 = 0 → fb_one i = FB.FizzBuzz := by
   intros i H3 H5 
   unfold fb_one
   split
-  ·
+  · rfl 
+  ·

Our goal, FB.Buzz = FB.FizzBuzz, should feel problematic to you: an axiom of inductive datatypes like FB is that they’re disjoint - every element of a given type can only be constructed in one way, so there’s no way for two different constructors of FB to produce “an equivalent” value. We should believe that the statement in the goal is false, and by simplifying our goal, we can see that Lean agrees with us.

 theorem mod_15_is_fizzbuzz : 
   ∀ (i : Nat), i % 3 = 0 → i % 5 = 0 → fb_one i = FB.FizzBuzz := by
   intros i H3 H5
   unfold fb_one
   split
   · rfl 
-  ·
+  · simp

Ooooook.

Our goal is to prove that the proposition False is true. That seems… actually impossible to do. This whole time we’ve been crunching the goal down until it’s something that Lean can reduce to True, after all!

There’s an escape hatch of sorts here: by the logical principle known, spectacularly, as both the principle of explosion and ex falso quodlibet. It states that any goal, no matter how absurd it seems, can be proven if you assume a contradiction.

Less philosophically-speaking: the purpose of this goal actually is to show that there’s a contradition somewhere in our context. It’s hard not to see where the contradiction is: we have both i % 3 = 0 and i % 5 = 0, but also ¬i % 15 = 0. If we can massage H3 and H5 into i % 15 = 0, then we’ve got our two contradictory statements.

Well, let’s do that, then: let’s state that, actually, i % 15 = 0:

 theorem mod_15_is_fizzbuzz : 
   ∀ (i : Nat), i % 3 = 0 → i % 5 = 0 → fb_one i = FB.FizzBuzz := by
   intros i H3 H5 
   unfold fb_one
   split
   · rfl 
   · simp
+    have H15: i % 15 = 0 := by

Notice that the goal has changed, because we’ve stated something with have that we are now going to prove.

I’ll admit that it’s been long enough since my undergrad math classes that I wasn’t entirely sure how to actually prove this! I knew it had something to do with 3 and 5 being coprime, and so probably falls out of the Fundamental Theorem of Arithmetic or something. I struggled for quite a while until I wondered whether, like Rocq, if Lean happens to have the lia tactic. Turns out it does!

lia is another automated tactic that comes with Mathlib, that specifically solves goals involving linear arithmetic, which is exactly what we’ve got on our hands here. The docs say it’s “inspired by modern SMT solvers” and uses cool research like e-matching. Automated tactics like lia bridge the gap between writing Lean proofs entirely by hand as we’ve been doing here, and handing off the entire proof burden to an automated solver, like we did with Dafny.

All right, with our lemma proven, the trap’s set:

 theorem mod_15_is_fizzbuzz : 
   ∀ (i : Nat), i % 3 = 0 → i % 5 = 0 → fb_one i = FB.FizzBuzz := by
   intros i H3 H5 
   unfold fb_one
   split
   · rfl 
   · simp
-    have H15: i % 15 = 0 := by
+    have H15 : i % 15 = 0 := by lia

Our contradiction is clear: certainly hypotheses H_15 and H+ can’t both be true, since one’s the negation of the other. So, we have two statements that contradict in our context! The only thing left to us is to close out the subgoal by marking it as a proof by contradiction:

 theorem mod_15_is_fizzbuzz : 
   ∀ (i : Nat), i % 3 = 0 → i % 5 = 0 → fb_one i = FB.FizzBuzz := by
   intros i H3 H5 
   unfold fb_one
   split
   · rfl 
   · simp 
     have H15 : i % 15 = 0 := by lia
+    contradiction

Simplifying our proofs with tactics combinators

Before proceeding, you should try and write the equivalent theorems for the other FB constructors.

This is what I got:

theorem mod_15_is_fizzbuzz :
  ∀ (i : Nat), i % 3 = 0 → i % 5 = 0 → fb_one i = FB.FizzBuzz := by
  intros i H3 H5
  unfold fb_one
  split 
  · rfl
  · simp
    have H15 : i % 15 = 0 := by lia
    contradiction


theorem mod_5_is_buzz : 
  ∀ (i : Nat), i % 3 ≠ 0 → i % 5 = 0 → fb_one i = FB.Buzz := by
  intros i H3 H5
  unfold fb_one 
  split 
  · simp
    have H3' : i % 3 = 0 := by lia
    contradiction
  · rfl 
  
theorem mod_3_is_fizz : 
  ∀ (i : Nat), i % 3 = 0 → i % 5 ≠ 0 → fb_one i = FB.Fizz := by
  intros i H3 H5
  unfold fb_one 
  split 
  · simp
    have H5' : i % 5 = 0 := by lia
    contradiction
  · rfl 

theorem else_is_num : 
  ∀ (i : Nat), i % 3 ≠ 0 → i % 5 ≠ 0 → fb_one i = FB.Num i := by
  intros i H3 H5
  unfold fb_one 
  split 
  · simp
    have H5' : i % 5 = 0 := by lia
    contradiction
  · simp

Just from a quick visual inspection you can see how similar the shape of these proofs look: Up until the split they’re identical, and the two subgoals involve a proof by contradiction and a straightforward “simplify and refl”. What differs between the proofs is which subgoal is trivally-solvable and which requires a proof by contradiction, and which requires introducing an additional hypothesis to direct the contradiction.

We can structure these four proofs around their commonalities by using tactics combinators (also called tacticals). Just like how a function combinator in functional programming is a higher-order function that combines, manipulates, or produces functions, a tactical lets us compose multiple tactics into new ones.

all_goals, any_goals, and <;> apply a tactic to all the open subgoals

Here’s a useful observation: For the contradictory subgoal, we always start by applying simp, and while we use rfl for the trivial subgoal, simp is smart enough to complete a proof that we have been using rfl for. This means that if we could tell Lean “run split, and then simp all the subgoals”, our proof becomes a lot shorter! We wouldn’t even have to manually prove the trivial subgoal, and we’d be able to skip manually simping the contradictory goal.

The all_goals tactical can do this for us: it takes as argument the name of a tactic to apply on all open subgoals. Let’s try this:

theorem mod_15_is_fizzbuzz :
  ∀ (i : Nat), i % 3 = 0 → i % 5 = 0 → fb_one i = FB.FizzBuzz := by
  intros i H3 H5
  unfold fb_one
  split
  all_goals simp

As we expected, the trivial goal has been eliminated completely, and we see the result of the simplification on the remaining contradictory goal. (This tactic is called all_goals because it will fail if, say, simp failed to apply in any of the goaisl. There’s also an any_goals t that applies t to all subgoals, but only fails if all subgoals rejected t.)

It’s common enough in proofs to “create a bunch of subgoals, and then apply the same tactics to all the goals” that there’s a nice shorthand for this: t1 <;> t2 first performs t1, and then for however many subgoals exist afterwards, t2 is applied to those. You can think of it as similar to t1 ; all_goals t2. (<;>'s a bit like xargs(1) in that sense.).

lia is magic, but opaque if you rely on it too much

This works for all our theorems here: the only thing that changes between them is which particular hypothesis we introduce to force the proof by contradiction:

But it turns out we’re being more explicit than we need to be: it’s actually sufficient to just apply lia after we unfold, avoiding all the case-splitting!

theorem mod_15_is_fizzbuzz :
  ∀ (i : Nat), i % 3 = 0 → i % 5 = 0 → fb_one i = FB.FizzBuzz := by
  intros i H3 H5
  unfold fb_one
  lia

theorem mod_5_is_buzz : 
  ∀ (i : Nat), i % 3 ≠ 0 → i % 5 = 0 → fb_one i = FB.Buzz := by
  intros i H3 H5
  unfold fb_one
  lia
  
theorem mod_3_is_fizz : 
  ∀ (i : Nat), i % 3 = 0 → i % 5 ≠ 0 → fb_one i = FB.Fizz := by
  intros i H3 H5
  unfold fb_one 
  lia 

theorem else_is_num : 
  ∀ (i : Nat), i % 3 ≠ 0 → i % 5 ≠ 0 → fb_one i = FB.Num i := by
  intros i H3 H5
  unfold fb_one 
  lia

Your turn: Connecting it back to fb_vec

OK, let’s see if we can quickly connect the proof we just wrote to one about fb_vec. Actually, you know what – by now you’ve seen enough proofs that I think you should give it a try before continuing!

Here’s the scaffholding to get you started:

theorem thm3 : 
  ∀ (i n : Nat) (H : i < n), 
    (i + 1) % 3 = 0 → (i + 1) % 5 = 0 → (fb_vec n)[i]'H = FB.FizzBuzz := by

This will involve both rewriting the goal based on fb_one_to_fb_vec and applying mod_15_is_fizzbuzz, the theorem we just proved.

(You might benefit from knowing about the exact tactic: exact p completes the proof if p is a proof of the goal. Don’t forget that a proof can take arguments.)

OK, here’s what I wrote:

theorem thm3 : 
  ∀ (i n : Nat) (H : i < n), 
    (i + 1) % 3 = 0 → (i + 1) % 5 = 0 → (fb_vec n)[i]'H = FB.FizzBuzz := by
    intros i n H H3 H5
    rw [fb_one_to_fb_vec]
    exact mod_15_is_fizzbuzz (i + 1) H3 H5

Completing the remaining three theorems is as simple as swapping out the appropriate final theorem to apply (e.g. mod_5_is_buzz for thm2, etc).

If you wanted, you could even go one step farther and just inline the definition of mod_15_is_fizzbuzz:

theorem thm3 : ∀ (i n : Nat) (H : i < n),
    (i + 1) % 3 = 0 → (i + 1) % 5 = 0 → (fb_vec n)[i]'H = FB.FizzBuzz := by
    intros i n H H3 H5
    rw [fb_one_to_fb_vec]
    unfold fb_one
    lia

At this point, lia handles all the particulars of the proof, so the body of all our four theorems would be identical! The amount of repetition here might now be driving you to distraction (though in many cases it’s okay to repeat yourself, actually), so you might argue that this is okay.

Short as it is, I actually like this proof a lot less, though, since it’s a lot more mysterious to the reader about why the proof is correct. Sure, it relies on some property of linear arithmetic, because we’re using lia, but beyond that, I don’t think I could explain this proof to another person; reading this proof without the context of us starting with the longer one feels like when you see a bit-twiddling hack - totally opaque as to how the programmer derived it! I likely could explain the longer one, though.

It all goes back to who your audience is! If you’re writing these proofs strictly for the typechecker, then fine, go with the shortest, most opaque proof; heck, at that point, run a code golf competition and see who on your team can find the shortest proof. If you want to use the specification as an implementation and debugging aid, though, maybe you want to optimise for something other than character count.

Adding our own definitions to simp’s proof search

That said, let’s see what we can do to reduce the amount of nontrivial boilerplate in our four main theorems. Those nontrivial things are twofold: we have to explicitly tell Lean that:

1. expanding the definition of fb_one simplifies our goal;
2. rewriting the equality in fb_one_to_fb_vec simplifies our goal.

The @[simp] annotation can be placed on definitions of functions and theorems; adding it means that definition will be added to Lean’s proof search database.

@[simp] -- NEW: `simp` will try to unfold this def'n when appropriate
def fb_one (i : Nat) :=
    if i % 15 = 0 then FB.FizzBuzz else
    if i % 5 = 0 then FB.Buzz else
    if i % 3 = 0 then FB.Fizz else
    FB.Num i

@[simp] --NEW: `simp` will attempt to apply this theorem when appropriate
theorem fb_one_to_fb_vec :
    ∀ (i n : Nat), (h : i < n) → (fb_vec n)[i]'h = fb_one (i + 1) := by
  intros i n h
  unfold fb_vec
  simp
  rw [Nat.add_comm]

(We say that fb_one_to_fb_vec, as well as the equality relating the function name fb_one to its body, have become “simp lemmas”.)

This really pays off in our four main specification theorem: the only thing they have to do now is introduce the relevant hypotheses, simplify the goal, and tidy up the linear arithmetic with lia!

theorem thm1 : ∀ (i n : Nat) (H : i < n),
    (i + 1) % 3 = 0 → (i + 1) % 5 ≠ 0 → (fb_vec n)[i]'H = FB.Fizz := by
    intros; 
    simp; -- NEW: No need to rewrite `fb_one_to_fb_vec` or unfold `fb_one`!
    lia

Could we have also, in order to simplify fb_one_to_fb_vec, tagged fb_vec as @[simp]? For sure. But, I like leaving the proofs that are more integral to the implementation explicit, and only relying on the mechanics of proof automation when there’s tonnes of repetition, or when being explicit doesn’t elucidate the thinking behind the proof.

Metaprogramming time: Extending Lean with custom tactics

One of my favourite languages, despite its dynamic typing, is Racket, not least because it’s so much fun, but also safe and well-defined, to extend the language’s syntax with the language itself. The mechanism Racket uses for letting programmers add new syntactic forms is called a syntax transformer: it takes in as input an AST and produces a new, transformed AST, before passing it along to to the rest of the interpreter. And, critically, the language in which one specifies language extensions is…the language itself!

Lean’s support for metaprogramming is similarly excellent too! Macros are used to define new tactics - we can have high-level syntax transformers that expand out into more tactics, or, at a lower-level, into the raw underlying monadic operations that implement tactics. The latter, which we call elaborator-based macros, is important if we really want to customise the behaviour of the tactic, like fiddling with the proof state to manipulate hypotheses. While it would be fun to write an elaborator-based macro for this, overkill as it would be, I don’t assume fluency with monads and do-notation, so let’s keep it simple today.

There’s an old canard that every problem in programming languages is a matter of syntax and of semantics. Defining a new tactic involves defining both to Lean.

Let’s create a new tactic called isl - this will simply run intros, followed by simp, concluding with lia. Since we only want to expand one tactic into three, let’s stick with the high level macro-based tactic definition for now. This means the only meaning that our new syntax has is purely in terms of existing syntax, which is fine for us.

Specifying new syntax with syntax

The syntax “command” lets us add isl to the tactic syntatic category thus:

-- This defines a new _syntax object_ - think of it as telling
-- the Lean parser that the `isl` token should be allowed wherever 
-- in the grammar a tactic is.
syntax (name := isl) "isl" : tactic

This is such a simple use of syntax that we are kind of doing it dirty here! It’s flexible enough to account for surrounding context like arguments, customised precedence levels, new operators (prefix, infix, or postfix!)… certainly anything I can think of.

Specifying new semantics with macro_rules

Okay, so we have defined isl syntactically. The Lean parser rejects us from using it in a context where a tactic isn’t expected, but we still can’t use it in a place where a tactic is expected:

theorem thm1 : ∀ (i n : Nat) (H : i < n),
    (i + 1) % 3 = 0 → (i + 1) % 5 ≠ 0 → (fb_vec n)[i]'H = FB.Fizz := by isl

Tactic `isl` has not been implemented

Sensible error message. Let’s go ahead and give isl meaning. macro_rules is a command that lets us specify syntactic patterns to match on, and to indicate how a matching pattern should be transformed. Our goal here is to transform isl into intros; simp; lia:

macro_rules
| `(tactic| isl) => `(tactic| intros; simp; lia)

If this reminds you of the syntax of a pattern match, I’m sure this is by design! With our macro defined, we can now change the proof of all our theorems into a single isl call.

theorem thm1 : ∀ (i n : Nat) (H : i < n),
    (i + 1) % 3 = 0 → (i + 1) % 5 ≠ 0 → (fb_vec n)[i]'H = FB.Fizz := by isl

Goals accomplished!

In a lot of ways this example undercuts how expressive metaprogramming in Lean can be; we really haven’t used a macro for anything fancier than what a #define could do for us. This is just scratching the surface; maybe in a future post we’ll find an opportunity to write a non-trivial elaborator-based macro.

Next time…

OK, here’s our implementation and specification:

import Mathlib.Data.Nat.Basic

inductive FB : Type where
  | Fizz
  | Buzz
  | FizzBuzz
  | Num (i : Nat)

instance : ToString FB where
  toString fb := match fb with
    | .Fizz => "Fizz"
    | .Buzz => "Buzz"
    | .FizzBuzz => "Fizzbuzz"
    | .Num i => toString i

@[simp] def fb_one (i : Nat) :=
  if i % 15 = 0 then FB.FizzBuzz else
  if i % 5 = 0 then FB.Buzz else
  if i % 3 = 0 then FB.Fizz else
  FB.Num i

def fb_vec (n : Nat) : Vector FB n :=
  Vector.range' 1 n |> Vector.map fb_one

-- specification

syntax (name := isl) "isl" : tactic

macro_rules
| `(tactic| isl) => `(tactic| intros; simp; lia)

@[simp] theorem fb_one_to_fb_vec :
    ∀ (i n : Nat), (h : i < n) → (fb_vec n)[i]'h = fb_one (i + 1) := by
  intros i n h
  unfold fb_vec
  simp
  rw [Nat.add_comm]
    
theorem thm1 : ∀ (i n : Nat) (H : i < n),
    (i + 1) % 3 = 0 → (i + 1) % 5 ≠ 0 → (fb_vec n)[i]'H = FB.Fizz := by isl

theorem thm2 : ∀ (i n : Nat) (H : i < n),
    (i + 1) % 3 ≠ 0 → (i + 1) % 5 = 0 → (fb_vec n)[i]'H = FB.Buzz := by isl

theorem thm3 : ∀ (i n : Nat) (H : i < n),
    (i + 1) % 3 = 0 → (i + 1) % 5 = 0 → (fb_vec n)[i]'H = FB.FizzBuzz := by isl

theorem thm4 : ∀ (i n : Nat) (H : i < n),
    (i + 1) % 3 ≠ 0 → (i + 1) % 5 ≠ 0 → (fb_vec n)[i]'H = FB.Num (i + 1) := by isl

This was actually a pretty productive session! We wrote our Fizzbuzz specification and our implementation successfully checks against it.

If you’ve been following along from the beginning, you might have said, “oh, I wouldn’t have implemented fb_one quite that way…” - maybe you would have chosen a match expression instead of an if-ladder, or maybe you would have done something altogether different. Next time, we’re going to have a brief diversion into other possible worlds and see how robust our specification is, and maybe even write a theorem to the effect of your_fb_one = my_fb_one…! See you then.