-- | Naming Module
module SSTG.Core.Language.Naming
    ( allNames
    , varName
    , nameOccStr
    , nameInt
    , freshStr
    , freshName
    , freshSeededName
    , freshNames
    , freshSeededNames
    ) where

import SSTG.Core.Language.Syntax

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

-- | Nameable typeclass.
class Nameable a where
    allNames :: a -> [Name]

-- | `Program` instance of `Nameable`.
instance Nameable Program where
    allNames (Program bindss) = concatMap allNames bindss

-- | `Binds` instance of `Nameable`
instance Nameable Binds where
    allNames (Binds _ kvs) = lhs ++ rhs
      where
        lhs = concatMap (allNames . fst) kvs
        rhs = concatMap (allNames . snd) kvs

-- | `Var` instance of `Nameable`
instance Nameable Var where
    allNames (Var name ty) = name : allNames ty

-- | `BindRhs` instance of `Nameable`
instance Nameable BindRhs where
    allNames (FunForm prms expr) = concatMap allNames prms ++ allNames expr
    allNames (ConForm dcon as) = allNames dcon ++ concatMap allNames as

-- | `Expr` instance of `Nameable`
instance Nameable Expr where
    allNames (Atom atom) = allNames atom
    allNames (Let binds expr) = allNames expr ++ allNames binds
    allNames (FunApp fun args) = allNames fun ++ concatMap allNames args
    allNames (PrimApp pfun args) = allNames pfun ++ concatMap allNames args
    allNames (ConApp dcon args) = allNames dcon ++ concatMap allNames args
    allNames (Case expr var alts) = allNames expr ++ concatMap allNames alts
                                                  ++ allNames var

-- | `Atom` instance of `Nameable`
instance Nameable Atom where
    allNames (LitAtom _) = []
    allNames (VarAtom var) = allNames var

-- | `PrimFun` instance of `Nameable`
instance Nameable PrimFun where
    allNames (PrimFun name ty) = name : allNames ty

-- | `DataCon` instance of `Nameable`
instance Nameable DataCon where
    allNames (DataCon name ty tys) = name : concatMap allNames (ty : tys)

-- | `Alt` instance of `Nameable`
instance Nameable Alt where
    allNames (Alt acon expr) = allNames acon ++ allNames expr

-- | `AltCon` instance of `Nameable`
instance Nameable AltCon where
    allNames (DataAlt dcon ps) = allNames dcon ++ concatMap allNames ps
    allNames _ = []

-- | `Type` instance of `Nameable`
instance Nameable Type where
    allNames (TyVarTy var) = allNames var
    allNames (AppTy ty1 ty2) = allNames ty1 ++ allNames ty2
    allNames (ForAllTy bndr ty) = allNames ty ++ allNames bndr
    allNames (FunTy ty1 ty2) = allNames ty1 ++ allNames ty2
    allNames (TyConApp tycon ty) = allNames tycon ++ concatMap allNames ty
    allNames (CoercionTy coer) = allNames coer
    allNames (CastTy ty coer) = allNames ty ++ allNames coer
    allNames (LitTy _) = []
    allNames (Bottom) = []

-- | `TyBinder` instance of `Nameable`
instance Nameable TyBinder where
    allNames (AnonTyBndr) = []
    allNames (NamedTyBndr name) = [name]

-- | `TyCon` instance of `Nameable`
instance Nameable TyCon where
    allNames (FamilyTyCon name params) = name : params
    allNames (SynonymTyCon name params) = name : params
    allNames (AlgTyCon name params rhs) = name : params ++ allNames rhs
    allNames (FunTyCon name bndrs) = name : concatMap allNames bndrs
    allNames (PrimTyCon name bndrs) = name : concatMap allNames bndrs
    allNames (Promoted name bndrs dcon) = name : concatMap allNames bndrs
                                              ++ allNames dcon

-- | `Coercion` instance of `Nameable`
instance Nameable Coercion where
    allNames (Coercion ty1 ty2) = allNames ty1 ++ allNames ty2

-- | `AlgTyRhs` instance of `Nameable`
instance Nameable AlgTyRhs where
    allNames (AbstractTyCon _) = []
    allNames (DataTyCon names) = names
    allNames (TupleTyCon name) = [name]
    allNames (NewTyCon name) = [name]

-- | A `Var`'s `Name`. Not to be confused with the other function.
varName :: Var -> Name
varName (Var name _) = name

-- | A `Name`'s occurrence string.
nameOccStr :: Name -> String
nameOccStr (Name occ _ _ _) = occ

-- | A `Name`'s unique int.
nameInt :: Name -> Int
nameInt (Name _ _ _ int) = int

-- | 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.
freshStr :: Int -> String -> S.Set String -> String
freshStr rand seed confs = if S.member seed confs
    then freshStr (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' = freshStr 1 occ (S.fromList alls)
    unq' = maxs + 1
    alls = map nameOccStr confs
    maxs = L.maximum (unq : map nameInt confs)

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

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