The Dutch National Flag Problem, inductively

Jeremy suggested that I look into Wouter Swierstra's Agda solution to the Dutch National Flag problem, which Shin has also written about in his Chinese blog using Dijkstra's guarded command language. There are two things that I'd like to improve:

1. In Agda every program has to be terminating, and one (perhaps the most obvious) way to show that is to make a program structurally recursive. Wouter did it by introducing an explicit tree-shaped call structure on which the program recurses, which is a general technique (the so-called Bove-Capretta method) for expressing well-founded general recursion in a total language. For this problem, however, it is obvious that the difference between the two indices j and k is decreasing, so we should be able to just do structural recursion on the difference. It would be a counterexample against using a total language if we could not express that directly and had to invent a tree structure that does not seem necessary.
2. Wouter then wrote additional proofs to verify the program in the traditional index-based style. That is, he worked in the externalist way. I would like to see an internalist solution, i.e., the proof being integrated into the program.

It turned out that there is an arguably cleaner solution — a few functional programs that use exclusively pattern matching!

First we look at the problem of expressing the program as structurally recursive. Wouter used the finite numbers as indices into vectors (naturally), and he defined a Difference relation on finite numbers as

data Difference : forall {n} -> (i j : Index n) -> Set where
  Same : forall {n} -> (i : Index n) -> Difference i i
  Step : forall {n} -> (i j : Index n) -> 
    Difference i j -> Difference (inj i) (Next j)

The reason that the program cannot be made structurally recursive using this type is that the section of unknown colour can shrink either leftwards or rightwards, i.e., the pair of indices inj i and Next j can proceed to either inj i and inj j or Next i and Next j. For the latter case, we need a separate function to tweak the difference, but then Agda cannot see that the tweaked difference is smaller than the given one. We do know, however, that the "size" of the difference, i.e., the number of Steps, is smaller, so all we need to do is expose that size, i.e., the underlying natural number, as an index in the type, using (ornamental-) algebraic ornamentation:

data Difference : forall {n} -> (i j : Index n) -> (m : Nat) -> Set where
  Same : forall {n} -> (i : Index n) -> Difference i i Zero
  Step : forall {n} -> (i j : Index n) -> forall {m} ->
    Difference i j m -> Difference (inj i) (Next j) (Succ m)

Now the program can just do structural recursion on m. No matter how the difference is tweaked, the size (reflected in the index of its type) can always be shown to be smaller, so Agda accepts the program as structurally recursive and thus terminating. This datatype will also appear in my program below (with a different syntax).

The second problem is whether we can make the program manifest its own correctness (so separate proofs are not needed) by making the types more precise. The crucial property we wish to express is the invariant that the array is always divided into four sections such that pebbles in the first/second/fourth section are all red/white/blue while those in the third section can be arbitrarily coloured (in red, white, or blue). It turned out that this invariant can be formulated inductively and baked into the list type. Assume there is a three-element datatype for the colours,

data Colour : Set where
  red white blue : Colour

and the type of pebbles is a type family Peb : Colour → Set. The Dutch vectors are defined by

data DVec : (i j k : ℕ) → Set where
  []   : DVec 0 0 0
  _R∷_ : (x : Peb red) →     ∀ {i j k} (xs : DVec i j k) → DVec (suc i) (suc j) (suc k)
  _W∷_ : (x : Peb white) →   ∀ {  j k} (xs : DVec 0 j k) → DVec 0       (suc j) (suc k)
  _A∷_ : ∀ {c} (x : Peb c) → ∀ {    k} (xs : DVec 0 0 k) → DVec 0       0       (suc k)
  _B∷_ : (x : Peb blue) →              (xs : DVec 0 0 0) → DVec 0       0       0 

It is easy to see from the indices that once we use the arbitrary cons _A∷_ it is no longer possible to use the blue cons _B∷_ (since the index k in the type of a Dutch vector constructed by _A∷_ is nonzero but _B∷_ expects k in the type of its tail to be 0), and once we use the white cons _W∷_ it is no longer possible to use _A∷_ and _B∷_, etc. Thus these constructors necessarily appear in the desired order. Moreover, the indices in the type of a Dutch vector are exactly the values of the variables indexing into the corresponding array satisfying the invariant. And interestingly, we can simply use natural numbers instead of finite numbers as the indices; in the type of a successfully constructed Dutch vector, the indices are necessarily within bounds.

Next we need to expose two more pieces of information in the indices, both of which are algebraic ornamentations. The first one is the termination measure, i.e., the difference of k and j. We redefine the datatype Difference as

data _-_≈_ : ℕ → ℕ → ℕ → Set where
  zero : ∀ {j} → j - j ≈ zero
  suc  : ∀ {j k m} → k - j ≈ m → suc k - j ≈ suc m

The difference can then be computed by the fold,

diff : ∀ {i j k} → DVec i j k → Σ ℕ (λ m → k - j ≈ m)
diff []        = _ , zero
diff (x R∷ xs) = _ , adjustLeft (suc (proj₂ (diff xs)))
diff (x W∷ xs) = _ , adjustLeft (suc (proj₂ (diff xs)))
diff (x A∷ xs) = _ , suc (proj₂ (diff xs))
diff (x B∷ xs) = _ , zero

where the function adjustLeft is defined by

adjustLeft : ∀ {j k m} (d : k - j ≈ suc m) → k - suc j ≈ m
adjustLeft {m = zero } (suc zero) = zero
adjustLeft {m = suc _} (suc d) = suc (adjustLeft d)

One can see that the definition is a bit tricky as it is structurally recursive on m instead of d. This will turn out to be useful later. After the algebraic ornamentation, the signature of the DVec type becomes

DVec : (i j k : ℕ) → ∀ {m} (d : k - j ≈ m) → Set

Another information we need is the colour of the first pebble in the arbitrarily coloured section, as it determines what kind of swap we will do. There may or may not be such a colour, depending on the size of the arbitrarily coloured section, which is m. So we define

MaybeColour : ℕ → Set
MaybeColour zero = ⊤
MaybeColour (suc _) = Colour

and perform algebraic ornamentation using the fold

firstArbitraryColour : ∀ {i j k m} {d : k - j ≈ m} → DVec i j k d → MaybeColour m
firstArbitraryColour []              = tt
firstArbitraryColour (x R∷ xs)       = firstArbitraryColour xs
firstArbitraryColour (x W∷ xs)       = firstArbitraryColour xs
firstArbitraryColour (_A∷_ {c} x xs) = c
firstArbitraryColour (x B∷ xs)       = tt

So the signature of the type of the Dutch vectors becomes

DVec : (i j k : ℕ) → ∀ {m} (d : k - j ≈ m) → MaybeColour m → Set

Now we are ready to write the swaps. First we consider the case where the first arbitrary coloured pebble is actually red, i.e., we need to complete the program

reduceRed :
  ∀ {i j k m} {d : k - j ≈ suc m} →
  (xs : DVec i j k d red) → DVec (suc i) (suc j) k (adjustLeft d) (nextColour xs)
reduceRed xs = ?

where nextColour computes the first colour of the arbitrarily coloured section after it is reduced. Note that in the type of reduceRed, the size of d is suc m, so we are allowed to specify that the first colour of xs is red. Asking Agda to perform case analysis, we get

reduceRed :
  ∀ {i j k m} {d : k - j ≈ suc m} →
  (xs : DVec i j k d red) → DVec (suc i) (suc j) k (adjustLeft d) (nextColour v)
reduceRed (x R∷ xs) = ?
reduceRed (x W∷ xs) = ?
reduceRed (x A∷ xs) = ?

In the red cons case, we should go past x and reduce the rest, so the goal is solved by x R∷ reduceRed xs. Here the typing works directly because of how we defined adjustLeft: Matching the input vector with red cons unifies d with adjustLeft (suc d') for some d' (of size suc m) appearing in the type of xs, so the difference in the goal type becomes adjustLeft (adjustLeft (suc d')), while in the type of x R∷ reduceRed xs it is adjustLeft (suc (adjustLeft d')), so we need the emphasised subterm in the goal type to compute further. We could have done case analysis on d' but that messes up the program; instructing adjustLeft to look at the size instead of the difference itself, however, directly makes the emphasised subterm compute.

reduceRed :
  ∀ {i j k m} {d : k - j ≈ suc m} →
  (xs : DVec i j k d red) → DVec (suc i) (suc j) k (adjustLeft d) (nextColour v)
reduceRed (x R∷ xs) = x R∷ reduceRed xs
reduceRed (x W∷ xs) = ?
reduceRed (x A∷ xs) = ?

In the white cons case, we need to swap x with the first arbitrarily coloured pebble, which we assume to be red, and change the constructor to a red cons, extending the red section. We thus need to define a function to lookup the pebble

firstPeb : ∀ {j k m} {d : k - j ≈ suc m} → DVec zero j k d red → Peb red
firstPeb (x W∷ xs) = firstPeb xs
firstPeb (x A∷ xs) = x

and a function that substitutes the white pebble x for that pebble in xs.

substWhite :
  ∀ {n j k m} {d : k - j ≈ suc m} →
  Peb white → (xs : DVec n zero j k d red) →
  DVec n zero (suc j) k (adjustLeft d) (nextColour xs)
substWhite x (y W∷ xs) = y W∷ substWhite x xs
substWhite x (y A∷ xs) = x W∷ xs

The case is then solved by firstPeb xs R∷ substWhite x xs.

reduceRed :
  ∀ {i j k m} {d : k - j ≈ suc m} →
  (xs : DVec i j k d red) → DVec (suc i) (suc j) k (adjustLeft d) (nextColour v)
reduceRed (x R∷ xs) = x R∷ reduceRed xs
reduceRed (x W∷ xs) = firstPeb xs R∷ substWhite x xs
reduceRed (x A∷ xs) = ?

For the last goal we simply replace the arbitrary cons with a red cons.

reduceRed :
  ∀ {i j k m} {d : k - j ≈ suc m} →
  (xs : DVec i j k d red) → DVec (suc i) (suc j) k (adjustLeft d) (nextColour v)
reduceRed (x R∷ xs) = x R∷ reduceRed xs
reduceRed (x W∷ xs) = firstPeb xs R∷ substWhite x xs
reduceRed (x A∷ xs) = x R∷ xs

Similarly, reduceWhite and reduceBlue can be defined by just pattern matching, although reduceBlue is slightly more complex. The three functions can be assembled into a tail-recursive function

reduce :
  ∀ {i j k m} {d : k - j ≈ m} {c} →
  DVec i j k d c → Σ[ i' ∶ ℕ ] Σ[ j' ∶ ℕ ] DVec i' j' j' zero tt
reduce {m = zero} {d = zero} v = _ , _ , v
reduce {m = suc _} {c = red  } v = reduce (reduceRed   v)
reduce {m = suc _} {c = white} v = reduce (reduceWhite v)
reduce {m = suc _} {c = blue } v = reduce (reduceBlue  v)

which is structurally recursive on m. We can package reduce so it works on ordinary lists:

fuel : ∀ {k} →  k - zero ≈ k
fuel {zero } = zero
fuel {suc n} = suc (fuel {n})

initialise :
  (xs : List (Σ Colour Peb)) →
  let n = length xs in Σ[ c ∶ MaybeColour n ] DVec zero zero n fuel c
initialise [] = tt , []
initialise ((c , x) ∷ xs) = c , x A∷ proj₂ (initialise xs)

forget : ∀ {i j k m} {d : k - j ≈ m} {c} → DVec i j k d c → List (Σ Colour Peb)
forget [] = []
forget (x R∷ xs) = (red   , x) ∷ forget xs
forget (x W∷ xs) = (white , x) ∷ forget xs
forget (x A∷ xs) = (_     , x) ∷ forget xs
forget (x B∷ xs) = (blue  , x) ∷ forget xs

dutchFlag : List (Σ Colour Peb) → List (Σ Colour Peb)
dutchFlag = forget ∘ (proj₂ ∘ proj₂ ∘ reduce) ∘ (proj₂ ∘ initialise)

The whole development is available here. I think the novelty of the approach is that the invariant is formulated inductively and integrated into the list type, so the proofs that certain swaps preserve the invariant can be written inductively in the form of simple list-manipulating functional programs. Also it demonstrates how algebraic ornamentation can be used to tidy up dependently typed programs.

--
Still five more blog posts in the queue, argh..