-- | 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 Names in 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

-- | Stack Names
stackNames :: Stack -> [Name]
stackNames (Stack [])     = []
stackNames (Stack (f:fs)) = frameNames f ++ stackNames (Stack fs)

-- | Frame Names
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

-- | Alt Names
altNames :: Alt -> [Name]
altNames (Alt _ vars expr) = (concatMap varNames vars) ++ exprNames expr

-- | Locals Names
localsNames :: Locals -> [Name]
localsNames (Locals lmap) = M.keys lmap

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

-- | Heap Object Names
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

-- | Symbol Names
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

-- | Lambda Form Names
bindRhsNames :: BindRhs -> [Name]
bindRhsNames (FunForm prms expr) = (concatMap varNames prms) ++ exprNames expr
bindRhsNames (ConForm dcon args) = dataNames dcon ++ concatMap atomNames args

-- | Var Names
varNames :: Var -> [Name]
varNames (Var n t) = n : typeNames t

-- | Atom Names
atomNames :: Atom -> [Name]
atomNames (VarAtom var) = varNames var
atomNames (LitAtom _)   = []

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

-- | Eval State Names
codeNames :: Code -> [Name]
codeNames (Return _)             = []
codeNames (Evaluate expr locals) = exprNames expr ++ localsNames locals

-- | Expression Names
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
-- | Type Names
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)          = []

-- | Prim Fun Names
pfunNames :: PrimFun -> [Name]
pfunNames (PrimFun n ty) = n : typeNames ty

-- | Data Constructor ID Names
conTagName :: ConTag -> Name
conTagName (ConTag n _) = n

-- | Data Constructor Names
dataNames :: DataCon -> [Name]
dataNames (DataCon tg ty tys) = conTagName tg : concatMap typeNames (ty : tys)

-- | Type Binder Names
tyBinderNames :: TyBinder -> [Name]
tyBinderNames (NamedTyBndr n ty) = n : typeNames ty
tyBinderNames (AnonTyBndr ty)    = typeNames ty

-- | Type Constructor Names
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

-- | Coercion Names
coercionNames :: Coercion -> [Name]
coercionNames (Coercion t1 t2) = typeNames t1 ++ typeNames t2

-- | Type Alg Rhs Names
algTyRhsNames :: AlgTyRhs -> [Name]
algTyRhsNames (AbstractTyCon _) = []
algTyRhsNames (DataTyCon tags)  = map conTagName tags
algTyRhsNames (TupleTyCon tag)  = [conTagName tag]
algTyRhsNames (NewTyCon tag)    = [conTagName tag]

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

-- | Path Constraint Names
pconsNames :: PathCons -> [Name]
pconsNames []     = []
pconsNames (c:cs) = pcondNames c ++ pconsNames cs

-- | Path Condition Names
pcondNames :: PathCond -> [Name]
pcondNames (PathCond (_, vars) expr locals _) = map varName vars ++
                                                exprNames expr   ++
                                                localsNames locals

-- | Symbolic Link Names
linksNames :: SymLinks -> [Name]
linksNames (SymLinks links) = concatMap (\(a, b) -> [a, b]) kvs
  where kvs = M.toList links

-- | Fresh String from Int Rand Seed
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 from Conflict List
freshName :: NameSpace -> [Name] -> Name
freshName nspace confs = freshSeededName seed confs
  where seed = Name "fs?" Nothing nspace 0

-- | Seeded Fresh Name from Conflict List
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)

-- | List of Fresh Names
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 Names
freshSeededNameList :: [Name] -> [Name] -> [Name]
freshSeededNameList []     _     = []
freshSeededNameList (n:ns) confs = name' : freshSeededNameList ns confs'
  where name'  = freshSeededName n confs
        confs' = name' : confs