TH.hs

-- SPDX-FileCopyrightText: 2020 Tocqueville Group
--
-- SPDX-License-Identifier: LicenseRef-MIT-TQ

module Util.TH
  ( deriveGADTNFData
  , lookupTypeNameOrFail
  ) where

import Language.Haskell.TH

-- | Generates an NFData instance for a GADT. /Note:/ This will not generate
-- additional constraints to the generated instance if those are required.
deriveGADTNFData :: Name -> Q [Dec]
deriveGADTNFData name = do
  (TyConI (DataD _ dataName vars _ cons _)) <- reify name
  let
    getNameFromVar (PlainTV n) = n
    getNameFromVar (KindedTV n _) = n
    convertTyVars orig = foldr (\a b -> AppT b . VarT $ getNameFromVar a) orig vars

    -- Unfolds multiple constructors of form "A, B, C :: A -> Stuff"
    -- into a list of tuples of constructor names and their data
    unfoldConstructor (GadtC cs bangs _) = map (,bangs) cs
    unfoldConstructor (ForallC _ _ c) = unfoldConstructor c
    unfoldConstructor _ = fail "Non GADT constructors are not supported."

    -- Constructs a clause "rnf (ConName a1 a2 ...) = rnf (a1, a2, ...)
    makeClauses (conName, bangs) = do
      varNames <- traverse (\_ -> newName "a") bangs
      let rnfExp e = AppE (VarE $ mkName "rnf") e
      return $
        (Clause
          [ConP conName $ map VarP varNames]
          (NormalB (rnfExp . TupE $ map VarE varNames))
          []
        )

    makeInstance clauses =
      InstanceD
        Nothing
        []
        (AppT (ConT $ mkName "NFData") (convertTyVars $ ConT dataName))
        [FunD (mkName "rnf") clauses]

  clauses <- traverse makeClauses $ cons >>= unfoldConstructor
  return [makeInstance clauses]

lookupTypeNameOrFail :: String -> Q Name
lookupTypeNameOrFail typeStr =
  lookupTypeName typeStr >>= \case
    Nothing -> fail $ "Failed type name lookup for: '" <> typeStr <> "'."
    Just tn -> pure tn