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
  seqQ <- [| seq |]
  unit <- [| () |]
  (TyConI (DataD _ dataName vars _ cons _)) <- reify name
  let
    getNameFromVar (PlainTV n) = n
    getNameFromVar (KindedTV n _) = n

    -- 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 `seq` rnf a2 `seq` rnf a3 `seq` ..."
    makeClauses (conName, bangs) = do
        varNames <- traverse (\_ -> newName "a") bangs
        let rnfVar = VarE 'rnf
        let rnfExp = AppE rnfVar . VarE
        let infixSeq e1 e2 = InfixE (Just e1) seqQ (Just e2)
        return $
          (Clause
            [ConP conName $ map VarP varNames]
            (NormalB $ foldl' infixSeq unit (map rnfExp varNames))
            []
          )

    nfDataT =
      AppT (ConT $ mkName "NFData") . foldl' AppT (ConT dataName) $
        map (VarT . getNameFromVar) vars

    makeInstance clauses =
      InstanceD Nothing [] nfDataT [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