ProvableQuicksort.idr

module ProvableSorting

import Data.So
import Data.List
import Control.Order
import Data.List.Elem
import Data.List.Quantifiers
import Data.Nat
import Data.Nat.Order
import Data.Nat.Views
import Data.Vect

-- Formally proved QuickSort
-- Consists of two parts: specification, and implementation. There are also extra comments on bottom of this file.
----------------------------------------
----------------------------------------
-- Specification below ↓. It's best read from bottom to top. 

data Sorted : List Nat -> Type where
    Empty : Sorted []
    One : (x: Nat) -> Sorted [x]
    Many : (x, y : Nat) -> (zs : List Nat) -> LTE x y -> Sorted (y :: zs) -> Sorted (x :: y :: zs)

data Permutation : (xs: List a) -> (xs': List a) -> Type where 
    PermNil : Permutation [] []  
    PermCons : Permutation xs xs' -> Permutation (x :: xs) (x :: xs')
    PermSwap : Permutation (x :: y :: xs) (y :: x :: xs)
    PermTrans : {ys : List a} -> Permutation xs ys -> Permutation ys zs -> Permutation xs zs  

-- What does it mean to have an element of type 'SortedList xs'?
-- It means that 
-- 1. We have a list (ws)
-- 2. ws is sorted (Sorted ws)
-- 3. ws has the same elements as xs (Permutation xs ws)
-- What is `Sorted`, or `Permutation`? Read above.
SortedList : (xs : List Nat) -> Type
SortedList xs = (ws : List Nat ** (Sorted ws, Permutation xs ws))

-- Defines the specification of our sorting function:
-- It takes a list of natural numbers(xs) and produces an element of SortedList xs. What is SortedList xs? Read above.
quicksort : (xs : List Nat) -> SortedList xs

-- Specification above ↑
----------------------------------------
----------------------------------------
-- Implementation below ↓


permutationRefl : {xs: List a} -> Permutation xs xs
permutationRefl {xs = []} = PermNil 
permutationRefl {xs = (x :: xs)} = PermCons permutationRefl {xs = xs}

permutationSym : Permutation xs ys -> Permutation ys xs
permutationSym PermNil = PermNil 
permutationSym (PermCons xs) = PermCons (permutationSym xs) 
permutationSym PermSwap = PermSwap 
permutationSym (PermTrans x y) = PermTrans (permutationSym y) (permutationSym x)

data Partition : (x: Nat) -> (xs: List Nat) -> Type where
    MkPartition : {ys, zs: List Nat}
      -> All (\x => LTE x p) ys -- p is bigger than all elements in ys
      -> All (\x => LTE p x) zs -- p is smaller than all elements in zs
      -> Permutation xs (ys ++ zs) -- if xs is a permutation of ys ++ zs, then
      -> Partition p xs -- p is the pivot, xs is the original list

-- if xs is a permutation of xs', then we can concatenate zs to both
permutationConcat : {xs, xs', zs: List a}
  -> Permutation xs xs' -> Permutation (xs ++ zs) (xs' ++ zs)
permutationConcat PermNil = permutationRefl 
permutationConcat (PermCons xs) = PermCons (permutationConcat xs) 
permutationConcat PermSwap = PermSwap 
permutationConcat (PermTrans p1 p2) = PermTrans (permutationConcat p1) (permutationConcat p2)


-- from two lists, we can concatenate them in two ways. They're permutations of each other
permutationComm : (xs, ys: List a)
  -> Permutation (xs ++ ys) (ys ++ xs)
permutationComm {xs = xs} {ys = []} = rewrite appendNilRightNeutral xs
                                      in permutationRefl 
permutationComm {xs = []} {ys = (y :: ys)} = rewrite appendNilRightNeutral ys
                                             in permutationRefl 
permutationComm {xs = (x :: xs)} {ys = (y :: ys)} = PermTrans p1 p2 where
    ih1 = permutationComm xs ys
    ih2 = permutationComm (x::xs) ys
    ih3 = permutationComm xs (y::ys)
    p1 = PermCons $ PermTrans ih3 (PermCons (permutationSym ih1))
    p2 = PermTrans PermSwap (PermCons ih2)

permutationConcatReverse : {xs, xs', zs: List a} -> Permutation xs xs' -> Permutation (zs ++ xs) (zs ++ xs')
permutationConcatReverse px = 
    let p1 = permutationConcat px {zs=zs}
        p2 = permutationComm {xs=xs} {ys=zs}
        p3 = PermTrans (permutationSym p1) p2
    in PermTrans (permutationSym p3) (permutationComm {xs=xs'} {ys=zs})

permutationConcatTrans : {xs, xs', ys, ys' : List a}
    -> Permutation xs xs'
    -> Permutation ys ys'
    -> Permutation (xs ++ ys) (xs' ++ ys')
permutationConcatTrans PermNil py = py 
permutationConcatTrans (PermCons px) py = PermCons (permutationConcatTrans px py) 
permutationConcatTrans PermSwap py = 
    let p1 = permutationConcatReverse py
        p2 = PermSwap
        p3 = PermCons (PermCons p1)
    in PermTrans p2 p3
permutationConcatTrans (PermTrans x y) py = 
    let ih1 = permutationConcatTrans x py
        ih2 = permutationConcatTrans y py
        p1 = permutationSym (permutationConcatReverse py)
    in PermTrans (PermTrans ih1 p1) ih2

permutationLemma : {x: a} -> {xs, ys, zs: List a}
  -> Permutation xs (ys ++ zs)
  -> Permutation (x :: xs) (ys ++ (x :: zs))
permutationLemma perm = PermTrans (PermCons {x=x} $ PermTrans perm (permutationComm {xs=ys} {ys=zs})) (permutationComm {xs=(x::zs)} {ys=ys})

partition : (p: Nat) -> (xs: List Nat) -> Partition p xs 
partition p [] = MkPartition [] [] PermNil   
partition p (x :: xs) = case lte x p of
    Yes lte => let (MkPartition allLte allGte perm) = partition p xs
               in MkPartition (lte::allLte) allGte (PermCons perm)
    No pf => let (MkPartition allLte allGte perm) = partition p xs
                 gte = lteSuccLeft $ notLTEImpliesGT pf
             in MkPartition allLte (gte::allGte) (permutationLemma perm) 

permutationAll : All p xs -> Permutation xs ws -> All p ws 
permutationAll (x :: y :: xs) PermSwap = y :: x :: xs  
permutationAll all PermNil = [] 
permutationAll (x :: xs) (PermCons c) = x :: permutationAll xs c
permutationAll all (PermTrans a b) = permutationAll (permutationAll all a) b


-- If all elements in a list satisfy a property, then the last element also satisfies it
lastAll : {p : a -> Type} -> {xs: List a} -> {auto pf: NonEmpty xs}
  -> All p xs -> p (last xs)
lastAll [] impossible
lastAll [pTerm] = pTerm
lastAll (pTerm :: pTerms) = case pTerms of
    [] => pTerm
    (q :: qs) => lastAll (q :: qs)

lteIfExists : (xs : List Nat) -> (n : Nat) -> Type
lteIfExists [] n = ()
lteIfExists (x :: xs) n = LTE (last (x :: xs)) n

-- appends a number `p` to the sorted list `xs` (and requires a proof that `p` is greater than all elements in `xs` if `xs isn't empty) and produces another sorted list
appendSorted : {xs : List Nat}
            -> Sorted xs
            -> (p : Nat)
            -> lteIfExists xs p
            -> Sorted (xs ++ [p])
appendSorted Empty p () = One p
appendSorted (One x) p proofXsmallerP = Many x p [] proofXsmallerP (One p)
appendSorted (Many x y zs proofXsmallerY sortedyzs) p proofLastSmallerP = 
    let rs = appendSorted sortedyzs p proofLastSmallerP
    in Many x y (zs ++ [p]) proofXsmallerY rs

sortingLemma : {xs: List Nat}
            -> Sorted xs
            -> (z : Nat)
            -> (zs : List Nat)
            -> Sorted (z::zs)
            -> lteIfExists xs z
            -> Sorted (xs ++ (z::zs))
sortingLemma Empty z zs zzs () = zzs
sortingLemma (One x) z zs zzs proofXsmallerZ = Many x z zs proofXsmallerZ zzs
sortingLemma (Many x x' xxs proofXsmallerX' s) z zs zzs proofXsmallerZ
  = let restSorted = sortingLemma s z zs zzs proofXsmallerZ
    in Many x x' (xxs ++ (z::zs)) proofXsmallerX' restSorted


quicksort [] = ([] ** (Empty, PermNil))
quicksort (p :: xs) = 
    let (MkPartition {ys=ys} {zs=zs} ltp gtp perm) = partition p xs 
        (lhs ** (stdLhs, permLhs)) = quicksort ys 
        (rhs ** (stdRhs, permRhs)) = quicksort zs
        lteLhs : All (\x => LTE x p) lhs
        lteLhs = permutationAll ltp permLhs
        gteRhs : All (\x => LTE p x) rhs
        gteRhs = permutationAll gtp permRhs

        rightSideSorted : Sorted (p :: rhs)
        rightSideSorted = case stdRhs of   
            Empty => One p
            One r => let [pfGte] = gteRhs
                     in Many p r [] pfGte (One r)
            rs@(Many rr hh sss _ _) =>
                let (pfGte :: _) = gteRhs
                in Many p rr (hh::sss) pfGte rs
        permFinal : Permutation (p :: xs) (lhs ++ (p :: rhs))
        permFinal = permutationLemma {x=p} (PermTrans perm (permutationConcatTrans permLhs permRhs))
        fullSorted : Sorted (lhs ++ (p :: rhs))
        fullSorted = case lteLhs of 
            [] => rightSideSorted
            (_::_) =>
              let lastLhsOld = lastAll lteLhs
              in sortingLemma stdLhs p rhs rightSideSorted lastLhsOld
    in (lhs ++ (p :: rhs) ** (fullSorted, permFinal))

----------------------------------------
----------------------------------------
-- Extra comments ↓
{-

If this code is ran, and tested on, say, the list [3,0,1] we get the following output:
ListSortingWorkingSimple> quicksort [3,0,1]
MkSortedList [0,
              1,
              3] (NonEmpty (Many LTEZero (Many (LTESucc LTEZero) One))) (PermTrans (PermCons (PermTrans (PermTrans (PermCons (PermCons PermNil)) (PermTrans (PermTrans (PermCons (PermTrans (PermTrans (PermTrans (PermTrans (PermTrans (PermTrans (PermCons (PermTrans (PermTrans PermNil (PermTrans PermNil (PermTrans PermNil PermNil))) PermNil)) (PermTrans (PermCons PermNil) (PermTrans (PermCons PermNil) (PermCons PermNil)))) (PermCons PermNil)) (PermTrans (PermCons PermNil) (PermTrans (PermCons PermNil) (PermCons PermNil)))) (PermTrans (PermTrans (PermCons (PermTrans (PermTrans (PermTrans (PermTrans PermNil (PermTrans PermNil (PermTrans PermNil PermNil))) PermNil) (PermTrans PermNil (PermTrans PermNil PermNil))) PermNil)) (PermTrans (PermCons PermNil) (PermTrans (PermCons PermNil) (PermCons PermNil)))) (PermCons PermNil))) (PermTrans (PermCons PermNil) (PermTrans (PermCons PermNil) (PermCons PermNil)))) (PermCons PermNil))) (PermTrans (PermCons (PermCons PermNil)) (PermTrans (PermCons (PermCons PermNil)) (PermCons (PermCons PermNil))))) (PermCons (PermCons PermNil)))) (PermCons (PermCons PermNil)))) (PermTrans (PermCons (PermTrans (PermCons (PermCons PermNil)) (PermCons (PermCons PermNil)))) (PermTrans PermSwap (PermCons (PermTrans (PermCons (PermTrans (PermCons PermNil) (PermCons PermNil))) (PermTrans PermSwap (PermCons (PermCons PermNil))))))))

Here the second component is a proof that the output is sorted.
This is machine checkable - we can automatically verify whether that output is of the type we expect.

But it turns out, we don't even need to check it! 
The fact that we were able to compile the code means that this function does the right thing for *every* input (and not just the one we happened to check).

Compilation = entering a building with a security badge (you need to show it)
Running the code = walking inside the building (you don't need to keep showing the badge to people)

-}