{-# OPTIONS -fglasgow-exts -cpp #-}

{-|

The scope analysis. Translates from concrete to abstract syntax.

-}
module Syntax.Scope
    ( -- Types
      ScopeInfo(..)
    , Modules
    , KindOfName(..)
    , ModuleScope(..)
    , DefinedName(..)
    , NameSpace(..)
    , ResolvedName(..)
    , ScopeM
      -- Functions
    , emptyScopeInfo
    , getScopeInfo
    , setScopeInfo
    , modScopeInfo
    , modScopeInfoM
    , getModule
    , getFixity
    , setContext
    , getContextPrecedence
    , resolvePatternNameM
    , resolveName
    , abstractName
    , notInScope
    , currentModuleScope
    , currentNameSpace
    , insideModule
    , defName
    , defineName
    , bindVariable
    , bindVariables
    , defineModule
    , openModule
    , importModule
    , implicitModule
    ) where

import Control.Exception
import Control.Monad.Reader
import Control.Monad.State
import Data.Generics (Data,Typeable)
import Data.Monoid
import Data.Typeable
import Data.Map as Map
import Data.List as List
import Data.Tree

import Syntax.Position
import Syntax.Common
import Syntax.Concrete
import Syntax.Concrete.Name as CName
import Syntax.Abstract.Name as AName
import Syntax.Fixity
import Syntax.ScopeInfo

import TypeChecking.Monad.Base
import TypeChecking.Monad.Options

import Utils.Monad
import Utils.Maybe
import Utils.Fresh
import Utils.Map

#include "../undefined.h"

---------------------------------------------------------------------------
-- * Types
---------------------------------------------------------------------------

-- | To simplify interaction between scope checking and type checking (in
--   particular when chasing imports) we use the same monad.
type ScopeM = TCM


---------------------------------------------------------------------------
-- * Instances
---------------------------------------------------------------------------

instance HasRange DefinedName where
    getRange = getRange . theName

-- | Names generated by the scope monad have odd ids. Names generated by the
--   type checking monad should have even ids.
abstractName :: CName.Name -> ScopeM AName.Name
abstractName x =
    do	i <- fresh
	return $ AName.Name i x

---------------------------------------------------------------------------
-- * Exceptions
---------------------------------------------------------------------------

notInScope :: CName.QName -> ScopeM a
notInScope x = typeError $ NotInScope [x]

clashingImport :: CName.Name -> AName.QName -> ScopeM a
clashingImport x x' = typeError $ ClashingImport x x'

clashingModuleImport :: CName.Name -> AName.ModuleName -> ScopeM a
clashingModuleImport x x' = typeError $ ClashingModuleImport x x'

noSuchModule :: CName.QName -> ScopeM a
noSuchModule x = typeError $ NoSuchModule x

uninstantiatedModule :: CName.QName -> ScopeM a
uninstantiatedModule x = typeError $ UninstantiatedModule x

clashingModule :: AName.ModuleName -> AName.ModuleName -> ScopeM a
clashingModule x y = typeError $ ClashingModule x y

clashingDefinition :: CName.Name -> AName.QName -> ScopeM a
clashingDefinition x y = typeError $ ClashingDefinition x y

---------------------------------------------------------------------------
-- * Updating name spaces
---------------------------------------------------------------------------

-- | The names of the name spaces should be the same. Assumes there
--   are no clashes.
plusNameSpace :: NameSpace -> NameSpace -> NameSpace
plusNameSpace (NSpace name ds1 m1) (NSpace _ ds2 m2) =
	NSpace name (Map.unionWith __IMPOSSIBLE__ ds1 ds2)
		    (Map.unionWith __IMPOSSIBLE__ m1 m2)

-- | Same as 'plusNameSpace' but allows the modules to overlap.
plusNameSpace_ :: NameSpace -> NameSpace -> ScopeM NameSpace
plusNameSpace_ (NSpace name ds1 m1) (NSpace _ ds2 m2) =
	NSpace name (Map.unionWith __IMPOSSIBLE__ ds1 ds2)
		<$> unionWithM plusModules m1 m2

-- | Merges two modules.
plusModules :: ModuleScope -> ModuleScope -> ScopeM ModuleScope
plusModules mi1 mi2
    | moduleAccess mi1 /= moduleAccess mi2 || moduleArity mi1 /= moduleArity mi2
	= clashingModule (moduleName $ moduleContents mi1)
			 (moduleName $ moduleContents mi2)    -- TODO: different exception?
    | otherwise	= do
	mc <- moduleContents mi1 `plusNameSpace_` moduleContents mi2
	return mi1 { moduleContents = mc }

topLevelNameSpace :: Modules -> NameSpace
topLevelNameSpace ms =
    (emptyNameSpace $ mkModuleName $ CName.QName CName.noName_)
	{ modules = ms }

-- | Throws an exception if the name exists.
addName :: CName.Name -> DefinedName -> NameSpace -> ScopeM NameSpace
addName x qx ns = do
	defs <- insertWithKeyM clash x qx $ definedNames ns
	return ns { definedNames = defs }
    where
	clash _ _ old = clashingImport x (theName old)

addModule :: CName.Name -> ModuleScope -> NameSpace -> ScopeM NameSpace
addModule x mi ns = do
	mods <- insertWithKeyM clash x mi $ modules ns
	return ns { modules = mods }
    where
	clash _ _ old = clashingModuleImport x (moduleName $ moduleContents old)

-- | Allows the module to already be defined. Used when adding modules
--   corresponding to directories (from imports).
addModule_ :: CName.Name -> ModuleScope -> NameSpace -> ScopeM NameSpace
addModule_ x mi ns = do
	mods <- insertWithKeyM clash x mi $ modules ns
	return ns { modules = mods }
    where
	clash x' mi mi' = plusModules mi mi'
	    -- TODO: we should only allow overlap of the pseudo-modules!

addQModule :: CName.QName -> ModuleScope -> Modules -> ScopeM Modules
addQModule x mi ms = modules <$> addQ [] x mi ns
    where
	ns = topLevelNameSpace ms

	addQ _  (CName.QName x) mi ns = addModule x mi ns
	addQ ms (Qual m x)	mi ns = do
	    mc <- addQ (ms ++ [m]) x mi (emptyNameSpace $ mkQual ms m)
	    let mi' = ModuleScope { moduleAccess   = PublicAccess
				  , moduleArity	   = 0
				  , moduleContents = mc
				  }
	    addModule_ m mi' ns
	    where
		mkQual ms x = mkModuleName $ foldr Qual (CName.QName x) ms

	next (CName.QName x) = x
	next (Qual m x)	     = m

addModules :: [(CName.Name, ModuleScope)] -> NameSpace -> ScopeM NameSpace
addModules ms ns = foldr (=<<) (return ns) $ List.map (uncurry addModule) ms

addNames :: [(CName.Name, DefinedName)] -> NameSpace -> ScopeM NameSpace
addNames ds ns = foldr (=<<) (return ns) $ List.map (uncurry addName) ds

-- | Add the names from the first name space to the second. Throws an
--   exception on clashes.
addNameSpace :: NameSpace -> NameSpace -> ScopeM NameSpace
addNameSpace ns0 ns =
    addNames (Map.assocs $ definedNames ns0)
    =<< addModules (Map.assocs $ modules ns0) ns


-- | Recompute canonical names. All mappings @x -> M1.M'.x@ will be replaced by @x -> M2.M'.x@
--   after @makeFreshCanonicalNames M1 M2@.  Recursively renames sub-modules.
makeFreshCanonicalNames :: ModuleName -> ModuleName -> NameSpace -> NameSpace
makeFreshCanonicalNames m1 m2 ns =
    ns { definedNames = Map.map newName	 $ definedNames ns
       , modules      = Map.map newModule $ modules ns
       , moduleName   = newModuleName $ moduleName ns
       }
    where
	-- TODO: not quite right?
	newModuleName m = fromId $ subst (mnameId m1) (mnameId m2) (mnameId m)
	    where
		subst [] m2 m' = m2 ++ m'
		subst (x:m1) m2 (y:m)
		    | x == y	= subst m1 m2 m
		subst _ _ _ = mnameId m

	fromId mid = MName mid $ mkQName mid
	    where
		mkQName []     = __IMPOSSIBLE__
		mkQName [x]    = CName.QName x
		mkQName (x:xs) = CName.Qual x $ mkQName xs

	newName d = d { theName = (theName d) { qnameModule = newModuleName $ qnameModule $ theName d } }
	newModule mi =
	    mi { moduleContents = makeFreshCanonicalNames m1 m2
				  $ (moduleContents mi)
				    { moduleName = newModuleName $ moduleName $ moduleContents mi }
	       }

---------------------------------------------------------------------------
-- * Updating the scope
---------------------------------------------------------------------------

updateNameSpace :: Access -> (NameSpace -> NameSpace) -> ScopeInfo -> ScopeInfo
updateNameSpace PublicAccess f si =
    si { publicNameSpace = f $ publicNameSpace si }
updateNameSpace PrivateAccess f si =
    si { privateNameSpace = f $ privateNameSpace si }

updateNameSpaceM :: Access -> (NameSpace -> ScopeM NameSpace) ->
		   ScopeInfo -> ScopeM ScopeInfo
updateNameSpaceM PublicAccess f si = do
    ns <- f $ publicNameSpace si
    return si { publicNameSpace = ns }
updateNameSpaceM PrivateAccess f si = do
    ns <- f $ privateNameSpace si
    return si { privateNameSpace = ns }

updateImports :: (Modules -> Modules) -> ScopeInfo -> ScopeInfo
updateImports f si = si { importedModules = f $ importedModules si }

defName :: Access -> KindOfName -> Fixity -> AName.Name ->
	   ScopeInfo -> ScopeM ScopeInfo
defName a k fx x si = do
    reportLn 7 $ "defined name " ++ concat (intersperse "." $ List.map show $ mnameId m) ++ "." ++ show x
    updateNameSpaceM a (addName (nameConcrete x) qx) si
    where
	m   = moduleName $ publicNameSpace si
	qx  = DefinedName a k fx (AName.qualify m x)

-- | Assumes that the name in the 'ModuleScope' fully qualified.
defModule :: CName.Name -> ModuleScope -> ScopeInfo -> ScopeInfo
defModule x mi = updateNameSpace (moduleAccess mi) f
    where
	f ns = ns { modules = Map.insert x mi $ modules ns }

bindVar, shadowVar :: AName.Name -> ScopeInfo -> ScopeInfo
bindVar x si	= si { localVariables = Map.insert (nameConcrete x) x (localVariables si) }
shadowVar x si	= si { localVariables = Map.delete (nameConcrete x)   (localVariables si) }

---------------------------------------------------------------------------
-- * Resolving names
---------------------------------------------------------------------------

setConcreteName :: CName.Name -> AName.Name -> AName.Name
setConcreteName c a = a { nameConcrete = c }

setConcreteQName :: CName.QName -> DefinedName -> DefinedName
setConcreteQName c d = d { theName = (theName d) { qnameConcrete = c } }

-- | Resolve a qualified name. Peals off name spaces until it gets
--   to an unqualified name and then applies the first argument.
resolve :: a -> (LocalVariables -> NameSpace -> CName.Name -> ScopeM a) ->
	   CName.QName -> ScopeInfo -> ScopeM a
resolve def f x si = res x vs (ns `plusNameSpace` ns' `plusNameSpace` nsi)
    where
	vs  = localVariables si
	ns  = publicNameSpace si
	ns' = privateNameSpace si
	nsi = topLevelNameSpace $ importedModules si

	res (CName.QName x) vs ns = f vs ns x
	res (Qual m x) vs ns =
	    case Map.lookup m $ modules ns of
		Nothing			   -> return $ def
		Just (ModuleScope 0 _ ns') -> res x empty ns'
		Just _			   -> uninstantiatedModule (CName.QName m)

-- | Figure out what a qualified name refers to.
resolveName :: CName.QName -> ScopeM ResolvedName
resolveName q = resolve UnknownName r q =<< getScopeInfo
    where
	r vs ns x =
	    return $ fromMaybe UnknownName $ mconcat
	    [ VarName . setConcreteName x  <$> Map.lookup x vs
	    , DefName . setConcreteQName q <$> Map.lookup x (definedNames ns)
	    ]

-- | This function doesn't bind 'VarName's, the caller has that responsibility.
resolvePatternNameM :: CName.QName -> ScopeM ResolvedName
resolvePatternNameM x =
    do	scope <- getScopeInfo
	resolve UnknownName r x scope
    where
	r vs ns x =
	    case Map.lookup x $ definedNames ns of
		Just c@(DefinedName _ ConName _ _) -> return $ DefName c
		Just c@(DefinedName _ FunName _ _) -> return $ DefName c
		_   -> VarName <$> abstractName x

-- | Figure out what module a qualified name refers to.
resolveModule :: CName.QName -> ScopeM ResolvedModule
resolveModule x = resolve UnknownModule r x =<< getScopeInfo
    where
	r _ ns x = return $ fromMaybe UnknownModule $
		    ModuleName <$> Map.lookup x (modules ns)


{--------------------------------------------------------------------------
    Wrappers for the resolve functions
 --------------------------------------------------------------------------}

-- | Make sure that a module hasn't been defined.
noModuleClash :: CName.QName -> ScopeM ()
noModuleClash x = do
    m <- resolveModule x
    case m of
	UnknownModule -> return ()
	ModuleName m  -> clashingModule (mkModuleName x) $ moduleName $ moduleContents m

-- | Get the module referred to by a name. Throws an exception if the module
--   doesn't exist.
getModule :: CName.QName -> ScopeM ModuleScope
getModule x = do
    m <- resolveModule x
    case m of
	UnknownModule -> noSuchModule x
	ModuleName m  -> return m

---------------------------------------------------------------------------
-- * Import directives
---------------------------------------------------------------------------

{-

Where should we check that the import directives are well-formed? I.e. that
    - you only refer to things that exist
    - you don't rename to something that's also imported
	using (x), renaming (y to x)
	renaming (y to x)   -- where x already exists
	hiding (x), renaming (y to x), should be ok though
Idea:
    - start with all names in the module
    - check that mentioned names exist
    - check that there are no internal clashes

import and open preserve canonical names
implicit modules create new canonical names

for implicit modules we can change the canonical names afterwards

-}

-- | Check that all names referred to in the import directive is exported by
--   the module.
invalidImportDirective :: NameSpace -> ImportDirective -> Maybe TypeError
invalidImportDirective ns i =
    case badNames ++ badModules of
	[]  -> Nothing
	xs  -> Just $ ModuleDoesntExport (moduleName ns) xs
    where
	referredNames = names (usingOrHiding i) ++ List.map fst (renaming i)
	badNames    = [ x | x@(ImportedName x') <- referredNames
			  , Nothing <- [Map.lookup x' $ definedNames ns]
		      ]
	badModules  = [ x | x@(ImportedModule x') <- referredNames
			  , Nothing <- [Map.lookup x' $ modules ns]
		      ]

	names (Using xs)    = xs
	names (Hiding xs)   = xs

{-| Figure out how an import directive affects a name.

  - @applyDirective d x == Just y@, if @d@ contains a renaming @x to y@.

  - else @applyDirective d x == Just x@, if @d@ doesn't mention @x@ in a hiding clause and
    if @d@ has a @using@ clause, it mentions @x@.

  - @applyDirective d x == Nothing@, otherwise.

-}
applyDirective :: ImportDirective -> ImportedName -> Maybe CName.Name
applyDirective d x
    | renamed   = just_renamed
    | hidden    = Nothing
    | otherwise = Just $ importedName x
    where
        renamed      = isJust just_renamed
        just_renamed = List.lookup x (renaming d)
        hidden       =
            case usingOrHiding d of
                Hiding xs   ->    elem x xs
                Using  xs   -> notElem x xs

-- | Compute the imported names from a module. When importing canonical names doesn't
--   change. For instance, if the module @A@ contains a function @f@ and we say
--
--   > import A as B, renaming (f to g)
--
--   the canonical form of @B.g@ is @A.f@. Compare this to implicit module definitions
--   which creates new canonical names.
importedNames :: ModuleScope -> ImportDirective -> ScopeM NameSpace
importedNames m i =
    case invalidImportDirective ns i of
	Just e	-> typeError e
	Nothing	-> addModules newModules =<< addNames newNames (emptyNameSpace name)
    where
	name = moduleName ns
	ns = moduleContents m
	newNames = [ (x',qx)
		   | (x,qx)  <- Map.assocs (definedNames ns)
		   , Just x' <- [applyDirective i $ ImportedName x]
		   ]
	newModules = [ (x',m)
		     | (x,m)   <- Map.assocs (modules ns)
		     , Just x' <- [applyDirective i $ ImportedModule x]
		     ]

---------------------------------------------------------------------------
-- * Utility functions
---------------------------------------------------------------------------

-- | Get the current 'ScopeInfo'.
getScopeInfo :: ScopeM ScopeInfo
getScopeInfo = gets stScopeInfo

setScopeInfo :: ScopeInfo -> ScopeM ()
setScopeInfo scope = modify $ \s -> s { stScopeInfo = scope }

modScopeInfoM :: (ScopeInfo -> ScopeM ScopeInfo) -> ScopeM a -> ScopeM a
modScopeInfoM f ret = do
    scope <- getScopeInfo
    setScopeInfo =<< f scope
    x <- ret
    setScopeInfo scope
    return x

modScopeInfo :: (ScopeInfo -> ScopeInfo) -> ScopeM a -> ScopeM a
modScopeInfo f = modScopeInfoM (return . f)

-- | Get the name of the current module.
getCurrentModule :: ScopeM ModuleName
getCurrentModule = moduleName . publicNameSpace <$> getScopeInfo


-- | Get a function that returns the operator version of a name.
getFixity :: CName.QName -> ScopeM Fixity
getFixity x = do
    d <- resolveName x
    case d of
	VarName x -> return defaultFixity
	DefName d -> return $ fixity d
	_	  -> notInScope x

-- | Get the current (public) name space.
currentNameSpace :: ScopeM NameSpace
currentNameSpace = publicNameSpace <$> getScopeInfo

-- | Extract the 'ModuleScope' of the current module.
currentModuleScope :: ScopeInfo -> ModuleScope
currentModuleScope scope = ModuleScope
    { moduleArity    = Map.size $ localVariables scope	-- TODO: Hack (but should work)
    , moduleAccess   = PublicAccess
    , moduleContents = publicNameSpace scope
    }

---------------------------------------------------------------------------
-- * Top-level functions
---------------------------------------------------------------------------

-- | Set the precedence of the current context. It's important to remember
--   to do this everywhere. It would be nice to have something that ensures
--   this.
setContext :: Precedence -> ScopeM a -> ScopeM a
setContext p =
    modScopeInfo $ \s -> s { contextPrecedence = p }

getContextPrecedence :: ScopeM Precedence
getContextPrecedence = contextPrecedence <$> getScopeInfo

-- | Work inside a module. This means moving everything in the
--   'publicNameSpace' to the 'privateNameSpace' and updating the names of
--   the both name spaces.
insideModule :: CName.QName -> ScopeM a -> ScopeM a
insideModule qx = modScopeInfoM upd
    where
	upd si = do
	    reportLn 5 $ "entering module " ++ show (mnameId m)
	    return $ si { publicNameSpace  = emptyNameSpace m
			, privateNameSpace = plusNameSpace pri pub
			}
	    where
		pub = publicNameSpace si
		pri = (privateNameSpace si) { moduleName = m }
		m   = qualifyModule (moduleName pub) (mkModuleName qx)

-- | Add a defined name to the current scope.
defineName :: Access -> KindOfName -> Fixity -> CName.Name -> (AName.Name -> ScopeM a) -> ScopeM a
defineName a k f x cont = do
    d  <- resolveName (CName.QName x)
    case d of
	UnknownName -> do
	    x' <- abstractName x
	    modScopeInfoM (defName a k f x') $ cont x'
	VarName _ -> do
	    x' <- abstractName x
	    modScopeInfoM (defName a k f x' . shadowVar x') $ cont x'
	DefName y   -> clashingDefinition x (theName y)


{-| If a variable shadows a defined name we still keep the defined name.  The
    reason for this is in patterns, where constructors should take precedence
    over variables (and that it would be some work to remove the defined name).
-}
bindVariable :: AName.Name -> ScopeM a -> ScopeM a
bindVariable x = modScopeInfo (bindVar x)

-- | Bind multiple variables.
bindVariables :: [AName.Name] -> ScopeM a -> ScopeM a
bindVariables = foldr (.) id . List.map bindVariable

-- | Defining a module. For explicit modules this should be done after scope
--   checking the module.
defineModule :: CName.Name -> ModuleScope -> ScopeM a -> ScopeM a
defineModule x mi cont =
    do	noModuleClash (CName.QName x)
	modScopeInfo (defModule x mi) cont

-- | Opening a module.
openModule :: CName.QName -> ImportDirective -> ScopeM a -> ScopeM a
openModule x i cont =
    do	m <- getModule x
	ns <- importedNames m i
	let access = if publicOpen i then PublicAccess else PrivateAccess
	modScopeInfoM (updateNameSpaceM access (addNameSpace ns)) cont

-- | Importing a module. The first argument is the name the module is imported
--   /as/. If there is no /as/ clause it should be the name of the module.
importModule :: CName.QName -> ModuleScope -> ImportDirective -> ScopeM a -> ScopeM a
importModule x m dir cont =
    do	noModuleClash x
	imps <- addQModule x m . importedModules =<< getScopeInfo
	modScopeInfo ( updateImports (const imps) ) cont

-- | Implicit module declaration.
implicitModule :: CName.Name -> Access -> Arity -> CName.QName -> ImportDirective ->
		  ScopeM a -> ScopeM a
implicitModule x ac ar x' i cont =
    do	noModuleClash (CName.QName x)
	m    <- getModule x'
	this <- getCurrentModule
	ns'  <- importedNames m i
	let newname = AName.qualifyModule' this x
	    ns	    = makeFreshCanonicalNames (moduleName ns') newname ns'
	    m' = ModuleScope { moduleAccess   = ac
			     , moduleArity    = ar
			     , moduleContents = ns
			     }
	modScopeInfoM (updateNameSpaceM ac (addModule x m')) cont

---------------------------------------------------------------------------
-- * Debugging
---------------------------------------------------------------------------

scopeTree :: NameSpace -> Tree String
scopeTree ns =
    Node (show $ moduleName ns)
	 $  (List.map leaf $ Map.assocs $ definedNames ns)
	 ++ (List.map (scopeTree . moduleContents) $ Map.elems $ modules ns)
    where
	leaf (x,d) = Node (unwords [show x,"-->",show $ theName d]) []

instance Show ScopeInfo where
    show si =
	unlines [ "Public name space:"
		, drawTree $ scopeTree $ publicNameSpace si
		, "Private name space:"
		, drawTree $ scopeTree $ privateNameSpace si
		, "Local variables: " ++ show (localVariables si)
		, "Precedence: " ++ show (contextPrecedence si)
		]