-- Copyright 2022 Google LLC
--
-- Use of this source code is governed by a BSD-style
-- license that can be found in the LICENSE file or at
-- https://developers.google.com/open-source/licenses/bsd

module QueryType (module QueryType, module QueryTypePure, toAtomVar) where

import Control.Category ((>>>))
import Control.Monad
import Control.Applicative
import Data.List (elemIndex)
import Data.Maybe (fromJust)
import Data.Functor ((<&>))

import Types.Primitives
import Types.Core
import Types.Source
import Types.Top
import Types.Imp
import IRVariants
import Core
import Err
import Name hiding (withFreshM)
import Subst
import Util
import PPrint
import QueryTypePure
import CheapReduction


-- === Exposed helpers for querying types and effects ===

caseAltsBinderTys :: (EnvReader m,  IRRep r) => Type r n -> m n [Type r n]
caseAltsBinderTys ty = case ty of
  TyCon (SumType types) -> return types -- need this case?
  TyCon (NewtypeTyCon t) -> case t of
    UserADTType _ defName params -> do
      def <- lookupTyCon defName
      ~(ADTCons cons) <- instantiateTyConDef def params
      return [repTy | DataConDef _ _ repTy _ <- cons]
    _ -> error msg
  _ -> error msg
  where msg = "Case analysis only supported on ADTs, not on " ++ pprint ty

extendEffect :: IRRep r => Effect r n -> EffectRow r n -> EffectRow r n
extendEffect eff (EffectRow effs t) = EffectRow (effs <> eSetSingleton eff) t

piTypeWithoutDest :: PiType SimpIR n -> PiType SimpIR n
piTypeWithoutDest (PiType bsRefB _) =
  case popNest bsRefB of
    Just (PairB bs (_:>RawRefTy ansTy)) -> do
      PiType bs $ EffTy Pure ansTy  -- XXX: we ignore the effects here
    _ -> error "expected trailing dest binder"

typeOfTabApp :: (IRRep r, EnvReader m) => Type r n -> Atom r n -> m n (Type r n)
typeOfTabApp (TyCon (TabPi tabTy)) i = instantiate tabTy [i]
typeOfTabApp ty _ = error $ "expected a table type. Got: " ++ pprint ty

typeOfApplyMethod :: EnvReader m => CDict n -> Int -> [CAtom n] -> m n (EffTy CoreIR n)
typeOfApplyMethod d i args = do
  ty <- toType <$> getMethodType d i
  appEffTy ty args

typeOfTopApp :: EnvReader m => TopFunName n -> [SAtom n] -> m n (EffTy SimpIR n)
typeOfTopApp f xs = do
  piTy <- getTypeTopFun f
  instantiate piTy xs

typeOfIndexRef :: (EnvReader m, Fallible1 m, IRRep r) => Type r n -> Atom r n -> m n (Type r n)
typeOfIndexRef (TyCon (RefType h s)) i = do
  TyCon (TabPi tabPi) <- return s
  eltTy <- instantiate tabPi [i]
  return $ toType $ RefType h eltTy
typeOfIndexRef _ _ = error "expected a ref type"

typeOfProjRef :: EnvReader m => Type r n -> Projection -> m n (Type r n)
typeOfProjRef (TyCon (RefType h s)) p = do
  toType . RefType h <$> case p of
    ProjectProduct i -> do
      ~(TyCon (ProdType tys)) <- return s
      return $ tys !! i
    UnwrapNewtype -> do
      case s of
        TyCon (NewtypeTyCon tc) -> snd <$> unwrapNewtypeType tc
        _ -> error "expected a newtype"
typeOfProjRef _ _ = error "expected a reference"

appEffTy  :: (IRRep r, EnvReader m) => Type r n -> [Atom r n] -> m n (EffTy r n)
appEffTy (TyCon (Pi piTy)) xs = instantiate piTy xs
appEffTy t _ = error $ "expected a pi type, got: " ++ pprint t

partialAppType  :: (IRRep r, EnvReader m) => Type r n -> [Atom r n] -> m n (Type r n)
partialAppType (TyCon (Pi (CorePiType appExpl expls bs effTy))) xs = do
  (_, expls2) <- return $ splitAt (length xs) expls
  PairB bs1 bs2 <- return $ splitNestAt (length xs) bs
  instantiate (Abs bs1 (toType $ CorePiType appExpl expls2 bs2 effTy)) xs
partialAppType _ _ = error "expected a pi type"

effTyOfHof :: (EnvReader m, IRRep r) => Hof r n -> m n (EffTy r n)
effTyOfHof hof = EffTy <$> hofEffects hof <*> typeOfHof hof

typeOfHof :: (EnvReader m, IRRep r) => Hof r n -> m n (Type r n)
typeOfHof = \case
  For _ ixTy f -> getLamExprType f >>= \case
    PiType (UnaryNest b) (EffTy _ eltTy) -> return $ TabTy (ixTypeDict ixTy) b eltTy
    _ -> error "expected a unary pi type"
  While _ -> return UnitTy
  Linearize f _ -> getLamExprType f >>= \case
    PiType (UnaryNest (binder:>a)) (EffTy Pure b) -> do
      let b' = ignoreHoistFailure $ hoist binder b
      let fLinTy = toType $ nonDepPiType [a] Pure b'
      return $ PairTy b' fLinTy
    _ -> error "expected a unary pi type"
  Transpose f _ -> getLamExprType f >>= \case
    PiType (UnaryNest (_:>a)) _ -> return a
    _ -> error "expected a unary pi type"
  RunReader _ f -> do
    (resultTy, _) <- getTypeRWSAction f
    return resultTy
  RunWriter _ _ f -> uncurry PairTy <$> getTypeRWSAction f
  RunState _ _ f -> do
    (resultTy, stateTy) <- getTypeRWSAction f
    return $ PairTy resultTy stateTy
  RunIO f -> return $ getType f
  RunInit f -> return $ getType f
  CatchException ty _ -> return ty

hofEffects :: (EnvReader m, IRRep r) => Hof r n -> m n (EffectRow r n)
hofEffects = \case
  For _ _ f     -> functionEffs f
  While body    -> return $ getEffects body
  Linearize _ _ -> return Pure  -- Body has to be a pure function
  Transpose _ _ -> return Pure  -- Body has to be a pure function
  RunReader _ f -> rwsFunEffects Reader f
  RunWriter d _ f -> maybeInit d <$> rwsFunEffects Writer f
  RunState  d _ f -> maybeInit d <$> rwsFunEffects State  f
  RunIO            f -> return $ deleteEff IOEffect        $ getEffects f
  RunInit          f -> return $ deleteEff InitEffect      $ getEffects f
  CatchException _ f -> return $ deleteEff ExceptionEffect $ getEffects f
  where maybeInit :: IRRep r => Maybe (Atom r i) -> (EffectRow r o -> EffectRow r o)
        maybeInit d = case d of Just _ -> (<>OneEffect InitEffect); Nothing -> id

deleteEff :: IRRep r => Effect r n -> EffectRow r n -> EffectRow r n
deleteEff eff (EffectRow effs t) = EffectRow (effs `eSetDifference` eSetSingleton eff) t

getMethodIndex :: EnvReader m => ClassName n -> SourceName -> m n Int
getMethodIndex className methodSourceName = do
  ClassDef _ _ methodNames _ _ _ _ _ <- lookupClassDef className
  case elemIndex methodSourceName methodNames of
    Nothing -> error $ pprint methodSourceName ++ " is not a method of " ++ pprint className
    Just i -> return i
{-# INLINE getMethodIndex #-}

getUVarType :: EnvReader m => UVar n -> m n (CType n)
getUVarType = \case
  UAtomVar v -> getType <$> toAtomVar v
  UTyConVar   v -> getTyConNameType v
  UDataConVar v -> getDataConNameType v
  UPunVar     v -> getStructDataConType v
  UClassVar v -> do
    ClassDef _ _ _ _ roleExpls bs _ _ <- lookupClassDef v
    return $ toType $ CorePiType ExplicitApp (map snd roleExpls) bs $ EffTy Pure TyKind
  UMethodVar  v -> getMethodNameType v

getMethodNameType :: EnvReader m => MethodName n -> m n (CType n)
getMethodNameType v = liftEnvReaderM $ lookupEnv v >>= \case
  MethodBinding className i -> do
    ClassDef _ _ _ paramNames _ paramBs scBinders methodTys <- lookupClassDef className
    refreshAbs (Abs paramBs $ Abs scBinders (methodTys !! i)) \paramBs' absPiTy -> do
      let params = toAtom <$> bindersVars paramBs'
      dictTy <- toType <$> dictType (sink className) params
      withFreshBinder noHint dictTy \dictB -> do
        scDicts <- getSuperclassDicts (toDict $ binderVar dictB)
        CorePiType appExpl methodExpls methodBs effTy <- instantiate (sink absPiTy) scDicts
        let paramExpls = paramNames <&> \name -> Inferred name Unify
        let expls = paramExpls <> [Inferred Nothing (Synth $ Partial $ succ i)] <> methodExpls
        return $ toType $ CorePiType appExpl expls (paramBs' >>> UnaryNest dictB >>> methodBs) effTy

getMethodType :: EnvReader m => CDict n -> Int -> m n (CorePiType n)
getMethodType dict i = do
  ~(TyCon (DictTy dictTy)) <- return $ getType dict
  case dictTy of
    DictType _ className params -> liftEnvReaderM $ withSubstReaderT do
      superclassDicts <- getSuperclassDicts dict
      classDef <- lookupClassDef className
      withInstantiated classDef params \ab -> do
        withInstantiated ab superclassDicts \(ListE methodTys) ->
          substM $ methodTys !! i
    IxDictType ixTy -> liftEnvReaderM case i of
      0 -> mkCorePiType []      NatTy -- size' : () -> Nat
      1 -> mkCorePiType [ixTy]  NatTy -- ordinal : (n) -> Nat
      2 -> mkCorePiType [NatTy] ixTy  -- unsafe_from_ordinal : (Nat) -> n
      _ -> error "Ix only has three methods"
    DataDictType _ -> error "Data class has no methods"

mkCorePiType :: EnvReader m => [CType n] -> CType n -> m n (CorePiType n)
mkCorePiType argTys resultTy = liftEnvReaderM $ withFreshBinders argTys \bs _ -> do
  expls <- return $ nestToList (const Explicit) bs
  return $ CorePiType ExplicitApp expls bs (EffTy Pure (sink resultTy))

getTyConNameType :: EnvReader m => TyConName n -> m n (Type CoreIR n)
getTyConNameType v = do
  TyConDef _ expls bs _ <- lookupTyCon v
  case bs of
    Empty -> return TyKind
    _ -> return $ toType $ CorePiType ExplicitApp (snd <$> expls) bs $ EffTy Pure TyKind

getDataConNameType :: EnvReader m => DataConName n -> m n (Type CoreIR n)
getDataConNameType dataCon = liftEnvReaderM $ withSubstReaderT do
  (tyCon, i) <- lookupDataCon dataCon
  tyConDef <- lookupTyCon tyCon
  buildDataConType tyConDef \expls paramBs' paramVs params -> do
    withInstantiatedNames tyConDef paramVs \(ADTCons dataCons) -> do
      DataConDef _ ab _ _ <- renameM (dataCons !! i)
      refreshAbs ab \dataBs UnitE -> do
        let appExpl = case dataBs of Empty -> ImplicitApp
                                     _     -> ExplicitApp
        let resultTy = toType $ UserADTType (getSourceName tyConDef) (sink tyCon) (sink params)
        let dataExpls = nestToList (const $ Explicit) dataBs
        return $ toType $ CorePiType appExpl (expls <> dataExpls) (paramBs' >>> dataBs) (EffTy Pure resultTy)

getStructDataConType :: EnvReader m => TyConName n -> m n (CType n)
getStructDataConType tyCon = liftEnvReaderM $ withSubstReaderT do
  tyConDef <- lookupTyCon tyCon
  buildDataConType tyConDef \expls paramBs' paramVs params -> do
    withInstantiatedNames tyConDef paramVs \(StructFields fields) -> do
      fieldTys <- forM fields \(_, t) -> renameM t
      let resultTy = toType $ UserADTType (getSourceName tyConDef) (sink tyCon) params
      Abs dataBs resultTy' <- return $ typesAsBinderNest fieldTys resultTy
      let dataExpls = nestToList (const Explicit) dataBs
      return $ toType $ CorePiType ExplicitApp (expls <> dataExpls) (paramBs' >>> dataBs) (EffTy Pure resultTy')

buildDataConType
  :: (EnvReader m, EnvExtender m)
  => TyConDef n
  -> (forall l. DExt n l => [Explicitness] -> Nest CBinder n l -> [CAtomName l] -> TyConParams l -> m l a)
  -> m n a
buildDataConType (TyConDef _ roleExpls bs _) cont = do
  let expls = snd <$> roleExpls
  expls' <- forM expls \case
    Explicit -> return $ Inferred Nothing Unify
    expl     -> return $ expl
  refreshAbs (Abs bs UnitE) \bs' UnitE -> do
    let vs = nestToNames bs'
    vs' <- mapM toAtomVar vs
    cont expls' bs' vs $ TyConParams expls (toAtom <$> vs')

makeTyConParams :: EnvReader m => TyConName n -> [CAtom n] -> m n (TyConParams n)
makeTyConParams tc params = do
  TyConDef _ expls _ _ <- lookupTyCon tc
  return $ TyConParams (map snd expls) params

dictType :: EnvReader m => ClassName n -> [CAtom n] -> m n (DictType n)
dictType className params = do
  ClassDef sourceName builtinName _ _ _ _ _ _ <- lookupClassDef className
  return case builtinName of
    Just Ix   -> IxDictType   singleTyParam
    Just Data -> DataDictType singleTyParam
    Nothing   -> DictType sourceName className params
    where singleTyParam = case params of
            [p] -> fromJust $ toMaybeType p
            _ -> error "not a single type param"

makePreludeMaybeTy :: EnvReader m => CType n -> m n (CType n)
makePreludeMaybeTy ty = do
  ~(Just (UTyConVar tyConName)) <- lookupSourceMap "Maybe"
  let params = TyConParams [Explicit] [toAtom ty]
  return $ toType $ UserADTType "Maybe" tyConName params

-- === computing effects ===

functionEffs :: (IRRep r, EnvReader m) => LamExpr r n -> m n (EffectRow r n)
functionEffs f = getLamExprType f >>= \case
  PiType b (EffTy effs _) -> return $ ignoreHoistFailure $ hoist b effs

rwsFunEffects :: (IRRep r, EnvReader m) => RWS -> LamExpr r n -> m n (EffectRow r n)
rwsFunEffects rws f = getLamExprType f >>= \case
   PiType (BinaryNest h ref) et -> do
     let effs' = ignoreHoistFailure $ hoist ref (etEff et)
     let hVal = toAtom $ AtomVar (binderName h) (TyCon HeapType)
     let effs'' = deleteEff (RWSEffect rws hVal) effs'
     return $ ignoreHoistFailure $ hoist h effs''
   _ -> error "Expected a binary function type"

getLamExprType :: (IRRep r, EnvReader m) => LamExpr r n -> m n (PiType r n)
getLamExprType (LamExpr bs body) =
   return $ PiType bs $ EffTy (getEffects body) (getType body)

getTypeRWSAction :: (IRRep r, EnvReader m) => LamExpr r n -> m n (Type r n, Type r n)
getTypeRWSAction f = getLamExprType f >>= \case
  PiType (BinaryNest regionBinder refBinder) (EffTy _ resultTy) -> do
    case binderType refBinder of
      RefTy _ referentTy -> do
        let referentTy' = ignoreHoistFailure $ hoist regionBinder referentTy
        let resultTy' = ignoreHoistFailure $ hoist (PairB regionBinder refBinder) resultTy
        return (resultTy', referentTy')
      _ -> error "expected a ref"
  _ -> error "expected a pi type"

getSuperclassDicts :: EnvReader m => CDict n -> m n ([CAtom n])
getSuperclassDicts dict = do
  case getType dict of
    TyCon (DictTy dTy) -> do
      ts <- getSuperclassTys dTy
      forM (enumerate ts) \(i, _) -> reduceSuperclassProj i dict
    _ -> error "expected a dict type"

getSuperclassTys :: EnvReader m => DictType n -> m n [CType n]
getSuperclassTys = \case
  DictType _ className params -> do
    ClassDef _ _ _ _ _ bs superclasses _ <- lookupClassDef className
    forM [0 .. nestLength superclasses - 1] \i -> do
      instantiate (Abs bs $ getSuperclassType REmpty superclasses i) params
  DataDictType _ -> return []
  IxDictType ty -> return [toType $ DataDictType ty]

getTypeTopFun :: EnvReader m => TopFunName n -> m n (PiType SimpIR n)
getTypeTopFun f = lookupTopFun f >>= \case
  DexTopFun _ (TopLam _ piTy _) _ -> return piTy
  FFITopFun _ iTy -> liftIFunType iTy

asTopLam :: (EnvReader m, IRRep r) => LamExpr r n -> m n (TopLam r n)
asTopLam lam = do
  piTy <- getLamExprType lam
  return $ TopLam False piTy lam

liftIFunType :: (IRRep r, EnvReader m) => IFunType -> m n (PiType r n)
liftIFunType (IFunType _ argTys resultTys) = liftEnvReaderM $ go argTys where
  go :: IRRep r => [BaseType] -> EnvReaderM n (PiType r n)
  go = \case
    [] -> return $ PiType Empty (EffTy (OneEffect IOEffect) resultTy)
      where resultTy = case resultTys of
              [] -> UnitTy
              [t] -> toType $ BaseType t
              [t1, t2] -> TyCon (ProdType [toType $ BaseType t1, toType $ BaseType t2])
              _ -> error $ "Not a valid FFI return type: " ++ pprint resultTys
    t:ts -> withFreshBinder noHint (toType $ BaseType t) \b -> do
      PiType bs effTy <- go ts
      return $ PiType (Nest b bs) effTy

-- === Data constraints ===

isData :: EnvReader m => Type CoreIR n -> m n Bool
isData ty = do
  result <- liftEnvReaderT $ withSubstReaderT $ go ty
  case result of
    Just () -> return True
    Nothing -> return False
  where
    go :: Type CoreIR i -> SubstReaderT Name (EnvReaderT Maybe) i o ()
    go = \case
      StuckTy _ _ -> notData
      TyCon con -> case con of
        TabPi (TabPiType _ b eltTy) -> renameBinders b \_ -> go eltTy
        DepPairTy (DepPairType _ b@(_:>l) r) -> go l >> renameBinders b \_ -> go r
        NewtypeTyCon nt -> do
          (_, ty') <- unwrapNewtypeType =<< renameM nt
          dropSubst $ go ty'
        BaseType _  -> return ()
        ProdType as -> mapM_ go as
        SumType  cs -> mapM_ go cs
        RefType _ _ -> return ()
        HeapType    -> return ()
        TypeKind -> notData
        DictTy _ -> notData
        Pi _     -> notData
      where notData = empty

checkExtends :: (Fallible m, IRRep r) => EffectRow r n -> EffectRow r n -> m ()
checkExtends allowed (EffectRow effs effTail) = do
  let (EffectRow allowedEffs allowedEffTail) = allowed
  case effTail of
    EffectRowTail _ -> assertEq allowedEffTail effTail ""
    NoTail -> return ()
  forM_ (eSetToList effs) \eff -> unless (eff `eSetMember` allowedEffs) $
    throwInternal $ "Unexpected effect: " ++ pprint eff ++
                       "\nAllowed: " ++ pprint allowed
{-# INLINE checkExtends #-}