module Vectorise.Utils.PADict (
  paDictArgType,
  paDictOfType,
  paMethod,
  prDictOfReprType,
  prDictOfPReprInstTyCon
) where

import GhcPrelude

import Vectorise.Monad
import Vectorise.Builtins
import Vectorise.Utils.Base

import CoreSyn
import CoreUtils
import FamInstEnv
import Coercion
import Type
import TyCoRep
import TyCon
import CoAxiom
import Var
import Outputable
import DynFlags
import FastString
import Util
import Control.Monad


-- |Construct the PA argument type for the tyvar. For the tyvar (v :: *) it's
-- just PA v. For (v :: (* -> *) -> *) it's
--
-- > forall (a :: * -> *). (forall (b :: *). PA b -> PA (a b)) -> PA (v a)
--
paDictArgType :: TyVar -> VM (Maybe Type)
paDictArgType tv = go (mkTyVarTy tv) (tyVarKind tv)
  where
    go ty (FunTy k1 k2)
      = do
          tv   <- if isCoercionType k1
                  then newCoVar (fsLit "c") k1
                  else newTyVar (fsLit "a") k1
          mty1 <- go (mkTyVarTy tv) k1
          case mty1 of
            Just ty1 -> do
                          mty2 <- go (mkAppTy ty (mkTyVarTy tv)) k2
                          return $ fmap (mkInvForAllTy tv . mkFunTy ty1) mty2
            Nothing  -> go ty k2

    go ty k
      | isLiftedTypeKind k
      = do
          pa_cls <- builtin paClass
          return $ Just $ mkClassPred pa_cls [ty]

    go _ _ = return Nothing


-- |Get the PA dictionary for some type
--
paDictOfType :: Type -> VM CoreExpr
paDictOfType ty
  = paDictOfTyApp ty_fn ty_args
  where
    (ty_fn, ty_args) = splitAppTys ty

    paDictOfTyApp :: Type -> [Type] -> VM CoreExpr
    paDictOfTyApp ty_fn ty_args
        | Just ty_fn' <- coreView ty_fn
        = paDictOfTyApp ty_fn' ty_args

    -- for type variables, look up the dfun and apply to the PA dictionaries
    -- of the type arguments
    paDictOfTyApp (TyVarTy tv) ty_args
      = do
        { dfun <- maybeCantVectoriseM "No PA dictionary for type variable"
                                      (ppr tv <+> text "in" <+> ppr ty)
                $ lookupTyVarPA tv
        ; dicts <- mapM paDictOfType ty_args
        ; return $ dfun `mkTyApps` ty_args `mkApps` dicts
        }

    -- for tycons, we also need to apply the dfun to the PR dictionary of
    -- the representation type if the tycon is polymorphic
    paDictOfTyApp (TyConApp tc []) ty_args
      = do
        { dfun <- maybeCantVectoriseM noPADictErr (ppr tc <+> text "in" <+> ppr ty)
                $ lookupTyConPA tc
        ; super <- super_dict tc ty_args
        ; dicts <- mapM paDictOfType ty_args
        ; return $ Var dfun `mkTyApps` ty_args `mkApps` super `mkApps` dicts
        }
      where
        noPADictErr = "No PA dictionary for type constructor (did you import 'Data.Array.Parallel'?)"

        super_dict _     []      = return []
        super_dict tycon ty_args
          = do
            { pr <- prDictOfPReprInst (TyConApp tycon ty_args)
            ; return [pr]
            }

    paDictOfTyApp _ _ = getDynFlags >>= failure

    failure dflags = cantVectorise dflags "Can't construct PA dictionary for type" (ppr ty)

-- |Produce code that refers to a method of the 'PA' class.
--
paMethod :: (Builtins -> Var) -> (TyCon -> Builtins -> Var) -> Type -> VM CoreExpr
paMethod _ query ty
  | Just tycon <- splitPrimTyCon ty             -- Is 'ty' from 'GHC.Prim' (e.g., 'Int#')?
  = liftM Var $ builtin (query tycon)
paMethod method _ ty
  = do
    { fn   <- builtin method
    ; dict <- paDictOfType ty
    ; return $ mkApps (Var fn) [Type ty, dict]
    }

-- |Given a type @ty@, return the PR dictionary for @PRepr ty@.
--
prDictOfPReprInst :: Type -> VM CoreExpr
prDictOfPReprInst ty
  = do
    { (FamInstMatch { fim_instance = prepr_fam, fim_tys = prepr_args })
          <- preprFamInst ty
    ; prDictOfPReprInstTyCon ty (famInstAxiom prepr_fam) prepr_args
    }

-- |Given a type @ty@, its PRepr synonym tycon and its type arguments,
-- return the PR @PRepr ty@. Suppose we have:
--
-- > type instance PRepr (T a1 ... an) = t
--
-- which is internally translated into
--
-- > type :R:PRepr a1 ... an = t
--
-- and the corresponding coercion. Then,
--
-- > prDictOfPReprInstTyCon (T a1 ... an) :R:PRepr u1 ... un = PR (T u1 ... un)
--
-- Note that @ty@ is only used for error messages
--
prDictOfPReprInstTyCon :: Type -> CoAxiom Unbranched -> [Type] -> VM CoreExpr
prDictOfPReprInstTyCon _ty prepr_ax prepr_args
  = do
      let rhs = mkUnbranchedAxInstRHS prepr_ax prepr_args []
      dict <- prDictOfReprType' rhs
      pr_co <- mkBuiltinCo prTyCon
      let co = mkAppCo pr_co
             $ mkSymCo
             $ mkUnbranchedAxInstCo Nominal prepr_ax prepr_args []
      return $ mkCast dict co

-- |Get the PR dictionary for a type. The argument must be a representation
-- type.
--
prDictOfReprType :: Type -> VM CoreExpr
prDictOfReprType ty
  | Just (tycon, tyargs) <- splitTyConApp_maybe ty
    = do
        prepr <- builtin preprTyCon
        if tycon == prepr
          then do
                 let [ty'] = tyargs
                 pa <- paDictOfType ty'
                 sel <- builtin paPRSel
                 return $ Var sel `App` Type ty' `App` pa
          else do
                 -- a representation tycon must have a PR instance
                 dfun <- maybeV (text "look up PR dictionary for" <+> ppr tycon) $
                           lookupTyConPR tycon
                 prDFunApply dfun tyargs

  | otherwise
    = do
        -- it is a tyvar or an application of a tyvar
        -- determine the PR dictionary from its PA dictionary
        --
        -- NOTE: This assumes that PRepr t ~ t is for all representation types
        -- t
        --
        -- FIXME: This doesn't work for kinds other than * at the moment. We'd
        -- have to simply abstract the term over the missing type arguments.
        pa    <- paDictOfType ty
        prsel <- builtin paPRSel
        return $ Var prsel `mkApps` [Type ty, pa]

prDictOfReprType' :: Type -> VM CoreExpr
prDictOfReprType' ty = prDictOfReprType ty `orElseV`
                       do dflags <- getDynFlags
                          cantVectorise dflags "No PR dictionary for representation type"
                                        (ppr ty)

-- | Apply a tycon's PR dfun to dictionary arguments (PR or PA) corresponding
-- to the argument types.
prDFunApply :: Var -> [Type] -> VM CoreExpr
prDFunApply dfun tys
  | Just [] <- ctxs    -- PR (a :-> b) doesn't have a context
  = return $ Var dfun `mkTyApps` tys

  | Just tycons <- ctxs
  , tycons `equalLength` tys
  = do
      pa <- builtin paTyCon
      pr <- builtin prTyCon
      dflags <- getDynFlags
      args <- zipWithM (dictionary dflags pa pr) tys tycons
      return $ Var dfun `mkTyApps` tys `mkApps` args

  | otherwise = do dflags <- getDynFlags
                   invalid dflags
  where
    -- the dfun's contexts - if its type is (PA a, PR b) => PR (C a b) then
    -- ctxs is Just [PA, PR]
    ctxs = fmap (map fst)
         $ sequence
         $ map splitTyConApp_maybe
         $ fst
         $ splitFunTys
         $ snd
         $ splitForAllTys
         $ varType dfun

    dictionary dflags pa pr ty tycon
      | tycon == pa = paDictOfType ty
      | tycon == pr = prDictOfReprType ty
      | otherwise   = invalid dflags

    invalid dflags = cantVectorise dflags "Invalid PR dfun type" (ppr (varType dfun) <+> ppr tys)