content: {-# Language BlockArguments, ImplicitParams #-}
  module Cryptol.TypeCheck.Module (doFunctorInst) where
  
  import Data.List(partition,unzip4)
  import Data.Text(Text)
  import Data.Map(Map)
  import Data.Maybe (maybeToList)
  import qualified Data.Map as Map
  import qualified Data.Map.Merge.Strict as Map
  import Data.Set (Set)
  import qualified Data.Set as Set
  import Control.Monad(unless,forM_,mapAndUnzipM)
  
  
  import Cryptol.Utils.Panic(panic)
  import Cryptol.Utils.Ident(Ident,Namespace(..),ModPath,isInfixIdent)
  import Cryptol.Parser.Position (Range,Located(..), thing)
  import qualified Cryptol.Parser.AST as P
  import Cryptol.ModuleSystem.Binds(newFunctorInst)
  import Cryptol.ModuleSystem.Name(nameIdent)
  import Cryptol.ModuleSystem.NamingEnv
            (NamingEnv(..), modParamNamingEnv, shadowing, without)
  import Cryptol.ModuleSystem.Interface
            ( IfaceG(..), IfaceDecls(..), IfaceNames(..), IfaceDecl(..)
            , filterIfaceDecls
            )
  import Cryptol.TypeCheck.AST
  import Cryptol.TypeCheck.Error
  import Cryptol.TypeCheck.Subst(Subst,listParamSubst,apSubst,mergeDistinctSubst)
  import Cryptol.TypeCheck.Solve(proveImplication)
  import Cryptol.TypeCheck.Monad
  import Cryptol.TypeCheck.Instantiate(instantiateWith)
  import Cryptol.TypeCheck.ModuleInstance
  import Cryptol.TypeCheck.ModuleBacktickInstance(MBQual, doBacktickInstance)
  
  doFunctorInst ::
    Located (P.ImpName Name)    {- ^ Name for the new module -} ->
    Located (P.ImpName Name)    {- ^ Functor being instantiated -} ->
    P.ModuleInstanceArgs Name   {- ^ Instance arguments -} ->
    Map Name Name
    {- ^ Instantitation.  These is the renaming for the functor that arises from
         generativity (i.e., it is something that will make the names "fresh").
    -} ->
    NamingEnv
    {- ^ Names in the enclosing scope of the instantiated module -} ->
    Maybe Text                  {- ^ Documentation for the module being generated -} ->
    InferM (Maybe TCTopEntity)
  doFunctorInst m f as instMap0 enclosingInScope doc =
    inRange (srcRange m)
    do mf    <- lookupFunctor (thing f)
       argIs <- checkArity (srcRange f) mf as
       m2 <- do let mpath = P.impNameModPath (thing m)
                as2 <- mapM (checkArg mpath) argIs
                let (tySus,paramTySyns,decls,paramInstMaps) =
                      unzip4 [ (su,ts,ds,im) | DefinedInst su ts ds im <- as2 ]
                instMap <- addMissingTySyns mpath mf instMap0
                let ?tVarSu = mergeDistinctSubst tySus
                    ?nameSu = instMap <> mconcat paramInstMaps
                let m1   = moduleInstance mf
                    m2   = m1 { mName             = m
                              , mDoc              = mempty
                              , mParamTypes       = mempty
                              , mParamFuns        = mempty
                              , mParamConstraints = mempty
                              , mParams           = mempty
                              , mTySyns = mconcat paramTySyns <> mTySyns m1
                              , mDecls = map NonRecursive (concat decls) ++
                                        mDecls m1
                              }
                let (tps,tcs,vps) =
                        unzip3 [ (xs,cs,fs) | ParamInst xs cs fs <- as2 ]
                    tpSet  = Set.unions tps
                    tpSet' = Set.map snd (Set.unions tps)
                    emit p = Set.null (freeParams (thing p)
                                                  `Set.intersection` tpSet')
  
                    (emitPs,delayPs) = partition emit (mParamConstraints m1)
  
                forM_ emitPs \lp ->
                  newGoals (CtModuleInstance (srcRange lp)) [thing lp]
  
                doBacktickInstance tpSet
                                   (map thing delayPs ++ concat tcs)
                                   (Map.unions vps)
                                   m2
  
       -- An instantiation doesn't really have anything "in scope" per se, but
       -- here we compute what would be in scope as if you hand wrote the
       -- instantiation by copy-pasting the functor then substituting the
       -- parameters. That is, it would be whatever is in scope in the functor,
       -- together with any names in the enclosing scope if this is a nested
       -- module, with the functor's names taking precedence. This is used to
       -- determine what is in scope at the REPL when the instantiation is loaded
       -- and focused.
       --
       -- The exception is when instantiating with _, in which case we must delete
       -- the module parameters from the naming environment.
       let inScope0 = mInScope m2 `without`
             mconcat [ modParamNamingEnv mp | (_, mp, AddDeclParams) <- argIs ]
           inScope = inScope0 `shadowing` enclosingInScope
  
       -- Combine the docstrings of:
       -- * The functor being instantiated
       -- * The module being generated
       let newDoc = maybeToList doc <> mDoc mf
  
       case thing m of
         P.ImpTop mn    -> newModuleScope newDoc mn (mExports m2) inScope
         P.ImpNested mn -> newSubmoduleScope mn newDoc (mExports m2) inScope
  
       mapM_ addTySyn     (Map.elems (mTySyns m2))
       mapM_ addNominal   (Map.elems (mNominalTypes m2))
       addSignatures      (mSignatures m2)
       addSubmodules      (mSubmodules m2)
       setNested          (mNested m2)
       addFunctors        (mFunctors m2)
       mapM_ addDecls     (mDecls m2)
  
       case thing m of
         P.ImpTop {}    -> Just <$> endModule
         P.ImpNested {} -> endSubmodule >> pure Nothing
  
  
  data ActualArg =
      UseParameter ModParam     -- ^ Instantiate using this parameter
    | UseModule (IfaceG ())     -- ^ Instantiate using this module
    | AddDeclParams             -- ^ Instantiate by adding parameters
  
  
  
  
  {- | Validate a functor application, just checking the argument names.
  The result associates a module parameter with the concrete way it should
  be instantiated.
  -}
  checkArity ::
    Range             {- ^ Location for reporting errors -} ->
    ModuleG ()        {- ^ The functor being instantiated -} ->
    P.ModuleInstanceArgs Name {- ^ The arguments -} ->
    InferM [(Range, ModParam, ActualArg)]
    {- ^ Associates functor parameters with the interfaces of the
         instantiating modules -}
  checkArity r mf args =
    case args of
  
      P.DefaultInstArg arg ->
        let p0 = case Map.keys ps0 of
                   p':_ -> p'
                   [] -> panic "checkArity" ["functor with no parameters"]
            i = Located { srcRange = srcRange arg
                        , thing    = p0
                        }
        in checkArgs [] ps0 [ P.ModuleInstanceNamedArg i arg ]
  
      P.NamedInstArgs as -> checkArgs [] ps0 as
  
      P.DefaultInstAnonArg {} -> panic "checkArity" [ "DefaultInstAnonArg" ]
    where
    ps0 = mParams mf
  
    checkArgs done ps as =
      case as of
  
        [] -> do forM_ (Map.keys ps) \p ->
                   recordErrorLoc (Just r) (FunctorInstanceMissingArgument p)
                 pure done
  
        P.ModuleInstanceNamedArg ll lm : more ->
          case Map.lookup (thing ll) ps of
            Just i ->
              do arg <- case thing lm of
                          P.ModuleArg m -> Just . UseModule <$> lookupModule m
                          P.ParameterArg p ->
                             do mb <- lookupModParam p
                                case mb of
                                  Nothing ->
                                     do inRange (srcRange lm)
                                                (recordError (MissingModParam p))
                                        pure Nothing
                                  Just a -> pure (Just (UseParameter a))
                          P.AddParams -> pure (Just AddDeclParams)
                 let next = case arg of
                              Nothing -> done
                              Just a  -> (srcRange lm, i, a) : done
                 checkArgs next (Map.delete (thing ll) ps) more
  
            Nothing ->
              do recordErrorLoc (Just (srcRange ll))
                                (FunctorInstanceBadArgument (thing ll))
                 checkArgs done ps more
  
  
  data ArgInst = -- | Argument that defines the params
                 DefinedInst Subst
                   (Map Name TySyn)
                   -- ^ Type synonyms created from the functor's type parameters
                   [Decl]
                   -- ^ Bindings for value parameters
                   (Map Name Name)
                   -- ^ Map from the functor's parameter names to the new names
                   --   created for the instantiation
  
               | ParamInst (Set (MBQual TParam)) [Prop] (Map (MBQual Name) Type)
                 -- ^ Argument that add parameters
                 -- The type parameters are in their module type parameter
                 -- form (i.e., tpFlav is TPModParam)
  
  
  
  {- | Check the argument to a functor parameter.
  Returns:
  
    * A substitution which will replace the parameter types with
      the concrete types that were provided
  
    * Some declarations that define the parameters in terms of the provided
      values.
  
    * XXX: Extra parameters for instantiation by adding params
  -}
  checkArg ::
    ModPath -> (Range, ModParam, ActualArg) -> InferM ArgInst
  checkArg mpath (r,expect,actual') =
    case actual' of
      AddDeclParams   -> paramInst
      UseParameter {} -> definedInst
      UseModule {}    -> definedInst
  
    where
    paramInst =
      do let as = Set.fromList
                     (map (qual . mtpParam) (Map.elems (mpnTypes params)))
             cs = map thing (mpnConstraints params)
             check = checkSimpleParameterValue r (mpName expect)
             qual a = (mpQual expect, a)
         fs <- Map.mapMaybeWithKey (\_ v -> v) <$> mapM check (mpnFuns params)
         pure (ParamInst as cs (Map.mapKeys qual fs))
  
    definedInst =
      do (tRens, tSyns, tInstMaps) <- unzip3 <$>
           mapM (checkParamType mpath r tyMap) (Map.toList (mpnTypes params))
         let renSu = listParamSubst (concat tRens)
  
         {- Note: the constraints from the signature are already added to the
            constraints for the functor and they are checked all at once in
            doFunctorInst -}
  
  
         (vDecls, vInstMaps) <-
           mapAndUnzipM (checkParamValue mpath r vMap)
             [ s { mvpType = apSubst renSu (mvpType s) }
             | s <- Map.elems (mpnFuns params) ]
  
         pure $ DefinedInst renSu (mconcat tSyns)
           (concat vDecls) (mconcat tInstMaps <> mconcat vInstMaps)
  
  
    params = mpParameters expect
  
    -- Things provided by the argument module
    tyMap :: Map Ident (Kind, Type)
    vMap  :: Map Ident (Name, Schema)
    (tyMap,vMap) =
      case actual' of
        UseParameter mp ->
          ( nameMapToIdentMap fromTP (mpnTypes ps)
          , nameMapToIdentMap fromVP (mpnFuns ps)
          )
          where
          ps        = mpParameters mp
          fromTP tp = (mtpKind tp, TVar (TVBound (mtpParam tp)))
          fromVP vp = (mvpName vp, mvpType vp)
  
        UseModule actual ->
          ( Map.unions [ nameMapToIdentMap fromTS      (ifTySyns decls)
                       , nameMapToIdentMap fromNominal (ifNominalTypes decls)
                       ]
  
          , nameMapToIdentMap fromD (ifDecls decls)
          )
  
          where
          localNames      = ifsPublic (ifNames actual)
          isLocal x       = x `Set.member` localNames
  
          -- Things defined by the argument module
          decls           = filterIfaceDecls isLocal (ifDefines actual)
  
          fromD d         = (ifDeclName d, ifDeclSig d)
          fromTS ts       = (kindOf ts, tsDef ts)
          fromNominal nt  = (kindOf nt, TNominal nt [])
  
        AddDeclParams -> panic "checkArg" ["AddDeclParams"]
  
  
  
  nameMapToIdentMap :: (a -> b) -> Map Name a -> Map Ident b
  nameMapToIdentMap f m =
    Map.fromList [ (nameIdent n, f v) | (n,v) <- Map.toList m ]
  
  
  
  
  -- | Check a type parameter to a module.
  checkParamType ::
    ModPath                    {- ^ The new module we are creating -} ->
    Range                      {- ^ Location for error reporting -} ->
    Map Ident (Kind,Type)      {- ^ Actual types -} ->
    (Name,ModTParam)           {- ^ Type parameter -} ->
    InferM ([(TParam,Type)], Map Name TySyn, Map Name Name)
      {- ^ Mapping from parameter name to actual type (for type substitution),
           type synonym map from a fresh type name to the actual type
             (only so that the type can be referred to in the REPL;
              type synonyms are fully inlined into types at this point),
           and a map from the old type name to the fresh type name
             (for instantiation) -}
  checkParamType mpath r tyMap (name,mp) =
    let i       = nameIdent name
        expectK = mtpKind mp
    in
    case Map.lookup i tyMap of
      Nothing ->
        do recordErrorLoc (Just r) (FunctorInstanceMissingName NSType i)
           pure ([], Map.empty, Map.empty)
      Just (actualK,actualT) ->
        do unless (expectK == actualK)
             (recordErrorLoc (Just r)
                             (KindMismatch (Just (TVFromModParam name))
                                                    expectK actualK))
           name' <- newFunctorInst mpath name
           let tySyn = TySyn { tsName = name'
                             , tsParams = []
                             , tsConstraints = []
                             , tsDef = actualT
                             , tsDoc = mtpDoc mp }
           pure ( [(mtpParam mp, actualT)]
                , Map.singleton name' tySyn
                , Map.singleton name name'
                )
  
  -- | Check a value parameter to a module.
  checkParamValue ::
    ModPath                 {- ^ The new module we are creating -} ->
    Range                   {- ^ Location for error reporting -} ->
    Map Ident (Name,Schema) {- ^ Actual values -} ->
    ModVParam               {- ^ The parameter we are checking -} ->
    InferM ([Decl], Map Name Name)
    {- ^ Decl mapping a new name to the actual value,
         and a map from the value param name in the functor to the new name
           (for instantiation) -}
  checkParamValue mpath r vMap mp =
    let name     = mvpName mp
        i        = nameIdent name
        expectT  = mvpType mp
    in case Map.lookup i vMap of
         Nothing ->
           do recordErrorLoc (Just r) (FunctorInstanceMissingName NSValue i)
              pure ([], Map.empty)
         Just actual ->
           do e <- mkParamDef r (name,expectT) actual
              name' <- newFunctorInst mpath name
              let d = Decl { dName        = name'
                           , dSignature   = expectT
                           , dDefinition  = DExpr e
                           , dPragmas     = []
                           , dInfix       = isInfixIdent (nameIdent name')
                           , dFixity      = mvpFixity mp
                           , dDoc         = mvpDoc mp
                           }
  
              pure ([d], Map.singleton name name')
  
  
  
  checkSimpleParameterValue ::
    Range                       {- ^ Location for error reporting -} ->
    Ident                       {- ^ Name of functor parameter -} ->
    ModVParam                   {- ^ Module parameter -} ->
    InferM (Maybe Type)  {- ^ Type to add to things, `Nothing` on err -}
  checkSimpleParameterValue r i mp =
    case (sVars sch, sProps sch) of
      ([],[]) -> pure (Just (sType sch))
      _ ->
        do recordErrorLoc (Just r)
              (FunctorInstanceBadBacktick
                 (BIPolymorphicArgument i (nameIdent (mvpName mp))))
           pure Nothing
    where
    sch = mvpType mp
  
  
  {- | Make an "adaptor" that instantiates the paramter into the form expected
  by the functor.  If the actual type is:
  
  > {x} P => t
  
  and the provided type is:
  
  > f : {y} Q => s
  
  The result, if successful would be:
  
    /\x \{P}. f @a {Q}
  
  To do this we need to find types `a` to instantiate `y`, and prove that:
    {x} P => Q[a/y] /\ s = t
  -}
  
  mkParamDef ::
    Range           {- ^ Location of instantiation for error reporting -} ->
    (Name,Schema)   {- ^ Name and type of parameter -} ->
    (Name,Schema)   {- ^ Name and type of actual argument -} ->
    InferM Expr
  mkParamDef r (pname,wantedS) (arg,actualS) =
    do (e,todo) <- collectGoals
            $ withTParams (sVars wantedS)
              do (e,t) <- instantiateWith pname(EVar arg) actualS []
                 props <- unify WithSource { twsType   = sType wantedS
                                           , twsSource = TVFromModParam arg
                                           , twsRange  = Just r
                                           }
                          t
                 newGoals (CtModuleInstance r) props
                 pure e
       su <- proveImplication False
                              (Just pname)
                              (sVars wantedS)
                              (sProps wantedS)
                              todo
       let res  = foldr ETAbs     res1            (sVars wantedS)
           res1 = foldr EProofAbs (apSubst su e)  (sProps wantedS)
  
       applySubst res
  
  
  -- | The instMap we get from the renamer will not contain the fresh names for
  -- certain things in the functor generated in the typechecking stage, if we are
  -- instantiating a functor that is in the same file, since renaming and
  -- typechecking happens together with the instantiation. In particular, if the
  -- functor's interface has type synonyms, they will only get copied over into
  -- the functor in the typechecker, so the renamer will not see them. Here we
  -- make the fresh names for those missing type synonyms and add them to the
  -- instMap.
  addMissingTySyns ::
    ModPath                  {- ^ The new module we are creating -} ->
    ModuleG ()               {- ^ The functor -} ->
    Map Name Name            {- ^ instMap we get from renamer -} ->
    InferM (Map Name Name)   {- ^ the complete instMap -}
  addMissingTySyns mpath f = Map.mergeA
    (Map.traverseMissing \name _ -> newFunctorInst mpath name)
    Map.preserveMissing
    (Map.zipWithMatched \_ _ name' -> name')
    (mTySyns f)
  