```-----------------------------------------------------------------------------
-- |
-- Module    : Documentation.SBV.Examples.WeakestPreconditions.Fib
-- Copyright : (c) Levent Erkok
-- Maintainer: erkokl@gmail.com
-- Stability : experimental
--
-- Proof of correctness of an imperative fibonacci algorithm, using weakest
-- preconditions. Note that due to the recursive nature of fibonacci, we
-- cannot write the spec directly, so we use an uninterpreted function
-- and proper axioms to complete the proof.
-----------------------------------------------------------------------------

{-# LANGUAGE DeriveAnyClass        #-}
{-# LANGUAGE DeriveFoldable        #-}
{-# LANGUAGE DeriveGeneric         #-}
{-# LANGUAGE DeriveTraversable     #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE NamedFieldPuns        #-}

module Documentation.SBV.Examples.WeakestPreconditions.Fib where

import Data.SBV
import Data.SBV.Control

import Data.SBV.Tools.WeakestPreconditions

import GHC.Generics (Generic)

-- * Program state

-- | The state for the sum program, parameterized over a base type @a@.
data FibS a = FibS { n :: a    -- ^ The input value
, i :: a    -- ^ Loop counter
, k :: a    -- ^ tracks @fib (i+1)@
, m :: a    -- ^ tracks @fib i@
}
deriving (Show, Generic, Mergeable, Functor, Foldable, Traversable)

-- | Show instance for 'FibS'. The above deriving clause would work just as well,
-- but we want it to be a little prettier here, and hence the @OVERLAPS@ directive.
instance {-# OVERLAPS #-} (SymVal a, Show a) => Show (FibS (SBV a)) where
show (FibS n i k m) = "{n = " ++ sh n ++ ", i = " ++ sh i ++ ", k = " ++ sh k ++ ", m = " ++ sh m ++ "}"
where sh v = case unliteral v of
Nothing -> "<symbolic>"
Just l  -> show l

-- | 'Fresh' instance for the program state
instance (SymVal a, SMTValue a) => Fresh IO (FibS (SBV a)) where
fresh = FibS <\$> freshVar_  <*> freshVar_  <*> freshVar_ <*> freshVar_

-- | Helper type synonym
type F = FibS SInteger

-- * The algorithm

-- | The imperative fibonacci algorithm:
--
-- @
--     i = 0
--     k = 1
--     m = 0
--     while i < n:
--        m, k = k, m + k
--        i++
-- @
--
-- When the loop terminates, @m@ contains @fib(n)@.
algorithm :: Stmt F
algorithm = Seq [ Assign \$ \st -> st{i = 0, k = 1, m = 0}
, assert "n >= 0" \$ \FibS{n} -> n .>= 0
, While "i < n"
(\FibS{n, i, k, m} -> i .<= n .&& k .== fib (i+1) .&& m .== fib i)
(Just (\FibS{n, i} -> [n-i]))
(\FibS{n, i} -> i .< n)
\$ Seq [ Assign \$ \st@FibS{m, k} -> st{m = k, k = m + k}
, Assign \$ \st@FibS{i}    -> st{i = i+1}
]
]

-- | Symbolic fibonacci as our specification. Note that we cannot
-- really implement the fibonacci function since it is not
-- symbolically terminating.  So, we instead uninterpret and
-- axiomatize it below.
--
-- NB. The concrete part of the definition is only used in calls to 'traceExecution'
-- and is not needed for the proof. If you don't need to call 'traceExecution', you
-- can simply ignore that part and directly uninterpret.
fib :: SInteger -> SInteger
fib x
| isSymbolic x = uninterpret "fib" x
| True         = go x
where go i = ite (i .== 0) 0
\$ ite (i .== 1) 1
\$ go (i-1) + go (i-2)

-- | Constraints and axioms we need to state explicitly to tell
-- the SMT solver about our specification for fibonacci.
axiomatizeFib :: Symbolic ()
axiomatizeFib = do -- Base cases.
-- Note that we write these in forms of implications,
-- instead of the more direct:
--
--    constrain \$ fib 0 .== 0
--    constrain \$ fib 1 .== 1
--
-- As otherwise they would be concretely evaluated and
-- would not be sent to the SMT solver!

x <- sInteger_
constrain \$ x .== 0 .=> fib x .== 0
constrain \$ x .== 1 .=> fib x .== 1

-- The inductive case. Unfortunately; SBV does not support
-- adding quantified constraints in the query mode. So we
-- have to write this axiom directly in SMT-Lib. Note also how
-- carefully we've chosen this axiom to work with our proof!
addAxiom "fib_n" [ "(assert (forall ((x Int))"
, "                (= (fib (+ x 2)) (+ (fib (+ x 1)) (fib x)))))"
]

-- | Precondition for our program: @n@ must be non-negative.
pre :: F -> SBool
pre FibS{n} = n .>= 0

-- | Postcondition for our program: @m = fib n@
post :: F -> SBool
post FibS{n, m} = m .== fib n

-- | Stability condition: Program must leave @n@ unchanged.
noChange :: Stable F
noChange = [stable "n" n]

-- | A program is the algorithm, together with its pre- and post-conditions.
imperativeFib :: Program F
imperativeFib = Program { setup         = axiomatizeFib
, precondition  = pre
, program       = algorithm
, postcondition = post
, stability     = noChange
}

-- * Correctness

-- | With the axioms in place, it is trivial to establish correctness:
--
-- >>> correctness
-- Total correctness is established.
-- Q.E.D.
--
-- Note that I found this proof to be quite fragile: If you do not get the algorithm right
-- or the axioms aren't in place, z3 simply goes to an infinite loop, instead of providing
-- counter-examples. Of course, this is to be expected with the quantifiers present.
correctness :: IO (ProofResult (FibS Integer))
correctness = wpProveWith defaultWPCfg{wpVerbose=True} imperativeFib

-- * Concrete execution
-- \$concreteExec

{- \$concreteExec

Example concrete run. As we mentioned in the definition for 'fib', the concrete-execution
function cannot deal with uninterpreted functions and axioms for obvious reasons. In those
cases we revert to the concrete definition. Here's an example run:

>>> traceExecution imperativeFib \$ FibS {n = 3, i = 0, k = 0, m = 0}
*** Precondition holds, starting execution:
{n = 3, i = 0, k = 0, m = 0}
===> [1.1] Assign
{n = 3, i = 0, k = 1, m = 0}
===> [1.2] Conditional, taking the "then" branch
{n = 3, i = 0, k = 1, m = 0}
===> [1.2.1] Skip
{n = 3, i = 0, k = 1, m = 0}
===> [1.3] Loop "i < n": condition holds, executing the body
{n = 3, i = 0, k = 1, m = 0}
===> [1.3.{1}.1] Assign
{n = 3, i = 0, k = 1, m = 1}
===> [1.3.{1}.2] Assign
{n = 3, i = 1, k = 1, m = 1}
===> [1.3] Loop "i < n": condition holds, executing the body
{n = 3, i = 1, k = 1, m = 1}
===> [1.3.{2}.1] Assign
{n = 3, i = 1, k = 2, m = 1}
===> [1.3.{2}.2] Assign
{n = 3, i = 2, k = 2, m = 1}
===> [1.3] Loop "i < n": condition holds, executing the body
{n = 3, i = 2, k = 2, m = 1}
===> [1.3.{3}.1] Assign
{n = 3, i = 2, k = 3, m = 2}
===> [1.3.{3}.2] Assign
{n = 3, i = 3, k = 3, m = 2}
===> [1.3] Loop "i < n": condition fails, terminating
{n = 3, i = 3, k = 3, m = 2}
*** Program successfully terminated, post condition holds of the final state:
{n = 3, i = 3, k = 3, m = 2}
Program terminated successfully. Final state:
{n = 3, i = 3, k = 3, m = 2}

As expected, @fib 3@ is @2@.
-}
```