-- | Simplified `derivingUnbox`.
module MCSP.Data.String.TH (
    EnumLike,
    derivingUnboxVia,
) where

import Control.Applicative (pure)
import Control.Monad (fail, (>>=))
import Data.Bool (otherwise, (&&))
import Data.Char (toLower)
import Data.Data (Proxy (..), Typeable, typeRep)
import Data.Function (($), (.))
import Data.Int (Int)
import Data.List ((++))
import Data.Ord (max, min, (<=))
import Data.String (String)
import Data.Vector.Unboxed.Deriving (derivingUnbox)
import GHC.Enum (Bounded (maxBound, minBound), Enum (fromEnum, toEnum))
import GHC.Err (error)
import Text.Show (show)

import Language.Haskell.TH (
    Dec,
    DecsQ,
    Name,
    Q,
    Type (..),
    TypeQ,
    nameBase,
    newName,
    sigD,
    varE,
    varP,
 )

-- | Makes the first character lowercase.
uncapitalize :: String -> String
uncapitalize :: [Char] -> [Char]
uncapitalize [] = []
uncapitalize (Char
ch : [Char]
rest) = Char -> Char
toLower Char
ch Char -> [Char] -> [Char]
forall a. a -> [a] -> [a]
: [Char]
rest

-- | Extract the name of a data type.
getName :: Type -> Q String
getName :: Type -> Q [Char]
getName (ConT Name
name) = [Char] -> Q [Char]
forall a. a -> Q a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([Char] -> Q [Char]) -> [Char] -> Q [Char]
forall a b. (a -> b) -> a -> b
$ Name -> [Char]
nameBase Name
name
getName Type
typ = [Char] -> Q [Char]
forall a. [Char] -> Q a
forall (m :: * -> *) a. MonadFail m => [Char] -> m a
fail ([Char] -> Q [Char]) -> [Char] -> Q [Char]
forall a b. (a -> b) -> a -> b
$ [Char]
"invalid type " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Type -> [Char]
forall a. Show a => a -> [Char]
show Type
typ

--- | Defines the signature and name for a cast function from @src@ to @dst@.
mkCastFn :: Type -> Type -> Q (Dec, Name)
mkCastFn :: Type -> Type -> Q (Dec, Name)
mkCastFn Type
src Type
dst = do
    [Char]
srcName <- Type -> Q [Char]
getName Type
src
    [Char]
dstName <- Type -> Q [Char]
getName Type
dst
    Name
fnName <- [Char] -> Q Name
forall (m :: * -> *). Quote m => [Char] -> m Name
newName ([Char] -> [Char]
uncapitalize [Char]
srcName [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"To" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
dstName)
    Dec
fnSig <- Name -> Q Type -> Q Dec
forall (m :: * -> *). Quote m => Name -> m Type -> m Dec
sigD Name
fnName [t|$(Type -> Q Type
forall a. a -> Q a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Type
src) -> $(Type -> Q Type
forall a. a -> Q a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Type
dst)|]
    pure (Dec
fnSig, Name
fnName)

-- | Creates a prefix name for deriving `Data.Vector.Unboxed.Unbox`.
mkPrefix :: Type -> Type -> Q String
mkPrefix :: Type -> Type -> Q [Char]
mkPrefix Type
typ Type
rep = do
    [Char]
typName <- Type -> Q [Char]
getName Type
typ
    [Char]
repName <- Type -> Q [Char]
getName Type
rep
    pure $ [Char]
typName [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"Via" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
repName

-- | @T@ and @U@ from @T -> U@.
splitTypRep :: Type -> Q (Type, Type)
splitTypRep :: Type -> Q (Type, Type)
splitTypRep (Type
ArrowT `AppT` Type
typ `AppT` Type
rep) = (Type, Type) -> Q (Type, Type)
forall a. a -> Q a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Type
typ, Type
rep)
splitTypRep Type
typ = [Char] -> Q (Type, Type)
forall a. [Char] -> Q a
forall (m :: * -> *) a. MonadFail m => [Char] -> m a
fail ([Char] -> Q (Type, Type)) -> [Char] -> Q (Type, Type)
forall a b. (a -> b) -> a -> b
$ [Char]
"invalid deriving rule " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Type -> [Char]
forall a. Show a => a -> [Char]
show Type
typ

-- | Given @T -> U@, derives `Data.Vector.Unboxed.Unbox` @T@ by casting it to @U@.
--
-- >>> import Data.Word (Word8)
-- >>> data DNA = A | C | G | T
-- >>> -- derivingUnboxVia [t|DNA -> Word8|]
derivingUnboxVia :: TypeQ -> DecsQ
derivingUnboxVia :: Q Type -> DecsQ
derivingUnboxVia Q Type
rule = do
    (Type
typ, Type
rep) <- Q Type
rule Q Type -> (Type -> Q (Type, Type)) -> Q (Type, Type)
forall a b. Q a -> (a -> Q b) -> Q b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Type -> Q (Type, Type)
splitTypRep
    (Dec
abSig, Name
aToB) <- Type -> Type -> Q (Dec, Name)
mkCastFn Type
typ Type
rep
    (Dec
baSig, Name
bToA) <- Type -> Type -> Q (Dec, Name)
mkCastFn Type
rep Type
typ
    [Dec]
decs <- [d|($(Name -> Q Pat
forall (m :: * -> *). Quote m => Name -> m Pat
varP Name
aToB), $(Name -> Q Pat
forall (m :: * -> *). Quote m => Name -> m Pat
varP Name
bToA)) = safeCasts|]
    [Char]
prefix <- Type -> Type -> Q [Char]
mkPrefix Type
typ Type
rep
    [Dec]
derive <- [Char] -> Q Type -> ExpQ -> ExpQ -> DecsQ
derivingUnbox [Char]
prefix [t|$(Type -> Q Type
forall a. a -> Q a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Type
typ) -> $(Type -> Q Type
forall a. a -> Q a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Type
rep)|] (Name -> ExpQ
forall (m :: * -> *). Quote m => Name -> m Exp
varE Name
aToB) (Name -> ExpQ
forall (m :: * -> *). Quote m => Name -> m Exp
varE Name
bToA)
    pure $ [Dec
Item [Dec]
abSig, Dec
Item [Dec]
baSig] [Dec] -> [Dec] -> [Dec]
forall a. [a] -> [a] -> [a]
++ [Dec]
decs [Dec] -> [Dec] -> [Dec]
forall a. [a] -> [a] -> [a]
++ [Dec]
derive

-- | Types convertible to a bounded integer.
type EnumLike a = (Enum a, Bounded a, Typeable a)

-- | Cast @a@ to @b@ via their integer representations.
cast :: (EnumLike a, EnumLike b) => a -> b
cast :: forall a b. (EnumLike a, EnumLike b) => a -> b
cast = Int -> b
forall a. Enum a => Int -> a
toEnum (Int -> b) -> (a -> Int) -> a -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Int
forall a. Enum a => a -> Int
fromEnum
{-# INLINE cast #-}

-- | Cast @b@ to @a@, clamping to valid values of @a@.
clamp :: forall a b. (EnumLike a, EnumLike b) => b -> a
clamp :: forall a b. (EnumLike a, EnumLike b) => b -> a
clamp = Int -> a
forall a. Enum a => Int -> a
toEnum (Int -> a) -> (b -> Int) -> b -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
minA (Int -> Int) -> (b -> Int) -> b -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
maxA (Int -> Int) -> (b -> Int) -> b -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. b -> Int
forall a. Enum a => a -> Int
fromEnum
  where
    minA :: Int
minA = forall a. Enum a => a -> Int
fromEnum @a a
forall a. Bounded a => a
minBound
    maxA :: Int
maxA = forall a. Enum a => a -> Int
fromEnum @a a
forall a. Bounded a => a
maxBound
{-# INLINE clamp #-}

-- | Holds information about an `EnumLike` @a@.
data Info a = Info {forall {k} (a :: k). Info a -> Int
minVal :: !Int, forall {k} (a :: k). Info a -> Int
maxVal :: !Int, forall {k} (a :: k). Info a -> [Char]
name :: !String}

-- | Extracts information about an `EnumLike` @a@.
info :: forall a. EnumLike a => Info a
info :: forall a. EnumLike a => Info a
info =
    Info
        { minVal :: Int
minVal = forall a. Enum a => a -> Int
fromEnum @a a
forall a. Bounded a => a
minBound,
          maxVal :: Int
maxVal = forall a. Enum a => a -> Int
fromEnum @a a
forall a. Bounded a => a
maxBound,
          name :: [Char]
name = TypeRep -> [Char]
forall a. Show a => a -> [Char]
show (TypeRep -> [Char]) -> TypeRep -> [Char]
forall a b. (a -> b) -> a -> b
$ Proxy a -> TypeRep
forall {k} (proxy :: k -> *) (a :: k).
Typeable a =>
proxy a -> TypeRep
typeRep (forall t. Proxy t
forall {k} (t :: k). Proxy t
Proxy @a)
        }

-- | Checks if @a@ is safely convertible to @b@ via `Enum` and `Bounded`, then returns the casts.
safeCasts :: forall a b. (EnumLike a, EnumLike b) => (a -> b, b -> a)
safeCasts :: forall a b. (EnumLike a, EnumLike b) => (a -> b, b -> a)
safeCasts
    | Info b -> Int
forall {k} (a :: k). Info a -> Int
minVal Info b
rep Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Info a -> Int
forall {k} (a :: k). Info a -> Int
minVal Info a
typ Bool -> Bool -> Bool
&& Info a -> Int
forall {k} (a :: k). Info a -> Int
maxVal Info a
typ Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Info b -> Int
forall {k} (a :: k). Info a -> Int
maxVal Info b
rep = (a -> b
forall a b. (EnumLike a, EnumLike b) => a -> b
cast, b -> a
forall a b. (EnumLike a, EnumLike b) => b -> a
clamp)
    | Bool
otherwise =
        [Char] -> (a -> b, b -> a)
forall a. HasCallStack => [Char] -> a
error ([Char] -> (a -> b, b -> a)) -> [Char] -> (a -> b, b -> a)
forall a b. (a -> b) -> a -> b
$ Info a -> [Char]
forall {k} (a :: k). Info a -> [Char]
name Info a
typ [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
" is larger than " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Info b -> [Char]
forall {k} (a :: k). Info a -> [Char]
name Info b
rep [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
" and cannot be safely converted"
  where
    typ :: Info a
typ = forall a. EnumLike a => Info a
info @a
    rep :: Info b
rep = forall a. EnumLike a => Info a
info @b