content: {-# LANGUAGE GADTs #-}
  
  -- | Translate Copilot Core expressions and operators to C99.
  module Copilot.Compile.C99.Expr
      ( transExpr
      , constArray
      )
    where
  
  -- External imports
  import           Control.Monad.State ( State, get, modify )
  import qualified Data.List.NonEmpty  as NonEmpty
  import qualified Language.C99.Simple as C
  
  -- Internal imports: Copilot
  import Copilot.Core ( Array, Expr (..), Field (..), Op1 (..), Op2 (..),
                        Op3 (..), Type (..), Value (..), accessorName,
                        arrayElems, toValues, typeLength, typeSize )
  
  -- Internal imports
  import Copilot.Compile.C99.Error ( impossible )
  import Copilot.Compile.C99.Name  ( exCpyName, streamAccessorName )
  import Copilot.Compile.C99.Type  ( transLocalVarDeclType, transType,
                                     transTypeName )
  
  -- | Translates a Copilot Core expression into a C99 expression.
  transExpr :: Expr a -> State FunEnv C.Expr
  transExpr (Const ty x) = return $ constTy ty x
  
  transExpr (Local ty1 _ name e1 e2) = do
    e1' <- transExpr e1
    let cTy1     = transLocalVarDeclType ty1
        initExpr = Just $ C.InitExpr e1'
  
    -- Add new decl to the tail of the fun env
    modify (\(i, x, y)
               -> (i, x ++ [C.VarDecln Nothing cTy1 name initExpr], y)
           )
  
    transExpr e2
  
  transExpr (Var _ n) = return $ C.Ident n
  
  transExpr (Drop _ amount sId) = do
    let accessVar = streamAccessorName sId
        index     = C.LitInt (fromIntegral amount)
    return $ funCall accessVar [index]
  
  transExpr (ExternVar _ name _) = return $ C.Ident (exCpyName name)
  
  transExpr (Label _ _ e) = transExpr e -- ignore label
  
  transExpr (Op1 op e) = do
    e' <- transExpr e
    return $ transOp1 op e'
  
  transExpr (Op2 (UpdateField ty1@(Struct _) ty2 f) e1 e2) = do
    -- Translating a struct update Op requires initializing a variable to the
    -- "old" value, updating the field, and returning the new value.
    e1' <- transExpr e1
    e2' <- transExpr e2
  
    -- Variable to hold the updated struct
    (i, _, _) <- get
    let varName = "_v" ++ show i
    modify (\(i, x, y) -> (i + 1, x, y))
  
    -- Add new var decl
    let initDecl = C.VarDecln Nothing cTy1 varName Nothing
        cTy1     = transLocalVarDeclType ty1
    modify (\(i, x, y) -> (i, x ++ [initDecl], y))
  
    -- Initialize the var to the same value as the original struct
    let initStmt = C.Expr
                 $ C.AssignOp
                     C.Assign
                     (C.Ident varName)
                     e1'
  
    -- Update field f with given value e2.
    let updateStmt = case ty2 of
          Array _ -> C.Expr $ memcpy dest e2' size
            where
              dest = C.Dot (C.Ident varName) (accessorName f)
              size = C.LitInt
                         (fromIntegral $ typeSize ty2)
                         C..* C.SizeOfType (C.TypeName (tyElemName ty2))
  
          _       -> C.Expr
                   $ C.AssignOp
                       C.Assign
                       (C.Dot (C.Ident varName) (accessorName f))
                       e2'
  
    modify (\(i, x, y) -> (i, x, y ++ [ initStmt, updateStmt ]))
  
    return $ C.Ident varName
  
  transExpr (Op2 op e1 e2) = do
    e1' <- transExpr e1
    e2' <- transExpr e2
    return $ transOp2 op e1' e2'
  
  transExpr e@(Op3 (UpdateArray arrTy@(Array ty2)) e1 e2 e3) = do
    e1' <- transExpr e1
    e2' <- transExpr e2
    e3' <- transExpr e3
  
    -- Variable to hold the updated array
    (i, _, _) <- get
    let varName = "_v" ++ show i
    modify (\(i, x, y) -> (i + 1, x, y))
  
    -- Add new var decl
    let initDecl = C.VarDecln Nothing cTy1 varName Nothing
        cTy1     = transType arrTy
    modify (\(i, x, y) -> (i, x ++ [initDecl], y))
  
    let size :: Type (Array n t) -> C.Expr
        size arrTy@(Array ty) = C.LitInt (fromIntegral $ typeLength arrTy)
                           C..* C.SizeOfType (C.TypeName $ transType ty)
  
    -- Initialize the var to the same value as the original array
    let initStmt = C.Expr $ memcpy (C.Ident varName) e1' (size arrTy)
  
    -- Update element of array
    let updateStmt = case ty2 of
          Array _ -> C.Expr $ memcpy dest e3' size
            where
              dest = C.Index (C.Ident varName) e2'
              size = C.LitInt
                         (fromIntegral $ typeSize ty2)
                         C..* C.SizeOfType (C.TypeName (tyElemName ty2))
  
          _ -> C.Expr
             $ C.AssignOp
                   C.Assign
                   (C.Index (C.Ident varName) e2')
                   e3'
  
    modify (\(i, x, y) -> (i, x, y ++ [ initStmt, updateStmt ]))
  
    return $ C.Ident varName
  
  transExpr (Op3 op e1 e2 e3) = do
    e1' <- transExpr e1
    e2' <- transExpr e2
    e3' <- transExpr e3
    return $ transOp3 op e1' e2' e3'
  
  -- | Translates a Copilot unary operator and its argument into a C99
  -- expression.
  transOp1 :: Op1 a b -> C.Expr -> C.Expr
  transOp1 op e =
    -- There are three types of ways in which a function in Copilot Core can be
    -- translated into C:
    --
    -- 1) Direct translation (perfect 1-to-1 mapping)
    -- 2) Type-directed translation (1-to-many mapping, choice based on type)
    -- 3) Desugaring/complex (expands to complex expression)
    case op of
      Not           -> (C..!) e
      Abs      ty   -> transAbs ty e
      Sign     ty   -> transSign ty e
      Recip    ty   -> constNumTy ty 1 C../ e
      Acos     ty   -> funCall (specializeMathFunName ty "acos") [e]
      Asin     ty   -> funCall (specializeMathFunName ty "asin") [e]
      Atan     ty   -> funCall (specializeMathFunName ty "atan") [e]
      Cos      ty   -> funCall (specializeMathFunName ty "cos") [e]
      Sin      ty   -> funCall (specializeMathFunName ty "sin") [e]
      Tan      ty   -> funCall (specializeMathFunName ty "tan") [e]
      Acosh    ty   -> funCall (specializeMathFunName ty "acosh") [e]
      Asinh    ty   -> funCall (specializeMathFunName ty "asinh") [e]
      Atanh    ty   -> funCall (specializeMathFunName ty "atanh") [e]
      Cosh     ty   -> funCall (specializeMathFunName ty "cosh") [e]
      Sinh     ty   -> funCall (specializeMathFunName ty "sinh") [e]
      Tanh     ty   -> funCall (specializeMathFunName ty "tanh") [e]
      Exp      ty   -> funCall (specializeMathFunName ty "exp") [e]
      Log      ty   -> funCall (specializeMathFunName ty "log") [e]
      Sqrt     ty   -> funCall (specializeMathFunName ty "sqrt") [e]
      Ceiling  ty   -> funCall (specializeMathFunName ty "ceil") [e]
      Floor    ty   -> funCall (specializeMathFunName ty "floor") [e]
      BwNot    _    -> (C..~) e
      Cast     _ ty -> C.Cast (transTypeName ty) e
      GetField (Struct _)  _ f -> C.Dot e (accessorName f)
  
  -- | Translates a Copilot binary operator and its arguments into a C99
  -- expression.
  --
  -- PRE: op is not a struct update operation (i.e., 'UpdateField').
  transOp2 :: Op2 a b c -> C.Expr -> C.Expr -> C.Expr
  transOp2 op e1 e2 = case op of
    And               -> e1 C..&& e2
    Or                -> e1 C..|| e2
    Add         _     -> e1 C..+  e2
    Sub         _     -> e1 C..-  e2
    Mul         _     -> e1 C..*  e2
    Mod         _     -> e1 C..%  e2
    Div         _     -> e1 C../  e2
    Fdiv        _     -> e1 C../  e2
    Pow         ty    -> funCall (specializeMathFunName ty "pow") [e1, e2]
    Logb        ty    -> funCall (specializeMathFunName ty "log") [e2] C../
                         funCall (specializeMathFunName ty "log") [e1]
    Atan2       ty    -> funCall (specializeMathFunName ty "atan2") [e1, e2]
    Eq          _     -> e1 C..== e2
    Ne          _     -> e1 C..!= e2
    Le          _     -> e1 C..<= e2
    Ge          _     -> e1 C..>= e2
    Lt          _     -> e1 C..<  e2
    Gt          _     -> e1 C..>  e2
    BwAnd       _     -> e1 C..&  e2
    BwOr        _     -> e1 C..|  e2
    BwXor       _     -> e1 C..^  e2
    BwShiftL    _ _   -> e1 C..<< e2
    BwShiftR    _ _   -> e1 C..>> e2
    Index       _     -> C.Index e1 e2
    UpdateField _ _ _ -> impossible "transOp2" "copilot-c99"
  
  -- | Translates a Copilot ternary operator and its arguments into a C99
  -- expression.
  transOp3 :: Op3 a b c d -> C.Expr -> C.Expr -> C.Expr -> C.Expr
  transOp3 op e1 e2 e3 = case op of
    Mux         _ -> C.Cond e1 e2 e3
    UpdateArray _ -> impossible "transOp3" "copilot-c99"
  
  -- | Translate @'Abs' e@ in Copilot Core into a C99 expression.
  --
  -- This function produces a portable implementation of abs in C99 that works
  -- for the type given, provided that the output fits in a variable of the same
  -- type (which may not be true, for example, for signed integers in the lower
  -- end of their type range). If the absolute value is out of range, the
  -- behavior is undefined.
  --
  -- PRE: The type given is a Num type (floating-point number, or a
  -- signed/unsigned integer of fixed size).
  transAbs :: Type a -> C.Expr -> C.Expr
  transAbs ty e
      -- Abs for floats/doubles is called fabs in C99's math.h.
      | typeIsFloating ty
      = funCall (specializeMathFunName ty "fabs") [e]
  
      -- C99 provides multiple implementations of abs, depending on the type of
      -- the arguments. For integers, it provides C99 abs, labs, and llabs, which
      -- take, respectively, an int, a long int, and a long long int.
      --
      -- However, the code produced by Copilot uses types with fixed width (e.g.,
      -- int16_t), and there is no guarantee that, for example, 32-bit int or
      -- 64-bit int will fit in a C int (only guaranteed to be 16 bits).
      -- Consequently, this function provides a portable version of abs for signed
      -- and unsigned ints implemented using shift and xor. For example, for a
      -- value x of type int32_t, the absolute value is:
      -- (x + (x >> sizeof(int32_t)-1)) ^ (x >> sizeof(int32_t)-1))
      | otherwise
      = (e C..+ (e C..>> tyBitSizeMinus1)) C..^ (e C..>> tyBitSizeMinus1)
    where
      -- Size of an integer type in bits, minus one. It's easier to hard-code
      -- them than to try and generate the right expressions in C using sizeof.
      --
      -- PRE: the type 'ty' is a signed or unsigned integer type.
      tyBitSizeMinus1 :: C.Expr
      tyBitSizeMinus1 = case ty of
        Int8   -> C.LitInt 7
        Int16  -> C.LitInt 15
        Int32  -> C.LitInt 31
        Int64  -> C.LitInt 63
        Word8  -> C.LitInt 7
        Word16 -> C.LitInt 15
        Word32 -> C.LitInt 31
        Word64 -> C.LitInt 63
        _      -> impossible
                    "transAbs"
                    "copilot-c99"
                    "Abs applied to unexpected types."
  
  -- | Translate @'Sign' e@ in Copilot Core into a C99 expression.
  --
  -- Sign is is translated as @e > 0 ? 1 : (e < 0 ? -1 : e)@, that is:
  --
  -- 1. If @e@ is positive, return @1@.
  --
  -- 2. If @e@ is negative, return @-1@.
  --
  -- 3. Otherwise, return @e@. This handles the case where @e@ is @0@ when the
  --    type is an integral type. If the type is a floating-point type, it also
  --    handles the cases where @e@ is @-0@ or @NaN@.
  --
  -- This implementation is modeled after how GHC implements 'signum'
  -- <https://gitlab.haskell.org/ghc/ghc/-/blob/aed98ddaf72cc38fb570d8415cac5de9d8888818/libraries/base/GHC/Float.hs#L523-L525 here>.
  transSign :: Type a -> C.Expr -> C.Expr
  transSign ty e = positiveCase $ negativeCase e
    where
      -- If @e@ is positive, return @1@, otherwise fall back to argument.
      --
      -- Produces the following code, where @<arg>@ is the argument to this
      -- function:
      -- @
      -- e > 0 ? 1 : <arg>
      -- @
      positiveCase :: C.Expr  -- ^ Value returned if @e@ is not positive.
                   -> C.Expr
      positiveCase =
        C.Cond (C.BinaryOp C.GT e (constNumTy ty 0)) (constNumTy ty 1)
  
      -- If @e@ is negative, return @1@, otherwise fall back to argument.
      --
      -- Produces the following code, where @<arg>@ is the argument to this
      -- function:
      -- @
      -- e < 0 ? -1 : <arg>
      -- @
      negativeCase :: C.Expr  -- ^ Value returned if @e@ is not negative.
                   -> C.Expr
      negativeCase =
        C.Cond (C.BinaryOp C.LT e (constNumTy ty 0)) (constNumTy ty (-1))
  
  -- | Transform a Copilot Core literal, based on its value and type, into a C99
  -- literal.
  constTy :: Type a -> a -> C.Expr
  constTy ty = case ty of
    Bool      -> C.LitBool
    Int8      -> explicitTy ty . C.LitInt . fromIntegral
    Int16     -> explicitTy ty . C.LitInt . fromIntegral
    Int32     -> explicitTy ty . C.LitInt . fromIntegral
    Int64     -> explicitTy ty . C.LitInt . fromIntegral
    Word8     -> explicitTy ty . C.LitInt . fromIntegral
    Word16    -> explicitTy ty . C.LitInt . fromIntegral
    Word32    -> explicitTy ty . C.LitInt . fromIntegral
    Word64    -> explicitTy ty . C.LitInt . fromIntegral
    Float     -> explicitTy ty . C.LitFloat
    Double    -> explicitTy ty . C.LitDouble
    Struct _  -> C.InitVal (transTypeName ty) . constStruct . toValues
    Array ty' -> C.InitVal (transTypeName ty) . constArray ty' . arrayElems
  
  -- | Transform a Copilot Core literal, based on its value and type, into a C99
  -- initializer.
  constInit :: Type a -> a -> C.Init
  constInit ty val = case ty of
    -- We include two special cases for Struct and Array to avoid using constTy
    -- on them.
    --
    -- In the default case (i.e., InitExpr (constTy ty val)), constant
    -- initializations are explicitly cast. However, doing so 1) may result in
    -- incorrect values for arrays, and 2) will be considered a non-constant
    -- expression in the case of arrays and structs, and thus not allowed as the
    -- initialization value for a global variable.
    --
    -- In particular, wrt. (1), for example, the nested array:
    --   [[0, 1], [2, 3]] :: Array 2 (Array 2 Int32)
    --
    -- with explicit casts, will be initialized in C as:
    --   { (int32_t[2]){(int32_t)(0), (int32_t)(1)},
    --     (int32_t[2]){(int32_t)(2), (int32_t)(3)} }
    --
    -- Due to the additional (int32_t[2]) casts, a C compiler will interpret the
    -- whole expression as an array of two int32_t's (as opposed to a nested
    -- array). This can either lead to compile-time errors (if you're lucky) or
    -- incorrect runtime semantics (if you're unlucky).
    Array ty' -> C.InitList $ constArray ty' $ arrayElems val
  
    -- We use InitArray to initialize a struct because the syntax used for
    -- initializing arrays and structs is compatible. For instance, {1, 2} works
    -- both for initializing an int array of length 2 as well as a struct with
    -- two int fields, although the two expressions are conceptually different
    -- (structs can also be initialized as { .a = 1, .b = 2}.
    Struct _  -> C.InitList $ constStruct (toValues val)
    _         -> C.InitExpr $ constTy ty val
  
  -- | Transform a Copilot Core struct field into a C99 initializer.
  constFieldInit :: Value a -> C.InitItem
  constFieldInit (Value ty (Field val)) = C.InitItem Nothing $ constInit ty val
  
  -- | Transform a Copilot Struct, based on the struct fields, into a list of C99
  -- initializer values.
  constStruct :: [Value a] -> NonEmpty.NonEmpty C.InitItem
  constStruct val = NonEmpty.fromList $ map constFieldInit val
  
  -- | Transform a Copilot Array, based on the element values and their type,
  -- into a list of C99 initializer values.
  constArray :: Type a -> [a] -> NonEmpty.NonEmpty C.InitItem
  constArray ty =
    NonEmpty.fromList . map (C.InitItem Nothing . constInit ty)
  
  -- | Explicitly cast a C99 value to a type.
  explicitTy :: Type a -> C.Expr -> C.Expr
  explicitTy ty = C.Cast (transTypeName ty)
  
  -- Translate a literal number of type @ty@ into a C99 literal.
  --
  -- PRE: The type of PRE is numeric (integer or floating-point), that
  -- is, not boolean, struct or array.
  constNumTy :: Type a -> Integer -> C.Expr
  constNumTy ty =
    case ty of
      Float  -> C.LitFloat . fromInteger
      Double -> C.LitDouble . fromInteger
      _      -> C.LitInt
  
  -- | Provide a specialized function name in C99 for a function given the type
  -- of its arguments, and its "family" name.
  --
  -- C99 provides multiple variants of the same conceptual function, based on the
  -- types. Depending on the function, common variants exist for signed/unsigned
  -- arguments, long or short types, float or double. The C99 standard uses the
  -- same mechanism to name most such functions: the default variant works for
  -- double, and there are additional variants for float and long double. For
  -- example, the sin function operates on double, while sinf operates on float,
  -- and sinl operates on long double.
  --
  -- This function only knows how to provide specialized names for functions in
  -- math.h that provide a default version for a double argument and vary for
  -- floats. It won't change the function name given if the variation is based on
  -- the return type, if the function is defined elsewhere, or for other types.
  specializeMathFunName :: Type a -> String -> String
  specializeMathFunName ty s
      -- The following function pattern matches based on the variants available
      -- for a specific function.
      --
      -- Do not assume that a function you need implemented follows the same
      -- standard as others: check whether it is present in the standard.
      | isMathFPArgs s
      , Float <- ty
      = s ++ "f"
  
      | otherwise
      = s
    where
      -- True if the function family name is part of math.h and follows the
      -- standard rule of providing multiple variants for floating point numbers
      -- based on the type of their arguments.
      --
      -- Note: nan is not in this list because the names of its variants are
      -- determined by the return type.
      --
      -- For details, see:
      -- "B.11 Mathematics <math.h>" in the C99 standard
      isMathFPArgs :: String -> Bool
      isMathFPArgs = flip elem
         [ "acos",   "asin",     "atan",      "atan2",      "cos",    "sin"
         , "tan",    "acosh",    "asinh",     "atanh",      "cosh",   "sinh"
         , "tanh",   "exp",      "exp2",      "expm1",      "frexp",  "ilogb"
         , "ldexp",  "log",      "log10",     "log1p",      "log2",   "logb"
         , "modf",   "scalbn",   "scalbln",   "cbrt",       "fabs",   "hypot"
         , "pow",    "sqrt",     "erf",       "erfc",       "lgamma", "tgamma"
         , "ceil",   "floor",    "nearbyint", "rint",       "lrint",  "llrint"
         , "round",  "lround",   "llround",   "trunc",      "fmod",   "remainder"
         , "remquo", "copysign", "nextafter", "nexttoward", "fdim"
         , "fmax",   "fmin",     "fma"
         ]
  
  -- * Auxiliary functions
  
  -- | True if the type given is a floating point number.
  typeIsFloating :: Type a -> Bool
  typeIsFloating Float  = True
  typeIsFloating Double = True
  typeIsFloating _      = False
  
  -- | Auxiliary type used to collect all the declarations of all the variables
  -- used in a function to be generated, since variable declarations are always
  -- listed first at the top of the function body.
  type FunEnv = (Int, [C.Decln], [C.Stmt])
  
  -- | Define a C expression that calls a function with arguments.
  funCall :: C.Ident   -- ^ Function name
          -> [C.Expr]  -- ^ Arguments
          -> C.Expr
  funCall name = C.Funcall (C.Ident name)
  
  -- Write a call to the memcpy function.
  memcpy :: C.Expr -> C.Expr -> C.Expr -> C.Expr
  memcpy dest src size = C.Funcall (C.Ident "memcpy") [dest, src, size]
  
  -- Translate a Copilot type to a C99 type, handling arrays especially.
  --
  -- If the given type is an array (including multi-dimensional arrays), the
  -- type is that of the elements in the array. Otherwise, it is just the
  -- equivalent representation of the given type in C.
  tyElemName :: Type a -> C.Type
  tyElemName ty = case ty of
    Array ty' -> tyElemName ty'
    _         -> transType ty
  