{-------------------------------------------------------------------------------

        Copyright:              Bernie Pope 2007

        Module:                 Rename

        Description:            Renames all lambda-bound identifiers apart in
                                expressions and programs. For example:

                                rename (\x -> plus 1 x) might produce
                                (\x_12 -> plus 1 x_12), assuming plus is not
                                lambda bound.

                                This is needed to avoid inadvertant name
                                capture during program evaluation. For example:

                                (\y -> \id -> (y, id)) id True

                                should reduce to:

                                (id, True)

                                and not:

                                (True, True)

        Primary Authors:        Bernie Pope

        Notes:

        We don't rename "top-level" identifiers (those defined in a declaration)
        because they are assumed to be unique. Baskell does not have local
        declarations so this does not cause any problems. It is a syntax error
        to have two top-level declarations with the same name, however Baskell
        might not check for this.

-------------------------------------------------------------------------------}

{-
    This file is part of baskell.

    baskell is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation; either version 2 of the License, or
    (at your option) any later version.

    baskell is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with baskell; if not, write to the Free Software
    Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
-}

module Rename
   ( renameExp
   , renameProgram
   )
   where

import AST
   ( Ident
   , Exp (..)
   , Decl (..)
   , Program (..)
   )

import qualified Data.Map as Map
   ( Map
   , empty
   , fromList
   , union
   , insert
   , lookup
   )

import Control.Monad
   ( liftM
   , liftM2
   )

import Control.Monad.State
   ( runStateT
   , evalStateT
   , get
   , put
   , StateT
   , gets
   , modify
   , execStateT
   )

import Control.Monad.Trans
   ( lift
   , liftIO
   )

import Control.Monad.Identity
   ( Identity
   , runIdentity
   )

import Control.Monad.Reader
   ( ReaderT
   , local
   , ask
   , runReaderT
   )

--------------------------------------------------------------------------------

type NameEnv = Map.Map Ident Ident
type RnState = Int
type Rn a = ReaderT NameEnv (StateT RnState Identity) a

renameProgram :: Program -> RnState -> (Program, RnState)
renameProgram prog state =
   runIdentity $ runStateT (runReaderT (renameProgramM prog) Map.empty) state

renameExp :: Exp -> RnState -> Exp
renameExp exp state =
   runIdentity $ evalStateT (runReaderT (renameExpM exp) Map.empty) state

extendEnv :: Ident -> Ident -> Rn a -> Rn a
extendEnv identOld identNew action =
   local (Map.insert identOld identNew) action

lookupIdent :: Ident -> Rn String
lookupIdent oldName = do
   env <- ask
   case Map.lookup oldName env of
      Nothing      -> return oldName
      Just newName -> return newName

-- renamed identifiers are appended with an underscore and a number
-- this is safe because underscores are not allowed by the parser
-- so there is no chance we will introduce an inadvertant name
-- clash during renaming
renameIdent :: Ident -> Rn Ident
renameIdent oldName = do
   count <- lift get
   lift $ put (count + 1)
   return $ oldName ++ "_" ++ show count

renameDecl :: Decl -> Rn Decl
renameDecl decl@(Sig {}) = return decl
renameDecl (Decl name body) =
   liftM (Decl name) $ renameExpM body

renameProgramM :: Program -> Rn Program
renameProgramM(Program decls) =
   liftM Program $ mapM renameDecl decls

renameExpM :: Exp -> Rn Exp
renameExpM (Var ident) =
   liftM Var $ lookupIdent ident

renameExpM (Lam ident body) = do
   newIdent <- renameIdent ident
   liftM (Lam newIdent) $ extendEnv ident newIdent $ renameExpM body

renameExpM (LamStrict ident body) = do
   newIdent <- renameIdent ident
   liftM (LamStrict newIdent) $ extendEnv ident newIdent $ renameExpM body

renameExpM (App e1 e2) =
   liftM2 App (renameExpM e1) (renameExpM e2)

renameExpM exp@(Literal lit) = return exp

renameExpM (Tuple exps) =
   liftM Tuple $ mapM renameExpM exps

renameExpM exp@(Prim _name _impl) =
   return exp