{-@ LIQUID "--reflection"     @-}
{-@ LIQUID "--ple"            @-}

module FingerTree where

import Language.Haskell.Liquid.ProofCombinators 

{-@ measure digitSize @-}
digitSize :: Digit a -> Int
digitSize (One {}) = 1
digitSize (Two {}) = 2
digitSize (Three {}) = 3
digitSize (Four {}) = 4

data Digit a
    = One a
    | Two a a
    | Three a a a
    | Four a a a a
    deriving (Show)

data Node a = Node2 {a2::a, b2::a} | Node3 {a3::a, b3::a, c3:: a}
    deriving (Show)

data FingerTree a
    = Empty
    | Single a
    | Deep (Digit a) (FingerTree (Node a)) (Digit a)
    deriving Show

{-@ reflect fingerTreeSize @-}
fingerTreeSize :: FingerTree a -> Int
fingerTreeSize t = size to1 t 

{-@ reflect size @-}
size :: (a -> Int) -> FingerTree a -> Int
size _ Empty        = 0
size f (Single a)   = f a
size f (Deep l m r) = digitS f l + size (nodeS f) m + digitS f r

{-@ reflect digitS @-} 
digitS :: (a -> Int) -> Digit a -> Int
digitS f (One a)        = f a
digitS f (Two a b)      = f a + f b
digitS f (Three a b c)  = f a + f b + f c
digitS f (Four a b c d) = f a + f b + f c + f d

{-@ reflect nodeS @-}
nodeS :: (a -> Int) -> Node a -> Int
nodeS f (Node2 a b)   = f a + f b
nodeS f (Node3 a b c) = f a + f b + f c

{-@ reflect consts @-}
consts a b = a

{-@ reflect to1 @-}
to1 :: a -> Int
to1 _ = 1

{-@ n2Int1 :: a:Int -> b:Int -> {n:Node Int | nodeS to1 n == 2} @-}
n2Int1 :: Int -> Int -> Node Int
n2Int1 a b = Node2 a b

-- {-@ n2Int :: a:Int -> b:Int -> {f:(Int -> Int) | f a == 1 && f b == 1} -> {n:Node Int | nodeS f n == 2} @-}
-- n2Int :: Int -> Int -> (Int -> Int) -> Node Int
-- n2Int a b f = Node2 a b

{-@ measure isEmpty @-}
isEmpty Empty      = True
isEmpty (Single _) = False
isEmpty Deep{}     = False

{-@ singleton :: v:Int -> {ft:FingerTree Int | isEmpty ft} @-}
singleton :: Int -> FingerTree Int
singleton a = Empty 

{-@ fromList :: xs:_ -> {t:_ | fingerTreeSize t == len xs} @-}
fromList :: [a] -> FingerTree a
fromList []     = Empty
fromList (x:xs) = add x (fromList xs)

{-@ infix <| @-}
-- | /O(1)/. Add an element to the left end of a sequence.
-- Mnemonic: a triangle with the single element at the pointy end.
{- (<|) :: a -> ft:FingerTree a -> {v : FingerTree a | fingerTreeSize v == fingerTreeSize ft + 1} @-}
{-@ reflect <| @-}
(<|) :: a -> FingerTree a -> FingerTree a
a <| Empty                    =  Single a
a <| Single b                 =  Deep (One a) Empty (One b)
a <| Deep (Four b c d e) m sf = {- m `seq` -} Deep (Two a b) (Node3 c d e <| m) sf
a <| Deep l m r               = Deep (consDigit a l) m r 

{-@ lem_add :: f:_ -> a:_ -> t:_ -> { size f (a <| t) == size f t + f a }  @-}
lem_add :: (a -> Int) -> a -> FingerTree a -> () 

lem_add f _ Empty       
  = () 

lem_add f _ (Single {}) 
  = () 

lem_add f a t@(Deep (Four b c d e) m sf) 
  =   size f t + f a 
  === digitS f (Four b c d e) + size (nodeS f) m + digitS f sf + f a
  === f a + f b + digitS f sf + f c + f d + f e + size (nodeS f) m 
  === f a + f b + digitS f sf + nodeS f (Node3 c d e) + size (nodeS f) m 
      ? lem_add (nodeS f) (Node3 c d e) m
  === f a + f b + digitS f sf + size (nodeS f) ((Node3 c d e) <| m)
  === size f (Deep (Two a b) ((Node3 c d e) <| m) sf)
  === size f (a <| (Deep (Four b c d e) m sf)) 
  === size f (a <| t)
  *** QED  

lem_add f a t@(Deep l m r)  
  = () 

{-@ thm_add :: a:_ -> t:_ -> { fingerTreeSize (a <| t) == 1 + fingerTreeSize t} @-}
thm_add :: a -> FingerTree a -> Proof --  FingerTree a 
thm_add a t 
  =   fingerTreeSize (a <| t) 
  === size to1 (a <| t) 
    ? lem_add to1 a t
  === size to1 t + to1 a 
  === fingerTreeSize t + 1 
  *** QED

{-@ add :: a:_ -> t:_ -> { v: _ | fingerTreeSize v == 1 + fingerTreeSize t} @-}
add :: a -> FingerTree a -> FingerTree a 
add a t = (a <| t) ? thm_add a t 

{-@ measure isFour @-}
isFour :: Digit a -> Bool 
isFour (Four {}) = True 
isFour _         = False 

{-@ reflect consDigit @-}
{-@ consDigit :: _ -> {d:_ | not (isFour d)} -> _ @-}
consDigit :: a -> Digit a -> Digit a
consDigit a (One b)        = Two a b
consDigit a (Two b c)      = Three a b c
consDigit a (Three b c d)  = Four a b c d

{-@ ft1 :: {v:_ | fingerTreeSize v == 10} @-}
ft1 :: FingerTree Int
ft1 = fromList [1,2,3,4,5,6,7,8,9,10]

{-@ ten :: {v:_ | len v = 10} @-}
ten :: [Int]
ten = [1,2,3,4,5,6,7,8,9,10]

ft2 :: FingerTree Int
ft2 = fromList [1..20000]

main :: IO ()
main = do
    print "test"
    print $ fingerTreeSize ft1 -- 10
    print $ fingerTreeSize ft2 -- 20000