Thursday, April 07, 2022

Using Agda's Induction/Recursion Library

module FastExp where
# Imports
open import Data.Nat
open import Data.Nat.Properties
open import Data.Nat.Induction hiding (rec)
open import Data.Product using (_×_; _,_; Σ; Σ-syntax; ; ∃-syntax; proj₁; proj₂)
open import Data.Sum using (_⊎_; inj₁; inj₂)
open import Induction
open import Relation.Binary.PropositionalEquality

open import Exponents
open import Parity

I've been putting off learning how to use Agda's `Induction` library
for some time because I knew it would take a serious effort. (The
documentation is incomplete.) However, I ran into yet another
situation that calls for it, so I finally decided to dive in!

The purpose of the `Induction` library is to provide alternate forms
of induction and recursion, such as complete induction, and it helps
you build your own forms of induction. (Induction and recursion are
the same thing in Agda.) Recall that Agda provides built-in support
for structural recursion, but sometimes you want to define a function
that doesn't fit into that mold. For example, suppose I wanted to
write down the fast exponentiation function. That is, I want to define
a function `fast-exp` such that

    fast-exp n x ≡ x ^ n

Here's a naive attempt to write the function in Agda. I'm using an
auxilliary function `parity` that determines whether a number `n` is
even (`n ≡ 2 * k`) or odd (`n ≡ 1 + 2 * k`).

    fast-exp : ℕ → ℕ → ℕ
    fast-exp zero x = 1
    fast-exp (suc n) x
        with parity n
    ... | inj₁ (k , refl) =  x * fast-exp k (x * x)
    fast-exp (suc n) x
        | inj₂ (k , refl) =  x * x * (fast-exp k (x * x))

Agda's termination checker rejects this program because it can't tell
that the argument `k` in the recursive call to `fast-exp` is smaller
than the input parameter `suc n`. We'll use the `Induction` library to
work around this problem.

There are two layers to the `Induction` library, there's a lower-level
layer that resides in the `src/Induction.agda` file of the Agda
standard library and there is a higher-level layer that is specific to
induction on natural numbers in the `src/Data/Nat/Induction.agda`
file. We'll start by using the higher-level layer to define `fast-exp`
and then we'll take a look at the lower-level layer and build a custom
induction principle for natural numbers and use it to redo our
definition of `fast-exp`.

The `Nat.Induction` library provides support for complete induction
(aka. strong induction), which allows you to make a recursive call
with any natural number smaller than the current one. The main
ingredient that you need to define is a "step" function. This function
looks a lot like the recursive function that you're trying to define,
but it takes an extra parameter, let's name it `rec`, that will give
you access to the recursive call. Here's a first (erroneous) attempt
to define such a step function for `fast-exp`. Notice how the
recursive calls to `fast-exp` above have been replaced by calls to `rec`.

    fe-step zero rec x = 1
    fe-step (suc n) rec x
        with parity n
    ... | inj₁ (k , refl) =  x * rec k (x * x)
    fe-step (suc n) rec x
        | inj₂ (k , refl) =  x * x * (rec k (x * x))

## The `CRec` Type Operator

The above doesn't quite work because the `rec` given to us by the
`Nat.Induction` library is not a function, it is a big tuple with one
element for every natural number smaller than the current one. The
type of this tuple is given by `CRec` in the `Nat.Induction`
library. The `CRec` type operator takes three parameters: a universe
level parameter (ignore that for now), a function that produces the
type for each element in the tuple (given its index from the back),
and the size of the tuple.

For the purposes of defining `fast-exp`, we want each element in the
tuple to be a function, in particular, the fast exponentiation
function that's been partially applied to its first parameter. So the
type of each element should be `ℕ → ℕ`. We define the following
`FERec` abbreviation for the use of `CRec` that matches our needs.

FERec :   Set
FERec n = CRec _  i    ) n

The next thing we'll need is a way to access the nth element of the
tuple. The Agda standard library has a `projₙ` function for this
purpose, but to use it we'd need to prove that the `CRec` type
operator produces a `Product` type. Instead we'll roll our own `projₙ`
function for `CRec`. It takes a tuple of length `suc k` (so that it's
non-empty), an index `n`, and a proof that `n` is less than `suc k`.

projₙ : ∀{ P k}  CRec  P (suc k)  (n : )  n ≤′ k  P n
projₙ {l} rec n ≤′-refl = proj₁ rec
projₙ {l} rec n (≤′-step n≤k) = projₙ (proj₂ rec) n n≤k

## A Step Function for Fast Exponentiation

Next we define the step function for fast exponentiation.  The type
for the `rec` parameter is `FERec n`.  To make the recursive call, we
use `projₙ` to access the appropriate partially-applied version of
fast exponentiation (for a smaller natural number) from the `rec`
tuple. However, to do so, we have to prove that `k` is less than `n`,
the length of the tuple.  (More about this below.)

fe-step : (n : )  FERec n    
fe-step zero rec x = 1
fe-step (suc n′) rec x
    with parity n′
... | inj₁ (k , refl) =  x * projₙ rec k lt (x * x)
      where lt : k ≤′ 2 * k
            lt = ≤⇒≤′ (m≤m+n k _)
fe-step (suc n′) rec x
    | inj₂ (k , refl) =  x * (x * (projₙ rec k lt (x * x)))
      where lt : k ≤′ 1 + (2 * k)
            lt = ≤⇒≤′ (≤-step (m≤m+n k _))

Regarding `fe-step`, the case for `n ≡ zero` is straightforward.  In
the case for `n ≡ suc n′`, we have two subcases to consider, when `n′`
is even (`n′ ≡ 2 * k`) and when `n′` is odd (`n′ ≡ 1 + 2 * k`).
For the even subcase, to call `projₙ` we need to show that
`k ≤ 2 * k`, which we do with the theorem `m≤m+n` from `Nat.Properties`.
For the odd subcase, to call `projₙ` we need to show that
`k ≤ 1 + 2 * k`, which we do using `≤-step` and then `m≤m+n`.

## Use `cRec` to Define Fast Exponentiation

The final step to defining `fast-exp` is to apply the `cRec` function
from `Nat.Induction` to our step function, `fe-step`. Similar to the
`CRec` type operator, the `cRec` function also needs to know the type
of the elements in the `rec` tuple, which is `ℕ → ℕ`.

fast-exp :     
fast-exp = cRec  _  (  )) fe-step

## Proof that Fast Exponentiation is Correct

Of course, the whole point of programming in a proof assistant like
Agda is to prove the correctness of our programs. Let's prove that

    fast-exp n x ≡ x ^ n

which will give us an opportunity to 1) use `Nat.Induction` in an
inductive proof and 2) reason about a function that was defined using

We're going to refer to the correctness condition many times, so we
define the following abbreviation for it.

fe-ok :   Set
fe-ok n =  x  fast-exp n x  x ^ n

When reasoning about `fast-exp`, we'll need to reason about the tuple
that gets passed to the `rec` parameter of `fe-step`.  It turns out
that the `cRec` function builds that tuple using an auxilliary
function named `cRecBuilder`. So we define the following abbreviation
named `fe-rec` for applying the `cRecBuilder` function to our step

fe-rec : (n : )  CRec _  _  (  )) n
fe-rec n = cRecBuilder  _  (  )) fe-step n

The nth function in the tuple produced by `fe-rec` is the `fast-exp n`

projₙ-fe :  n k x (lt : n ≤′ k)  projₙ (fe-rec (suc k)) n lt x  fast-exp n x
projₙ-fe n n x ≤′-refl = refl
projₙ-fe n (suc k) x (≤′-step lt) = projₙ-fe n k x lt

(This proof goes through easily because we do induction on `n ≤′ k`,
which is using the alternate form of less-than.  If we had instead
used the normal less-than `≤` in the definition of `projₙ`, this proof
would be more difficult.)

We prove that `fast-exp` is correct using complete induction. Since
induction and recursion are the same thing in Agda, this means the
proof is a (dependetly typed) recursive function defined using `cRec`.

Recall that the first argument to `cRec` is the type for the elements
of the `rec` tuple. However, because we are now doing induction, we
should instead think of the `rec` tuple as the induction hypothesis.
In our step function for the proof we will use the parameter name `IH`
instead of `rec`.  Furthermore, because this tuple serves as the
induction hypothesis, its elements will need to be proofs that
`fast-exp` is correct for particular (smaller) exponents. So the type
of the element at position `n` should be the proposition `fe-ok n`.
Also, recall that the `CRec` type operator produces the type of
the tuple. So for our current purposes, `CRec _ fe-ok n` should be the
type of `IH`.

The second argument to `cRec` is a step function, which in this case
needs to construct a proof that `fast-exp n x ≡ x ^ n` for an
arbitrary `n`, given the induction hypothesis `IH`. We define the
`step` function in the `where` clause of our theorem below. The `step`
function mimics the structure of the `fe-step` function, doing case
analysis on the result of `parity n`.  We then do the appropriate
equational reasoning.  The one unusual step is the first one, which
uses the `projₙ-fe` lemma to replace the "raw" recursive call via
`projₙ` (as it appears in the body of `fe-step`) with a call to

fast-exp-is-correct :  n x  fast-exp n x  x ^ n
fast-exp-is-correct = cRec fe-ok step
  step : (n : )  CRec _ fe-ok n  fe-ok n
  step zero IH x = refl
  step (suc n) IH x
      with parity n
  ... | inj₁ (k , refl) =
          x * projₙ (fe-rec (1 + 2 * k)) k lt (x * x)   ≡⟨ cong  X  x * X) (projₙ-fe k (k + (k + zero)) (x * x) lt) 
          x * fast-exp k (x * x)                        ≡⟨ cong  X  x * X) (projₙ IH k (≤⇒≤′ (m≤m+n _ _)) (x * x)) 
          x * ((x * x) ^ k)                             ≡⟨ cong  X  x * X) (*-distribˡ-^ x x k) 
          x * (x ^ k * x ^ k)                           ≡⟨ cong  X  x * (x ^ k * x ^ X)) (sym (+-identityʳ k)) 
          x * (x ^ k * x ^ (k + 0))                     ≡⟨ cong  X  x * X) (sym (^-distribˡ-+-* x k (k + zero))) 
          x * (x ^ (2 * k))
        open ≡-Reasoning
        lt = ≤⇒≤′ (m≤m+n k (k + zero))
  ... | inj₂ (k , refl) =
          x * (x * (projₙ (fe-rec (2 + (2 * k))) k lt (x * x))) ≡⟨ cong  X  x * (x * X)) (projₙ-fe k (1 + (2 * k)) (x * x) lt) 
          x * (x * (fast-exp k (x * x)))                        ≡⟨ cong  X  x * (x * X)) (projₙ IH k (≤⇒≤′ (≤-step (m≤m+n _ _))) (x * x))  
          x * (x * (x * x) ^ k)                                 ≡⟨ cong  X  x * (x * X)) (*-distribˡ-^ x x k) 
          x * (x * ((x ^ k) * (x ^ k)))                         ≡⟨ cong  X  x * (x * ((x ^ k) * (x ^ X))))(sym(+-identityʳ k)) 
          x * (x * (x ^ k * x ^ (k + 0)))                       ≡⟨ cong  X  x * (x * X)) (sym (^-distribˡ-+-* x k (k + zero))) 
          x * (x * (x ^ (2 * k)))
        open ≡-Reasoning
        lt = ≤⇒≤′ (≤-step (m≤m+n k (k + zero)))

## Diving Deeper, Definition Induction using the `Induction` Library

We have seen how to use the facilities in `Nat.Induction` for complete
induction. Next we explore how to use the lower-level library in
`src/Induction.agda` to build our own recursion/induction principle.
In particular, we'll build an alternative form of complete induction
that should be more familiar to everyone, one in which the `rec`
parameter is simply a function, not a tuple of functions.

## The `RecStruct` Type Operator and `SRec`

The `RecStruct` type operator in `src/Induction.agda` produces the
type for the operators like `CRec`. The parameter `A` of `RecStruct`
is for the things you're doing induction on (e.g. `ℕ`) and the
universe levels `ℓ₁` and `ℓ₂` can be ignored for now.

    RecStruct : ∀ {a} → Set a → (ℓ₁ ℓ₂ : Level) → Set _
    RecStruct A ℓ₁ ℓ₂ = Pred A ℓ₁ → Pred A ℓ₂

Recall that `Pred A ℓ₁` is a function from an element of `A` to a type
(an instance of `Set`). In `FERec` above, `(λ i → ℕ → ℕ)` is an
example of something of type `Pred ℕ _`, and we applied `CRec` to this

Let us take a look at the definition of `CRec`. It generates a
tuple type of length `n` where each element is of type `P i`
(for `i < n`) except for `0`. The last element is just unit.

    CRec : ∀ ℓ → RecStruct ℕ ℓ ℓ
    CRec ℓ P zero    = ⊤
    CRec ℓ P (suc n) = P n × CRec ℓ P n

For our alternative form of complete induction, we replace the tuple
with a function. In particular, a function that takes a number `k`, a
proof that `k` is smaller than the current `n`, and produces the
function's result for `k`. So the type of our `rec` parameter is given
by the following `SRec` operator.

SRec :    RecStruct   
SRec  P = λ n   k  k <′ n  P k

(We use `<′` instead of `<` to make it easier to define `sRecBuilder`
and prove `fe-rec-fast-exp₂` in the following.)

## Building the Arguments for the `rec` Parameter

The next step is to build a value of type `SRec` for every natural
number, so that we can pass these values into the `rec` parameter of
the client's step function. Thus, we have to define a function
analogous to the `cRecBuilder` we discussed above. Let us name our
builder function `sRecBuilder`. As a guide, the `Induction` library
defines the `RecursorBuilder` type operator that specifies the type
for builder functions. Its input parameter `Rec` has type `RecStruct`
(e.g. `SRec` is a valid argument to `RecursorBuilder`).

    RecursorBuilder : ∀ {a ℓ₁ ℓ₂} {A : Set a} → RecStruct A ℓ₁ ℓ₂ → Set _
    RecursorBuilder Rec = ∀ P → (Rec P ⊆′ P) → Universal (Rec P)

A builder function takes 1) a predicate `P` that specifies the result
type of the recursive function (just like the `P` in `SRec`), and 2) a
step function that produces a `P` given a `rec` parameter of type `Rec
P`.  The result of the builder function is a value of type `Rec P`,
that is, a value that can be passed into the `rec` parameter of the
client's step function.

So our `sRecBuilder` function has type `RecursorBuilder (SRec ℓ)`.  So
it takes a predicate `P`, a step function, and its output is of type
`SRec ℓ P`, so the output is a function with parameters `n` and `k`
and a proof of `k <′ n`. The function produces an element of `P k`. We
define `sRecBuilder` using induction on `k <′ n`. We discuss
the two cases below.

sRecBuilder :  {}  RecursorBuilder (SRec )
sRecBuilder P step .(suc k) k ≤′-refl = step k rec
  where rec = sRecBuilder P step k
sRecBuilder P step (suc n) k (≤′-step lt) = sRecBuilder P step n k lt

* In the case of `≤′-refl`, we have `n = suc k`. We need to produce `P k`,
  which we can do with the call `step k rec`, but we need to fill in
  the second argument `rec`. This we do by the recursive call to

* The case for `≤′-step` is even easier. We simply call `sRecBuilder`

## The `build` Function to Finish `sRec` 

The final step in creating our custom induction/recursion principle is
to invoke the `build` function in `src/Induction.agda` to produce
`sRec`. The `build` function takes a builder function, such as
`sRecBuilder`, and produces a `Recursor`, which is a function that,
given a step function, produces a recursive function.

sRec : ∀{}  Recursor (SRec )
sRec = build sRecBuilder

## Revisiting Fast Exponentiation with Strong Recursion/Induction

Analogous to `FERec`, we define the type for the `rec` parameter of
our step function with the below `FERec₂`, but this time use `SRec`
instead of `CRec`.

FERec₂ :   Set
FERec₂ n = SRec _  i    ) n

The step function `fe-step₂` is similar to `fe-step`, but this time
the `rec` parameter is easier to work with. It's a function that we
can call. It just requires an extra argument with the proof that the
argument `k` is less than `n`.

fe-step₂ : (n : )  FERec₂ n  (  )
fe-step₂ zero rec x = 1
fe-step₂ (suc n′) rec x
    with parity n′
... | inj₁ (k , refl) =  x * rec k lt (x * x)
      where lt : k <′ 1 + (2 * k)
            lt = ≤⇒≤′ (s≤s (m≤m+n k _))
... | inj₂ (k , refl) =  x * x * rec k lt (x * x)
      where lt : k <′ 2 + (2 * k)
            lt = ≤⇒≤′ (s≤s (≤-step (m≤m+n k _)))

We define our second fast exponentiation using `sRec` and our new step
function `fe-step₂`.

fast-exp₂ :     
fast-exp₂ = sRec  _  (  )) fe-step₂

## Revisiting the Proof that Fast Exponentiation is Correct

Let us see how this alternate definition of fast exponentiation
affects our proof of correctness. (TLDR: it only changes one lemma.)
We define an abbreviation for the correctness condition as follows.

fe-ok₂ :   Set
fe-ok₂ n =  x  fast-exp₂ n x  x ^ n

And another appreviation for calling `sRecBuilder` with the new step
function `fe-step₂`.

fe-rec₂ : (n : )  SRec _  _  (  )) n
fe-rec₂ n = sRecBuilder  _  (  )) fe-step₂ n

Last time we proved the lemma `projₙ-fe` to relate the "raw" recursive
call to a call to `fast-exp`. Here we need a similar lemma, but it is
somewhat simpler because we no longer need to use `projₙ`.  So we just
need to relate `fe-rec₂` to `fast-exp₂`.

fe-rec-fast-exp₂ :  k n x (lt : k <′ suc n)  fe-rec₂ (suc n) k lt x  fast-exp₂ k x
fe-rec-fast-exp₂ k .k x ≤′-refl = refl
fe-rec-fast-exp₂ k (suc n′) x (≤′-step lt) = fe-rec-fast-exp₂ k n′ x lt

We prove that `fast-exp₂` is correct, again with a proof by complete
induction. The only difference is that the first step in the equational
reasoning is to use `fe-rec-fast-exp₂` instead of `projₙ-fe`.

fast-exp₂-is-correct :  n x  fast-exp₂ n x  x ^ n
fast-exp₂-is-correct = cRec fe-ok₂ step
  step : (n : )  CRec _ fe-ok₂ n  fe-ok₂ n
  step zero IH x = refl
  step (suc n′) IH x
      with parity n′
  ... | inj₁ (k , refl) =
        x * fe-rec₂ (suc n′) k lt (x * x) ≡⟨ cong X  x * X) (fe-rec-fast-exp₂ k n′ (x * x) lt) 
        x * fast-exp₂ k (x * x)           ≡⟨ cong X  x * X) (projₙ IH k lt₂ (x * x)) 
        x * (x * x) ^ k                   ≡⟨ cong X  x * X) (*-distribˡ-^ x x k) 
        x * ((x ^ k) * (x ^ k))           ≡⟨ cong X  x * (x ^ k * x ^ X)) (sym(+-identityʳ k)) 
        x * (x ^ k * x ^ (k + 0))         ≡⟨ cong X  x * X) (sym (^-distribˡ-+-* x k (k + zero))) 
        x * x ^ (2 * k)
      open ≡-Reasoning
      lt = s≤′s (≤⇒≤′ (m≤m+n k (k + zero)))
      lt₂ = ≤⇒≤′ (m≤m+n k (k + zero))
  ... | inj₂ (k , refl) =
      (x * x) * fe-rec₂ (2 + 2 * k) k lt (x * x) ≡⟨ cong X  (x * x) * X)(fe-rec-fast-exp₂ k (1 + (2 * k)) (x * x) lt) 
      (x * x) * fast-exp₂ k (x * x)              ≡⟨ cong X  (x * x) * X)(projₙ IH k lt₂ (x * x)) 
      (x * x) * (x * x) ^ k                      ≡⟨ cong X  (x * x) * X) (*-distribˡ-^ x x k) 
      (x * x) * (x ^ k  *  x ^ k)                ≡⟨ cong X  (x * x) * (x ^ k  *  x ^ X)) (sym(+-identityʳ k)) 
      (x * x) * (x ^ k  *  x ^ (k + 0))          ≡⟨ cong X  (x * x) * X) (sym (^-distribˡ-+-* x k (k + zero))) 
      (x * x) * (x ^ (2 * k))                    ≡⟨ *-assoc x _ _ 
      x * (x * (x ^ (2 * k)))
      open ≡-Reasoning
      lt = s≤′s (≤⇒≤′ (≤-step (m≤m+n k (k + zero))))
      lt₂ = (≤⇒≤′ (≤-step (m≤m+n k (k + zero))))

## A Parting Thought

As I was fumbling around and bumping into deadends on my way to
writing the above definitions and proofs, perhaps the trickiest part
of using the `Induction` library was stating and proving the lemmas
that relate the "raw" recursive call to the recursive function, e.g.,
`projₙ-fe` and `fe-rec-fast-exp₂`. I wonder whether the `Induction`
library could somehow also provide a general lemma for that, perhaps
with help from the client.