TH.hs

{-# LANGUAGE CPP #-}
{-# LANGUAGE TemplateHaskellQuotes #-}

module Data.Recollections.TH
  ( mkCollection
  , mkIndices

  , mkDistribute

  , mkIndex
  , mkTabulate
  ) where

import Data.Char
import Data.Foldable
import Data.Traversable
import GHC.Generics (Generic, Generic1, Generically1)
import Language.Haskell.TH

mkCollection :: Name -> Q [Dec]
mkCollection tags = do
  let nothingsBanger = bang noSourceUnpackedness noSourceStrictness
  let collectionName = mkName "Collection"

  names <- tagNames tags
  let
    fields = do
      name <- names
      pure $ varBangType (mkName $ fieldNameOf name) $ (bangType nothingsBanger) (varT (mkName "a"))
  let constr = recC collectionName fields
  let
    derivs =
      [ derivClause (Just StockStrategy)
          [ conT ''Show
          , conT ''Eq
          , conT ''Generic
          , conT ''Generic1
          , conT ''Functor
          , conT ''Foldable
          , conT ''Traversable
          ]
      , derivClause (Just $ ViaStrategy $ ConT ''Generically1 `AppT` ConT collectionName)
        [ conT ''Applicative
        ]
      ]
  pure <$> dataD mempty collectionName [collectionTyVar] Nothing [constr] derivs

-- | The @a@ binder of @data Collection a@.
--
-- @template-haskell-2.21@ (GHC 9.8) changed the binder flag of 'dataD'
-- from @()@ to 'BndrVis'.
#if MIN_VERSION_template_haskell(2,21,0)
collectionTyVar :: TyVarBndrVis
collectionTyVar = PlainTV (mkName "a") BndrReq
#else
collectionTyVar :: TyVarBndrUnit
collectionTyVar = PlainTV (mkName "a") ()
#endif

mkIndices :: Name -> Q [Dec]
mkIndices tags = do
  let collectionName = mkName "Collection"
  let indicesName = mkName "indices"
  names <- tagNames tags
  sig <- sigD indicesName $ conT collectionName `appT` conT tags
  let body = foldl' appE (conE collectionName) (map (conE . mkName) names)
  fun <- funD indicesName [ clause [] (normalB body) [] ]
  pure [sig, fun]

-- * Distributive

{- |
@
distribute f = Collection
  { this      = this      <$> f
  , that      = that      <$> f
  , something = something <$> f
  , else'     = else'     <$> f
  , entirely  = entirely  <$> f
  }
@
-}
mkDistribute :: Name -> Q [Dec]
mkDistribute tags = do
  let collectionName = mkName "Collection"
  let distributeName = mkName "distribute"
  names <- tagNames tags
  let f = mkName "f"
  let a = mkName "a"
  {- The binder must not shadow any field selector the body refers to by
     'mkName' (those are resolved lexically, by occurrence name). Field names
     always start with a lowercase letter, so a leading underscore is safe;
     plain @f@ would break for a tag named @F@. -}
  arg <- newName "_f"
  sig <- sigD distributeName $
    forallT [] (cxt [conT ''Functor `appT` varT f]) $
      appT (appT arrowT (varT f `appT` (conT collectionName `appT` (varT a)))) $
        (conT collectionName `appT` (varT f `appT` varT a))
  let
    body =
      foldl' appE (conE collectionName) do
        name <- names
        let fieldSelector = varE . mkName $ fieldNameOf name
        pure $ appE (varE 'fmap) fieldSelector `appE` varE arg
  fun <- funD distributeName
    [ clause [varP arg] (normalB body) []
    ]
  inl <- inlineP distributeName
  pure [sig, fun, inl]

-- * Representable

{- |
@
index c = \case This -> this c; That -> that c; …
@
-}
mkIndex :: Name -> Q [Dec]
mkIndex tags = do
  let collectionName = mkName "Collection"
  let indexName = mkName "index"
  names <- tagNames tags
  let a = mkName "a"
  sig <- sigD indexName $
    appT (appT arrowT (conT collectionName `appT` varT a)) $
      appT (appT arrowT (conT tags)) (varT a)
  bounds <- for names \name ->
    (,) name <$> newName (fieldNameOf name)
  let
    body = lamCaseE do
      (name, bound) <- bounds
      pure $ match (conP (mkName name) []) (normalB $ varE bound) []

  fun <- funD indexName
    [ clause
        [ conP collectionName $ map (varP . snd) bounds
        ]
        (normalB body)
        []
    ]
  inl <- inlineP indexName
  pure [sig, fun, inl]

{- |
@
tabulate k = Collection (k This) (k That) (k Something) (k Else) (k Entirely)
@
-}
mkTabulate :: Name -> Q [Dec]
mkTabulate tags = do
  let collectionName = mkName "Collection"
  let tabulateName = mkName "tabulate"
  names <- tagNames tags
  let a = mkName "a"
  sig <- sigD tabulateName $
    appT (appT arrowT (appT (appT arrowT (conT tags)) (varT a))) $
      conT collectionName `appT` varT a
  f <- newName "_f"
  let
    body =
      foldl' appE (conE collectionName) do
        name <- names
        pure $ appE (varE f) $ conE (mkName name)
  fun <- funD tabulateName
    [ clause [varP f] (normalB body) []
    ]
  inl <- inlineP tabulateName
  pure [sig, fun, inl]

-- * Utils

tagNames :: Name -> Q [String]
tagNames tags =
 reify tags >>= \case
    TyConI (DataD _ _ _ _ constructors _) ->
      foldrM (flip extractTags) [] constructors
    _ ->
      fail "Expected a type constructor name"

extractTags :: [String] -> Con -> Q [String]
extractTags acc = \case
  NormalC name [] ->
    pure $ nameBase name : acc
  huh ->
    -- Skipping would produce a 'Collection' with fewer fields than there are
    -- tags, making the generated 'index' silently non-exhaustive.
    fail $ "Expected a nullary constructor, got: " <> show huh

-- | Lowercase the leading character and dodge reserved words.
fieldNameOf :: String -> String
fieldNameOf = \case
  [] -> []
  n : ame -> legalizeVariableName $ toLower n : ame

legalizeVariableName :: String -> String
legalizeVariableName name
  | name `elem` illegal = name ++ "'"
  | otherwise = name

-- | Haskell 2010 reserved words, which cannot be used as record field names.
illegal :: [String]
illegal =
  [ "case", "class", "data", "default", "deriving", "do"
  , "else", "foreign", "if", "import", "in"
  , "infix", "infixl", "infixr", "instance"
  , "let", "module", "newtype", "of"
  , "then", "type", "where", "_"
  ]

inlineP :: Name -> Q Dec
inlineP name = pragInlD name Inline FunLike AllPhases