-----------------------------------------------------------------------
-- |
-- Module           : Lang.Crucible.Simulator.BoundedRecursion
-- Description      : Support for bounding function recursion depth
-- Copyright        : (c) Galois, Inc 2019
-- License          : BSD3
-- Maintainer       : Rob Dockins <rdockins@galois.com>
-- Stability        : provisional
--
-- This module provides an execution feature for bounding recursion.
-- Essentially, we bound the number of times any particular function
-- is allowed to have active frames on the call stack.
------------------------------------------------------------------------

{-# 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.BoundedRecursion
  ( boundedRecursionFeature
  ) where

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

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

import           What4.Interface

import           Lang.Crucible.Backend
import           Lang.Crucible.CFG.Common
import           Lang.Crucible.FunctionHandle
import           Lang.Crucible.Panic
import           Lang.Crucible.Simulator.SimError
import           Lang.Crucible.Simulator.ExecutionTree
import           Lang.Crucible.Simulator.EvalStmt
import           Lang.Crucible.Simulator.Intrinsics
import           Lang.Crucible.Simulator.GlobalState
import           Lang.Crucible.Types

type BoundedRecursionMap = Map.Map SomeHandle Word64

instance IntrinsicClass sym "BoundedRecursionData" where
  type Intrinsic sym "BoundedRecursionData" ctx = [BoundedRecursionMap]
  muxIntrinsic _sym _iTypes _nm _ _p x _y = return x

type BoundedRecursionGlobal = GlobalVar (IntrinsicType "BoundedRecursionData" EmptyCtx)

-- | This execution feature allows users to place a bound on the number of
--   recursive calls that a function can execute.  Each time a function is
--   called, the number of activations of the functions is incremented, and
--   the path is aborted if the bound is exceeded.
--
--   The boolean argument indicates if we should generate proof obligations when
--   we cut off recursion.  If true, recursion cutoffs will generate proof obligations
--   which will be provable only if the function actually could not have executed that number
--   of times.  If false, the execution of recursive functions will be aborted without
--   generating side conditions.
boundedRecursionFeature ::
  (SomeHandle -> IO (Maybe Word64))
    {- ^ Action for computing what recursion depth to allow for the given function -}  ->
  Bool {- ^ Produce a proof obligation when resources are exhausted? -} ->
  IO (GenericExecutionFeature sym)

boundedRecursionFeature getRecursionBound generateSideConditions =
  do gvRef <- newIORef (error "Global variable for BoundedRecursionData not initialized")
     return $ GenericExecutionFeature $ onStep gvRef

 where
 popFrame ::
   IORef BoundedRecursionGlobal ->
   (SimState p sym ext rtp f args -> ExecState p sym ext rtp) ->
   SimState p sym ext rtp f args ->
   IO (ExecutionFeatureResult p sym ext rtp)
 popFrame gvRef mkSt st =
   do gv <- readIORef gvRef
      case lookupGlobal gv (st ^. stateGlobals) of
        Nothing -> panic "bounded recursion" ["global not defined!"]
        Just [] -> panic "bounded recursion" ["pop on empty stack!"]
        Just (_:xs) ->
          do let st' = st & stateGlobals %~ insertGlobal gv xs
             return (ExecutionFeatureModifiedState (mkSt st'))

 pushFrame ::
   IORef BoundedRecursionGlobal ->
   (BoundedRecursionMap -> BoundedRecursionMap -> [BoundedRecursionMap] -> [BoundedRecursionMap]) ->
   SomeHandle ->
   (SimState p sym ext rtp f args -> ExecState p sym ext rtp) ->
   SimState p sym ext rtp f args ->
   IO (ExecutionFeatureResult p sym ext rtp)
 pushFrame gvRef rebuildStack h mkSt st = stateSolverProof st $
     do let sym = st^.stateSymInterface
        let simCtx = st^.stateContext
        gv <- readIORef gvRef
        case lookupGlobal gv (st ^. stateGlobals) of
          Nothing -> panic "bounded recursion" ["global not defined!"]
          Just [] -> panic "bounded recursion" ["empty stack!"]
          Just (x:xs) ->
            do mb <- getRecursionBound h
               let v = 1 + fromMaybe 0 (Map.lookup h x)
               case mb of
                 Just b | v > b ->
                   do loc <- getCurrentProgramLoc sym
                      let msg = ("reached maximum number of recursive calls to function " ++ show h ++ " (" ++ show b ++ ")")
                      let err = SimError loc (ResourceExhausted msg)
                      when generateSideConditions $ withBackend simCtx $ \bak ->
                        addProofObligation bak (LabeledPred (falsePred sym) err)
                      return (ExecutionFeatureNewState (AbortState (AssertionFailure err) st))
                 _ ->
                   do let x'  = Map.insert h v x
                      let st' = st & stateGlobals %~ insertGlobal gv (rebuildStack x' x xs)
                      x' `seq` return (ExecutionFeatureModifiedState (mkSt st'))

 onStep ::
   IORef BoundedRecursionGlobal ->
   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 "BoundedRecursionData") knownRepr
        writeIORef gvRef gv
        let simctx'  = simctx{ ctxIntrinsicTypes = MapF.insert
                                   (knownSymbol @"BoundedRecursionData")
                                   IntrinsicMuxFn
                                   (ctxIntrinsicTypes simctx) }
        let globals' = insertGlobal gv [mempty] globals
        return (ExecutionFeatureModifiedState (InitialState simctx' globals' ah ret cont))

   CallState rh call st ->
     pushFrame gvRef (\a b xs -> a:b:xs) (resolvedCallHandle call) (CallState rh call) st

   TailCallState vfv call st ->
     pushFrame gvRef (\a _ xs -> a:xs) (resolvedCallHandle call) (TailCallState vfv call) st

   ReturnState nm vfv pr st ->
     popFrame gvRef (ReturnState nm vfv pr) st

   UnwindCallState vfv ar st ->
     popFrame gvRef (UnwindCallState vfv ar) st

   _ -> return ExecutionFeatureNoChange