content: {-# LANGUAGE TypeApplications #-}
  
  -- | Compilation to Haskell types.
  module Agda2Hs.Compile.Type where
  
  import Control.Arrow ( (>>>) )
  import Control.Monad ( forM, when, unless )
  import Control.Monad.Trans ( lift )
  import Control.Monad.Reader ( asks )
  import Data.List ( find )
  import Data.Maybe ( mapMaybe, isJust )
  import qualified Data.Set as Set ( singleton )
  
  import Agda.Compiler.Backend hiding ( Args )
  
  import Agda.Syntax.Common
  import Agda.Syntax.Internal
  import Agda.Syntax.Common.Pretty ( prettyShow )
  
  import Agda.TypeChecking.Pretty
  import Agda.TypeChecking.Reduce ( reduce, unfoldDefinitionStep, instantiate )
  import Agda.TypeChecking.Substitute
  import Agda.TypeChecking.Telescope
  
  import Agda.Utils.Impossible ( __IMPOSSIBLE__ )
  import Agda.Utils.List ( downFrom )
  import Agda.Utils.Maybe ( ifJustM, fromMaybe )
  import Agda.Utils.Monad ( ifM, whenM, unlessM, and2M, or2M )
  import Agda.Utils.Size ( Sized(size) )
  import Agda.Utils.Functor ( ($>) )
  
  import Agda2Hs.Compile.Name ( compileQName )
  import Agda2Hs.Compile.Types
  import Agda2Hs.Compile.Utils
  import Agda2Hs.Compile.Var
  import Agda2Hs.AgdaUtils
  import qualified Agda2Hs.Language.Haskell as Hs
  import Agda2Hs.Language.Haskell.Utils
    ( Strictness(Lazy), hsName, tApp, constrainType, qualifyType )
  
  
  -- | Type definitions from the prelude that get special translation rules.
  isSpecialType :: QName -> Maybe (Elims -> C (Hs.Type ()))
  isSpecialType = prettyShow >>> \case
    "Haskell.Prim.Tuple._×_"    -> Just tupleType
    "Haskell.Prim.Tuple._×_×_"  -> Just tupleType
    "Haskell.Extra.Sigma.Σ"     -> Just tupleType
    "Haskell.Extra.Erase.Erase" -> Just unitType
    "Haskell.Extra.Delay.Delay" -> Just delayType
    _                           -> Nothing
  
  
  -- | Compile all the elims into a n-uple.
  tupleType :: Elims -> C (Hs.Type ())
  tupleType es = do
    let Just as = allApplyElims es
    ts <- mapM (compileType . unArg) as
    return $ Hs.TyTuple () Hs.Boxed ts
  
  
  -- | Ignore arguments and return the unit type.
  unitType :: Elims -> C (Hs.Type ())
  unitType _ = return $ Hs.TyTuple () Hs.Boxed []
  
  
  -- | Compile fully applied Delay type as its only type argument.
  delayType :: Elims -> C (Hs.Type ())
  delayType (Apply a : _) = compileType (unArg a)
  delayType (_       : _) = __IMPOSSIBLE__
  delayType [] = agda2hsError "Cannot compile unapplied Delay type"
  
  
  -- | Compile an Agda term into a Haskell type, along with its strictness.
  compileTypeWithStrictness :: Term -> C (Strictness, Hs.Type ())
  compileTypeWithStrictness t = do
    s <- case t of
      Def f es -> fromMaybe Lazy <$> isUnboxRecord f
      _        -> return Lazy
    ty <- compileType t
    pure (s, ty)
  
  
  -- | Compile an Agda term into a Haskell type.
  compileType :: Term -> C (Hs.Type ())
  compileType t = do
  
    reportSDoc "agda2hs.compile.type" 12 $ text "Compiling type" <+> prettyTCM t
    reportSDoc "agda2hs.compile.type" 22 $ text "Compiling type" <+> pretty t
  
    whenM (isErasedBaseType t) fail
  
    instantiate t >>= \case
      Pi a b -> do
        reportSDoc "agda2hs.compile.type" 13 $ text "Compiling pi type (" <+> prettyTCM (absName b)
          <+> text ":" <+> prettyTCM a <+> text ") -> " <+> underAbstraction a b prettyTCM
        let compileB = underAbstraction a b (compileType . unEl)
        compileDomType (absName b) a >>= \case
          DomType _ hsA -> Hs.TyFun () hsA <$> compileB
          DomConstraint hsA -> constrainType hsA <$> compileB
          DomDropped -> compileB
          DomForall hsA -> qualifyType hsA <$> compileB
  
      Def f es -> maybeUnfoldCopy f es compileType $ \f es -> do
        def <- getConstInfo f
        if | not (usableModality def) ->
              agda2hsErrorM $
                    text "Cannot use erased definition" <+> prettyTCM f
                <+> text "in Haskell type"
           | Just semantics <- isSpecialType f -> setCurrentRange f $ semantics es
           | Just args <- allApplyElims es ->
               ifJustM (isUnboxRecord f) (\_ -> compileUnboxType f args) $
               ifJustM (isTupleRecord f) (\b -> compileTupleType f b args) $
               ifM (isTransparentFunction f) (compileTransparentType (defType def) args) $
               ifM (isInlinedFunction f) (compileInlineType f es) $ do
                 vs <- compileTypeArgs (defType def) args
                 f <- compileQName f
                 return $ tApp (Hs.TyCon () f) vs
           | otherwise -> fail
  
      Var x es | Just args <- allApplyElims es -> do
        CtxVar _ ti <- lookupBV x
        unless (usableModality ti) $ agda2hsErrorM $
              text "Cannot use erased variable" <+> prettyTCM (var x)
          <+> text "in Haskell type"
        vs <- compileTypeArgs (unDom ti) args
        x  <- hsName <$> compileDBVar x
        return $ tApp (Hs.TyVar () x) vs
  
      Sort s -> return (Hs.TyStar ())
  
      Lam argInfo restAbs -> do
        (body , x0) <- underAbstraction_ restAbs $ \b ->
          (,) <$> compileType b <*> (hsName <$> compileDBVar 0)
  
        -- TODO: we should also drop lambdas that can be erased based on their type
        -- (e.g. argument is of type Level/Size or in a Prop) but currently we do
        -- not have access to the type of the lambda here.
        if | hasQuantity0 argInfo -> return body
           -- Rewrite `\x -> (a -> x)` to `(->) a`
           | Hs.TyFun _ a (Hs.TyVar _ y) <- body
           , y == x0 -> return $ Hs.TyApp () (Hs.TyCon () $ Hs.Special () $ Hs.FunCon ()) a
           -- Rewrite `\x -> f x` to `f`
           | Hs.TyApp _ f (Hs.TyVar _ y) <- body
           , y == x0 -> return f
           | otherwise -> agda2hsErrorM $ text "Not supported: type-level lambda" <+> prettyTCM t
  
      _ -> fail
    where fail = agda2hsErrorM $ text "Bad Haskell type:" <?> prettyTCM t
  
  
  compileTypeArgs :: Type -> Args -> C [Hs.Type ()]
  compileTypeArgs ty [] = pure []
  compileTypeArgs ty (x:xs) = do
    (a, b) <- mustBePi ty
    reportSDoc "agda2hs.compile.type" 16 $ text "compileTypeArgs x =" <+> prettyTCM x
    reportSDoc "agda2hs.compile.type" 16 $ text "                a =" <+> prettyTCM a
    reportSDoc "agda2hs.compile.type" 16 $ text "         modality =" <+> prettyTCM (getModality a)
    let rest = compileTypeArgs (absApp b $ unArg x) xs
    let fail msg = agda2hsErrorM $ (text msg <> text ":") <+> parens (prettyTCM (absName b) <+> text ":" <+> prettyTCM (unDom a))
    compileDom a >>= \case
      DODropped -> rest
      DOInstance -> fail "Type-level instance argument not supported"
      DOType -> do
        (:) <$> compileType (unArg x) <*> rest
      DOTerm -> fail "Type-level term argument not supported"
  
  compileTel :: Telescope -> C [Hs.Type ()]
  compileTel EmptyTel = return []
  compileTel (ExtendTel a tel) = compileDom a >>= \case
    DODropped  -> underAbstraction a tel compileTel
    DOInstance -> __IMPOSSIBLE__
    DOType     -> __IMPOSSIBLE__
    DOTerm     -> (:) <$> compileType (unEl $ unDom a) <*> underAbstraction a tel compileTel
  
  -- Version of @compileTel@ that just computes the size,
  -- and avoids compiling the types themselves.
  compileTelSize :: Telescope -> C Int
  compileTelSize EmptyTel = return 0
  compileTelSize (ExtendTel a tel) = compileDom a >>= \case
    DODropped -> underAbstraction a tel compileTelSize
    DOInstance -> __IMPOSSIBLE__
    DOType -> __IMPOSSIBLE__
    DOTerm -> (1+) <$> underAbstraction a tel compileTelSize
  
  compileUnboxType :: QName -> Args -> C (Hs.Type ())
  compileUnboxType r pars = do
    def <- getConstInfo r
    let tel = recTel (theDef def) `apply` pars
    compileTel tel >>= \case
      [t] -> return t
      _   -> __IMPOSSIBLE__
  
  compileTupleType :: QName -> Hs.Boxed -> Args -> C (Hs.Type ())
  compileTupleType r b pars = do
    tellUnboxedTuples b
    def <- getConstInfo r
    let tel = recTel (theDef def) `apply` pars
    ts <- compileTel tel
    return $ Hs.TyTuple () b ts
  
  compileTransparentType :: Type -> Args -> C (Hs.Type ())
  compileTransparentType ty args = compileTypeArgs ty args >>= \case
    (v:vs) -> return $ v `tApp` vs
    []     -> __IMPOSSIBLE__
  
  
  compileInlineType :: QName -> Elims -> C (Hs.Type ())
  compileInlineType f args = do
    Function { funClauses = cs } <- theDef <$> getConstInfo f
  
    let [ Clause { namedClausePats = pats } ] = filter (isJust . clauseBody) cs
  
    when (length args < length pats) $ agda2hsErrorM $
      text "Cannot compile inlinable type alias" <+> prettyTCM f <+> text "as it must be fully applied."
  
    r <- liftReduce $ locallyReduceDefs (OnlyReduceDefs $ Set.singleton f)
                    $ unfoldDefinitionStep (Def f args) f args
  
    case r of
      YesReduction _ t -> compileType t
      _                -> agda2hsErrorM $ text "Could not reduce inline type alias " <+> prettyTCM f
  
  
  data DomOutput = DOInstance | DODropped | DOType | DOTerm
  
  compileDom :: Dom Type -> C DomOutput
  compileDom a = do
    isErasable <- pure (not $ usableModality a) `or2M` canErase (unDom a)
    isClassConstraint <- pure (isInstance a) `and2M` isClassType (unDom a)
    isType <- endsInSort (unDom a)
    return $ if
      | isErasable        -> DODropped
      | isClassConstraint -> DOInstance
      | isType            -> DOType
      | otherwise         -> DOTerm
  
  -- | Compile a function type domain.
  -- A domain can either be:
  --
  -- - dropped if the argument is erased.
  -- - added as a class constraint.
  -- - added as a type parameter
  -- - kept as a regular explicit argument.
  compileDomType :: ArgName -> Dom Type -> C CompiledDom
  compileDomType x a =
    compileDom a >>= \case
      DODropped  -> pure DomDropped
      DOInstance -> DomConstraint . Hs.TypeA () <$> compileType (unEl $ unDom a)
      DOType     -> do
        -- We compile (non-erased) type parameters to an explicit forall if they
        -- come from a module parameter or if we are in a nested position inside the type.
        reportSDoc "agda2hs.compile.type" 15 $ text "Compiling forall type:" <+> prettyTCM a
        isNested <- asks isNestedInType
        ctx <- getContextSize
        npars <- size <$> (lookupSection =<< currentModule)
        if
          | isNested -> do
            tellExtension Hs.RankNTypes
            -- tellExtension Hs.ExistentialQuantification
            return $ DomForall $ Hs.UnkindedVar () $ Hs.Ident () x
          | ctx < npars -> do
            tellExtension Hs.ScopedTypeVariables
            return $ DomForall $ Hs.UnkindedVar () $ Hs.Ident () x
          | otherwise -> return DomDropped
      DOTerm     -> fmap (uncurry DomType) . withNestedType . compileTypeWithStrictness . unEl $ unDom a
  
  compileTeleBinds :: Bool -> Telescope -> C [Hs.TyVarBind ()]
  compileTeleBinds kinded = go
    where
    go EmptyTel = return []
    go (ExtendTel a tel) = do
      reportSDoc "agda2hs.compile.type" 15 $ text "Compiling type parameter: " <+> prettyTCM (absName tel) <+> ":" <+> prettyTCM a
      let fail msg = agda2hsErrorM $ (text msg <> text ":") <+> parens (prettyTCM (absName tel) <+> text ":" <+> prettyTCM (unDom a))
      compileDom a >>= \case
        DODropped  -> underAbstraction a tel go
        DOType -> do
          ha <- compileKeptTeleBind kinded (hsName $ absName tel) (unDom a)
          (ha:) <$> underAbstraction a tel go
        DOInstance -> agda2hsError "Constraint in type parameter not supported"
        DOTerm -> agda2hsError "Term variable in type parameter not supported"
  
  compileKeptTeleBind :: Bool -> Hs.Name () -> Type -> C (Hs.TyVarBind ())
  compileKeptTeleBind kinded x t = do
    checkValidTyVarName x
    k <- compileKind t
    pure $ if kinded
      then Hs.KindedVar () x k
      else Hs.UnkindedVar () x
  
  compileKind :: Type -> C (Hs.Kind ())
  compileKind t = case unEl t of
    Sort (Type _) -> pure (Hs.TyStar ())
    Pi a b -> compileDom a >>= \case
      DODropped -> underAbstraction a b compileKind
      DOType -> Hs.TyFun () <$> compileKind (unDom a) <*> underAbstraction a b compileKind
      DOTerm -> err
      DOInstance -> err
    _ -> err
    where
      err = agda2hsErrorM $ text "Not a valid Haskell kind: " <+> prettyTCM t
  