content: {-# LANGUAGE ConstraintKinds        #-}
  {-# LANGUAGE DataKinds              #-}
  {-# LANGUAGE FlexibleContexts       #-}
  {-# LANGUAGE FlexibleInstances      #-}
  {-# LANGUAGE FunctionalDependencies #-}
  {-# LANGUAGE GADTs                  #-}
  {-# LANGUAGE KindSignatures         #-}
  {-# LANGUAGE LambdaCase             #-}
  {-# LANGUAGE MultiParamTypeClasses  #-}
  {-# LANGUAGE PolyKinds              #-}
  {-# LANGUAGE RankNTypes             #-}
  {-# LANGUAGE ScopedTypeVariables    #-}
  {-# LANGUAGE UndecidableInstances   #-}
  {-# LANGUAGE TypeApplications       #-}
  {-# LANGUAGE CPP                    #-}
  
  {-# OPTIONS_GHC -Wall #-}
  
  module Language.Fortran.Model.Op.Core.Eval
    ( MonadEvalFortran
    , evalCoreOp
    ) where
  
  import           Control.Applicative                  (liftA2)
  
  import           Data.SBV.Dynamic                     (SVal)
  import qualified Data.SBV.Dynamic                     as SBV
  
  #if MIN_VERSION_singletons(3,0,0)
  import           Data.List.Singletons
  import           GHC.TypeLits.Singletons
  #else
  import           Data.Singletons
  import           Data.Singletons.Prelude.List
  import           Data.Singletons.TypeLits
  #endif
  
  import           Data.Vinyl                           hiding (Field)
  import           Data.Vinyl.Curry
  
  import           Language.Fortran.Model.Op.Core.Core
  import           Language.Fortran.Model.Op.Core.Match
  import           Language.Fortran.Model.Op.Eval
  import           Language.Fortran.Model.Repr
  import           Language.Fortran.Model.Repr.Prim
  import           Language.Fortran.Model.Singletons
  import           Language.Fortran.Model.Types
  import           Language.Fortran.Model.Types.Match
  
  --------------------------------------------------------------------------------
  --  Evaluation
  --------------------------------------------------------------------------------
  
  evalCoreOp
    :: (MonadEvalFortran r m)
    => Op (Length args) ok -> OpSpec ok args result -> Rec CoreRepr args -> m (CoreRepr result)
  evalCoreOp op = \case
    OSLit px x -> \_ -> primFromVal px <$> primLit px x
  
    OSNum1 _ _ p2 ->
      primUnop True p2 (numUnop op)
    OSNum2 nk1 nk2 p1 p2 p3 ->
      primBinop True p1 p2 p3 (numBinop (nkBothInts nk1 nk2) op)
  
    OSLogical1 _ p2 -> primUnop True p2 (logicalUnop op)
    OSLogical2 p1 p2 p3 -> primBinop True p1 p2 p3 (logicalBinop op)
  
    OSEq cmp p1 p2 p3 -> primBinop False p1 p2 p3 (eqBinop cmp op)
    OSRel cmp p1 p2 p3 -> primBinop False p1 p2 p3 (relBinop cmp op)
  
    OSLookup _ -> return . runcurry lookupArr
  
    OSDeref _ s -> return . runcurry (derefData s)
  
  --------------------------------------------------------------------------------
  --  General
  --------------------------------------------------------------------------------
  
  primToVal :: CoreRepr (PrimS a) -> SVal
  primToVal = \case
    CRPrim (DPrim _) v -> v
  
  primFromVal :: Prim p k a -> SVal -> CoreRepr (PrimS a)
  primFromVal p v = CRPrim (DPrim p) v
  
  primUnop
    :: (MonadEvalFortran r m)
    => Bool
    -> Prim p2 k2 b -- ^ The target type
    -> (SVal -> SVal)
    -> Rec CoreRepr '[PrimS a] -> m (CoreRepr (PrimS b))
  primUnop shouldCoerce p2 f = runcurry $ fmap (primFromVal p2 . f) . maybeCoerce . primToVal
    where
      maybeCoerce
        | shouldCoerce = coercePrimSVal p2
        | otherwise = return
  
  primBinop
    :: (MonadEvalFortran r m)
    => Bool
    -- ^ True to coerce arguments to result value. False to coerce arguments to
    -- ceiling of each.
    -> Prim p1 k1 a -> Prim p2 k2 b -> Prim p3 k3 c
    -> (SVal -> SVal -> SVal)
    -> Rec CoreRepr '[PrimS a, PrimS b] -> m (CoreRepr (PrimS c))
  primBinop takesResultVal p1 p2 p3 (.*.) =
    fmap (primFromVal p3) .
    runcurry (\x y -> liftA2 (.*.) (coerceArg (primToVal x)) (coerceArg (primToVal y)))
  
    where
      coerceToCeil = case primCeil p1 p2 of
        Just (MakePrim pCeil) -> coercePrimSVal pCeil
        _                     -> return
  
      coerceArg
        | takesResultVal = coercePrimSVal p3
        | otherwise = coerceToCeil
  
  --------------------------------------------------------------------------------
  --  Numeric
  --------------------------------------------------------------------------------
  
  -- TODO: Worry about what happens when LHS and RHS have different types
  
  nkBothInts :: NumericBasicType k1 -> NumericBasicType k2 -> Bool
  nkBothInts NBTInt NBTInt = True
  nkBothInts _ _           = False
  
  numUnop :: Op 1 'OKNum -> SVal -> SVal
  numUnop = \case
    OpNeg -> SBV.svUNeg
    OpPos -> id
  
  numBinop :: Bool -> Op 2 'OKNum -> SVal -> SVal -> SVal
  numBinop isInt = \case
    OpAdd -> SBV.svPlus
    OpSub -> SBV.svMinus
    OpMul -> SBV.svTimes
    OpDiv -> if isInt then SBV.svQuot else SBV.svDivide
  
  --------------------------------------------------------------------------------
  --  Logical
  --------------------------------------------------------------------------------
  
  -- TODO: Always return single bits, special-case when inputs are not single bits.
  -- TODO: Worry about what happens when LHS and RHS have different types
  
  logicalUnop :: Op 1 'OKLogical -> SVal -> SVal
  logicalUnop = \case
    OpNot -> SBV.svNot
  
  logicalBinop :: Op 2 'OKLogical -> SVal -> SVal -> SVal
  logicalBinop = \case
    OpAnd -> SBV.svAnd
    OpOr -> SBV.svOr
    OpEquiv -> svEquiv
    OpNotEquiv -> svNotEquiv
    where
      svEquiv x y = (x `SBV.svAnd` y) `SBV.svOr` (SBV.svNot x `SBV.svAnd` SBV.svNot y)
      svNotEquiv x y = SBV.svNot (x `svEquiv` y)
  
  --------------------------------------------------------------------------------
  --  Equality
  --------------------------------------------------------------------------------
  
  -- TODO: Worry about what happens when LHS and RHS have different types
  
  eqBinop :: ComparableBasicTypes k1 k2 -> Op 2 'OKEq -> SVal -> SVal -> SVal
  eqBinop _ = \case
    OpEq -> SBV.svEqual
    OpNE -> SBV.svNotEqual
  
  --------------------------------------------------------------------------------
  --  Relational
  --------------------------------------------------------------------------------
  
  -- TODO: Worry about what happens when LHS and RHS have different types
  
  relBinop :: ComparableBasicTypes k1 k2 -> Op 2 'OKRel -> SVal -> SVal -> SVal
  relBinop _ = \case
    OpLT -> SBV.svLessThan
    OpLE -> SBV.svLessEq
    OpGT -> SBV.svGreaterThan
    OpGE -> SBV.svGreaterEq
  
  --------------------------------------------------------------------------------
  --  Lookup
  --------------------------------------------------------------------------------
  
  lookupArrRepr
    :: CoreRepr i
    -> D (Array i v)
    -> ArrRepr i v
    -> CoreRepr v
  lookupArrRepr ixRepr (DArray ixIndex@(Index _) valAV) arrRepr =
    case ixRepr of
      CRPrim _ ixVal -> case (valAV, arrRepr) of
        (ArrPrim _, ARPrim arr) ->
          CRPrim (dArrValue valAV) (SBV.readSArr arr ixVal)
        (ArrData _ dfields, ARData afields) ->
          let avD = (dArrValue valAV)
          in CRData avD (
            rzipWith
            (zipFieldsWith (lookupArrRepr ixRepr . DArray ixIndex))
            dfields
            afields)
  
  lookupArr
    :: CoreRepr (Array i v)
    -> CoreRepr i
    -> CoreRepr v
  lookupArr (CRArray arrD arrRepr) ixRepr = lookupArrRepr ixRepr arrD arrRepr
  
  --------------------------------------------------------------------------------
  --  Deref
  --------------------------------------------------------------------------------
  
  derefData
    :: forall a fname fields i rname. RElem '(fname, a) fields i
    => SSymbol fname
    -> CoreRepr (Record rname fields)
    -> CoreRepr a
  derefData _ (CRData _ dataRec) =
        case rget @'(fname, a) @fields dataRec of
      Field _ x -> x
  
  --------------------------------------------------------------------------------
  --  Equality of operators
  --------------------------------------------------------------------------------
  
  -- eqOpRArgs
  --   :: (forall x y. (x -> y -> Bool) -> f x -> g y -> Bool)
  --   -> OpSpec ok1 args1 r1
  --   -> OpSpec ok2 args2 r2
  --   -> Rec f args1
  --   -> Rec g args2
  --   -> Bool
  -- eqOpRArgs le opr1 opr2 =
  --   case (opr1, opr2) of
  --     (OSLit _ _, OSLit _ _) -> \_ _ -> True
  --     (OSNum1 _ px _, OSNum1 _ py _) -> runcurry $ runcurry . le (eqPrimS px py)
  --     (OSNum2 _ _ px1 px2 _, OSNum2 _ _ py1 py2 _) ->
  --       runcurry $ \x1 x2 -> runcurry $ \y1 y2 ->
  --       le (eqPrimS px1 py1) x1 y1 && le (eqPrimS px2 py2) x2 y2
  --     (OSLogical1 px _, OSLogical1 py _) -> runcurry $ runcurry . le (eqPrimS px py)
  --     (OSLogical2 px1 px2 _, OSLogical2 py1 py2 _) ->
  --       runcurry $ \x1 x2 -> runcurry $ \y1 y2 ->
  --       le (eqPrimS px1 py1) x1 y1 && le (eqPrimS px2 py2) x2 y2
  --     (OSEq _ px1 px2 _, OSEq _ py1 py2 _) ->
  --       runcurry $ \x1 x2 -> runcurry $ \y1 y2 ->
  --       le (eqPrimS px1 py1) x1 y1 && le (eqPrimS px2 py2) x2 y2
  --     (OSRel _ px1 px2 _, OSRel _ py1 py2 _) ->
  --       runcurry $ \x1 x2 -> runcurry $ \y1 y2 ->
  --       le (eqPrimS px1 py1) x1 y1 && le (eqPrimS px2 py2) x2 y2
  --     (OSLookup (DArray (Index pi1) _), OSLookup (DArray (Index pi2) _)) ->
  --       runcurry $ \arr1 i1 -> runcurry $ \arr2 i2 ->
  --       le _ arr1 arr2 && le (eqPrimS pi1 pi2) i1 i2
  --       -- le ()
  
  -- liftEqRec
  --   :: (forall a b. f a -> g b -> Bool)
  --   -> Rec f xs -> Rec g ys
  --   -> Bool
  -- liftEqRec f r1 r2 = case (r1, r2) of
  --   (RNil, RNil) -> True
  --   (h1 :& t1, h2 :& t2) -> f h1 h2 && liftEqRec f t1 t2
  --   _ -> False
  