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

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

import Control.Monad.Fix (mfix)
import Language.Haskell.TH
import Prelude hiding (Type)

{-# ANN module ("HLint: ignore Language.Haskell.TH should be imported post-qualified or with an explicit import list" :: Text) #-}

-- | Generates an 'NFData' instance for a GADT.
--
-- On superclass constraints for type arguments:
-- we use heuristics to guess for which type arguments
-- we need to add @NFData@ instance.
-- If this behaves not as you want, probably it's just worth
-- starting passing the necessary constraints to this function manually.
deriveGADTNFData :: Name -> Q [Dec]
deriveGADTNFData name = do
  seqQ <- [| seq |]
  unit <- [| () |]
  (TyConI (DataD _ dataName vars _ cons _)) <- reify name
  tyArgRoles <- reifyRoles name
  let
    nfDataC = ConT $ mkName "NFData"

    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 nfDataC . foldl' AppT (ConT dataName) $
        map (VarT . getNameFromVar) vars

    nfDataConstr = do
      (var, role) <- zip vars tyArgRoles

      -- Phantom type arguments do not require constraints
      case role of
        NominalR -> mzero
        RepresentationalR -> pass
        PhantomR -> mzero
        InferR -> error "unexpected InferR returned by reifyRole"

      -- Only types of 'Type' kind may require 'NFData' constraint
      varTy <- case var of
        PlainTV v _ -> pure v
        KindedTV v _ k -> do
          guard (k == StarT)
          pure v

      return $ nfDataC `AppT` VarT varTy

    makeInstance clauses =
      InstanceD Nothing nfDataConstr 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 -> TypeQ
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

-- | Given a list of types, produce the type of a tuple of
-- those types. This is analogous to 'tupE' and 'tupP'.
--
-- @
-- tupT [[t|Int|], [t|Char|], [t|Bool]] = [t| (Int, Char, Bool) |]
-- @
tupT :: [Q Type] -> Q Type
tupT ts = do
  -- We build the expression with a thunk inside that will be filled in with
  -- the length of the list once that's been determined. This works
  -- efficiently (in one pass) because TH.Type is rather lazy. Why isn't this
  -- just a left fold? A left fold will produce a big Q action that, when run,
  -- will produce the expression. We want to produce the expression incrementally
  -- as we run the Q action. foldM lets us do that, and mfix gives us the thunk
  -- for the tuple size. The irrefutable pattern is required as usual because the
  -- function passed to mfix must never force its argument.
  (res, !_n) <- mfix (\ ~(_res, n) -> foldM go (TupleT n, 0) ts)
  pure res
  where
    go (acc, !k) ty = do
      ty' <- ty
      pure (acc `AppT` ty', k + 1)