## Thursday, April 07, 2022

### Using Agda's Induction/Recursion Library

FastExp

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
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,


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
Nat.Induction.

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
function.


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
function.


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.


fast-exp-is-correct : ∀ n x → fast-exp n x ≡ x ^ n
fast-exp-is-correct = cRec fe-ok step
where
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) =
begin
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))
∎
where
open ≡-Reasoning
lt = ≤⇒≤′ (m≤m+n k (k + zero))
... | inj₂ (k , refl) =
begin
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)))
∎
where
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
predicate.

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
sRecBuilder.

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

## 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
where
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) =
begin
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)
∎
where
open ≡-Reasoning
lt = s≤′s (≤⇒≤′ (m≤m+n k (k + zero)))
lt₂ = ≤⇒≤′ (m≤m+n k (k + zero))
... | inj₂ (k , refl) =
begin
(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)))
∎
where
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.