content: {-# LANGUAGE BlockArguments #-}
  {-# LANGUAGE DerivingVia #-}
  {-# LANGUAGE GeneralisedNewtypeDeriving #-}
  {-# LANGUAGE RankNTypes #-}
  {-# LANGUAGE RecordWildCards #-}
  {-# LANGUAGE TypeFamilies #-}
  
  {-|
  Module      : Data.Input.Config
  Description : My horror at his crimes was lost in my admiration at his skill.
  Copyright   : (c) Tom Harding, 2020
  License     : MIT
  
  'MoriarT' is a monad transformer implementing the 'MonadCell' class with
  provenance and backtracking search. In other words, it can search large
  parameter spaces using different parameter configurations, looking for
  contradicting sets of parameters to prune out parts of the search tree. It does
  this by keeping track of which cells influence which results, and considering
  any influencers on a failure to be contradictory.
  
  In other words: if a write to cell @A@ fails, and the write was based on values
  from cells @B@ and @C@, any search branch in which @B@ and @C@ have these
  current values will be /pruned/ from the search, and we won't try them.
  
  (In practice, this isn't strictly true: we just abort any branch that ever
  produces any cell with any provenance that contains those values for @B@ and
  @C@. This is a "lazier" strategy, and doesn't involve evaluating the search
  space up front).
  -}
  module Control.Monad.MoriarT
    ( MoriarT (..)
  
    , runAll
    , runOne
    , solve
    , unsafeRead
    ) where
  
  import Control.Applicative (Alternative (..))
  import Control.Monad (MonadPlus, guard)
  import Control.Monad.Cell.Class (MonadCell (..))
  import qualified Control.Monad.Cell.Class as Cell
  import Control.Monad.IO.Class (MonadIO (..))
  import Control.Monad.Logic (MonadLogic, LogicT (..))
  import qualified Control.Monad.Logic as LogicT
  import Control.Monad.Primitive (PrimMonad (..))
  import Control.Monad.Reader.Class (MonadReader (..))
  import qualified Control.Monad.Reader.Class as Reader
  import Control.Monad.State.Class (MonadState (..))
  import qualified Control.Monad.State.Class as State
  import Control.Monad.Trans.Class (MonadTrans (..))
  import Control.Monad.Trans.Reader (ReaderT (..))
  import Control.Monad.Trans.State (StateT (..))
  import qualified Control.Monad.Trans.State as StateT
  import qualified Data.CDCL as CDCL
  import Data.Foldable (asum)
  import Data.Function ((&))
  import Data.Functor ((<&>))
  import Data.Input.Config (Config (..))
  import Data.JoinSemilattice.Class.Boolean (BooleanR (..))
  import Data.JoinSemilattice.Class.Eq (EqR (..))
  import Data.JoinSemilattice.Class.Merge (Merge (..), Result (..))
  import Data.Kind (Type)
  import Data.Maybe (listToMaybe)
  import Data.Monoid (Ap (..))
  import Data.Primitive.MutVar (MutVar)
  import qualified Data.Primitive.MutVar as MutVar
  import Data.Propagator (Prop)
  import qualified Data.Propagator as Prop
  import Type.Reflection (Typeable)
  
  -- | The constraint-solving monad transformer. We implement the current
  -- computation context with 'MonadReader', and the current "no goods" list with
  -- 'MonadState'.
  --
  -- This transformer exposes its internals through the 'MonadReader',
  -- 'MonadState', 'MonadLogic', and 'MonadIO' interfaces, and should therefore
  -- /not/ be used directly. The reason is simply that misuse of any of these
  -- will break the computation, so the library provides "Control.Monad.Holmes"
  -- and "Control.Monad.Watson", who do their best to thwart 'MoriarT'.
  newtype MoriarT (m :: Type -> Type) (x :: Type)
    = MoriarT
        { unMoriarT :: ReaderT CDCL.Rule (LogicT (StateT CDCL.Group m)) x
        }
     deriving newtype
      ( Functor
      , Applicative
      , Alternative
      , Monad
      , MonadIO
      , MonadLogic
      , MonadPlus
      , MonadReader CDCL.Rule
      , MonadState CDCL.Group
      )
    deriving (Semigroup, Monoid)
      via (Ap (MoriarT m) x)
  
  instance MonadTrans MoriarT where
    lift = MoriarT . lift . lift . lift
  
  instance PrimMonad m => PrimMonad (MoriarT m) where
    type PrimState (MoriarT m) = PrimState m
  
    primitive = lift . primitive
  
  instance PrimMonad m => MonadCell (MoriarT m) where
    newtype Cell (MoriarT m) (content :: Type)
      = Cell (MutVar (PrimState m) (CDCL.Rule, content, MoriarT m ()))
  
    discard = do
      context <- Reader.ask
      State.modify (CDCL.resolve context) -- Add this context to the "no goods" list.
      
      empty
  
    fill content = do
      context <- Reader.ask
      mutVar  <- MutVar.newMutVar (context, content, mempty)
      pure (Cell mutVar)
  
    watch cell@(Cell mutVar) propagator = do
      let next = with cell propagator
  
      before@(provenance, content, callbacks)
        <- MutVar.readMutVar mutVar
  
      MutVar.writeMutVar mutVar (provenance, content, callbacks *> next) *> next
        <|> MutVar.writeMutVar mutVar before *> empty -- Undo the action for the next branch.
  
    with (Cell mutVar) callback = do
      (provenance, content, _) <- MutVar.readMutVar mutVar
      Reader.local (<> provenance) (callback content)
  
    write (Cell mutVar) news = do
      context <- Reader.ask
      nogoods <- State.get
  
      before@(provenance, content, callbacks)
        <- MutVar.readMutVar mutVar
  
      let provenance' = context <> provenance
          content'    = content <<- news
  
      -- Skip this branch if the provenance is no good.
      guard (not (nogoods `CDCL.implies` provenance'))
  
      case content' of
        Changed update -> do
          MutVar.writeMutVar mutVar (provenance', update, callbacks) *> callbacks
            <|> MutVar.writeMutVar mutVar before *> empty
  
        Failure   -> Reader.local (<> context) discard
        Unchanged -> pure ()
  
  -- | Unsafely read from a cell. This operation is unsafe because it doesn't
  -- factor this cell into the provenance of any subsequent writes. If this value
  -- ends up causing a contradiction, we may end up removing branches of the
  -- search tree that are totally valid! This operation is safe as long as it is
  -- the __very last thing__ you do in a computation, and its value is __never__
  -- used to influence any writes in any way.
  unsafeRead :: PrimMonad m => Cell (MoriarT m) x -> MoriarT m x
  unsafeRead (Cell mutVar) = do
    (_, content, _) <- MutVar.readMutVar mutVar
  
    pure content
  
  -- | Run a 'MoriarT' computation and return the list of __all__ valid branches'
  -- results, in the order in which they were discovered.
  runAll :: Monad m => MoriarT m x -> m [ x ]
  runAll
    = flip StateT.evalStateT mempty
    . LogicT.observeAllT
    . flip runReaderT mempty
    . unMoriarT
  
  -- | Run a 'MoriarT' computation and return the /first/ valid branch's result.
  runOne :: Monad m => MoriarT m x -> m (Maybe x)
  runOne
    = fmap listToMaybe
    . flip StateT.evalStateT mempty
    . LogicT.observeManyT 1
    . flip runReaderT mempty
    . unMoriarT
  
  -- | Given an input configuration, and a predicate on those input variables,
  -- compute the configurations that satisfy the predicate. This result (or these
  -- results) can be extracted using 'runOne' or 'runAll'.
  solve
    :: ( PrimMonad m
       , EqR f
       , Merge (f x)
       , Typeable x
       )
    => Config (MoriarT m) (f x)
    -> (forall m'. MonadCell m' => [ Prop m' (f x) ] -> Prop m' (f Bool))
    -> MoriarT m [ f x ]
  solve Config{..} predicate = do
    inputs <- traverse Cell.fill initial
    output <- Prop.down (predicate (map Prop.up inputs))
    Cell.write output trueR
  
    _ <- zip [0 ..] inputs & traverse \(major, cell) -> do
      current     <- unsafeRead cell
      refinements <- refine current
  
      input <- asum $ CDCL.index major refinements <&> \(rule, content) ->
        fmap Cell (MutVar.newMutVar (rule, content, mempty))
  
      Cell.unify cell input
  
    traverse unsafeRead inputs
  