{-# LANGUAGE FlexibleInstances
           , GADTs
           , GeneralizedNewtypeDeriving
           , MultiParamTypeClasses
           , TypeFamilies
           , TypeOperators
           , UndecidableInstances #-}

module Algorithm.Pdr ( runpdr ) where

import Control.Monad
import Control.Monad.IO.Class
import Control.Monad.Trans.Class
import Control.Monad.Trans.Except
import Control.Monad.Trans.Reader
import Control.Monad.Trans.State.Strict

import Algorithm.Pdr.State
import Domain
import Monitor
import Utils.Miscellaneous

-- A wrapper of except monad transformer to distinguish from an arbitrary
-- except.
newtype PdrT mon a m b = PdrT { runPdrT :: ReaderT mon (ExceptT a m) b }
    deriving (Functor, Applicative, Monad)

-- The monad stack used to express the most of the algorithm.
type Pdr mon d a b = PdrT mon a (PdrStateT d) b

instance MonadTrans (PdrT mon a) where
    lift = PdrT . lift . lift

instance MonadTransFunctor (PdrT mon a) where
    mapT f m = PdrT . ReaderT $ mapT f . runReaderT (runPdrT m)

instance MonadDomain m => MonadIO (PdrT mon a m) where
    liftIO = PdrT . liftIO

-- Expose domain specific operations in the monad stack by delegating them to
-- the underlying domain (the base monad).
instance MonadDomain d => MonadDomain (PdrT mon a d) where
    type Element (PdrT mon a d) = Element d
    type Set (PdrT mon a d) = Set d
    type Relation (PdrT mon a d) = Relation d

    context              = lift1 context
    bottom               = lift  bottom
    top                  = lift  top
    bottomR              = lift  bottomR
    topR                 = lift  topR
    identityR            = lift  identityR
    singleton            = lift1 singleton
    enumerate            = lift1 enumerate
    relate               = lift2 relate
    select               = lift2 select
    project              = lift1 project
    compose              = lift2 compose
    isEmpty              = lift1 isEmpty
    intersect            = lift2 intersect
    union                = lift2 union
    complement           = lift1 complement
    exclude              = lift2 exclude
    intersectR           = lift2 intersectR
    unionR               = lift2 unionR
    complementR          = lift1 complementR
    excludeR             = lift2 excludeR
    reverseR             = lift1 reverseR
    subsetOf             = lift2 subsetOf
    elementOf            = lift2 elementOf
    equals               = lift2 equals
    refine               = lift2 refine
    generalise           = lift2 generalise
    stringifyElement     = lift1 stringifyElement
    stringifySet         = lift1 stringifySet
    stringifyRelation    = lift1 stringifyRelation
    stringifyAbstraction = lift1 stringifyAbstraction

instance FunctionAbstracting s => FunctionAbstracting (PdrT mon a s) where
    withFunctionAbstractionMap abs = mapT $ withFunctionAbstractionMap abs
    getFunctionAbstractionMap      = lift   getFunctionAbstractionMap
    mapToCalleeContext             = lift4  mapToCalleeContext

instance Encoding e s => Encoding e (PdrT mon a s) where
    encodeInstruction = lift5 encodeInstruction
    encodeTransition  = lift3 encodeTransition

instance MonadSystem e s => MonadSystem e (PdrT mon a s) where
    getSystem  = lift  getSystem
    transition = lift2 transition
    refute     = lift1 refute
    connect    = lift2 connect

instance MonadPdrState m => MonadPdrState (PdrT mon a m) where
    getInitial       = lift  getInitial
    getProperty      = lift  getProperty
    getCurrentFrame  = lift  getCurrentFrame
    getNextFrame     = lift  getNextFrame
    getPreviousFrame = lift  getPreviousFrame
    setCurrentFrame  = lift1 setCurrentFrame
    goFrameUp        = lift  goFrameUp
    goFrameDown      = lift  goFrameDown
    goToFirstFrame   = lift  goToFirstFrame
    goToLastFrame    = lift  goToLastFrame
    isTopFrame       = lift  isTopFrame
    isBotFrame       = lift  isBotFrame
    pushFrame        = lift1 pushFrame
    getFrameNum      = lift  getFrameNum
    getFramesNum     = lift  getFramesNum

instance (MonadDomain d, MonadPdrState d, Monitor mon)
    => Monitored (PdrT mon a d) where
        notifyBad e = do
            mon <- PdrT ask
            updateBad mon e
        notifyCounterexample c = do
            mon <- PdrT ask
            updateCounterexample mon c
        notifyInvariant i = do
            mon <- PdrT ask
            updateInvariant mon i

-- PDR [4] (IC3 [1-3]) computes overapproximation of forward-reachable sets of
-- states (frames) in increasing distance from initial states. In each
-- iteration it checks reachability of bad states which violate target
-- property.
--
-- Predecessors of bad states are traced back toward initial states and
-- blocked at respective level when possible, i.e. when there is no further
-- predecessor at previous level.
--
-- When unreachability of bad states is shown at the top level (including
-- recursive reachability checks), inductive facts from all frames are
-- propagated to higher levels. This is done in hope to arrive at an inductive
-- invariant (an overapproximation of reachable states closed under
-- post-transformer) that satisfies the target property.
--
-- When predecessors of bad states reach the initial states, there is a
-- counterexample to the target property.
--
-- [1] A.R.Bradley. SAT-Based Model Checking without Unrolling. VMCAI 2011.
-- [2] A.R.Bradley. Understanding IC3. SAT 2012.
-- [3] A.Cimatti, A.Griggio, S.Mover, S.Tonetta. IC3 Modulo Theories via
--     Implicit Abstraction. TACAS 2014.
-- [4] N.Een, A.Mishchenko, R.Brayton. Efficient Implementation of Property
--     Directed Reachability. FMCAD 2011.
--
-- pdr implements the algorithm:
--   1) init: check whether initial states satisfy the property
--   2) loop: do until counterexample or invariant is reported
--      2.1) bad: extract bad (predecessor of an error state) at current level
--      2.2) block: attempt to block it recursively (may report counterexample)
--      2.3) induct: propagate inductive facts (may report invariant)
--
pdr :: (MonadSystem e s, Monitor mon) => Pdr mon s (Proof s) a
pdr = init >> loop (bad >>= block ><> induct) where

    -- Check zero-length counterexample, i.e. incidence of the initial states
    -- and bad states which violate the desired property
    init :: (MonadSystem e s, Monitor mon) => Pdr mon s (Proof s) ()
    init = do
        i <- getInitial
        p <- getProperty

        -- Setup context of the domain
        context [i, p]

        -- Initial intersection not Property
        bad <- exclude i p
        empty <- isEmpty bad

        if empty
        then pushFrame i
        else do
            trace <- ((: []) . head) <$> enumerate bad
            mapL Left $ counterexample trace

    counterexample :: (MonadSystem e s, Monitor mon)
                   => [Element s]
                   -> Pdr mon s (Counterexample s) ()
    counterexample c = do
        notifyCounterexample c
        PdrT . lift . throwE . Counterexample $ c

    invariant :: (MonadSystem e s, Monitor mon)
              => Set s
              -> Pdr mon s (Invariant s) ()
    invariant i = do
        notifyInvariant i
        PdrT . lift . throwE . Invariant $ i

    -- Loop forever until either a counterexample or an invariant is found.
    loop :: Monad m => m a -> m b
    loop = forever

    -- Extract a bad state reachable in one step from current depth.
    bad :: (MonadSystem e s, Monitor mon) => Pdr mon s (Proof s) (Set s)
    bad = do
        c <- getCurrentFrame
        p <- getProperty

        -- Assume:
        --   current frame
        --   property satisfied by the pre-state
        -- Require:
        --   property satisfied by any post-state
        --
        -- i.e.:
        --   s in Frame implies post(s) in Property
        --   post(Frame) subset of Property
        --
        -- Dually inspect for emptiness:
        --   post(Frame) intersection not Property
        --
        p' <- complement p
        c' <- post c

        c' `intersect` p'

    -- Blocks predecessors of a candidate bad state if possible. Violation of
    -- the property is reachable from bad. We may either succeed in showing
    -- that there is no further predecessor at lower frame or we need to prove
    -- unreachability of those predecessors recursively. If we hit the initial
    -- frame, there is an abstract counterexample.
    block :: (MonadSystem e s, Monitor mon)
          => Set s
          -> Pdr mon s (Counterexample s) ()
    block bad = catch (block' bad) analyseCounterexample where

        analyseCounterexample :: (MonadSystem e s, Monitor mon)
                              => Counterexample s
                              -> Pdr mon s (Counterexample s) ()
        analyseCounterexample (Counterexample trace) = do
                goToLastFrame
                r <- refute trace
                case r of
                    Just trace -> counterexample trace
                    _ -> return ()

        block' :: (MonadSystem e s, Monitor mon)
               => Set s
               -> Pdr mon s (Counterexample s) ()
        block' bad = mapM_ (block'' []) =<< enumerate bad

        -- Maintain a trace to witness reachability of the violation.
        block'' :: (MonadSystem e s, Monitor mon)
               => [Element s]
               -> Element s
               -> Pdr mon s (Counterexample s) ()
        block'' trace bad = do
            let trace' = bad : trace

            notifyBad bad

            bot    <- isBotFrame
            c      <- getCurrentFrame
            preBad <- intersect c =<< pre bad

            -- As long as the set of predecessors is realized in
            -- the current frame we keep searching for further
            -- predecessors.
            unlessM (isEmpty preBad) $
                if bot
                then do
                    -- At this point we have proven abstract reachability
                    -- of a violation of the property.
                    --
                    -- Finalise the witness.
                    trace'' <- (: trace') . head <$> enumerate preBad

                    counterexample trace''
                else mapM_ (block''' trace') =<< enumerate preBad

        -- Recurse and block the candidate in the higher frame.
        block''' :: (MonadSystem e s, Monitor mon)
                => [Element s]
                -> Element s
                -> Pdr mon s (Counterexample s) ()
        block''' trace bad = do
            goFrameDown
            block'' trace bad
            goFrameUp
            c' <- post =<< getPreviousFrame
            s  <- singleton =<< generalise c' bad
            modifyFrame (`exclude` s)

    -- Moves facts to higher frames to reduce deep blocking recursions in
    -- future as well as improve chances of finding an inductive invariant.
    induct :: (MonadSystem e s, Monitor mon) => Pdr mon s (Invariant s) ()
    induct = do
        b <- bottom
        t <- top
        p <- getProperty

        pushFrame p
        goToFirstFrame
        induct' b t where

            -- Carry inductive facts toward the top frame.
            induct' :: (MonadSystem e s, Monitor mon)
                    => Set s
                    -> Set s
                    -> Pdr mon s (Invariant s) ()
            induct' prev ind = do
                -- Strengthen current frame with inductive facts.
                modifyFrame (intersect ind)

                c <- getCurrentFrame

                -- Test for fixpoint.
                whenM (c `subsetOf` prev) $ invariant c

                -- Collect all states unreachable in k steps (those excluded
                -- from current frame) that cannot be reached from other states
                -- reachable in k steps (the current frame) in one step. Such
                -- states (unreachable in k+1 steps) can be excluded from the
                -- next frame. However, only states that are not yet excluded
                -- from the next frame need to be examined, wasteful
                -- examination could defeat the efficiency of the algorithm.
                unlessM isTopFrame $ do
                    t   <- top
                    c'  <- post c
                    c'' <- c `union` c'

                    -- blocked in next frame
                    next <- getNextFrame

                    -- blocked in current frame, but not next frame
                    excluded <- enumerate =<< intersect next =<< complement c

                    ind' <- flip filterM excluded $ \e -> do
                        s <- singleton e

                        -- reachable in one additional step
                        out <- s `intersect` c'

                        isEmpty out

                    -- collect all blocked cubes generalised w.r.t. the current and next frames
                    ind'' <- foldM exclude t =<< mapM (singleton <=< generalise c'') ind'

                    -- Send inductive facts to higher frames.
                    goFrameUp
                    induct' c ind''

    -- Auxiliary branching to handle occasional `e` or fallback
    (><>) :: (MonadDomain d, Monitor mon) => (Set d -> Pdr mon d a c)
          -> Pdr mon d b c
          -> Set d
          -> Pdr mon d (Either a b) c
    (><>) l r e = do
        empty <- isEmpty e
        if empty
        then mapL Right r
        else mapL Left (l e)

    -- Auxiliary merge of branches
    mapL :: (MonadDomain d, Monitor mon)
         => (a -> c)
         -> Pdr mon d a b
         -> Pdr mon d c b
    mapL f p = PdrT
             $ ReaderT $ \r ->
                ExceptT
                . fmap (either (Left . f) Right)
                . runExceptT
                . flip runReaderT r
                . runPdrT

                $ p

    catch :: (MonadDomain d, Monitor mon)
          => Pdr mon d e ()
          -> (e -> Pdr mon d e ())
          -> Pdr mon d e ()
    catch a h = PdrT $ ReaderT $ \r ->
        let a'   = flip runReaderT r . runPdrT $ a
            h' e = flip runReaderT r . runPdrT $ h e in catchE a' h'

-- Wrapper of the generic algorithm that feeds the initial configuration in and
-- retrieves the results.
runpdr :: (MonadDomain m, MonadSystem e m, Monitor mon)
       => mon -> Set m -> Set m -> m (Proof m)
runpdr m i p = do
    r <- runPdrStateT i p        -- unpack PdrState
       . fmap (either id id)     -- reconcile termination and non-termination
       . runExceptT              -- extract proof or do not terminate
       . flip runReaderT m       -- attach monitor
       . runPdrT                 -- unpack Pdr

       $ pdr                     -- perform Pdr

    waitFor m                    -- some monitors want to hang around

    return r
