content: {-|
  Module      : Lang.Crucible.Backend.Assumptions
  Copyright   : (c) Galois, Inc 2014-2024
  License     : BSD3
  Maintainer  : Langston Barrett <langston@galois.com>
  -}
  
  {-# LANGUAGE ConstraintKinds #-}
  {-# LANGUAGE DataKinds #-}
  {-# LANGUAGE DeriveFunctor #-}
  {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE FunctionalDependencies #-}
  {-# LANGUAGE GADTs #-}
  {-# LANGUAGE LambdaCase #-}
  {-# LANGUAGE MultiParamTypeClasses #-}
  {-# LANGUAGE PolyKinds #-}
  {-# LANGUAGE RankNTypes #-}
  {-# LANGUAGE OverloadedStrings #-}
  {-# LANGUAGE TypeFamilies #-}
  {-# LANGUAGE TypeOperators #-}
  {-# LANGUAGE ViewPatterns #-}
  
  module Lang.Crucible.Backend.Assumptions
    ( CrucibleAssumption(..)
    , CrucibleEvent(..)
    , CrucibleAssumptions(..)
    , Assumption
    , Assumptions
  
    , concretizeEvents
    , ppEvent
    , singleEvent
    , singleAssumption
    , trivialAssumption
    , ppAssumption
    , assumptionLoc
    , eventLoc
    , mergeAssumptions
    , assumptionPred
    , forgetAssumption
    , assumptionsPred
    , flattenAssumptions
    , assumptionsTopLevelLocs
    , ppAssumptions'
    , ppAssumptions
    ) where
  
  
  import           Control.Lens (Traversal, folded)
  import           Data.Kind (Type)
  import qualified Data.Foldable as F
  import           Data.Functor.Identity
  import           Data.Functor.Const
  import qualified Data.Parameterized.TraversableF as TF
  import qualified Data.Sequence as Seq
  import           Data.Sequence (Seq)
  import qualified Prettyprinter as PP
  
  import           What4.Expr.Builder
  import           What4.Interface
  import           What4.ProgramLoc
  import           What4.Expr (GroundValue, GroundValueWrapper(..))
  
  import           Lang.Crucible.Simulator.SimError
  
  type Assumption sym = CrucibleAssumption (SymExpr sym)
  type Assumptions sym = CrucibleAssumptions (SymExpr sym)
  
  -- | This type describes assumptions made at some point during program execution.
  data CrucibleAssumption (e :: BaseType -> Type)
    = GenericAssumption ProgramLoc String (e BaseBoolType)
      -- ^ An unstructured description of the source of an assumption.
  
    | BranchCondition ProgramLoc (Maybe ProgramLoc) (e BaseBoolType)
      -- ^ This arose because we want to explore a specific path.
      -- The first location is the location of the branch predicate.
      -- The second one is the location of the branch target.
  
    | AssumingNoError SimError (e BaseBoolType)
      -- ^ An assumption justified by a proof of the impossibility of
      -- a certain simulator error.
  
  instance TF.FunctorF CrucibleAssumption where
    fmapF = TF.fmapFDefault
  instance TF.FoldableF CrucibleAssumption where
    foldMapF = TF.foldMapFDefault
  instance TF.TraversableF CrucibleAssumption where
    traverseF = traverseAssumption
  
  -- | This type describes events we can track during program execution.
  data CrucibleEvent (e :: BaseType -> Type) where
    -- | This event describes the creation of a symbolic variable.
    CreateVariableEvent ::
      ProgramLoc {- ^ location where the variable was created -} ->
      String {- ^ user-provided name for the variable -} ->
      BaseTypeRepr tp {- ^ type of the variable -} ->
      e tp {- ^ the variable expression -} ->
      CrucibleEvent e
  
    -- | This event describes reaching a particular program location.
    LocationReachedEvent ::
      ProgramLoc ->
      CrucibleEvent e
  
  instance TF.FunctorF CrucibleEvent where
    fmapF = TF.fmapFDefault
  instance TF.FoldableF CrucibleEvent where
    foldMapF = TF.foldMapFDefault
  instance TF.TraversableF CrucibleEvent where
    traverseF = traverseEvent
  
  -- | Pretty print an event
  ppEvent' :: (forall t. e t -> PP.Doc ann) -> CrucibleEvent e -> PP.Doc ann
  ppEvent' ppExp (CreateVariableEvent loc nm _tpr v) =
    "create var" PP.<+> PP.pretty nm PP.<+> "=" PP.<+> ppExp v PP.<+> "at" PP.<+> PP.pretty (plSourceLoc loc)
  ppEvent' _ppExp (LocationReachedEvent loc) =
    "reached" PP.<+> PP.pretty (plSourceLoc loc) PP.<+> "in" PP.<+> PP.pretty (plFunction loc)
  
  -- | Pretty print an event
  ppEvent :: IsExpr e => CrucibleEvent e -> PP.Doc ann
  ppEvent = ppEvent' printSymExpr
  
  -- | Return the program location associated with an event
  eventLoc :: CrucibleEvent e -> ProgramLoc
  eventLoc (CreateVariableEvent loc _ _ _) = loc
  eventLoc (LocationReachedEvent loc) = loc
  
  -- | Return the program location associated with an assumption
  assumptionLoc :: CrucibleAssumption e -> ProgramLoc
  assumptionLoc r =
    case r of
      GenericAssumption l _ _ -> l
      BranchCondition  l _ _   -> l
      AssumingNoError s _    -> simErrorLoc s
  
  -- | Get the predicate associated with this assumption
  assumptionPred :: CrucibleAssumption e -> e BaseBoolType
  assumptionPred (AssumingNoError _ p) = p
  assumptionPred (BranchCondition _ _ p) = p
  assumptionPred (GenericAssumption _ _ p) = p
  
  forgetAssumption :: CrucibleAssumption e -> CrucibleAssumption (Const ())
  forgetAssumption = runIdentity . traverseAssumption (\_ -> Identity (Const ()))
  
  -- | Check if an assumption is trivial (always true)
  trivialAssumption :: IsExpr e => CrucibleAssumption e -> Bool
  trivialAssumption a = asConstantPred (assumptionPred a) == Just True
  
  traverseAssumption :: Traversal (CrucibleAssumption e) (CrucibleAssumption e') (e BaseBoolType) (e' BaseBoolType)
  traverseAssumption f = \case
    GenericAssumption loc msg p -> GenericAssumption loc msg <$> f p
    BranchCondition l t p -> BranchCondition l t <$> f p
    AssumingNoError err p -> AssumingNoError err <$> f p
  
  -- | This type tracks both logical assumptions and program events
  --   that are relevant when evaluating proof obligations arising
  --   from simulation.
  data CrucibleAssumptions (e :: BaseType -> Type) where
    SingleAssumption :: CrucibleAssumption e -> CrucibleAssumptions e
    SingleEvent      :: CrucibleEvent e -> CrucibleAssumptions e
    ManyAssumptions  :: Seq (CrucibleAssumptions e) -> CrucibleAssumptions e
    MergeAssumptions ::
      e BaseBoolType {- ^ branch condition -} ->
      CrucibleAssumptions e {- ^ "then" assumptions -} ->
      CrucibleAssumptions e {- ^ "else" assumptions -} ->
      CrucibleAssumptions e
  
  instance Semigroup (CrucibleAssumptions e) where
    ManyAssumptions xs <> ManyAssumptions ys = ManyAssumptions (xs <> ys)
    ManyAssumptions xs <> y = ManyAssumptions (xs Seq.|> y)
    x <> ManyAssumptions ys = ManyAssumptions (x Seq.<| ys)
    x <> y = ManyAssumptions (Seq.fromList [x,y])
  
  instance Monoid (CrucibleAssumptions e) where
    mempty = ManyAssumptions mempty
  
  instance TF.FunctorF CrucibleAssumptions where
    fmapF = TF.fmapFDefault
  instance TF.FoldableF CrucibleAssumptions where
    foldMapF = TF.foldMapFDefault
  instance TF.TraversableF CrucibleAssumptions where
    traverseF f = \case
      SingleAssumption a ->
        SingleAssumption <$> TF.traverseF f a
      SingleEvent e ->
        SingleEvent <$> TF.traverseF f e
      ManyAssumptions xs ->
        ManyAssumptions <$> traverse (TF.traverseF f) xs
      MergeAssumptions c xs ys ->
        MergeAssumptions <$> f c <*> TF.traverseF f xs <*> TF.traverseF f ys
  
  singleAssumption :: CrucibleAssumption e -> CrucibleAssumptions e
  singleAssumption x = SingleAssumption x
  
  singleEvent :: CrucibleEvent e -> CrucibleAssumptions e
  singleEvent x = SingleEvent x
  
  -- | Collect the program locations of all assumptions and
  --   events that did not occur in the context of a symbolic branch.
  --   These are locations that every program path represented by
  --   this @CrucibleAssumptions@ structure must have passed through.
  assumptionsTopLevelLocs :: CrucibleAssumptions e -> [ProgramLoc]
  assumptionsTopLevelLocs (SingleEvent e)      = [eventLoc e]
  assumptionsTopLevelLocs (SingleAssumption a) = [assumptionLoc a]
  assumptionsTopLevelLocs (ManyAssumptions as) = concatMap assumptionsTopLevelLocs as
  assumptionsTopLevelLocs MergeAssumptions{}   = []
  
  -- | Compute the logical predicate corresponding to this collection of assumptions.
  assumptionsPred :: IsExprBuilder sym => sym -> Assumptions sym -> IO (Pred sym)
  assumptionsPred sym (SingleEvent _) =
    return (truePred sym)
  assumptionsPred _sym (SingleAssumption a) =
    return (assumptionPred a)
  assumptionsPred sym (ManyAssumptions xs) =
    andAllOf sym folded =<< traverse (assumptionsPred sym) xs
  assumptionsPred sym (MergeAssumptions c xs ys) =
    do xs' <- assumptionsPred sym xs
       ys' <- assumptionsPred sym ys
       itePred sym c xs' ys'
  
  traverseEvent :: Applicative m =>
    (forall tp. e tp -> m (e' tp)) ->
    CrucibleEvent e -> m (CrucibleEvent e')
  traverseEvent f (CreateVariableEvent loc nm tpr v) = CreateVariableEvent loc nm tpr <$> f v
  traverseEvent _ (LocationReachedEvent loc) = pure (LocationReachedEvent loc)
  
  -- | Given a ground evaluation function, compute a linear, ground-valued
  --   sequence of events corresponding to this program run.
  concretizeEvents ::
    IsExpr e =>
    (forall tp. e tp -> IO (GroundValue tp)) ->
    CrucibleAssumptions e ->
    IO [CrucibleEvent GroundValueWrapper]
  concretizeEvents f = loop
    where
      loop (SingleEvent e) =
        do e' <- traverseEvent (\v -> GVW <$> f v) e
           return [e']
      loop (SingleAssumption _) = return []
      loop (ManyAssumptions as) = concat <$> traverse loop as
      loop (MergeAssumptions p xs ys) =
        do b <- f p
           if b then loop xs else loop ys
  
  -- | Given a @CrucibleAssumptions@ structure, flatten all the muxed assumptions into
  --   a flat sequence of assumptions that have been appropriately weakened.
  --   Note, once these assumptions have been flattened, their order might no longer
  --   strictly correspond to any concrete program run.
  flattenAssumptions :: IsExprBuilder sym => sym -> Assumptions sym -> IO [Assumption sym]
  flattenAssumptions sym = loop Nothing
    where
      loop _mz (SingleEvent _) = return []
      loop mz (SingleAssumption a) =
        do a' <- maybe (pure a) (\z -> traverseAssumption (impliesPred sym z) a) mz
           if trivialAssumption a' then return [] else return [a']
      loop mz (ManyAssumptions as) =
        concat <$> traverse (loop mz) as
      loop mz (MergeAssumptions p xs ys) =
        do pnot <- notPred sym p
           px <- maybe (pure p) (andPred sym p) mz
           py <- maybe (pure pnot) (andPred sym pnot) mz
           xs' <- loop (Just px) xs
           ys' <- loop (Just py) ys
           return (xs' <> ys')
  
  -- | Merge the assumptions collected from the branches of a conditional.
  mergeAssumptions ::
    IsExprBuilder sym =>
    sym ->
    Pred sym ->
    Assumptions sym ->
    Assumptions sym ->
    IO (Assumptions sym)
  mergeAssumptions _sym p thens elses =
    return (MergeAssumptions p thens elses)
  
  ppAssumption :: (forall tp. e tp -> PP.Doc ann) -> CrucibleAssumption e -> PP.Doc ann
  ppAssumption ppDoc e =
    -- TODO(lb): These should really all be `align`ed, but that breaks a bunch
    -- of tests.
    case e of
      GenericAssumption l msg p ->
        PP.vsep [ ppLocated l (PP.pretty msg)
                , ppDoc p
                ]
      BranchCondition l Nothing p ->
        PP.vsep [ "The branch in" PP.<+> ppFn l PP.<+> "at" PP.<+> ppLoc l
                , ppDoc p
                ]
      BranchCondition l (Just t) p ->
        PP.vsep [ "The branch in" PP.<+> ppFn l PP.<+> "from" PP.<+> ppLoc l PP.<+> "to" PP.<+> ppLoc t
                , ppDoc p
                ]
      AssumingNoError simErr p ->
        PP.vsep [ "Assuming the following error does not occur:"
                , PP.indent 2 (ppSimError simErr)
                , ppDoc p
                ]
    where
      ppLocated :: ProgramLoc -> PP.Doc ann -> PP.Doc ann
      ppLocated l x = "in" PP.<+> ppFn l PP.<+> ppLoc l PP.<> ":" PP.<+> x
  
      ppFn :: ProgramLoc -> PP.Doc ann
      ppFn l = PP.pretty (plFunction l)
  
      ppLoc :: ProgramLoc -> PP.Doc ann
      ppLoc l = PP.pretty (plSourceLoc l)
  
  -- | Pretty-print 'CrucibleAssumptions'.
  ppAssumptions' ::
    -- | How to print expressions. If @'IsExpr' e@ holds, then see 'ppAssumptions'
    -- for a version that uses 'printSymExpr'.
    (forall tp. e tp -> PP.Doc ann) ->
    CrucibleAssumptions e ->
    PP.Doc ann
  ppAssumptions' ppExp =
    \case
      SingleAssumption asmp -> ppAssumption ppExp asmp
      SingleEvent e -> ppEvent' ppExp e
      ManyAssumptions asmps -> PP.list (map (ppAssumptions' ppExp) (F.toList asmps))
      MergeAssumptions b thn els ->
        PP.align $
          PP.vcat
          [ "if " <> PP.align (ppExp b)
          , "then " <> PP.align (ppAssumptions' ppExp thn)
          , "else " <> PP.align (ppAssumptions' ppExp els)
          ]
  
  -- | @'ppAssumptions' = `ppAssumptions'` 'printSymExpr'@
  ppAssumptions :: IsExpr e => CrucibleAssumptions e -> PP.Doc ann
  ppAssumptions = ppAssumptions' printSymExpr
  