content: {-# LANGUAGE GADTs               #-}
  {-# LANGUAGE LambdaCase          #-}
  {-# LANGUAGE NamedFieldPuns      #-}
  {-# LANGUAGE RankNTypes          #-}
  {-# LANGUAGE Safe                #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  
  -- | Translate Copilot specifications into IL specifications.
  module Copilot.Theorem.IL.Translate ( translate, translateWithBounds ) where
  
  import Copilot.Theorem.IL.Spec
  
  import qualified Copilot.Core as C
  
  import qualified Data.Map.Strict as Map
  
  import Control.Monad       (forM, liftM2, when)
  import Control.Monad.State
  
  import Data.Char
  import Data.List (find)
  
  import Text.Printf
  
  import GHC.Float (float2Double)
  
  import Data.Typeable (Typeable)
  
  -- 'nc' stands for naming convention.
  ncSeq :: C.Id -> SeqId
  ncSeq = printf "s%d"
  
  -- We assume all local variables have distinct names whatever their scopes.
  ncLocal :: C.Name -> SeqId
  ncLocal s = "l" ++ dropWhile (not . isNumber) s
  
  ncExternVar :: C.Name -> SeqId
  ncExternVar n = "ext_" ++ n
  
  ncUnhandledOp :: String -> String
  ncUnhandledOp = id
  
  ncMux :: Integer -> SeqId
  ncMux n = "mux" ++ show n
  
  -- | Translate a Copilot specification to an IL specification.
  translate :: C.Spec -> IL
  translate = translate' False
  
  -- | Translate a Copilot specification to an IL specification, adding
  -- constraints for limiting the values of numeric expressions to known bounds
  -- based on their specific types (only for integers or natural numbers).
  translateWithBounds :: C.Spec -> IL
  translateWithBounds = translate' True
  
  translate' :: Bool -> C.Spec -> IL
  translate' b (C.Spec {C.specStreams, C.specProperties}) = runTrans b $ do
  
    let modelInit = concatMap streamInit specStreams
  
    mainConstraints <- mapM streamRec specStreams
  
    localConstraints <- popLocalConstraints
    properties <- Map.fromList <$>
      forM specProperties
        (\(C.Property {C.propertyName, C.propertyProp}) -> do
          -- Soundness note: it is OK to call `extractProp` here to drop the
          -- quantifier from the proposition `propertyProp`. This is because we
          -- IL translation always occurs within the context of a function that
          -- returns a `Proof`, and these `Proof` functions are always careful to
          -- use `Prover`s that respect the propositions's quantifier.
          e' <- expr (C.extractProp propertyProp)
          propConds <- popLocalConstraints
          return (propertyName, (propConds, e')))
  
    return IL
      { modelInit
      , modelRec = mainConstraints ++ localConstraints
      , properties
      , inductive = not $ null specStreams
      }
  
  bound :: Expr -> C.Type a -> Trans ()
  bound s t = case t of
    C.Int8    -> bound' C.Int8
    C.Int16   -> bound' C.Int16
    C.Int32   -> bound' C.Int32
    C.Int64   -> bound' C.Int64
    C.Word8   -> bound' C.Word8
    C.Word16  -> bound' C.Word16
    C.Word32  -> bound' C.Word32
    C.Word64  -> bound' C.Word64
    _         -> return ()
    where
      bound' :: (Bounded a, Integral a) => C.Type a -> Trans ()
      bound' t = do
        b <- addBounds <$> get
        when b $ localConstraint (Op2 Bool And
          (Op2 Bool Le (trConst t minBound) s)
          (Op2 Bool Ge (trConst t maxBound) s))
  
  streamInit :: C.Stream -> [Expr]
  streamInit (C.Stream { C.streamId       = id
                       , C.streamBuffer   = b :: [val]
                       , C.streamExprType = t }) =
    zipWith initConstraint [0..] b
    where
      initConstraint :: Integer -> val -> Expr
      initConstraint p v = Op2 Bool Eq
        (SVal (trType t) (ncSeq id) (Fixed p))
        $ trConst t v
  
  streamRec :: C.Stream -> Trans Expr
  streamRec (C.Stream { C.streamId       = id
                      , C.streamExpr     = e
                      , C.streamBuffer   = b
                      , C.streamExprType = t })
    = do
    let s = SVal (trType t) (ncSeq id) (_n_plus $ length b)
    bound s t
    e' <- expr e
    return $ Op2 Bool Eq s e'
  
  expr :: Typeable a => C.Expr a -> Trans Expr
  
  expr (C.Const t v) = return $ trConst t v
  
  expr (C.Label _ _ e) = expr e
  
  expr (C.Drop t k id) = return $ SVal (trType t) (ncSeq id) (_n_plus k)
  
  expr (C.Local ta _ name ea eb) = do
    ea' <- expr ea
    localConstraint (Op2 Bool Eq (SVal (trType ta) (ncLocal name) _n_) ea')
    expr eb
  
  expr (C.Var t name) = return $ SVal (trType t) (ncLocal name) _n_
  
  expr (C.ExternVar t name _) = bound s t >> return s
    where
      s = SVal (trType t) (ncExternVar name) _n_
  
  expr (C.Op1 (C.Sign ta) e) = case ta of
    C.Int8   -> trSign ta e
    C.Int16  -> trSign ta e
    C.Int32  -> trSign ta e
    C.Int64  -> trSign ta e
    C.Float  -> trSign ta e
    C.Double -> trSign ta e
    _        -> expr $ C.Const ta 1
    where
      trSign :: (Typeable a, Ord a, Num a) => C.Type a -> C.Expr a -> Trans Expr
      trSign ta e =
        expr (C.Op3 (C.Mux ta)
          (C.Op2 (C.Lt ta) e (C.Const ta 0))
          (C.Const ta (-1))
          (C.Op3 (C.Mux ta)
            (C.Op2 (C.Gt ta) e (C.Const ta 0))
            (C.Const ta 1)
            (C.Const ta 0)))
  expr (C.Op1 (C.Sqrt _) e) = do
    e' <- expr e
    return $ Op2 Real Pow e' (ConstR 0.5)
  expr (C.Op1 (C.Cast _ _) e) = expr e
  expr (C.Op1 op e) = do
    e' <- expr e
    return $ Op1 t' op' e'
    where
      (op', t') = trOp1 op
  
  expr (C.Op2 (C.Ne t) e1 e2) = do
    e1' <- expr e1
    e2' <- expr e2
    return $ Op1 Bool Not (Op2 t' Eq e1' e2')
    where
      t' = trType t
  
  expr (C.Op2 op e1 e2) = do
    e1' <- expr e1
    e2' <- expr e2
    return $ Op2 t' op' e1' e2'
    where
      (op', t') = trOp2 op
  
  expr (C.Op3 (C.Mux t) cond e1 e2) = do
    cond' <- expr cond
    e1'   <- expr e1
    e2'   <- expr e2
    newMux cond' (trType t) e1' e2'
  
  trConst :: C.Type a -> a -> Expr
  trConst t v = case t of
    C.Bool   -> ConstB v
    C.Float  -> negifyR (float2Double v)
    C.Double -> negifyR v
    t@C.Int8   -> negifyI v (trType t)
    t@C.Int16  -> negifyI v (trType t)
    t@C.Int32  -> negifyI v (trType t)
    t@C.Int64  -> negifyI v (trType t)
    t@C.Word8  -> negifyI v (trType t)
    t@C.Word16 -> negifyI v (trType t)
    t@C.Word32 -> negifyI v (trType t)
    t@C.Word64 -> negifyI v (trType t)
    where
      negifyR :: Double -> Expr
      negifyR v
        | v >= 0    = ConstR v
        | otherwise = Op1 Real Neg $ ConstR $ negate $ v
      negifyI :: Integral a => a -> Type -> Expr
      negifyI v t
        | v >= 0    = ConstI t $ toInteger v
        | otherwise = Op1 t Neg $ ConstI t $ negate $ toInteger v
  
  trOp1 :: C.Op1 a b -> (Op1, Type)
  trOp1 = \case
    C.Not     -> (Not, Bool)
    C.Abs t   -> (Abs, trType t)
    -- C.Sign t  ->
    -- C.Recip t ->
    C.Exp t   -> (Exp, trType t)
    -- C.Sqrt t  ->
    C.Log t   -> (Log, trType t)
    C.Sin t   -> (Sin, trType t)
    C.Tan t   -> (Tan, trType t)
    C.Cos t   -> (Cos, trType t)
    C.Asin t  -> (Asin, trType t)
    C.Atan t  -> (Atan, trType t)
    C.Acos t  -> (Acos, trType t)
    C.Sinh t  -> (Sinh, trType t)
    C.Tanh t  -> (Tanh, trType t)
    C.Cosh t  -> (Cosh, trType t)
    C.Asinh t -> (Asinh, trType t)
    C.Atanh t -> (Atanh, trType t)
    C.Acosh t -> (Acosh, trType t)
    -- C.BwNot t ->
    -- C.Cast t  ->
    _ -> error "Unsupported unary operator in input."
  
  trOp2 :: C.Op2 a b c -> (Op2, Type)
  trOp2 = \case
    C.And          -> (And, Bool)
    C.Or           -> (Or, Bool)
  
    C.Add t        -> (Add, trType t)
    C.Sub t        -> (Sub, trType t)
    C.Mul t        -> (Mul, trType t)
  
    C.Mod t        -> (Mod, trType t)
    -- C.Div t        ->
  
    C.Fdiv t       -> (Fdiv, trType t)
  
    C.Pow t        -> (Pow, trType t)
    -- C.Logb t       ->
  
    C.Eq _         -> (Eq, Bool)
    -- C.Ne t         ->
  
    C.Le t         -> (Le, trType t)
    C.Ge t         -> (Ge, trType t)
    C.Lt t         -> (Lt, trType t)
    C.Gt t         -> (Gt, trType t)
  
    -- C.BwAnd t      ->
    -- C.BwOr t       ->
    -- C.BwXor t      ->
    -- C.BwShiftL t _ ->
    -- C.BwShiftR t _ ->
  
    _ -> error "Unsupported binary operator in input."
  
  trType :: C.Type a -> Type
  trType = \case
    C.Bool   -> Bool
    C.Int8   -> SBV8
    C.Int16  -> SBV16
    C.Int32  -> SBV32
    C.Int64  -> SBV64
    C.Word8  -> BV8
    C.Word16 -> BV16
    C.Word32 -> BV32
    C.Word64 -> BV64
    C.Float  -> Real
    C.Double -> Real
  
  -- | Translation state.
  data TransST = TransST
    { localConstraints :: [Expr]
    , muxes            :: [(Expr, (Expr, Type, Expr, Expr))]
    , nextFresh        :: Integer
    , addBounds        :: Bool
    }
  
  newMux :: Expr -> Type -> Expr -> Expr -> Trans Expr
  newMux c t e1 e2 = do
    ms <- muxes <$> get
    case find ((==mux) . snd) ms of
      Nothing -> do
        f <- fresh
        let v = SVal t (ncMux f) _n_
        modify $ \st -> st { muxes = (v, mux) : ms }
        return v
      Just (v, _) -> return v
    where
      mux = (c, t, e1, e2)
  
  getMuxes :: Trans [Expr]
  getMuxes = muxes <$> get >>= return . concat . (map toConstraints)
    where
      toConstraints (v, (c, _, e1, e2)) =
        [ Op2 Bool Or (Op1 Bool Not c) (Op2 Bool Eq v e1)
        , Op2 Bool Or c (Op2 Bool Eq v e2)
        ]
  
  -- | A state monad over the translation state ('TransST').
  type Trans = State TransST
  
  fresh :: Trans Integer
  fresh = do
    modify $ \st -> st {nextFresh = nextFresh st + 1}
    nextFresh <$> get
  
  localConstraint :: Expr -> Trans ()
  localConstraint c =
    modify $ \st -> st {localConstraints = c : localConstraints st}
  
  popLocalConstraints :: Trans [Expr]
  popLocalConstraints = liftM2 (++) (localConstraints <$> get) getMuxes
    <* (modify $ \st -> st {localConstraints = [], muxes = []})
  
  runTrans :: Bool -> Trans a -> a
  runTrans b m = evalState m $ TransST [] [] 0 b
  