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

import SSTG.Core.Language.Support
import SSTG.Core.Language.Syntax

import qualified Data.List as L
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)
    acc_ns = stack_ns ++ heap_ns ++ glbls_ns ++ expr_ns ++ pcons_ns

-- | `Name`s in a `Stack`.
stackNames :: Stack -> [Name]
stackNames stack = concatMap frameNames (stackToList stack)

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

-- | `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 = map fst (localsToList locals)

-- | `Name`s in the `Heap`.
heapNames :: Heap -> [Name]
heapNames heap = concatMap (heapObjNames . snd) (heapToList heap)

-- | `Name`s in a `HeapObj`.
heapObjNames :: HeapObj -> [Name]
heapObjNames (AddrObj _) = []
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_ns
  where
    scls_ns = 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) = concatMap atomNames args ++ dataNames dcon

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

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

-- | `Name`s in `Globals`.
globalsNames :: Globals -> [Name]
globalsNames globals = map fst (globalsToList globals)

-- | `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 (Let bnd expr) = exprNames expr ++ bindNames bnd
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 (Case expr var alts) = exprNames expr ++ concatMap altNames alts
                                                ++ varNames var
-- | `Name`s in a `Type`.
typeNames :: Type -> [Name]
typeNames (TyVarTy n ty) = n : typeNames ty
typeNames (CoercionTy coer) = coercionNames coer
typeNames (AppTy t1 t2) = typeNames t1 ++ typeNames t2
typeNames (CastTy ty coer) = typeNames ty ++ coercionNames coer
typeNames (ForAllTy bnd ty) = typeNames ty ++ tyBinderNames bnd
typeNames (FunTy t1 t2) = typeNames t1 ++ typeNames t2
typeNames (TyConApp tc ty) = tyConNames tc ++ concatMap typeNames ty
typeNames (LitTy _) = []
typeNames (Bottom) = []

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

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

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

-- | `Name`s in a `TyCon`.
tyConNames :: TyCon -> [Name]
tyConNames (FamilyTyCon n ns) = n : ns
tyConNames (SynonymTyCon n ns) = n : ns
tyConNames (AlgTyCon n ns r) = n : ns ++ algTyRhsNames r
tyConNames (FunTyCon n bs) = n : concatMap tyBinderNames bs
tyConNames (PrimTyCon n bs) = n : concatMap tyBinderNames bs
tyConNames (Promoted n bs dc) = n : concatMap tyBinderNames bs ++ dataNames dc

-- | `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 ns) = ns
algTyRhsNames (TupleTyCon n) = [n]
algTyRhsNames (NewTyCon n) = [n]

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

-- | `Name`s in a `PathCons`.
pconsNames :: PathCons -> [Name]
pconsNames pathcons = concatMap constraintNames (pathconsToList pathcons)

-- | `Name`s in a `PathCons`.
constraintNames :: Constraint -> [Name]
constraintNames (Constraint (_, vs) e locs _) = exprNames e ++ localsNames locs
                                                            ++ map varName vs

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