content: {-# LANGUAGE DataKinds #-}
  {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE FlexibleInstances #-}
  {-# LANGUAGE GADTs #-}
  {-# LANGUAGE LambdaCase #-}
  {-# LANGUAGE MultiParamTypeClasses #-}
  {-# LANGUAGE OverloadedStrings #-}
  {-# LANGUAGE PolyKinds #-}
  {-# LANGUAGE RankNTypes #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-# LANGUAGE TupleSections #-}
  {-# LANGUAGE TypeApplications #-}
  {-# LANGUAGE TypeFamilies #-}
  {-# LANGUAGE TypeOperators #-}
  {-# LANGUAGE UndecidableInstances #-}
  {-# LANGUAGE ViewPatterns #-}
  {-# LANGUAGE MultiWayIf #-}
  
  -- | Normalization and equivalence checking for expressions
  module What4.Serialize.Normalize
    ( normSymFn
    , normExpr
    , testEquivSymFn
    , testEquivExpr
    , ExprEquivResult(..)
    ) where
  
  import qualified Data.Parameterized.Context as Ctx
  import qualified Data.Parameterized.TraversableFC as FC
  
  import qualified What4.Interface as S
  import qualified What4.Expr as S
  import qualified What4.Expr.Builder as B
  import qualified What4.Expr.WeightedSum as WSum
  import           Data.Parameterized.Classes
  
  -- | Apply some normalizations to make function call arguments more readable.  Examples include:
  --
  --  * Avoid wrapping single literals in a 'B.SemiRingLiteral' and just represent them as a bare integer literals
  --  * Attempt to reduce function calls with constant arguments where possible
  normSymFn :: forall sym st fs t args ret. sym ~ B.ExprBuilder t st fs
            => sym
            -> B.ExprSymFn t args ret
            -> Ctx.Assignment (S.Expr t) args
            -> IO (S.Expr t ret)
  normSymFn sym symFn argEs = case B.symFnInfo symFn of
    B.DefinedFnInfo argBVs expr _ -> do
      argEs' <- FC.traverseFC (normExpr sym) argEs
      expr' <- B.evalBoundVars sym expr argBVs argEs'
      normExpr sym expr'
    _ -> S.applySymFn sym symFn argEs
  
  
  normExpr :: forall sym st fs t tp
            . sym ~ B.ExprBuilder t st fs
           => sym
           -> B.Expr t tp -> IO (B.Expr t tp)
  normExpr sym e = go e
    where go :: B.Expr t tp -> IO (B.Expr t tp)
          go (B.SemiRingLiteral S.SemiRingIntegerRepr val _) = S.intLit sym val
          go (B.AppExpr appExpr) = normAppExpr sym appExpr
          go x@(B.NonceAppExpr nae) =
            case B.nonceExprApp nae of
              B.FnApp fn args -> normSymFn sym fn args
              _ -> return x
          go x = return x
  
  -- | Normalize an expression by passing it back through the builder
  --
  -- NOTE: We may want to audit the cases here for completeness
  normAppExpr :: forall sym st fs t tp
               . sym ~ S.ExprBuilder t st fs
              => sym
              -> S.AppExpr t tp
              -> IO (S.Expr t tp)
  normAppExpr sym ae = do
    e' <- go (S.appExprApp ae)
    B.sbMakeExpr sym e'
    where norm2 :: forall tp' tp'' tp'''
                 . (S.Expr t tp' -> S.Expr t tp'' -> IO (S.Expr t tp'''))
                -> S.Expr t tp' -> S.Expr t tp'' -> IO (S.Expr t tp''')
          norm2 f e1 e2 = do
            e1' <- normExpr sym e1
            e2' <- normExpr sym e2
            f e1' e2'
  
          go :: forall tp'. S.App (S.Expr t) tp' -> IO (S.App (S.Expr t) tp')
          go (S.BaseIte _ _ test then_ else_) = do
            test' <- normExpr sym test
            then' <- normExpr sym then_
            else' <- normExpr sym else_
            Just sm' <- B.asApp <$> S.baseTypeIte sym test' then' else'
            return sm'
          go x@(S.SemiRingSum sm) =
            case WSum.sumRepr sm of
              S.SemiRingIntegerRepr -> do
                let
                  smul si i = do
                      i' <- normExpr sym i
                      si' <- S.intLit sym si
                      S.intMul sym si' i'
                Just sm' <- B.asApp <$> WSum.evalM (norm2 $ S.intAdd sym) smul (S.intLit sym) sm
                return sm'
              _ -> return x
          go x@(S.SemiRingProd pd) =
            case WSum.prodRepr pd of
              S.SemiRingIntegerRepr -> do
                  maybeS <- WSum.prodEvalM (norm2 $ S.intMul sym) return pd
                  case maybeS of
                    Just s | Just sm' <- B.asApp s -> return sm'
                    _ -> return x
              _ -> return x
          go x@(S.SemiRingLe sr e1 e2) = do
            case sr of
              S.OrderedSemiRingIntegerRepr -> do
                Just sm' <- B.asApp <$> (norm2 $ S.intLe sym) e1 e2
                return sm'
              _ -> return x
          go x = return x
  
  
  
  data ExprEquivResult = ExprEquivalent | ExprNormEquivalent | ExprUnequal
  
  testEquivExpr :: forall sym st fs t tp tp'. sym ~ S.ExprBuilder t st fs => sym -> B.Expr t tp -> B.Expr t tp' -> IO (ExprEquivResult)
  testEquivExpr sym e1 e2 = case testEquality e1 e2 of
    Just Refl -> return ExprEquivalent
    _ -> do
      e1' <- normExpr sym e1
      e2' <- normExpr sym e2
      case testEquality e1' e2' of
        Just Refl -> return ExprNormEquivalent
        _ -> return ExprUnequal
  
  testEquivSymFn :: forall sym st fs t args ret args' ret'. sym ~ S.ExprBuilder t st fs => sym -> S.SymFn sym args ret -> S.SymFn sym args' ret' -> IO (ExprEquivResult)
  testEquivSymFn sym fn1 fn2 =
    let
      argTypes1 = S.fnArgTypes fn1
      argTypes2 = S.fnArgTypes fn2
      retType1 = S.fnReturnType fn1
      retType2 = S.fnReturnType fn2
    in if | Just Refl <- testEquality argTypes1 argTypes2
          , Just Refl <- testEquality retType1 retType2
          , B.symFnName fn1 == B.symFnName fn2 ->
            case (S.symFnInfo fn1, S.symFnInfo fn2) of
             (S.DefinedFnInfo argBVs1 efn1 _, S.DefinedFnInfo argBVs2 efn2 _) -> do
               args <- FC.traverseFC (\bv -> S.freshConstant sym (S.bvarName bv) (B.bvarType bv)) argBVs1
               expr1 <- B.evalBoundVars sym efn1 argBVs1 args
               expr2 <- B.evalBoundVars sym efn2 argBVs2 args
               case testEquality expr1 expr2 of
                 Just Refl -> return ExprEquivalent
                 Nothing -> do
                   expr1' <- normExpr sym expr1
                   expr2' <- normExpr sym expr2
                   case testEquality expr1' expr2' of
                     Just Refl -> return ExprNormEquivalent
                     Nothing -> return ExprUnequal
             (S.UninterpFnInfo _ _, S.UninterpFnInfo _ _) -> return ExprEquivalent
             (S.MatlabSolverFnInfo _ _ _, _) -> fail "Unsupported function type for equivalence check."
             (_, S.MatlabSolverFnInfo _ _ _) -> fail "Unsupported function type for equivalence check."
             (_, _) -> return ExprUnequal
          | otherwise -> return ExprUnequal
  