content: {-|
  Module      : Lang.Crucible.Backend.Goals
  Copyright   : (c) Galois, Inc 2025
  License     : BSD3
  
  This module defines a data strucutre for storing a collection of
  proof obligations, and the current state of assumptions.
  -}
  
  {-# LANGUAGE LambdaCase #-}
  {-# LANGUAGE RankNTypes #-}
  {-# LANGUAGE TupleSections #-}
  
  module Lang.Crucible.Backend.Goals
    ( ProofGoal(..)
    , Goals(..)
    , ppGoals
    , goalsToList
    , assuming
    , proveAll
    , goalsConj
      -- * Traversals
    , traverseGoals
    , traverseOnlyGoals
    , traverseGoalsWithAssumptions
    , traverseGoalsSeq
    )
    where
  
  import           Control.Monad.Reader (ReaderT(..), withReaderT)
  import           Data.Functor.Const (Const(..))
  import           Data.Sequence (Seq)
  import qualified Data.Sequence as Seq
  import qualified Prettyprinter as PP
  
  -- | A proof goal consists of a collection of assumptions
  --   that were in scope when an assertion was made, together
  --   with the given assertion.
  data ProofGoal asmp goal =
    ProofGoal
    { proofAssumptions :: asmp
    , proofGoal        :: goal
    }
  
  -- | A collection of goals, which can represent shared assumptions.
  data Goals asmp goal =
      -- | Make an assumption that is in context for all the
      --   contained goals.
      Assuming asmp !(Goals asmp goal)
  
      -- | A proof obligation, to be proved in the context of
      --   all previously-made assumptions.
    | Prove goal
  
      -- | A conjunction of two goals.
    | ProveConj !(Goals asmp goal) !(Goals asmp goal)
      deriving Show
  
  -- | Intended for debugging, this is not generally a user-facing datatype.
  ppGoals ::
    (asmp -> PP.Doc ann) ->
    (goal -> PP.Doc ann) ->
    Goals asmp goal ->
    PP.Doc ann
  ppGoals ppAsmp ppGoal =
    \case
      Assuming asmp gls ->
        PP.align $
          PP.vcat
          [ PP.pretty "Assuming:"
          , PP.indent 2 (ppAsmp asmp)
          , PP.pretty "Prove:"
          , PP.indent 2 (ppGoals ppAsmp ppGoal gls)
          ]
      Prove gl -> ppGoal gl
      ProveConj gls gls' ->
        PP.align $
          PP.vcat
          [ PP.pretty "Prove both:"
          , PP.indent 2 (ppGoals ppAsmp ppGoal gls)
          , PP.indent 2 (ppGoals ppAsmp ppGoal gls')
          ]
  
  -- | Intended for debugging, this is not generally a user-facing datatype.
  instance (PP.Pretty asmp, PP.Pretty goal) => PP.Pretty (Goals asmp goal) where
    pretty = ppGoals PP.pretty PP.pretty
  
  -- | Construct a goal that first assumes a collection of
  --   assumptions and then states a goal.
  assuming :: Monoid asmp => asmp -> Goals asmp goal -> Goals asmp goal
  assuming as (Assuming bs g) = assuming (as <> bs) g
  assuming as g = Assuming as g
  
  -- | Construct a 'Goals' object from a collection of subgoals, all of
  --   which are to be proved.  This yields 'Nothing' if the collection
  --   of goals is empty, and otherwise builds a conjunction of all the
  --   goals.  Note that there is no new sharing in the resulting structure.
  proveAll :: Foldable t => t (Goals asmp goal) -> Maybe (Goals asmp goal)
  proveAll = foldr f Nothing
   where
   f x Nothing  = Just $! x
   f x (Just y) = Just $! ProveConj x y
  
  -- | Helper to conjoin two possibly trivial 'Goals' objects.
  goalsConj :: Maybe (Goals asmp goal) -> Maybe (Goals asmp goal) -> Maybe (Goals asmp goal)
  goalsConj Nothing y = y
  goalsConj x Nothing = x
  goalsConj (Just x) (Just y) = Just (ProveConj x y)
  
  -- | Render the tree of goals as a list instead, duplicating
  --   shared assumptions over each goal as necessary.
  goalsToList :: Monoid asmp => Goals asmp goal -> [ProofGoal asmp goal]
  goalsToList =
    getConst . traverseGoalsWithAssumptions
      (\as g -> Const [ProofGoal as g])
  
  -- | Traverse the structure of a 'Goals' data structure.  The function for
  --   visiting goals my decide to remove the goal from the structure.  If
  --   no goals remain after the traversal, the resulting value will be a 'Nothing'.
  --
  --   In a call to 'traverseGoals assumeAction transformer goals', the
  --   arguments are used as follows:
  --
  --   * 'traverseGoals' is an action is called every time we encounter
  --     an 'Assuming' constructor.  The first argument is the original
  --     sequence of assumptions.  The second argument is a continuation
  --     action.  The result is a sequence of transformed assumptions
  --     and the result of the continuation action.
  --
  --   * 'assumeAction' is a transformer action on goals.  Return
  --     'Nothing' if you wish to remove the goal from the overall tree.
  traverseGoals :: (Applicative f, Monoid asmp') =>
                   (forall a. asmp -> f a -> f (asmp', a))
                -> (goal -> f (Maybe goal'))
                -> Goals asmp goal
                -> f (Maybe (Goals asmp' goal'))
  traverseGoals fas fgl = go
    where
    go (Prove gl)        = fmap Prove <$> fgl gl
    go (Assuming as gl)  = assuming' <$> fas as (go gl)
    go (ProveConj g1 g2) = goalsConj <$> go g1 <*> go g2
  
    assuming' (_, Nothing) = Nothing
    assuming' (as, Just g) = Just $! assuming as g
  
  
  traverseOnlyGoals :: (Applicative f, Monoid asmp) =>
    (goal -> f (Maybe goal')) ->
    Goals asmp goal -> f (Maybe (Goals asmp goal'))
  traverseOnlyGoals f = traverseGoals (\as m -> (as,) <$> m) f
  
  -- | Traverse a sequence of 'Goals' data structures.  See 'traverseGoals'
  --   for an explanation of the action arguments.  The resulting sequence
  --   may be shorter than the original if some 'Goals' become trivial.
  traverseGoalsSeq :: (Applicative f, Monoid asmp') =>
    (forall a. asmp -> f a -> f (asmp', a)) ->
    (goal -> f (Maybe goal')) ->
    Seq (Goals asmp goal) -> f (Seq (Goals asmp' goal'))
  traverseGoalsSeq fas fgl = go
    where
    go Seq.Empty      = pure Seq.Empty
    go (g Seq.:<| gs) = combine <$> traverseGoals fas fgl g <*> go gs
  
    combine Nothing gs  = gs
    combine (Just g) gs = g Seq.<| gs
  
  -- | Visit every goal in a 'Goals' structure, remembering the sequence of
  --   assumptions along the way to that goal.
  traverseGoalsWithAssumptions :: (Applicative f, Monoid asmp) =>
    (asmp -> goal -> f (Maybe goal')) ->
    Goals asmp goal -> f (Maybe (Goals asmp goal'))
  
  traverseGoalsWithAssumptions f gls =
     runReaderT (traverseGoals fas fgl gls) mempty
    where
    fas a m = (a,) <$> withReaderT (<> a) m
    fgl gl  = ReaderT $ \as -> f as gl
  
  