content: {-# LANGUAGE GADTs #-}
  {-# LANGUAGE LambdaCase #-}
  {-# LANGUAGE TypeApplications #-}
  {-# LANGUAGE ViewPatterns #-}
  
  module Bool where
  
  import Control.Monad (unless, when)
  import Control.Monad.IO.Class (MonadIO, liftIO)
  import qualified Control.Monad.State.Strict as State
  import Control.Monad.Trans (lift)
  import Data.Coerce (coerce)
  import Data.Either (isRight)
  import Data.Foldable (traverse_)
  import qualified Data.Map as Map
  import qualified Data.Parameterized.Map as MapF
  import Data.Parameterized.Nonce (newIONonceGenerator)
  import Data.Parameterized.Some (Some(Some))
  import Hedgehog (GenT)
  import qualified Hedgehog.Gen as Gen
  import qualified Hedgehog.Internal.Gen as HG
  import qualified Hedgehog.Internal.Property as HG
  import qualified Test.Tasty.Hedgehog as THG
  import qualified Test.Tasty as T
  import qualified What4.Expr.BoolMap as BM
  import What4.Expr.Builder
  import What4.Expr (EmptyExprBuilderState(EmptyExprBuilderState))
  import What4.Interface
  
  -- | A tree of API calls to 'IsExprBuilder' methods.
  --
  -- Instances may be \"interpreted\" into 'IsExprBuilder' calls via 'toSymExpr'.
  -- Data flows from children to parents.
  --
  -- Given a means to evaluate variables to 'Bool's, these expressions can
  -- also be evaluated directly (via 'eval'), in order to compare the result to
  -- 'asConstantPred'.
  data BExpr var
    = -- 0-ary
      -- | 'falsePred', 'truePred'
      Lit !Bool
      -- | 'freshConstant'
    | Var !var
      -- unary
      -- | 'notPred'
    | Not !(BExpr var)
      -- binary
      -- | 'andPred'
    | And !(BExpr var) !(BExpr var)
      -- | 'eqPred'
    | Eq !(BExpr var) !(BExpr var)
      -- | 'orPred'
    | Or !(BExpr var) !(BExpr var)
      -- | 'xorPred'
    | Xor !(BExpr var) !(BExpr var)
      -- tertiary
      -- | 'itePred'
    | Ite !(BExpr var) !(BExpr var) !(BExpr var)
    deriving Show
  
  genBExpr :: HG.MonadGen m => m var -> m (BExpr var)
  genBExpr var =
    Gen.recursive
      Gen.choice
      [ -- 0-ary
        Lit <$> Gen.bool
      , Var <$> var
      ]
      [ -- unary
        Not <$> genBExpr var
        -- binary
      , And <$> genBExpr var <*> genBExpr var
      -- TODO: Generate Eq, Xor.
      --
      -- This would require updating 'isNormal' to take these into account.
      --
      -- , Eq <$> genBExpr var <*> genBExpr var
      , Or <$> genBExpr var <*> genBExpr var
      -- , Xor <$> genBExpr var <*> genBExpr var
      , Ite <$> genBExpr var <*> genBExpr var <*> genBExpr var
      ]
  
  newtype Valuation t
    = Valuation { getValuation :: Map.Map (ExprBoundVar t BaseBoolType) Bool }
    deriving Show
  
  getValue :: ExprBoundVar t BaseBoolType -> Valuation t -> Bool
  getValue v vs =
    case Map.lookup v (getValuation vs) of
      Nothing -> error "getValue: bad variable"
      Just b -> b
  
  genFreshVar ::
    (HG.MonadGen m, MonadIO m) =>
    ExprBuilder t st fs ->
    State.StateT (Valuation t) m (ExprBoundVar t BaseBoolType)
  genFreshVar sym = do
    v <- lift (liftIO (freshConstant sym (safeSymbol "b") BaseBoolRepr))
    case v of
      BoundVarExpr v' -> do
        b <- Gen.bool
        State.modify (coerce (Map.insert v' b))
        pure v'
      _ -> error "Not a bound variable?"
  
  -- | Generate a new variable ('genFreshVar') or reuse an existing one
  genVar ::
    (HG.MonadGen m, MonadIO m) =>
    ExprBuilder t st fs ->
    State.StateT (Valuation t) m (ExprBoundVar t BaseBoolType)
  genVar sym = do
    b <- Gen.bool
    if b
      then genFreshVar sym
      else do
        vs <- State.gets (Map.toList . getValuation)
        case vs of
          [] -> genFreshVar sym
          _ -> Gen.choice (map (pure . fst) vs)
  
  doGenExpr ::
    ExprBuilder t st fs ->
    GenT IO (BExpr (ExprBoundVar t BaseBoolType), Valuation t)
  doGenExpr sym =
    let vars0 = Valuation Map.empty in
    State.runStateT (genBExpr @(State.StateT _ (GenT IO)) (genVar @(GenT IO) sym)) vars0
  
  toSymExpr ::
    IsExprBuilder sym =>
    sym ->
    -- | How to handle variables
    (var -> IO (SymExpr sym BaseBoolType)) ->
    BExpr var ->
    IO (SymExpr sym BaseBoolType)
  toSymExpr sym var = go
    where
    go =
      \case
        Lit True -> pure (truePred sym)
        Lit False -> pure (falsePred sym)
        Var v -> var v
        Not e -> notPred sym =<< go e
        And l r -> do
          l' <- go l
          r' <- go r
          andPred sym l' r'
        Eq l r -> do
          l' <- go l
          r' <- go r
          eqPred sym l' r'
        Or l r -> do
          l' <- go l
          r' <- go r
          orPred sym l' r'
        Xor l r -> do
          l' <- go l
          r' <- go r
          xorPred sym l' r'
        Ite c l r -> do
          c' <- go c
          l' <- go l
          r' <- go r
          itePred sym c' l' r'
  
  -- | For use with 'toSymExpr', to leave variables uninterpreted
  uninterpVar :: ExprBoundVar t BaseBoolType -> Expr t BaseBoolType
  uninterpVar = BoundVarExpr
  
  eval :: Applicative f => (var -> f Bool) -> BExpr var -> f Bool
  eval var = go
    where
    ite c l r = if c then l else r
    go =
      \case
        Lit True -> pure True
        Lit False -> pure False
        Var v -> var v
        Not e -> not <$> go e
        And l r -> (&&) <$> go l <*> go r
        Eq l r -> (==) <$> go l <*> go r
        Or l r -> (||) <$> go l <*> go r
        Xor l r -> (/=) <$> go l <*> go r
        Ite c l r -> ite <$> go c <*> go l <*> go r
  
  -- | For use with 'eval', to interpret variables
  getVar :: ExprBoundVar t BaseBoolType -> State.State (Valuation t) Bool
  getVar v = State.gets (getValue v)
  
  isNot :: Expr t BaseBoolType -> Bool
  isNot e =
    case e of
      AppExpr ae ->
        case appExprApp ae of
          NotPred {} -> True
          _ -> False
      _ -> False
  
  isNormalIte ::
    ExprBuilder t st fs ->
    Expr t BaseBoolType -> 
    Expr t BaseBoolType -> 
    Expr t BaseBoolType -> 
    Either String ()
  isNormalIte sym c l r = do
    isNormal sym c
    isNormal sym l
    isNormal sym r
    when (isNot c) (Left "negated ite condition")
    when (c == l) (Left "ite cond == LHS")
    when (c == r) (Left "ite cond == RHS")
    when (c == truePred sym) (Left "ite cond == true")
    when (c == falsePred sym) (Left "ite cond == false")
  
  isNormalConjunct ::
    ExprBuilder t st fs ->
    Expr t BaseBoolType ->
    BM.Polarity ->
    Either String ()
  isNormalConjunct sym expr pol =
    case expr of
      BoolExpr {} -> Left "boolean literal inside conjunction"
      BoundVarExpr {} -> Right ()
      AppExpr ae ->
        case appExprApp ae of
          NotPred {} -> Left "not should be expressed via polarity"
          -- This must be an OR, if it is an AND it should be combined with
          -- its parent
          ConjPred cm' -> do
            when (pol == BM.Positive) (Left "and inside and")
            -- Note that it is possible to have ORs inside ORs, e.g., if the outer
            -- OR used to be an AND but was negated.
            isNormalMap sym cm'
          BaseIte BaseBoolRepr _sz c l r -> isNormalIte sym c l r
          _ -> Left "non-normal app in conjunct"
      _ -> Left "non-normal expr in conjunct"
  
  isNormalMap :: ExprBuilder t st fs -> BM.ConjMap (Expr t) -> Either String ()
  isNormalMap sym cm =
    case BM.viewConjMap cm of
      BM.ConjTrue -> Left "empty conjunction map"
      BM.ConjFalse -> Left "inconsistent conjunction map"
      BM.Conjuncts conjs -> traverse_ (uncurry (isNormalConjunct sym)) conjs
  
  -- | Is this boolean expression sufficiently normalized?
  isNormal :: ExprBuilder t st fs -> Expr t BaseBoolType -> Either String ()
  isNormal sym =
    \case
      BoolExpr {} -> Right ()
      BoundVarExpr {} -> Right ()
      AppExpr ae ->
        case appExprApp ae of
          NotPred (asApp -> Just NotPred {}) -> Left "double negation"
          NotPred e -> isNormal sym e
          ConjPred cm -> isNormalMap sym cm
          BaseIte BaseBoolRepr _sz c l r -> isNormalIte sym c l r
          _ -> Left "non-normal app"
      _ -> Left "non-normal expr"
  
  boolTests :: T.TestTree
  boolTests =
    T.testGroup
      "boolean normalization tests"
      [ -- Test that the rewrite rules rewrite expressions into a sufficiently
        -- normal form (defined by 'isNormal').
        THG.testProperty "boolean rewrites normalize" $
          HG.property $ do
            Some ng <- liftIO newIONonceGenerator
            sym <- liftIO (newExprBuilder FloatIEEERepr EmptyExprBuilderState ng)
            (e, _vars) <- HG.forAllT (doGenExpr sym)
            e' <- liftIO (toSymExpr sym (pure . uninterpVar) e)
            let ok = isNormal sym e'
            unless (isRight ok) $
              liftIO (putStrLn ("Not normalized:\n" ++ show (printSymExpr e')))
            ok HG.=== Right ()
      , THG.testProperty "boolean rewrites preserve semantics" $
          HG.property $ do
            Some ng <- liftIO newIONonceGenerator
            sym <- liftIO (newExprBuilder FloatIEEERepr EmptyExprBuilderState ng)
            (e, vars) <- HG.forAllT (doGenExpr sym)
            -- Concretely evaluate the `BExpr` to get the expected semantics.
            let expected = State.evalState (eval getVar e) vars
            -- Generate a `Expr` with uninterpreted variables. It is important to
            -- not interpret the variables into `truePred` and `falsePred` here,
            -- to avoid only hitting the `asConstantPred` cases in the rewrites.
            e' <- liftIO (toSymExpr sym (pure . uninterpVar) e)
            -- Finally, substitute values in for the variables, simplifying the
            -- `Expr` along the way until we get a concrete boolean.
            let vs = Map.toList (getValuation vars)
            let substs = foldr (\(v, b) -> MapF.insert v (if b then truePred sym else falsePred sym)) MapF.empty vs
            e'' <- liftIO (substituteBoundVars sym substs e')
            -- Check that the `BExpr` and `Expr` agreed on the semantics.
            case asConstantPred e'' of
              Just actual -> actual HG.=== expected
              Nothing -> HG.failure
      ]
  