-----------------------------------------------------------------------
-- |
-- Module           : Lang.Crucible.Simulator.RegValue
-- Description      : Runtime representation of CFG registers
-- Copyright        : (c) Galois, Inc 2014
-- License          : BSD3
-- Maintainer       : Joe Hendrix <jhendrix@galois.com>
-- Stability        : provisional
--
-- RegValue is a type family that defines the runtime representation
-- of crucible types.
------------------------------------------------------------------------
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE PatternGuards #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeOperators #-}
module Lang.Crucible.Simulator.RegValue
  ( RegValue
  , CanMux(..)
  , RegValue'(..)
  , MuxFn

    -- * Register values
  , AnyValue(..)
  , FnVal(..)
  , fnValType
  , RolledType(..)
  , SymSequence(..)

  , VariantBranch(..)
  , injectVariant

    -- * Value mux functions
  , ValMuxFn
  , eqMergeFn
  , mergePartExpr
  , muxRecursive
  , muxStringMap
  , muxStruct
  , muxVariant
  , muxVector
  , muxSymSequence
  , muxHandle
  ) where

import           Control.Monad
import           Control.Monad.Trans.Class
import           Data.Kind
import           Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map
import           Data.Proxy
import qualified Data.Set as Set
import           Data.Text (Text)
import qualified Data.Vector as V
import           Data.Word
import           GHC.TypeNats (KnownNat)

import qualified Data.Parameterized.Context as Ctx

import           What4.FunctionName
import           What4.Interface
import           What4.InterpretedFloatingPoint
import           What4.Partial
import           What4.WordMap

import           Lang.Crucible.FunctionHandle
import           Lang.Crucible.Simulator.Intrinsics
import           Lang.Crucible.Simulator.SymSequence
import           Lang.Crucible.Types
import           Lang.Crucible.Utils.MuxTree
import           Lang.Crucible.Backend

type MuxFn p v = p -> v -> v -> IO v

-- | Maps register types to the runtime representation.
type family RegValue (sym :: Type) (tp :: CrucibleType) :: Type where
  RegValue sym (BaseToType bt) = SymExpr sym bt
  RegValue sym (FloatType fi) = SymInterpretedFloat sym fi
  RegValue sym AnyType = AnyValue sym
  RegValue sym UnitType = ()
  RegValue sym NatType = SymNat sym
  RegValue sym CharType = Word16
  RegValue sym (FunctionHandleType a r) = FnVal sym a r
  RegValue sym (MaybeType tp) = PartExpr (Pred sym) (RegValue sym tp)
  RegValue sym (VectorType tp) = V.Vector (RegValue sym tp)
  RegValue sym (SequenceType tp) = SymSequence sym (RegValue sym tp)
  RegValue sym (StructType ctx) = Ctx.Assignment (RegValue' sym) ctx
  RegValue sym (VariantType ctx) = Ctx.Assignment (VariantBranch sym) ctx
  RegValue sym (ReferenceType tp) = MuxTree sym (RefCell tp)
  RegValue sym (WordMapType w tp) = WordMap sym w tp
  RegValue sym (RecursiveType nm ctx) = RolledType sym nm ctx
  RegValue sym (IntrinsicType nm ctx) = Intrinsic sym nm ctx
  RegValue sym (StringMapType tp) = Map Text (PartExpr (Pred sym) (RegValue sym tp))

-- | A newtype wrapper around RegValue.  This is wrapper necessary because
--   RegValue is a type family and, as such, cannot be partially applied.
newtype RegValue' sym tp = RV { unRV :: RegValue sym tp }

------------------------------------------------------------------------
-- FnVal

-- | Represents a function closure.
data FnVal (sym :: Type) (args :: Ctx CrucibleType) (res :: CrucibleType) where
  ClosureFnVal ::
    !(FnVal sym (args ::> tp) ret) ->
    !(TypeRepr tp) ->
    !(RegValue sym tp) ->
    FnVal sym args ret

  VarargsFnVal ::
    !(FnHandle (args ::> VectorType AnyType) ret) ->
    !(CtxRepr addlArgs) ->
    FnVal sym (args <+> addlArgs) ret

  HandleFnVal ::
    !(FnHandle a r) ->
    FnVal sym a r


closureFunctionName :: FnVal sym args res -> FunctionName
closureFunctionName (ClosureFnVal c _ _) = closureFunctionName c
closureFunctionName (HandleFnVal h) = handleName h
closureFunctionName (VarargsFnVal h _) = handleName h

-- | Extract the runtime representation of the type of the given 'FnVal'
fnValType :: FnVal sym args res -> TypeRepr (FunctionHandleType args res)
fnValType (HandleFnVal h) = FunctionHandleRepr (handleArgTypes h) (handleReturnType h)
fnValType (VarargsFnVal h addlArgs) =
  case handleArgTypes h of
    args Ctx.:> _ -> FunctionHandleRepr (args Ctx.<++> addlArgs) (handleReturnType h)
fnValType (ClosureFnVal fn _ _) =
  case fnValType fn of
    FunctionHandleRepr allArgs r ->
      case allArgs of
        args Ctx.:> _ -> FunctionHandleRepr args r

instance Show (FnVal sym a r) where
  show = show . closureFunctionName

-- | Version of 'MuxFn' specialized to 'RegValue'
type ValMuxFn sym tp = MuxFn (Pred sym) (RegValue sym tp)

------------------------------------------------------------------------
-- CanMux

-- | A class for 'CrucibleType's that have a
--   mux function.
class CanMux sym (tp :: CrucibleType) where
   muxReg :: sym
          -> p tp          -- ^ Unused type to identify what is being merged.
          -> ValMuxFn sym tp

-- | Merge function that checks if two values are equal, and
-- fails if they are not.
{-# INLINE eqMergeFn #-}
eqMergeFn :: (IsExprBuilder sym, Eq v) => sym -> String -> MuxFn p v
eqMergeFn sym nm = \_ x y ->
  if x == y then
    return x
  else
    throwUnsupported sym $ "Cannot merge dissimilar " ++ nm ++ "."

------------------------------------------------------------------------
-- RegValue AnyType instance

data AnyValue sym where
  AnyValue :: TypeRepr tp -> RegValue sym tp -> AnyValue sym

------------------------------------------------------------------------
-- RegValue () instance

instance CanMux sym UnitType where
  muxReg _ = \_ _ x _y -> return x

------------------------------------------------------------------------
-- RegValue instance for base types

instance IsExprBuilder sym => CanMux sym BoolType where
  {-# INLINE muxReg #-}
  muxReg s = const $ itePred s

instance IsExprBuilder sym => CanMux sym NatType where
  {-# INLINE muxReg #-}
  muxReg s = \_ -> natIte s

instance IsExprBuilder sym => CanMux sym IntegerType where
  {-# INLINE muxReg #-}
  muxReg s = \_ -> intIte s

instance IsExprBuilder sym => CanMux sym RealValType where
  {-# INLINE muxReg #-}
  muxReg s = \_ -> realIte s

instance IsInterpretedFloatExprBuilder sym => CanMux sym (FloatType fi) where
  {-# INLINE muxReg #-}
  muxReg s = \_ -> iFloatIte @sym @fi s

instance IsExprBuilder sym => CanMux sym ComplexRealType where
  {-# INLINE muxReg #-}
  muxReg s = \_ -> cplxIte s

instance IsExprBuilder sym => CanMux sym (StringType si) where
  {-# INLINE muxReg #-}
  muxReg s = \_ -> stringIte s

instance IsExprBuilder sym => CanMux sym (IEEEFloatType fpp) where
  muxReg s = \_ -> floatIte s

------------------------------------------------------------------------
-- RegValue Vector instance

{-# INLINE muxVector #-}
muxVector :: IsExprBuilder sym =>
             sym -> MuxFn p e -> MuxFn p (V.Vector e)
muxVector sym f p x y
  | V.length x == V.length y = V.zipWithM (f p) x y
  | otherwise =
      throwUnsupported sym "Cannot merge vectors with different dimensions."

instance (IsSymInterface sym, CanMux sym tp) => CanMux sym (VectorType tp) where
  {-# INLINE muxReg #-}
  muxReg s _ = muxVector s (muxReg s (Proxy :: Proxy tp))

------------------------------------------------------------------------
-- RegValue WordMap instance

instance (IsExprBuilder sym, KnownNat w, KnownRepr BaseTypeRepr tp)
  => CanMux sym (WordMapType w tp) where
  {-# INLINE muxReg #-}
  muxReg s _ p = muxWordMap s knownNat knownRepr p

------------------------------------------------------------------------
-- RegValue MatlabChar instance

instance IsSymInterface sym => CanMux sym CharType where
  {-# INLINE muxReg #-}
  muxReg s = \_ -> eqMergeFn s "characters"

------------------------------------------------------------------------
-- RegValue Maybe instance

mergePartExpr :: IsExprBuilder sym
              => sym
              -> (Pred sym -> v -> v -> IO v)
              -> Pred sym
              -> PartExpr (Pred sym) v
              -> PartExpr (Pred sym) v
              -> IO (PartExpr (Pred sym) v)
mergePartExpr sym fn = mergePartial sym (\c a b -> lift (fn c a b))

instance (IsExprBuilder sym, CanMux sym tp) => CanMux sym (MaybeType tp) where
  {-# INLINE muxReg #-}
  muxReg s = \_ -> do
    let f = muxReg s (Proxy :: Proxy tp)
     in mergePartExpr s f

------------------------------------------------------------------------
-- RegValue FunctionHandleType instance

-- TODO: Figure out how to actually compare these.
{-# INLINE muxHandle #-}
muxHandle :: IsExpr (SymExpr sym)
          => sym
          -> Pred sym
          -> FnVal sym a r
          -> FnVal sym a r
          -> IO (FnVal sym a r)
muxHandle _ c x y
  | Just b <- asConstantPred c = pure $! if b then x else y
  | otherwise = return x

instance IsExprBuilder sym => CanMux sym (FunctionHandleType a r) where
  {-# INLINE muxReg #-}
  muxReg s = \_ c x y -> do
    muxHandle s c x y

------------------------------------------------------------------------
-- RegValue IdentValueMap instance

-- | Merge to string maps together.
{-# INLINE muxStringMap #-}
muxStringMap :: IsExprBuilder sym
             => sym
             -> MuxFn (Pred sym) e
             -> MuxFn (Pred sym) (Map Text (PartExpr (Pred sym) e))
muxStringMap sym = \f c x y -> do
  let keys = Set.toList $ Set.union (Map.keysSet x) (Map.keysSet y)
  fmap Map.fromList $ forM keys $ \k -> do
    let vx = joinMaybePE (Map.lookup k x)
    let vy = joinMaybePE (Map.lookup k y)
    r <- mergePartExpr sym f c vx vy
    return (k,r)

------------------------------------------------------------------------
-- RegValue Recursive instance

newtype RolledType sym nm ctx = RolledType { unroll :: RegValue sym (UnrollType nm ctx) }


{-# INLINE muxRecursive #-}
muxRecursive
   :: IsRecursiveType nm
   => (forall tp. TypeRepr tp -> ValMuxFn sym tp)
   -> SymbolRepr nm
   -> CtxRepr ctx
   -> ValMuxFn sym (RecursiveType nm ctx)
muxRecursive recf = \nm ctx p x y -> do
   RolledType <$> recf (unrollType nm ctx) p (unroll x) (unroll y)

------------------------------------------------------------------------
-- RegValue Struct instance

{-# INLINE muxStruct #-}
muxStruct
   :: (forall tp. TypeRepr tp -> ValMuxFn sym tp)
   -> CtxRepr ctx
   -> ValMuxFn sym (StructType ctx)
muxStruct recf ctx = \p x y ->
  Ctx.generateM (Ctx.size ctx) $ \i -> do
    RV <$> recf (ctx Ctx.! i) p (unRV $ x Ctx.! i) (unRV $ y Ctx.! i)

------------------------------------------------------------------------
-- RegValue Variant instance

newtype VariantBranch sym tp = VB { unVB :: PartExpr (Pred sym) (RegValue sym tp) }

-- | Construct a 'VariantType' value by identifying which branch of
--   the variant to construct, and providing a value of the correct type.
injectVariant ::
  IsExprBuilder sym =>
  sym {- ^ symbolic backend -} ->
  CtxRepr ctx {- ^ Types of the variant branches -} ->
  Ctx.Index ctx tp {- ^ Which branch -} ->
  RegValue sym tp  {- ^ The value to inject -} ->
  RegValue sym (VariantType ctx)
injectVariant sym ctxRepr idx val =
  Ctx.generate (Ctx.size ctxRepr) $ \j ->
    case testEquality j idx of
      Just Refl -> VB (PE (truePred sym) val)
      Nothing -> VB Unassigned


{-# INLINE muxVariant #-}
muxVariant
   :: IsExprBuilder sym
   => sym
   -> (forall tp. TypeRepr tp -> ValMuxFn sym tp)
   -> CtxRepr ctx
   -> ValMuxFn sym (VariantType ctx)
muxVariant sym recf ctx = \p x y ->
  Ctx.generateM (Ctx.size ctx) $ \i ->
     VB <$> mergePartExpr sym
                (recf (ctx Ctx.! i))
                p
                (unVB (x Ctx.! i))
                (unVB (y Ctx.! i))