Challenge 3 of the Great Theorem Prover Showdown

Fulcrum.agda

-- Problem 3 of @Hillelogram's Great Theorem Prover Showdown.
-- https://www.hillelwayne.com/post/theorem-prover-showdown
--
-- Implemented using current stable-2.5 branch of Agda
-- (https://github.com/agda/agda/tree/08664ac) and the agda-prelude
-- library (https://github.com/UlfNorell/agda-prelude/tree/d193a94).

module Fulcrum where

open import Prelude hiding (_≤?_)
open import Prelude.Int.Properties
open import Container.List renaming (forgetAny to ixToNat)
open import Container.List.Properties
open import Tactic.Nat
open import Tactic.Cong

-- Refinement type. Σ with irrelevant second component. We could use Σ-types
-- but there would be more work making sure that the proofs don't get in the
-- way of performance. With this type the second component is completely
-- erased at runtime.
syntax Refine A (λ x → B) = [ x ∈ A ∣ B ]   -- let's have some nice syntax
record Refine {a b} (A : Set a) (B : A → Set b) : Set (a ⊔ b) where
  constructor _,_
  field value  : A
        .proof : B value
open Refine using (value)

-- The library compare gives you strict inequalities, but for our purposes
-- it's more convenient to work with ≤.
_≤?_ : (a b : Nat) → Either (a ≤ b) (b ≤ a)
a ≤? b with compare a b
... | less    lt = left  (lt-to-leq {A = Nat} lt)
... | equal   eq = left  (eq-to-leq eq)
... | greater gt = right (lt-to-leq {A = Nat} gt)

-- A bunch of lemmas about lists that really ought to go in a library somewhere.
module _ {a} {A : Set a} where

  Nonempty : List A → Set
  Nonempty []      = ⊥
  Nonempty (_ ∷ _) = ⊤

  nonempty? : (xs : List A) → Dec (Nonempty xs)
  nonempty? []      = no λ()
  nonempty? (_ ∷ _) = yes _

  drop-suc : ∀ (xs : List A) {y ys} i → drop i xs ≡ y ∷ ys → drop (suc i) xs ≡ ys
  drop-suc [] zero    ()
  drop-suc [] (suc i) ()
  drop-suc (x ∷ xs) zero    refl = refl
  drop-suc (x ∷ xs) (suc i) eq   = drop-suc xs i eq

  take-suc : ∀ (xs : List A) {y ys} i → drop i xs ≡ y ∷ ys → take (suc i) xs ≡ take i xs ++ [ y ]
  take-suc [] zero    ()
  take-suc [] (suc i) ()
  take-suc (x ∷ xs) zero    refl = refl
  take-suc (x ∷ xs) (suc i) eq   = x ∷_ $≡ take-suc xs i eq

  take-all : ∀ (xs : List A) i → i ≥ length xs → take i xs ≡ xs
  take-all [] zero    _ = refl
  take-all [] (suc i) _ = refl
  take-all (x ∷ xs) zero    geq = refute geq
  take-all (x ∷ xs) (suc i) geq = x ∷_ $≡ take-all xs i (by geq)

  drop-all : ∀ (xs : List A) i → i ≥ length xs → drop i xs ≡ []
  drop-all [] zero    _ = refl
  drop-all [] (suc i) _ = refl
  drop-all (x ∷ xs) zero    geq = refute geq
  drop-all (x ∷ xs) (suc i) geq = drop-all xs i (by geq)

  append-assoc : ∀ (xs : List A) {ys zs} → (xs ++ ys) ++ zs ≡ xs ++ (ys ++ zs)
  append-assoc []       = refl
  append-assoc (x ∷ xs) = x ∷_ $≡ append-assoc xs

private
  sum-append : ∀ (xs : List Int) {ys} → sumR (xs ++ ys) ≡ sumR xs + sumR ys
  sum-append []       = sym (addInt-zero-l _)
  sum-append (x ∷ xs) = x +_ $≡ sum-append xs ⟨≡⟩ addInt-assoc x _ _

module _ {a} {A : Set a} where

  -- Packing up an element of a list with its index (proof that it's in the list).
  Elem : List A → Set a
  Elem xs = Σ A (_∈ xs)

  -- We can look up an element using a number if it's small enough. Returns
  --   - the element and its well-typed index as an Elem xs
  --   - a proof that the index is the number we started with
  getElem : (xs : List A) (i : Nat) → i < length xs → Σ (Elem xs) λ e → ixToNat (snd e) ≡ i
  getElem [] i lt = refute lt
  getElem (x ∷ xs) zero    lt = (x , zero!) , refl
  getElem (x ∷ xs) (suc i) lt with getElem xs i (by lt)
  ... | (z , ix) , ix=i = (z , suc ix) , (suc $≡ ix=i)

-- Being the smallest value of a type wrt to some measure.
Minimal : ∀ {a} {A : Set a} → (A → Nat) → A → Set a
Minimal μ x = ∀ y → μ x ≤ μ y

-- Given a nonempty list and a measure we can find its minimal Elem.

module _ {a} {A : Set a} (μ : A → Nat) where

  -- It's hard to get the certified findMin to run without performance
  -- penalties. For efficiency the best thing to do would be to define
  -- a simply typed findMin and prove it correct separately, but that
  -- means duplicating the logic, so we'll just take the performance hit.
  -- (around +20%).
  findMin : (xs : List A) ⦃ _ : Nonempty xs ⦄ → [ e ∈ Elem xs ∣ Minimal (μ ∘ fst) e ]
  findMin [] ⦃()⦄
  findMin xs@(x ∷ xs′) = go x xs′ zero! suc λ where
                            w zero!   → left (diff! 0)
                            w (suc i) → right i
    where
      -- As we traverse the list, z is the best value so far and ys is the
      -- part of the list we haven't looked at yet. We also keep track of
      -- three invariants:
      --   (1) z is in the original list xs
      --   (2) ys is a sublist of xs, needed for (1) when we find a new minimum
      --   (3) any element in the original list falls either in the part we've
      --       already looked at, in which case it's bigger than the current
      --       minimum, or it's in ys. This is needed to establish minimality at
      --       the end.
      go : ∀ z ys →
             z ∈ xs →                                        -- (1)
             (∀ {y} → y ∈ ys → y ∈ xs) →                    -- (2)
             (∀ w → w ∈ xs → Either (μ z ≤ μ w) (w ∈ ys)) → -- (3)
             [ e ∈ Elem xs ∣ Minimal (μ ∘ fst) e ]
      go z [] z∈xs emb split = (z , z∈xs) , λ w →
         case uncurry split w of λ where
           (left z≤w) → z≤w
           (right ())
      go z (y ∷ ys) z∈xs emb split with μ z ≤? μ y
      ...   | left  z≤y = go z ys z∈xs (emb ∘ suc) λ w iw →
                              case split w iw of λ where
                                (left z≤w)      → left z≤w
                                (right zero!)   → left z≤y
                                (right (suc i)) → right i
      ...   | right y≤z = go y ys (emb zero!) (emb ∘ suc) λ w iw →
                              case split w iw of λ where
                                (left z≤w) → left (leq-trans {A = Nat} y≤z z≤w)
                                (right zero!)     → left (diff! 0)
                                (right (suc i))  → right i

-- Up to this point we've done nothing specific to the problem at hand. All of that could
-- conceivably go in a library. Now it's fulcrum time!

-- The fulcrum value to minimize. This is part of the specification and won't execute at
-- runtime, so we use sumR = foldr _+_ 0. This is easier to reason about, but has a bad
-- space leak.
fv : List Int → Nat → Nat
fv xs i = abs (sumR (take i xs) - sumR (drop i xs))

-- To build fulcrum values incrementally we need to keep track of the left and right sums.
record State (xs : List Int) : Set where
  constructor mkState
  field
    i        : Nat
    leftSum  : Int
    rightSum : Int
    leftPrf  : leftSum  ≡ sumR (take i xs)
    rightPrf : rightSum ≡ sumR (drop i xs)

-- It's easy to compute the fulcrum value from a state.
fvFromState : ∀ {xs} → State xs → Nat
fvFromState (mkState _ l r _ _) = abs (l - r)

private
  sumR=sum : (xs : List Int) → sum xs ≡ sumR xs
  sumR=sum xs = foldl!-foldl _+_ 0 xs ⟨≡⟩
                foldl-foldr  _+_ 0 addInt-assoc addInt-zero-l addInt-zero-r xs

-- The starting state. We actually need to evaluate this sum at runtime, so we
-- use the strict foldl sum instead of the easier-to-reason-about sumR.
zeroState : ∀ {xs} → State xs
zeroState {xs} = mkState 0 0 (sum xs) refl (sumR=sum xs)

-- In the main loop we recurse over the list, so we'll keep track of the
-- invariant that the state and the current tail match up.
HasTail : ∀ {xs} → List Int → State xs → Set
HasTail {xs} tl s = drop (State.i s) xs ≡ tl

-- The step function for the state. Some boring work to prove that the
-- invariants are maintained.
sucState : ∀ {x xs tl} → Σ (State xs) (HasTail (x ∷ tl)) → Σ (State xs) (HasTail tl)
sucState {x} {xs} {tl} (mkState i l r lprf rprf , tlprf) =
  mkState (suc i) (l + x) (r - x) lprf′ rprf′ , drop-suc xs i tlprf
  where
    lprf′ : l + x ≡ sumR (take (suc i) xs)
    lprf′ rewrite take-suc xs i tlprf | lprf = sym $
      sumR (take i xs ++ [ x ])     ≡⟨ sum-append (take i xs) ⟩
      sumR (take i xs) + sumR [ x ] ≡⟨ by-cong (addInt-zero-r x) ⟩
      sumR (take i xs) + x ∎

    rprf′ : r - x ≡ sumR (drop (suc i) xs)
    rprf′ rewrite drop-suc xs i tlprf | rprf | tlprf =
      x + sumR tl - x          ≡⟨ addInt-assoc x _ _ ⟩ʳ
      x + (sumR tl - x)        ≡⟨ x +_ $≡ addInt-commute (sumR tl) _ ⟩
      x + (negate x + sumR tl) ≡⟨ addInt-assoc x _ _ ⟩
      x - x + sumR tl          ≡⟨ _+ sumR tl $≡ subInt-equal x x refl ⟩
      0 + sumR tl              ≡⟨ addInt-zero-l _ ⟩
      sumR tl ∎   -- TODO: integer tactics!!!

-- We split the algorithm in two parts:
--   - computing the list of all fulcrum states,
--   - and finding the minimal one

-- Computing the list of fulcrum states is straight-forward using the step
-- function above. Note that we iterate over a state whose index matches up
-- with the current tail.
fulcrums′ : ∀ {xs} tl → Σ (State xs) (HasTail tl) → List (State xs)
fulcrums′ []       s = fst s ∷ []
fulcrums′ (x ∷ tl) s = fst s ∷ fulcrums′ tl (sucState s)

fulcrums : ∀ xs → List (State xs)
fulcrums xs = fulcrums′ xs (zeroState , refl)

instance
  nonemptyFulcrums : ∀ {xs} → Nonempty (fulcrums xs)
  nonemptyFulcrums {[]}     = _
  nonemptyFulcrums {x ∷ xs} = _

-- Now, we can easily find the minimal state in the list, but that's
-- not enough for two reasons:
--   - there might be smaller values not in the list
--   - we haven't proved that fvFromState is the same as fv
-- This is the mainLemma below.

fulcrum′ : (xs : List Int) → [ e ∈ Elem (fulcrums xs) ∣ Minimal (fvFromState ∘ fst) e ]
fulcrum′ xs = findMin fvFromState (fulcrums xs)

-- First let's prove that the only thing interesting about the fulcrum
-- state is the index:

fulcrum-state : ∀ {xs} → Nat → State xs
fulcrum-state {xs} i = mkState i (sumR (take i xs)) (sumR (drop i xs)) refl refl

unique-state : ∀ {xs} (s : State xs) i → State.i s ≡ i → s ≡ fulcrum-state i
unique-state (mkState i l r refl refl) i refl = refl

-- The ith element in fulcrums xs is fulcrum-state i.
fulcrums-spec : ∀ {xs s} (ix : s ∈ fulcrums xs) i → ixToNat ix ≡ i → s ≡ fulcrum-state i
fulcrums-spec {xs} {s} ix i refl = unique-state s i (by (go xs (zeroState , refl) ix))
  where go : ∀ tl (s₀ : Σ (State xs) (HasTail tl)) (ix : s ∈ fulcrums′ tl s₀) →
               State.i s ≡ ixToNat ix + State.i (fst s₀)
        go []       _ zero! = refl
        go (x ∷ tl) _ zero! = refl
        go [] _       (suc ())
        go (x ∷ tl) s (suc i) = by (go tl (sucState s) i)

-- We're also going to need a boring lemma about the number of fulcrums states.
fulcrums-length : ∀ xs → suc (length xs) ≡ length (fulcrums xs)
fulcrums-length xs = go xs (zeroState , refl)
  where go : ∀ tl s → suc (length tl) ≡ length (fulcrums′ {xs} tl s)
        go []       s = refl
        go (x ∷ tl) s = suc $≡ go tl (sucState s)

-- The specification says we only need to minimize over indices (0 .. length xs), but
-- the take/drop specification works for any natural:
fv-out-of-bounds : ∀ xs j → j ≥ length xs → fv xs j ≡ fv xs (length xs)
fv-out-of-bounds xs j geq rewrite take-all xs j geq | drop-all xs j geq
                                | take-all xs (length xs) (diff! 0)
                                | drop-all xs (length xs) (diff! 0) = refl

-- Finally we can show the main lemma: the minimal element in our list of
-- fulcrum states has the minimal index with respect to `fv`.
mainLemma : ∀ {xs} (e : Elem (fulcrums xs)) i → State.i (fst e) ≡ i →
              Minimal (fvFromState ∘ fst) e → Minimal (fv xs) i
mainLemma {xs} (s , ix) i ix=i isMin j rewrite unique-state s i ix=i =
    either (in-bounds j) (out-of-bounds j) (j ≤? length xs)
  where
    in-bounds : ∀ j → j ≤ length xs → fv xs i ≤ fv xs j
    in-bounds j leq rewrite fulcrums-length xs with getElem (fulcrums xs) j leq
    ... | (js , jx) , js=j =
      case fulcrums-spec jx j js=j of λ { refl → isMin (fulcrum-state j , jx) }

    out-of-bounds : ∀ j → j ≥ length xs → fv xs i ≤ fv xs j
    out-of-bounds j geq rewrite fv-out-of-bounds xs j geq =
      in-bounds (length xs) (diff! 0)

-- Using the main lemma we can patch up the result of fulcrum′ to get the final
-- result. Here spending the extra effort of proving the out-of-bounds case pays
-- off in a very clean specification:
--   you get the `i ∈ Nat` that minimizes `fv xs i`

fulcrum : (xs : List Int) → [ i ∈ Nat ∣ Minimal (fv xs) i ]
fulcrum xs = case fulcrum′ xs of λ where
  (e , isMin) → State.i (fst e) , mainLemma e _ refl isMin

-- Performance testing ------

-- We should probably make sure that it does indeed run in O(n) time and space...

-- Let's have some random numbers
blumBlumShub : (x₀ M len : Nat) ⦃ _ : NonZero M ⦄ → List Int
blumBlumShub x₀ M len = go (step x₀) len
  where offs = fromNat (M div 2)
        step = λ x → x * x mod M
        mkInt = λ x → fromNat x - offs
        go : Nat → Nat → List Int
        go x zero    = mkInt x ∷ []
        go x (suc n) = mkInt x ∷ go (step x) n

input : Nat → List Int
input n = blumBlumShub 3 209 n

test : Nat → Nat
test n = value (fulcrum (input n))

inputLength : IO Nat
inputLength = getArgs >>= λ where
  []      → pure 10
  (s ∷ _) → pure $ maybe 10 id (parseNat s)

main : IO ⊤
main = do
  n ← inputLength
  print (test n)

-- Benchmarks on my machine (2013 MacBook Pro)
--
--         n    time  memory
-- 1,000,000    0.5s     40M
-- 2,000,000    1.0s     80M
-- 4,000,000    2.1s    160M
-- 8,000,000    4.3s    320M
--
-- That certainly looks like O(n).
--
-- For reference the (insultingly short) unverified Haskell version:
--
--   index (i, _, _) = i
--   fv (_, l, r) = abs (l - r)
--   fulcrums xs = go xs (0, 0, sum xs)
--     where go [] s = [s]
--           go (x : xs) s@(i, l, r) = s : go xs (i + 1, l + x, r - x)
--   fulcrum = index . minimumBy (compare `on` fv) . fulcrums
--
-- performs as follows
--
--         n    time  memory
-- 1,000,000    0.4s     39M
-- 2,000,000    0.8s     80M
-- 4,000,000    1.6s    148M
-- 8,000,000    3.1s    306M
--
-- So, for the small price of 250 lines of code, a few hours of dependently
-- typed fun, and a 40% performance hit we get a fully verified fulcrum
-- implementation!