In the previous part we ended with a nifty way of computing the last bit of a mask using the prefix sum. The key of the approch is based on the fact that the last element of the prefix sum of a mask is either the population count or population count minus one.

Let’s prove that using Lean.

What we want to prove

Given the following two definitions, we want to prove the presented Theorem.

This is pretty obvious after thinking about it for a bit but we will prove it anyways.

Lean

Lean is a functional programming language built on top of a powerful formalism called calculus of constructions and supporting a thing called inductive types (types built on top of other types in a systematic way). But the most interesting part of Lean is its capability as a proof assistant.

There is ample documentation of the Lean language covering the different users and use cases.

Representation

While Lean supports Array and Vector types, it will make everything a bit simpler if we prove it using plain old List types. Functional languages have excellent support for list-like types due to its inductive way of constructing them (a list is either the empty list, or the resulting list of prepending an element in front of an existing list). In fact, in Lean, sees Array (proof-wise) as a wrapper to a List (where the size is cached in the data type itself, rather than being recomputed as in the length of a List) and a Vector is a wrapper of an Array where the size is a parameter of the type making it a dependent type. Vector being a dependent type means that some proofs get a bit hairy specially if we intend to use induction (that we will because makes things easier).

It makes sense to represent the 1 and 0 above by the true and false terms of type Bool. For this reason we will define an alias Mask to be a List Bool.

abbrev Mask := List Bool

Population count

Now we can define popc using the List.count function that basically counts the number of ocurrences that are true in a given mask.

def popc (m: Mask) : Nat := List.count true m

We can already define some very simple theorems as a warm up for popc. They are very simple so the rfl tactic is enough.

theorem popc_of_empty_is_zero : popc [] = 0 := by rfl

theorem popc_of_singleton_false : popc [false] = 0 := by rfl
theorem popc_of_singleton_true : popc [true] = 1 := by rfl

We can generalise the last two theorems into a single one.

theorem popc_of_singleton : popc [b] = if b then 1 else 0 := by
  cases b with
  | true => rw[popc_of_singleton_true]; rfl
  | false => rw[popc_of_singleton_false]; rfl

Now we will prove that popc of a list that has been prepended an element is the popcount of that list plus whether that element is true or false. Before that we need a small helper

theorem count_true_is_boolean (x: Bool) : List.count true [x] = if x then 1 else 0 := by grind

In some of our proofs we will slightly (ab)use the grind tactic. It feels a bit like cheating but my colleagues with more knowledge than me on these topics tell me that it is fine.

Now we can prove that popc of prepending a mask with another element is basically the popc of that list and adding one if the new element prepended is true.

theorem popc_of_cons : popc (a :: b) = (if a then 1 else 0) + popc b := by
  induction b with
  | nil => rw [popc_of_empty_is_zero]
           rw [popc_of_singleton]
           simp
  | cons x b1 ih => unfold popc
                    have count_associates_1 : List.count true (a :: x :: b1) = List.count true [a] + List.count true [x] + List.count true b1 := by grind
                    have count_associates_2 : List.count true (x :: b1) = List.count true [x] + List.count true b1  := by grind
                    rw [count_associates_1]
                    rw [count_associates_2]
                    rw [count_true_is_boolean x]
                    rw [count_true_is_boolean a]
                    ring

Using a List is convenient because we can prove using induction. The base case is prepending to an empty list. We will be using rw a lot as we are basically proving equalities. Then we prove it for the inductive case and this needs to rewrite the inductive case until we get something simple enough.

Finally we want to prove that popc distributes over kind concatenation (I called it associates in the proof but I don’t think this is the right terminology).

theorem popc_associates : popc (a ++ b) = popc a + popc b := by
 induction a with
  | nil => rw [popc_of_empty_is_zero]; simp
  | cons x a1 ih => unfold popc
                    have count_associates_1: List.count true (x :: a1 ++ b) = List.count true [x] + List.count true a1 + List.count true b := by grind
                    have count_associates_2: List.count true (x :: a1) = List.count true [x] + List.count true a1 := by grind
                    rw [count_associates_1]
                    rw [count_associates_2]

I keep abusing the grind tactic for the subtheorems introduced using have.

Prefix sum

Now we can define the prefix sum.

def prefixsum (m: Mask) : List Nat := List.mapIdx (fun idx _ => popc (List.take idx m)) m

Note that this definition puts the “least significant bits” at the beginning of the list. This holds:

#eval prefixsum [true, false, true, true, false]
                [0,    1,     1,    2,    3    ]
#eval prefixsum [false, false, true, true, false]
                [0,     0,     0,    1,    2    ]

Before we can prove our theorem, we will need some helpers: non-empty lists have a last element, the emptiness of a list is the same as the emptiness of that list after we have mapped and how to obtain that last element we have appended to a list.

theorem non_empty_lists_have_a_last (l : List α) (h : l ≠ []) : ∃ (l': List α) (last: α), l = l' ++ [last] := by
  match l with
  | [x] => use []; use x; rfl
  | x :: y :: xs => obtain ⟨hl, hlast, p⟩ := non_empty_lists_have_a_last (y::xs) (List.cons_ne_nil y xs)
                    use (x :: hl); use hlast; rw [p]; simp

theorem map_ne_nil_iff: List.map f l ≠ [] ↔ l ≠ [] := by
  simp

theorem get_last_of_concat_to_singleton (h : l ++ [x] ≠ []) : List.getLast (l ++ [x]) h = x := by simp

Our proof gets a bit annoying because we want to use List.getLast but it needs a proof that the list is non empty. We didn’t do explicitly in our proof but it should be possible to deduce that m ≠ [] implies prefixsum m ≠ []. We simply passed the two proofs as a parameter. It should be possible to have a theorem with only the first assumption and prove the second one before rewriting with our presented theorem.

Part of the tedium of this proof is making sure we can do the rewrites. Every operand of List.getLast needs a proof that the list is not empty. But if we substitue the operand then the proof may not apply and Lean will diagnose a type mismatch. So we need to slowly replace everything. This leads to some gnarly equivalences in the successive eq1, eq2, etc. A lot of the subtheorems (h3, h4, …) are just there to be able to carry the proof of the non-emptyness prior to the rewrite.

Finally once we have managed to replace everything, we can pull the last item of the List.getLast using get_last_of_concat_to_singleton. Witht this the proof is a relatively simple one where we prove by cases of the last element.

theorem last_element_of_prefixsum_is_popc_or_popc_minus_one (m: Mask) (h : m ≠ []) (h2 : prefixsum m ≠ []) : 
   List.getLast (prefixsum m) h2 = popc m ∨  List.getLast (prefixsum m) h2 = (popc m) - 1 := by

   obtain ⟨prefix_bools, last_bool, lh⟩ := non_empty_lists_have_a_last m h
   have h3 : prefix_bools ++ [last_bool] ≠ [] :=
     by simp

   have h2 : prefixsum m ≠ [] := by
      unfold prefixsum
      apply Iff.mpr List.mapIdx_ne_nil_iff h
   have h4 : prefixsum (prefix_bools ++ [last_bool]) ≠ [] := by
     unfold prefixsum
     apply Iff.mpr List.mapIdx_ne_nil_iff h3

   let last2 := List.getLast (prefixsum (prefix_bools ++ [last_bool])) h4

   have eq1 : List.getLast (prefixsum m) h2 = List.getLast (prefixsum (prefix_bools ++ [last_bool])) h4 := by
      unfold prefixsum
      grind

   have h5 : [last_bool] ≠ [] := by simp
   have h6 : prefixsum [last_bool] ≠ [] := by
     unfold prefixsum
     apply Iff.mpr List.mapIdx_ne_nil_iff h5

   rw [eq1]

   have h7 : (prefixsum prefix_bools) ++ List.map (fun x => x + popc prefix_bools) (prefixsum [last_bool]) ≠ [] := by
    apply List.append_ne_nil_of_right_ne_nil
    apply Iff.mpr map_ne_nil_iff h6 

   have eq2 : List.getLast (prefixsum (prefix_bools ++ [last_bool])) h4 = 
              List.getLast ((prefixsum prefix_bools) ++ List.map (fun x => x + popc prefix_bools) (prefixsum [last_bool])) h7 := by 
              unfold prefixsum
              aesop

   rw [eq2]

   have h8 : (prefixsum prefix_bools) ++ [popc prefix_bools] ≠ [] := by
    apply List.append_ne_nil_of_right_ne_nil
    apply Iff.mpr map_ne_nil_iff h6 

   have eq3: List.getLast ((prefixsum prefix_bools) ++ List.map (fun x => x + popc prefix_bools) (prefixsum [last_bool])) h7 =
            List.getLast ((prefixsum prefix_bools) ++ [popc prefix_bools]) h8 := by
            unfold prefixsum
            aesop

   rw [eq3]

   rw [get_last_of_concat_to_singleton]

   -- I am sure this can be simplified
   match last_bool with
   | true => right
             rw [lh]
             rw [popc_associates]
             rw [popc_of_singleton]
             simp
   | false => left
              rw [lh]
              rw [popc_associates]
              rw [popc_of_singleton]
              simp

Again I have abused of grind and aesop in the subtheorems.

I’m clearly not even a rookie proving theorems but it was a fun exercise. I stay convinced of two things: 1) my proofs can be improved a lot 2) being able to prove things with the help of a proof assistant such as Lean is very useful to gain confidence in our algorithms.

It would have been cool to prove that vlast.m as implemented using the prefix sum in Part 1 computes indeed the index of the last non-zero element. I’m happy enough to have proved the key insight that allows us to use prefix sum to do that. I’ll leave that as an exercise for the reader 😉.