BoundedExec.hs

-----------------------------------------------------------------------
-- |
-- Module           : Lang.Crucible.Simulator.BoundedExec
-- Description      : Support for bounding loop depth
-- Copyright        : (c) Galois, Inc 2018
-- License          : BSD3
-- Maintainer       : Rob Dockins <rdockins@galois.com>
-- Stability        : provisional
--
-- This module provides an execution feature for bounding the
-- number of iterations that a loop will execute in the simulator.
------------------------------------------------------------------------
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}

{-# OPTIONS_GHC -Wno-orphans #-}
module Lang.Crucible.Simulator.BoundedExec
  ( boundedExecFeature
  ) where

import           Control.Lens ( (^.), to, (&), (%~), (.~) )
import           Control.Monad ( when )
import           Data.IORef
import           Data.Map (Map)
import qualified Data.Map as Map
import           Data.Maybe (fromMaybe)
import           Data.Sequence (Seq)
import qualified Data.Sequence as Seq
import qualified Data.Text as Text
import           Data.Word


import qualified Data.Parameterized.Context as Ctx
import qualified Data.Parameterized.Map as MapF

import           Lang.Crucible.Analysis.Fixpoint.Components
import           Lang.Crucible.Backend
import           Lang.Crucible.CFG.Core
import           Lang.Crucible.FunctionHandle
import           Lang.Crucible.Panic
import           Lang.Crucible.Simulator.CallFrame
import           Lang.Crucible.Simulator.ExecutionTree
import           Lang.Crucible.Simulator.GlobalState
import           Lang.Crucible.Simulator.Intrinsics
import           Lang.Crucible.Simulator.EvalStmt
import           Lang.Crucible.Simulator.SimError

import           What4.FunctionName
import           What4.Interface

data FrameBoundData =
  forall args ret.
    FrameBoundData
    { frameBoundHandle :: !(FnHandle args ret)
    , frameBoundLimit :: !Word64
    , frameWtoMap :: !(Map Int (Int,Int))
    , frameBoundCounts :: Seq Word64
    }

-- | This function takes weak topological order data and computes
--   a mapping from block ID number to (position, depth) pair.  The
--   position value indicates which the position in the WTO listing
--   in which the block ID appears, and the depth indicates the number
--   of nested components the block ID appears in.  Loop backedges
--   occur exactly at places where control flows from a higher position
--   number to a lower position number.  Jumps that exit inner loops
--   to the next iteration of an outer loop are identified by backedges
--   that pass from higher depths to lower depths.
buildWTOMap :: [WTOComponent (Some (BlockID blocks))] -> Map Int (Int,Int)
buildWTOMap = snd . go 0 0 Map.empty
 where
 go :: Int -> Int -> Map Int (Int,Int) -> [WTOComponent (Some (BlockID blocks))] -> (Int, Map Int (Int,Int))
 go !x !_ m [] = (x,m)
 go !x !d m (Vertex (Some bid) : cs) =
    let m' = Map.insert (Ctx.indexVal (blockIDIndex bid)) (x,d) m
     in go (x+1) d m' cs
 go !x !d m (SCC scc : cs) =
    let m'  = viewSome (\hd -> Map.insert (Ctx.indexVal (blockIDIndex hd)) (x,d+1) m) (wtoHead scc)
        (x',m'') = go (x+1) (d+1) m' $ wtoComps scc
     in go x' d m'' cs


-- | This function updates the loop bound count at the given depth.
--   Any loop bounds deeper than this are discarded.  If the given
--   sequence is too short to accommodate the given depth, the sequence
--   is extended with 0 counters to the correct depth.
incrementBoundCount :: Seq Word64 -> Int -> (Seq Word64, Word64)
incrementBoundCount cs depth =
  case Seq.lookup depth cs of
     Just n ->
       do let n' = n+1
          let cs' = Seq.update depth n' $ Seq.take (depth+1) cs
          n' `seq` cs' `seq` (cs', n')
     Nothing ->
       do let cs' = cs <> Seq.replicate (depth - Seq.length cs) 0 <> Seq.singleton 1
          cs' `seq` (cs', 1)

instance IntrinsicClass sym "BoundedExecFrameData" where
  type Intrinsic sym "BoundedExecFrameData" ctx = [Either FunctionName FrameBoundData]

  muxIntrinsic _sym _iTypes _nm _ _p fd1 fd2 = combineFrameBoundData fd1 fd2

mergeCounts :: Seq Word64 -> Seq Word64 -> Seq Word64
mergeCounts cx cy =
  Seq.fromFunction
    (max (Seq.length cx) (Seq.length cy))
    (\i -> max (fromMaybe 0 $ Seq.lookup i cx)
               (fromMaybe 0 $ Seq.lookup i cy))

mergeFBD ::
  FrameBoundData ->
  FrameBoundData ->
  IO FrameBoundData
mergeFBD x@FrameBoundData{ frameBoundHandle = hx } y@FrameBoundData{ frameBoundHandle = hy }
  | Just _ <- testEquality (handleID hx) (handleID hy) =
       return x{ frameBoundCounts = mergeCounts (frameBoundCounts x) (frameBoundCounts y) }

  | otherwise =
       panic "BoundedExec.mergeFBD"
       [ "Attempted to merge frame bound data from different function activations: "
       , " ** " ++ show hx
       , " ** " ++ show hy
       ]


combineFrameBoundData ::
  [Either FunctionName FrameBoundData] ->
  [Either FunctionName FrameBoundData] ->
  IO [Either FunctionName FrameBoundData]
combineFrameBoundData [] [] = return []

combineFrameBoundData (Left nmx:xs) (Left nmy : _) | nmx == nmy
  = return (Left nmx : xs)

combineFrameBoundData (Right x:xs) (Right y:_)
  = (\x' -> Right x' : xs) <$> mergeFBD x y

combineFrameBoundData xs ys
  = panic "BoundedExec.combineFrameBoundData"
      [ "Attempt to combine incompatible frame bound data: stack shape mismatch:"
      , " *** " ++ show (printStack xs)
      , " *** " ++ show (printStack ys)
      ]

printStack :: [Either FunctionName FrameBoundData] -> [String]
printStack [] = []
printStack (Left nm :xs) = show nm : printStack xs
printStack (Right FrameBoundData{ frameBoundHandle = h } : xs) = show h : printStack xs


type BoundedExecGlobal = GlobalVar (IntrinsicType "BoundedExecFrameData" EmptyCtx)


-- | This execution feature allows users to place a bound on the number
--   of iterations that a loop will execute.  Each time a function is called,
--   the included action is called to determine if the loops in that function
--   should be bounded, and what their iteration bound should be.
--
--   The boolean argument indicates if we should generate proof obligations when
--   we cut off loop execution.  If true, loop cutoffs will generate proof obligations
--   which will be provable only if the loop actually could not have executed that number
--   of iterations.  If false, the execution of loops will be aborted without generating
--   side conditions.
--
--   Note that we compute a weak topological ordering on control flow graphs
--   to determine loop heads and loop nesting structure.  Loop bounds for inner
--   loops are reset on every iteration through an outer loop.
boundedExecFeature ::
  (SomeHandle -> IO (Maybe Word64))
    {- ^ Action for computing loop bounds for functions when they are called -} ->
  Bool {- ^ Produce a proof obligation when resources are exhausted? -} ->
  IO (GenericExecutionFeature sym)
boundedExecFeature getLoopBounds generateSideConditions =
  do gvRef <- newIORef (error "Global variable for BoundedExecFrameData not initialized")
     return $ GenericExecutionFeature $ onStep gvRef

 where
 buildFrameData :: ResolvedCall p sym ext ret -> IO (Either FunctionName FrameBoundData)
 buildFrameData (OverrideCall ov _) = return (Left (overrideName ov))
 buildFrameData (CrucibleCall _entry CallFrame{ _frameCFG = g }) =
   do let wtoMap = buildWTOMap (cfgWeakTopologicalOrdering g)
      mn <- getLoopBounds (SomeHandle (cfgHandle g))
      case mn of
        Nothing -> return $ Left  $ handleName (cfgHandle g)
        Just n  -> return $ Right $ FrameBoundData
                       { frameBoundHandle = cfgHandle g
                       , frameBoundLimit  = n
                       , frameWtoMap      = wtoMap
                       , frameBoundCounts = mempty
                       }

 checkBackedge ::
   IORef BoundedExecGlobal ->
   Some (BlockID blocks) ->
   BlockID blocks tgt_args ->
   SymGlobalState sym ->
   IO (SymGlobalState sym, Maybe Word64)
 checkBackedge gvRef (Some bid_curr) bid_tgt globals =
   do gv <- readIORef gvRef
      case fromMaybe [] (lookupGlobal gv globals) of
        ( Right fbd : rest ) ->
          do let id_curr = Ctx.indexVal (blockIDIndex bid_curr)
             let id_tgt  = Ctx.indexVal (blockIDIndex bid_tgt)
             let m = frameWtoMap fbd
             case (Map.lookup id_curr m, Map.lookup id_tgt m) of
               (Just (cx, _cd), Just (tx, td)) | tx <= cx ->
                  do let cs       = frameBoundCounts fbd
                     let (cs', q) = incrementBoundCount cs td
                     let fbd'     = fbd{ frameBoundCounts = cs' }
                     let globals' = insertGlobal gv (Right fbd' : rest) globals
                     if q > frameBoundLimit fbd then
                       return (globals', Just (frameBoundLimit fbd))
                     else
                       return (globals', Nothing)

               _ -> return (globals, Nothing)
        _ -> return (globals, Nothing)

 modifyStackState ::
   IORef BoundedExecGlobal ->
   (SimState p sym ext rtp f args -> ExecState p sym ext rtp) ->
   SimState p sym ext rtp f args ->
   ([Either FunctionName FrameBoundData] -> [Either FunctionName FrameBoundData]) ->
   IO (ExecutionFeatureResult p sym ext rtp)
 modifyStackState gvRef mkSt st f =
   do gv <- readIORef gvRef
      let xs = case lookupGlobal gv (st ^. stateGlobals) of
                 Nothing -> error "bounded execution global not defined!"
                 Just v  -> v
      let st' = st & stateGlobals %~ insertGlobal gv (f xs)
      return (ExecutionFeatureModifiedState (mkSt st'))

 onTransition ::
   IORef BoundedExecGlobal ->
   BlockID blocks tgt_args ->
   ControlResumption p sym ext rtp (CrucibleLang blocks ret) ->
   SimState p sym ext rtp (CrucibleLang blocks ret) ('Just a) ->
   IO (ExecutionFeatureResult p sym ext rtp)
 onTransition gvRef tgt_id res st = stateSolverProof st $
  do let sym = st^.stateSymInterface
     let simCtx = st^.stateContext
     (globals', overLimit) <- checkBackedge gvRef (st^.stateCrucibleFrame.frameBlockID) tgt_id (st^.stateGlobals)
     let st' = st & stateGlobals .~ globals'
     case overLimit of
       Just n ->
         do let msg = "reached maximum number of loop iterations (" ++ show n ++ ")"
            let loc = st^.stateCrucibleFrame.to frameProgramLoc
            let err = SimError loc (ResourceExhausted msg)
            when generateSideConditions $ withBackend simCtx $ \bak ->
              addProofObligation bak (LabeledPred (falsePred sym) err)
            return (ExecutionFeatureNewState (AbortState (AssertionFailure err) st'))
       Nothing -> return (ExecutionFeatureModifiedState (ControlTransferState res st'))

 onStep ::
   IORef BoundedExecGlobal ->
   ExecState p sym ext rtp ->
   IO (ExecutionFeatureResult p sym ext rtp)

 onStep gvRef = \case
   InitialState simctx globals ah ret cont ->
     do let halloc = simHandleAllocator simctx
        gv <- freshGlobalVar halloc (Text.pack "BoundedExecFrameData") knownRepr
        writeIORef gvRef gv
        let globals' = insertGlobal gv [Left "_init"] globals
        let simctx' = simctx{ ctxIntrinsicTypes = MapF.insert (knownSymbol @"BoundedExecFrameData") IntrinsicMuxFn (ctxIntrinsicTypes simctx) }
        return (ExecutionFeatureModifiedState (InitialState simctx' globals' ah ret cont))

   CallState rh call st ->
     do boundData <- buildFrameData call
        modifyStackState gvRef (CallState rh call) st (boundData:)

   TailCallState vfv call st ->
     do boundData <- buildFrameData call
        modifyStackState gvRef (TailCallState vfv call) st ((boundData:) . drop 1)

   ReturnState nm vfv pr st ->
        modifyStackState gvRef (ReturnState nm vfv pr) st (drop 1)

   UnwindCallState vfv ar st ->
        modifyStackState gvRef (UnwindCallState vfv ar) st (drop 1)

   ControlTransferState res st ->
     case res of
       ContinueResumption (ResolvedJump tgt_id _)  ->  onTransition gvRef tgt_id res st
       CheckMergeResumption (ResolvedJump tgt_id _) -> onTransition gvRef tgt_id res st
       _ -> return ExecutionFeatureNoChange

   _ -> return ExecutionFeatureNoChange