-- | Naming Module
module SSTG.Core.Execution.Naming
    ( allNames
    , freshString
    , freshName
    , freshSeededName
    , freshNameList
    , freshSeededNameList
    ) where

import SSTG.Core.Syntax
import SSTG.Core.Execution.Models

import qualified Data.List as L
import qualified Data.Map  as M
import qualified Data.Set  as S

-- | All `Name`s in a `State`.
allNames :: State -> [Name]
allNames state = L.nub acc_ns
  where stack_ns = stackNames   (state_stack   state)
        heap_ns  = heapNames    (state_heap    state)
        glbls_ns = globalsNames (state_globals state)
        expr_ns  = codeNames    (state_code    state)
        pcons_ns = pconsNames   (state_paths   state)
        links_ns = linksNames   (state_links   state)
        acc_ns   = stack_ns ++ heap_ns  ++ glbls_ns ++
                   expr_ns  ++ pcons_ns ++ links_ns

-- | `Name`s in a `Stack`.
stackNames :: Stack -> [Name]
stackNames (Stack [])     = []
stackNames (Stack (f:fs)) = frameNames f ++ stackNames (Stack fs)

-- | `Name`s in a `Frame`.
frameNames :: Frame -> [Name]
frameNames (UpdateFrame _)          = []
frameNames (ApplyFrame as lcs)      = concatMap atomNames as ++ localsNames lcs
frameNames (CaseFrame var alts lcs) = varNames var ++ (concatMap altNames alts)
                                                   ++ localsNames lcs

-- | `Name`s in an `Alt`.
altNames :: Alt -> [Name]
altNames (Alt _ vars expr) = (concatMap varNames vars) ++ exprNames expr

-- | `Name`s in the `Locals`
localsNames :: Locals -> [Name]
localsNames (Locals lmap) = M.keys lmap

-- | `Name`s in the `Heap`.
heapNames :: Heap -> [Name]
heapNames (Heap heap _) = concatMap (heapObjNames . snd) kvs
  where kvs = M.toList heap

-- | `Name`s in a `HeapObj`.
heapObjNames :: HeapObj -> [Name]
heapObjNames (Blackhole)           = []
heapObjNames (LitObj _)            = []
heapObjNames (SymObj sym)          = symbolNames sym
heapObjNames (ConObj dcon _)       = dataNames dcon
heapObjNames (FunObj ps expr locs) = exprNames expr ++ localsNames locs
                                                    ++ concatMap varNames ps

-- | `Name`s in a `Symbol`.
symbolNames :: Symbol -> [Name]
symbolNames (Symbol sym mb_scls) = varNames sym ++ scls_names
  where scls_names = case mb_scls of
                         Nothing     -> []
                         Just (e, l) -> exprNames e ++ localsNames l

-- | `Name`s in a `BindRhs`.
bindRhsNames :: BindRhs -> [Name]
bindRhsNames (FunForm prms expr) = (concatMap varNames prms) ++ exprNames expr
bindRhsNames (ConForm dcon args) = dataNames dcon ++ concatMap atomNames args

-- | `Name`s in a `Var`.
varNames :: Var -> [Name]
varNames (Var n t) = n : typeNames t

-- | `Name`s in an `Atom`.
atomNames :: Atom -> [Name]
atomNames (VarAtom var) = varNames var
atomNames (LitAtom _)   = []

-- | `Name`s in `Globals`.
globalsNames :: Globals -> [Name]
globalsNames (Globals gmap) = M.keys gmap

-- | `Name`s in the current evaluation `Code`.
codeNames :: Code -> [Name]
codeNames (Return _)             = []
codeNames (Evaluate expr locals) = exprNames expr ++ localsNames locals

-- | `Name`s in an `Expr`.
exprNames :: Expr -> [Name]
exprNames (Atom atom)          = atomNames atom
exprNames (FunApp fun args)    = varNames  fun  ++ concatMap atomNames args
exprNames (PrimApp prim args)  = pfunNames prim ++ concatMap atomNames args
exprNames (ConApp dcon args)   = dataNames dcon ++ concatMap atomNames args
exprNames (Let binds expr)     = bindingNames binds ++ exprNames expr
exprNames (Case expr var alts) = varNames var ++ exprNames expr
                                              ++ concatMap altNames alts
-- | `Name`s in a `Type`.
typeNames :: Type -> [Name]
typeNames (TyVarTy n ty)    = n : typeNames ty
typeNames (AppTy t1 t2)     = typeNames t1  ++ typeNames t2
typeNames (ForAllTy bnd ty) = tyBinderNames bnd ++ typeNames ty
typeNames (CastTy ty coer)  = typeNames ty  ++ coercionNames coer
typeNames (TyConApp tc ty)  = tyConNames tc ++ concatMap typeNames ty
typeNames (CoercionTy coer) = coercionNames coer
typeNames (LitTy _)         = []
typeNames (FunTy t1 t2)     = typeNames t1 ++ typeNames t2
typeNames (Bottom)          = []

-- | `Name`s in a `PrimFun`.
pfunNames :: PrimFun -> [Name]
pfunNames (PrimFun n ty) = n : typeNames ty

-- | `Name`s in a `ConTag`.
conTagName :: ConTag -> Name
conTagName (ConTag n _) = n

-- | `Name`s in a `DataCon`.
dataNames :: DataCon -> [Name]
dataNames (DataCon tg ty tys) = conTagName tg : concatMap typeNames (ty : tys)

-- | `Name`s in a `TyBinder`.
tyBinderNames :: TyBinder -> [Name]
tyBinderNames (NamedTyBndr n ty) = n : typeNames ty
tyBinderNames (AnonTyBndr ty)    = typeNames ty

-- | `Name`s in a `TyCon`.
tyConNames :: TyCon -> [Name]
tyConNames (FunTyCon n)      = [n]
tyConNames (AlgTyCon n r)    = n : algTyRhsNames r
tyConNames (SynonymTyCon n)  = [n]
tyConNames (FamilyTyCon n)   = [n]
tyConNames (PrimTyCon n)     = [n]
tyConNames (TcTyCon n)       = [n]
tyConNames (Promoted n dcon) = n : dataNames dcon

-- | `Name`s in a `Coercion`.
coercionNames :: Coercion -> [Name]
coercionNames (Coercion t1 t2) = typeNames t1 ++ typeNames t2

-- | `Name`s in a `AlgTyRhs`.
algTyRhsNames :: AlgTyRhs -> [Name]
algTyRhsNames (AbstractTyCon _) = []
algTyRhsNames (DataTyCon tags)  = map conTagName tags
algTyRhsNames (TupleTyCon tag)  = [conTagName tag]
algTyRhsNames (NewTyCon tag)    = [conTagName tag]

-- | `Name`s in a `Binding`.
bindingNames :: Binding -> [Name]
bindingNames (Binding _ bnd) = lhs ++ rhs
  where lhs = concatMap (varNames . fst) bnd
        rhs = concatMap (bindRhsNames . snd) bnd

-- | `Name`s in a `PathCons`.
pconsNames :: PathCons -> [Name]
pconsNames []     = []
pconsNames (c:cs) = pcondNames c ++ pconsNames cs

-- | `Name`s in a `PathCond`.
pcondNames :: PathCond -> [Name]
pcondNames (PathCond (_, vars) expr locals _) = map varName vars ++
                                                exprNames expr   ++
                                                localsNames locals

-- | `Name`s in a `SymLinks`.
linksNames :: SymLinks -> [Name]
linksNames (SymLinks links) = concatMap (\(a, b) -> [a, b]) kvs
  where kvs = M.toList links

-- | Create a fresh seed given any `Int`, a `String` seed, and a `Set` of
-- `String`s that we do not want our new `String` to conflict with. The sole
-- purpose of the `Int` seed is to allow us tell us how much to multiply some
-- prime number to "orbit" an index around a fixed list of acceptable `Char`s.
freshString :: Int -> String -> S.Set String -> String
freshString rand seed confs = if S.member seed confs
    then freshString (rand + 1) (seed ++ [pick]) confs
    else seed
  where pick  = bank !! index
        index = raw_i `mod` (length bank)
        raw_i = (abs rand) * prime
        prime = 151  -- The original? :)
        bank  = lower ++ upper ++ nums
        lower = "abcdefghijlkmnopqrstuvwxyz"
        upper = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
        nums  = "1234567890"

-- | Fresh `Name` given a list of `Name`s that acts as conflicts. The fresh
-- `Name`s generated in this manner are prefixed with @"fs?"@, which is not a
-- valid identifier in Haskell, but okay in SSTG. we also specify the
-- `NameSpace` under which the `Name` will be generated. This will generally
-- be `VarNSpace` in actual usage.
freshName :: NameSpace -> [Name] -> Name
freshName nspace confs = freshSeededName seed confs
  where seed = Name "fs?" Nothing nspace 0

-- | A fresh `Name` generated from a seed `Name`, which will act as the prefix
-- of the new `Name`. We ues the same `NameSpace` as the seed `Name` when
-- generating this way.
freshSeededName :: Name -> [Name] -> Name
freshSeededName seed confs = Name occ' mdl ns unq'
  where Name occ mdl ns unq = seed
        occ' = freshString 1 occ (S.fromList alls)
        unq' = maxs + 1
        alls = map nameOccStr confs
        maxs = L.maximum (unq : map nameUnique confs)

-- | Generate a list of `Name`s, each corresponding to the appropriate element
-- of the `NameSpace` list.
freshNameList :: [NameSpace] -> [Name] -> [Name]
freshNameList []           _     = []
freshNameList (nspace:nss) confs = name' : freshNameList nss confs'
  where name'  = freshName nspace confs
        confs' = name' : confs

-- | List of seeded fresh `Name`s.
freshSeededNameList :: [Name] -> [Name] -> [Name]
freshSeededNameList []     _     = []
freshSeededNameList (n:ns) confs = name' : freshSeededNameList ns confs'
  where name'  = freshSeededName n confs
        confs' = name' : confs