TH.hs

-- SPDX-FileCopyrightText: 2021 Oxhead Alpha
-- SPDX-License-Identifier: LicenseRef-MIT-OA

module Morley.Util.TH
  ( deriveGADTNFData
  , lookupTypeNameOrFail
  , isTypeAlias
  , addTypeVariables
  ) where

import Language.Haskell.TH as 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

-- | Check if name is a @type@
isTypeAlias :: Name -> Q Bool
isTypeAlias typeName = reify typeName <&> \case
  TyConI (TySynD {}) -> True
  _ -> False

-- | Accepts a type constructor and fills it with variables until
-- getting a type of kind @*@.
addTypeVariables :: Name -> Q TH.Type
addTypeVariables tyCtor = do
  tyVarBindrs <- reify tyCtor >>= \case
    TyConI (DataD _ _ tyVarBindrs _ _ _) -> pure tyVarBindrs
    TyConI (NewtypeD _ _ tyVarBindrs _ _ _) -> pure tyVarBindrs
    _ -> fail "Expected a plain datatype"
  let vars = tyVarBindrs <&> \case
        PlainTV vName -> vName
        KindedTV vName _ -> vName
  return $ foldl (\acc var -> acc `AppT` VarT var) (ConT tyCtor) vars