{-# OPTIONS_GHC -fglasgow-exts #-}

module BinarySearch where

{- 
This module implements binary search. 
The core algoritm statically checks the index calculations so
the minimum number of run-time bounds checks are necessary.
Any performance advantage of this may be lost in the unary encoding
of natural numbers or the nature of Haskell itself.
-}


-- Write only arrays of size n
data Array a n where
    Array :: Nat n ->
             (forall i . Nat i -> Less i n -> a) ->
             Array a n

-- A family of singleton types, isomorphic to the natural numbers
-- defined by "data Nat = Zero | Succ Nat"

data Z
data S a

data Nat n where
    Zero :: Nat Z
    Succ :: Nat n -> Nat (S n)

data Less a b where
    Base  :: Nat b -> Less Z (S b)
    Induc :: Less a b -> Less (S a) (S b)

data Equal a b where
    Equal :: Equal a a

type LTE a b = Either (Less a b) (Equal a b)

{-
GHC 6.6 doesn't have type functions, so we have this type relation.
The burden is upon the programmer to prove that Middle a b c realy
proves that b is directly between a and c. This burden could sit 
quietly unproven, or the programmer could provide a proof of the
type:
middle :: Middle a b c -> Diff b a low -> Diff c b high -> 
   Either (Equal low high) Either ((Equal (S low) high) (Equal low (S high)))
after providing a relation Diff
-}
data Middle a b c where
    Empty    :: Nat a -> Middle a a a
    One      :: Nat a -> Middle a a (S a)
    Divide   :: Middle a b c -> Middle (S a) (S b) (S c)
    AddRight :: Middle Z b c -> Middle Z (S b) (S (S c))

{-
Here is an example of a proof of the well-kindedness of the 
Middle relation. This term is not used anywhere else - it's
sole pupose is as a compiler-checked comment.
-}
wellKindedMiddle :: Middle a b c -> (Nat a, Nat b, Nat c)
wellKindedMiddle (Empty x) = (x,x,x)
wellKindedMiddle (One x) = (x, x, Succ x)
wellKindedMiddle (Divide r) = 
    let (x,y,z) = wellKindedMiddle r
    in (Succ x, Succ y, Succ z)
wellKindedMiddle (AddRight r) = 
    let (x,y,z) = wellKindedMiddle r
    in (x, Succ y, Succ $ Succ z)

-- A proof of the lower inequality implied by Middle
low :: Middle a b c -> Nat b -> LTE a b
low (Empty x) _ = Right Equal
low (One x) _ = Right Equal
low (Divide x) (Succ y) = 
    case low x y of
      Left p -> Left $ Induc p
      Right Equal -> Right Equal
low (AddRight x) (Succ m) =
    case low x m of
      Left p -> Left $ oneMoreLess p
      Right Equal -> Left $ oneMore m

-- A proof of the higher inequality implied by Middle
high :: Middle a b c -> Nat b -> LTE b c
high (Empty x) _ = Right Equal
high (One x) _ = Left $ oneMore x
high (Divide x) (Succ m) = 
    case high x m of
      Left q -> Left $ Induc q
      Right Equal -> Right Equal
high (AddRight x) (Succ m) =
    case high x m of
      Left q -> Left $ Induc $ oneMoreLess q
      Right Equal -> Left $ Induc $ oneMore m

data MiddleType a b where
    MiddleType :: Middle a c b -> Nat c -> MiddleType a b

-- This function actually calculates the middle value of two naturals
makeMiddle :: Nat a -> LTE a b -> Nat b -> MiddleType a b
makeMiddle Zero _ Zero = MiddleType (Empty Zero) Zero
makeMiddle Zero _ (Succ Zero) = MiddleType (One Zero) Zero
makeMiddle Zero _ (Succ (Succ Zero)) = MiddleType (AddRight $ Empty Zero) (Succ Zero)
makeMiddle Zero _ (Succ (Succ n@(Succ m))) =
    case makeMiddle Zero (Left $ Base m) n of
      MiddleType x l -> MiddleType (AddRight x) (Succ l)
makeMiddle (Succ n) p (Succ m) =
    case makeMiddle n (dropOneLTE p) m of
      MiddleType x l -> MiddleType (Divide x) (Succ l)

--Below are several three-line auxiliary functions used in constructing proofs
dropOneLTE :: LTE (S a) (S b) -> LTE a b
dropOneLTE (Left (Induc x)) = Left x
dropOneLTE (Right Equal) = Right Equal

oneMore :: Nat n -> Less n (S n)
oneMore Zero = Base Zero
oneMore (Succ n) = Induc (oneMore n)

oneMoreLess :: Less a b -> Less a (S b)
oneMoreLess (Base x) = Base (Succ x)
oneMoreLess (Induc x) = Induc (oneMoreLess x)

transitiveELL :: LTE a b -> Less b c -> Less a c
transitiveELL (Left d) e = transitiveLLL d e
transitiveELL (Right Equal) e = e

transitiveLLL :: Less a b -> Less b c -> Less a c
transitiveLLL (Base d) (Induc e) = 
    let (_,f) = wellKindedLess e 
    in Base f
transitiveLLL (Induc d) (Induc e) = Induc $ transitiveLLL d e

wellKindedLess :: Less a b -> (Nat a, Nat b)
wellKindedLess (Base d) = (Zero, Succ d)
wellKindedLess (Induc d) = let (e,f) = wellKindedLess d 
                           in (Succ e, Succ f)

bsearch :: (Ord t) => t -> Array t n -> Maybe t
bsearch a e@(Array b c) = case b of
                            Zero -> Nothing
                            Succ Zero -> bsearch1 a (Zero, Right Equal, Zero) e (Base Zero)
                            Succ (Succ n) -> bsearch1 a (Zero, Left $ Base n, Succ n) e (oneMore $ Succ n)
bsearch1 :: (Ord t) => t -> (Nat a, LTE a b, Nat b) -> Array t n -> Less b n -> Maybe t
bsearch1 c (d,e,f) p@(Array _ h) j = 
    case makeMiddle d e f of
      MiddleType k l -> 
          let 
              o = low k l
              m = high k l
              n = transitiveELL m j
              ans = h l n
          in 
            case compare c ans of 
              LT -> case k of
                      Empty _ -> Nothing
                      _ -> bsearch1 c (d,o,l) p n
              EQ -> Just ans
              GT -> case k of
                      Empty _ -> Nothing
                      _ -> bsearch1 c (l,m,f) p j