-- |
-- Module      :   Grisette.Internal.TH.Ctor.Common
-- Copyright   :   (c) Sirui Lu 2024
-- License     :   BSD-3-Clause (see the LICENSE file)
--
-- Maintainer  :   siruilu@cs.washington.edu
-- Stability   :   Experimental
-- Portability :   GHC only
module Grisette.Internal.TH.Ctor.Common
  ( withNameTransformer,
    prefixTransformer,
    decapitalizeTransformer,
  )
where

import Control.Monad (unless)
import Data.Char (isAlphaNum, toLower)
import Data.Foldable (traverse_)
import Language.Haskell.TH (Dec, Name, Q, nameBase)
import Language.Haskell.TH.Datatype
  ( ConstructorInfo (constructorName),
    DatatypeInfo (datatypeCons),
    reifyDatatype,
  )

checkName :: String -> Q ()
checkName :: String -> Q ()
checkName String
name =
  Bool -> Q () -> Q ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless ((Char -> Bool) -> String -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (\Char
x -> Char -> Bool
isAlphaNum Char
x Bool -> Bool -> Bool
|| Char
x Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== Char
'\'' Bool -> Bool -> Bool
|| Char
x Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== Char
'_') String
name) (Q () -> Q ()) -> Q () -> Q ()
forall a b. (a -> b) -> a -> b
$
    String -> Q ()
forall a. String -> Q a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail
      ( String
"Constructor name contain invalid characters, consider providing a "
          String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"custom name: "
          String -> String -> String
forall a. [a] -> [a] -> [a]
++ String -> String
forall a. Show a => a -> String
show String
name
      )

-- | Generate smart constructor given a type name, using a name transformer
-- to transform constructor names.
withNameTransformer ::
  -- | A function that generates decs given a list of constructor names and a
  -- type name
  ([String] -> Name -> Q [Dec]) ->
  -- | A function that transforms constructor names
  (String -> String) ->
  -- | The type to generate the wrappers for
  Name ->
  Q [Dec]
withNameTransformer :: ([String] -> Name -> Q [Dec])
-> (String -> String) -> Name -> Q [Dec]
withNameTransformer [String] -> Name -> Q [Dec]
namedGen String -> String
nameTransformer Name
typName = do
  d <- Name -> Q DatatypeInfo
reifyDatatype Name
typName
  let constructorNames = Name -> String
nameBase (Name -> String)
-> (ConstructorInfo -> Name) -> ConstructorInfo -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ConstructorInfo -> Name
constructorName (ConstructorInfo -> String) -> [ConstructorInfo] -> [String]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> DatatypeInfo -> [ConstructorInfo]
datatypeCons DatatypeInfo
d
  let transformedNames = String -> String
nameTransformer (String -> String) -> [String] -> [String]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [String]
constructorNames
  traverse_ checkName transformedNames
  namedGen transformedNames typName

-- | A name transformer that prefixes a string to the constructor name
prefixTransformer :: String -> String -> String
prefixTransformer :: String -> String -> String
prefixTransformer = String -> String -> String
forall a. [a] -> [a] -> [a]
(++)

-- | A name transformer that converts the first character to lowercase
decapitalizeTransformer :: String -> String
decapitalizeTransformer :: String -> String
decapitalizeTransformer (Char
x : String
xs) = Char -> Char
toLower Char
x Char -> String -> String
forall a. a -> [a] -> [a]
: String
xs
decapitalizeTransformer [] = []