{-# LANGUAGE DataKinds                 #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE FlexibleContexts          #-}
{-# LANGUAGE FlexibleInstances         #-}
{-# LANGUAGE GADTs                     #-}
{-# LANGUAGE KindSignatures            #-}
{-# LANGUAGE ScopedTypeVariables       #-}
{-# LANGUAGE StandaloneDeriving        #-}
{-# LANGUAGE Trustworthy               #-}
{-# LANGUAGE TypeApplications          #-}
{-# LANGUAGE TypeOperators             #-}
{-# LANGUAGE UndecidableInstances      #-}
-- |
-- Description: Typing for Core.
-- Copyright:   (c) 2011 National Institute of Aerospace / Galois, Inc.
--
-- All expressions and streams in Core are accompanied by a representation of
-- the types of the underlying expressions used or carried by the streams.
-- This information is needed by the compiler to generate code, since it must
-- initialize variables and equivalent representations for those types in
-- the target languages.
module Copilot.Core.Type
    ( Type (..)
    , Typed (..)
    , typeOfDefault
    , UType (..)
    , SimpleType (..)

    , typeSize
    , typeLength

    , Value (..)
    , toValues
    , toValuesDefault
    , Field (..)
    , typeName
    , typeNameDefault

    , Struct
    , fieldName
    , accessorName
    , updateField
    , updateFieldDefault
    )
  where

-- External imports
import Data.Char          (isLower, isUpper, toLower)
import Data.Coerce        (coerce)
import Data.Int           (Int16, Int32, Int64, Int8)
import Data.List          (intercalate)
import Data.Proxy         (Proxy (..))
import Data.Type.Equality as DE
import Data.Typeable      (Typeable, eqT, typeRep)
import Data.Word          (Word16, Word32, Word64, Word8)
import GHC.Generics       (Datatype (..), D1, Generic (..), K1 (..), M1 (..),
                           U1 (..), (:*:) (..))
import GHC.TypeLits       (KnownNat, KnownSymbol, Symbol, natVal, sameNat,
                           sameSymbol, symbolVal)

-- Internal imports
import Copilot.Core.Type.Array (Array)

-- | The value of that is a product or struct, defined as a constructor with
-- several fields.
class Struct a where
  -- | Returns the name of struct in the target language.
  typeName :: a -> String

  -- | Transforms all the struct's fields into a list of values.
  toValues :: a -> [Value a]

  -- | Update the value of a struct field. This is only used by the Copilot
  -- interpreter.
  --
  -- If you do not plan to use the interpreter, you can omit an implementation
  -- of this method. If you do so, it is recommended that you derive a 'Generic'
  -- instance for the struct data type. This is because in a future release, the
  -- default implementation of 'updateField' (which will be picked if there is
  -- not a manually written implementation) will be changed to require a
  -- 'Generic' instance.
  --
  -- In order to implement 'updateField', pick one of the following approaches:
  --
  -- * Derive a 'Generic' instance for the struct data type and then define
  --   @'updateField' = 'updateFieldDefault'@ in the 'Struct' instance.
  --
  -- * Manually implement 'updateField' by doing the following for each 'Field'
  --   in a struct:
  --
  --   1. Check that the name of the 'Field' matches the name of the supplied
  --      'Value' (using 'GHC.TypeLits.sameSymbol').
  --
  --   2. Check that the type of the 'Field' matches the 'Type' of the supplied
  --      'Value' (using 'DE.testEquality').
  --
  --   3. If both (1) and (2) succeed, update the corresponding struct field
  --      using a record update.
  --
  --   For a complete end-to-end example that demonstrates how to manually
  --   implement 'updateField' and use it in the Copilot interpreter, see the
  --   @examples/StructsUpdateField.hs@ example in the @copilot@ library.
  updateField :: a -> Value t -> a
  updateField = error $ unlines
    [ "Field updates not supported for this type."
    , "(Perhaps you need to implement 'updateField' for a 'Struct' instance?)"
    ]

-- | The field of a struct, together with a representation of its type.
data Value a =
  forall s t . (Typeable t, KnownSymbol s, Show t) => Value (Type t) (Field s t)

-- | A field in a struct. The name of the field is a literal at the type
-- level.
data Field (s :: Symbol) t = Field t

-- | Extract the name of a field.
fieldName :: forall s t . KnownSymbol s => Field s t -> String
fieldName _ = symbolVal (Proxy :: Proxy s)

-- | Extract the name of an accessor (a function that returns a field of a
-- struct).
accessorName :: forall a s t . (Struct a, KnownSymbol s)
             => (a -> Field s t) -> String
accessorName _ = symbolVal (Proxy :: Proxy s)

instance (KnownSymbol s, Show t) => Show (Field s t) where
  show f@(Field v) = fieldName f ++ ":" ++ show v

instance {-# OVERLAPPABLE #-} (Typed t, Struct t) => Show t where
  show t = "<" ++ fields ++ ">"
    where
      fields = intercalate "," $ map showfield (toValues t)
      showfield (Value _ field) = show field

-- | A Type representing the types of expressions or values handled by
-- Copilot Core.
--
-- Note that both arrays and structs use dependently typed features. In the
-- former, the length of the array is part of the type. In the latter, the
-- names of the fields are part of the type.
data Type :: * -> * where
  Bool   :: Type Bool
  Int8   :: Type Int8
  Int16  :: Type Int16
  Int32  :: Type Int32
  Int64  :: Type Int64
  Word8  :: Type Word8
  Word16 :: Type Word16
  Word32 :: Type Word32
  Word64 :: Type Word64
  Float  :: Type Float
  Double :: Type Double
  Array  :: forall n t . ( KnownNat n
                         , Typed t
                         ) => Type t -> Type (Array n t)
  Struct :: (Typed a, Struct a) => a -> Type a
deriving instance Show (Type a)

-- | Return the length of an array from its type
typeLength :: forall n t . KnownNat n => Type (Array n t) -> Int
typeLength _ = fromIntegral $ natVal (Proxy :: Proxy n)

-- | Return the total (nested) size of an array from its type
typeSize :: forall n t . KnownNat n => Type (Array n t) -> Int
typeSize ty@(Array ty'@(Array _)) = typeLength ty * typeSize ty'
typeSize ty@(Array _            ) = typeLength ty

instance TestEquality Type where
  testEquality Bool   Bool   = Just DE.Refl
  testEquality Int8   Int8   = Just DE.Refl
  testEquality Int16  Int16  = Just DE.Refl
  testEquality Int32  Int32  = Just DE.Refl
  testEquality Int64  Int64  = Just DE.Refl
  testEquality Word8  Word8  = Just DE.Refl
  testEquality Word16 Word16 = Just DE.Refl
  testEquality Word32 Word32 = Just DE.Refl
  testEquality Word64 Word64 = Just DE.Refl
  testEquality Float  Float  = Just DE.Refl
  testEquality Double Double = Just DE.Refl
  testEquality (Array t1) (Array t2) =
      testArrayEquality t1 t2
    where
      testArrayEquality :: forall n1 a1 n2 a2.
                           (KnownNat n1, KnownNat n2)
                        => Type a1
                        -> Type a2
                        -> Maybe (Array n1 a1 :~: Array n2 a2)
      testArrayEquality ty1 ty2
        | Just DE.Refl <- sameNat (Proxy :: Proxy n1) (Proxy :: Proxy n2)
        , Just DE.Refl <- testEquality ty1 ty2
        = Just DE.Refl
        | otherwise
        = Nothing
  testEquality (Struct _) (Struct _) = eqT
  testEquality _ _ = Nothing

-- | A simple, monomorphic representation of types that facilitates putting
-- variables in heterogeneous lists and environments in spite of their types
-- being different.
data SimpleType where
  SBool   :: SimpleType
  SInt8   :: SimpleType
  SInt16  :: SimpleType
  SInt32  :: SimpleType
  SInt64  :: SimpleType
  SWord8  :: SimpleType
  SWord16 :: SimpleType
  SWord32 :: SimpleType
  SWord64 :: SimpleType
  SFloat  :: SimpleType
  SDouble :: SimpleType
  SArray  :: Type t -> SimpleType
  SStruct :: SimpleType

-- | Type equality, used to help type inference.

-- This instance is necessary, otherwise the type of SArray can't be inferred.
instance Eq SimpleType where
  SBool   == SBool   = True
  SInt8   == SInt8   = True
  SInt16  == SInt16  = True
  SInt32  == SInt32  = True
  SInt64  == SInt64  = True
  SWord8  == SWord8  = True
  SWord16 == SWord16 = True
  SWord32 == SWord32 = True
  SWord64 == SWord64 = True
  SFloat  == SFloat  = True
  SDouble == SDouble = True
  (SArray t1) == (SArray t2) | Just DE.Refl <- testEquality t1 t2 = True
                             | otherwise                          = False
  SStruct == SStruct = True
  _ == _ = False

-- | A typed expression, from which we can obtain the two type representations
-- used by Copilot: 'Type' and 'SimpleType'.
class (Show a, Typeable a) => Typed a where
  typeOf     :: Type a
  simpleType :: Type a -> SimpleType
  simpleType _ = SStruct

instance Typed Bool where
  typeOf       = Bool
  simpleType _ = SBool

instance Typed Int8 where
  typeOf       = Int8
  simpleType _ = SInt8

instance Typed Int16 where
  typeOf       = Int16
  simpleType _ = SInt16

instance Typed Int32 where
  typeOf       = Int32
  simpleType _ = SInt32

instance Typed Int64 where
  typeOf       = Int64
  simpleType _ = SInt64

instance Typed Word8 where
  typeOf       = Word8
  simpleType _ = SWord8

instance Typed Word16 where
  typeOf       = Word16
  simpleType _ = SWord16

instance Typed Word32 where
  typeOf       = Word32
  simpleType _ = SWord32

instance Typed Word64 where
  typeOf       = Word64
  simpleType _ = SWord64

instance Typed Float where
  typeOf       = Float
  simpleType _ = SFloat

instance Typed Double where
  typeOf       = Double
  simpleType _ = SDouble

instance (Typeable t, Typed t, KnownNat n) => Typed (Array n t) where
  typeOf               = Array typeOf
  simpleType (Array t) = SArray t

-- | A untyped type (no phantom type).
data UType = forall a . Typeable a => UType (Type a)

instance Eq UType where
  UType ty1 == UType ty2 = typeRep ty1 == typeRep ty2

-- * GHC.Generics-based defaults

-- | A default implementation of 'typeName' that leverages 'Generic'. In order
-- to use this, make sure you derive a 'Generic' instance for your data type and
-- then define @'typeName' = 'typeNameDefault'@ in its 'Struct' instance.
--
-- This generates a struct name that consists of the name of the original
-- Haskell data type, but converted to snake_case.
typeNameDefault :: (Generic a, GDatatype (Rep a)) => a -> String
typeNameDefault = convert . gTypeName . from
  where
    convert :: String -> String
    convert = convert' True True

    convert' :: Bool   -- ^ Is this the first letter
             -> Bool   -- ^ Was the previous letter capital
             -> String -- ^ Remainder of the string
             -> String
    convert' _ _ []    = []
    convert' _ v [x]
      | v && isUpper x = toLower x : []
      | isUpper x      = '_' : toLower x : []
      | otherwise      = x : []
    convert' b v (x1:x2:xs)
      | not b && isUpper x1 && (isLower x2 || not v)
      = '_' : toLower x1 : convert' False (isUpper x1) (x2:xs)
      | otherwise
      = toLower x1 : convert' False (isUpper x1) (x2:xs)

-- | A default implementation of 'toValues' that leverages 'Generic'. In order
-- to use this, make sure you derive a 'Generic' instance for your data type and
-- then define @'toValues' = 'toValuesDefault'@ in its 'Struct' instance.
toValuesDefault :: (Generic a, GStruct (Rep a)) => a -> [Value a]
toValuesDefault x = coerce $ gToValues $ from x

-- | A default implementation of 'updateField' that leverages 'Generic'. In
-- order to use this, make sure you derive a 'Generic' instance for your data
-- type and then define @'updateField' = 'updateFieldDefault'@ in its 'Struct'
-- instance.
updateFieldDefault :: (Generic a, GStruct (Rep a)) => a -> Value t -> a
updateFieldDefault a v@(Value _ field)
    | updated   = to a'
    | otherwise = error $ "Unexpected field: " ++ show field
  where
    (a', updated) = gUpdateField (from a) v

-- | A default implementation of 'typeOf' that leverages 'Generic'. In order to
-- use this, make sure you derive a 'Generic' instance for your data type and
-- then define @'typeOf' = 'typeOfDefault'@ in its 'Typed' instance.
typeOfDefault ::
  forall a. (Typed a, Struct a, Generic a, GTypedStruct (Rep a)) => Type a
typeOfDefault = Struct $ to $ gStructPlaceholder @(Rep a) @()

-- ** Generic-based classes (not exported)

-- | Capture the name of a Haskell data type from its 'Generic' metadata.
class GDatatype f where
  -- | Returns the name of a Haskell data type. (Note that this differs from
  -- 'typeName', which is expected to return the name of the struct in the
  -- /target/ language).
  gTypeName :: f p -> String

-- | The only 'GDatatype' instance we need is for 'D1', which describes
-- 'Datatype' metadata (@d@). We ignore all other metadata (@_f@).
instance Datatype d => GDatatype (D1 d _f) where
  gTypeName = datatypeName

-- | Perform struct-related operations on 'Generic' representation types.
class GStruct f where
  -- | Transforms all the struct representation's fields into a list of values.
  gToValues :: f p -> [Value (f p)]

  -- | Update the value of a struct representation's field. This returns two
  -- things:
  --
  -- 1. @f p@: The struct representation, but with the field updated.
  --
  -- 2. 'Bool': This is 'True' if the field was successfully updated and 'False'
  --    otherwise. If this returns 'False', it is the responsibility of the
  --    caller to raise an error.
  gUpdateField :: f p -> Value t -> (f p, Bool)

-- | 'U1' represents a data constructor with no fields. As such, 'gToValues'
-- returns an empty list of 'Value's, and 'gUpdateField' does not update
-- anything.
instance GStruct U1 where
  gToValues U1 = []
  gUpdateField u _ = (u, False)

-- | 'M1' is only used to store metadata, which the 'GStruct' class does not
-- make use of. As such, this instance discards the metadata and recurses.
instance GStruct f => GStruct (M1 _i _c f) where
  gToValues (M1 x) = coerce (gToValues x)
  gUpdateField (M1 x) v = (M1 x', updated)
    where
      (x', updated) = gUpdateField x v

-- | @(':*:')@ represents a data constructor with multiple fields.
instance (GStruct f, GStruct g) => GStruct (f :*: g) where
  -- Recursively compute two lists of Values and append them.
  gToValues (f :*: g) = coerce (gToValues f) ++ coerce (gToValues g)
  -- Recursively attempt to update the field in both branches and combine the
  -- updated branches. We will have successfully updated the field if either
  -- branch was successfully updated.
  gUpdateField (f :*: g) v = (f' :*: g', updatedF || updatedG)
    where
      (f', updatedF) = gUpdateField f v
      (g', updatedG) = gUpdateField g v

-- | 'K1' represents a single field in a data constructor. This is the base
-- case.
instance (KnownSymbol name, Typed ty, c ~ Field name ty) =>
    GStruct (K1 _i c) where
  -- Now that we have the field, return it in a singleton list.
  gToValues (K1 field) = [Value typeOf field]
  -- In order to update the field, we check that the field names and types
  -- match. If they do, return the updated field and declare the update as
  -- successful. Otherwise, return the old field and declare the update as
  -- unsuccessful.
  gUpdateField (K1 oldField) (Value newTy (newField :: Field newName newTy)) =
    case (sameSymbol pName pNewName, testEquality ty newTy) of
      (Just Refl, Just Refl) -> (K1 newField, True)
      _                      -> (K1 oldField, False)
    where
      pName    = Proxy @name
      pNewName = Proxy @newName
      ty       = typeOf @ty

-- | Compute a 'Generic' placeholder value to use for a struct type.
class GTypedStruct f where
  -- | A placeholder value to supply to use in a generic implementation of
  -- 'typeOf' for a struct type.
  gStructPlaceholder :: f p

-- | 'U1' represents a data constructor with no fields. As such,
-- 'gStructPlaceholder' simply returns the data constructor with no other
-- information.
instance GTypedStruct U1 where
  gStructPlaceholder = U1

-- | 'M1' is only used to store metadata, which the 'GTypedStruct' class does
-- not make use of. As such, this instance recursively computes a placeholder
-- value without inspecting the metadata.
instance GTypedStruct f => GTypedStruct (M1 _i _c f) where
  gStructPlaceholder = M1 gStructPlaceholder

-- | @(':*:')@ represents a data constructor with multiple fields. As such,
-- 'gStructPlaceholder' recursively computes placeholders for each field and
-- combines them into the overall data constructor.
instance (GTypedStruct f, GTypedStruct g) => GTypedStruct (f :*: g) where
  gStructPlaceholder = gStructPlaceholder :*: gStructPlaceholder

-- | 'K1' represents a single field in a data constructor. This is the base
-- case. This instance computes a placeholder value that works for any field of
-- any type.
instance (c ~ Field name ty) => GTypedStruct (K1 _i c) where
  -- We use 'undefined' as the actual value for the 'Field' because Copilot
  -- never inspects the value.
  gStructPlaceholder = K1 $ Field undefined
