module Wingman.Naming where

import           Control.Monad.State.Strict
import           Data.Bool (bool)
import           Data.Char
import           Data.Map (Map)
import qualified Data.Map as M
import           Data.Set (Set)
import qualified Data.Set as S
import           Data.Traversable
import           Name
import           TcType
import           TyCon
import           Type
import           TysWiredIn (listTyCon, pairTyCon, unitTyCon)


------------------------------------------------------------------------------
-- | Use type information to create a reasonable name.
mkTyName :: Type -> String
-- eg. mkTyName (a -> B) = "fab"
mkTyName :: Type -> String
mkTyName (Type -> ([Type], Type)
tcSplitFunTys -> ([a :: Type
a@(Type -> Bool
isFunTy -> Bool
False)], Type
b))
  = String
"f" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Type -> String
mkTyName Type
a String -> String -> String
forall a. [a] -> [a] -> [a]
++ Type -> String
mkTyName Type
b
-- eg. mkTyName (a -> b -> C) = "f_C"
mkTyName (Type -> ([Type], Type)
tcSplitFunTys -> (Type
_:[Type]
_, Type
b))
  = String
"f_" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Type -> String
mkTyName Type
b
-- eg. mkTyName (Either A B) = "eab"
mkTyName (HasDebugCallStack => Type -> Maybe (TyCon, [Type])
Type -> Maybe (TyCon, [Type])
splitTyConApp_maybe -> Just (TyCon
c, [Type]
args))
  = TyCon -> String
mkTyConName TyCon
c String -> String -> String
forall a. [a] -> [a] -> [a]
++ (Type -> String) -> [Type] -> String
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Type -> String
mkTyName [Type]
args
-- eg. mkTyName (f a) = "fa"
mkTyName (Type -> (Type, [Type])
tcSplitAppTys -> (Type
t, args :: [Type]
args@(Type
_:[Type]
_)))
  = Type -> String
mkTyName Type
t String -> String -> String
forall a. [a] -> [a] -> [a]
++ (Type -> String) -> [Type] -> String
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Type -> String
mkTyName [Type]
args
-- eg. mkTyName a = "a"
mkTyName (Type -> Maybe TyVar
getTyVar_maybe -> Just TyVar
tv)
  = OccName -> String
occNameString (OccName -> String) -> OccName -> String
forall a b. (a -> b) -> a -> b
$ TyVar -> OccName
forall name. HasOccName name => name -> OccName
occName TyVar
tv
-- eg. mkTyName (forall x. y) = "y"
mkTyName (Type -> ([TyVar], [Type], Type)
tcSplitSigmaTy -> (TyVar
_:[TyVar]
_, [Type]
_, Type
t))
  = Type -> String
mkTyName Type
t
mkTyName Type
_ = String
"x"


------------------------------------------------------------------------------
-- | Get a good name for a type constructor.
mkTyConName :: TyCon -> String
mkTyConName :: TyCon -> String
mkTyConName TyCon
tc
  | TyCon
tc TyCon -> TyCon -> Bool
forall a. Eq a => a -> a -> Bool
== TyCon
listTyCon = String
"l_"
  | TyCon
tc TyCon -> TyCon -> Bool
forall a. Eq a => a -> a -> Bool
== TyCon
pairTyCon = String
"p_"
  | TyCon
tc TyCon -> TyCon -> Bool
forall a. Eq a => a -> a -> Bool
== TyCon
unitTyCon = String
"unit"
  | Bool
otherwise
      = Int -> String -> String
forall a. Int -> [a] -> [a]
take Int
1
      (String -> String) -> (OccName -> String) -> OccName -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Char -> Char) -> String -> String
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Char -> Char
toLower
      (String -> String) -> (OccName -> String) -> OccName -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Char -> Bool) -> Char -> String -> String
forall a. (a -> Bool) -> a -> [a] -> [a]
filterReplace Char -> Bool
isSymbol      Char
's'
      (String -> String) -> (OccName -> String) -> OccName -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Char -> Bool) -> Char -> String -> String
forall a. (a -> Bool) -> a -> [a] -> [a]
filterReplace Char -> Bool
isPunctuation Char
'p'
      (String -> String) -> (OccName -> String) -> OccName -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. OccName -> String
occNameString
      (OccName -> String) -> OccName -> String
forall a b. (a -> b) -> a -> b
$ TyCon -> OccName
forall a. NamedThing a => a -> OccName
getOccName TyCon
tc


------------------------------------------------------------------------------
-- | Maybe replace an element in the list if the predicate matches
filterReplace :: (a -> Bool) -> a -> [a] -> [a]
filterReplace :: (a -> Bool) -> a -> [a] -> [a]
filterReplace a -> Bool
f a
r = (a -> a) -> [a] -> [a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\a
a -> a -> a -> Bool -> a
forall a. a -> a -> Bool -> a
bool a
a a
r (Bool -> a) -> Bool -> a
forall a b. (a -> b) -> a -> b
$ a -> Bool
f a
a)


------------------------------------------------------------------------------
-- | Produce a unique, good name for a type.
mkGoodName
    :: Set OccName  -- ^ Bindings in scope; used to ensure we don't shadow anything
    -> Type       -- ^ The type to produce a name for
    -> OccName
mkGoodName :: Set OccName -> Type -> OccName
mkGoodName Set OccName
in_scope Type
t =
  let tn :: String
tn = Type -> String
mkTyName Type
t
   in String -> OccName
mkVarOcc (String -> OccName) -> String -> OccName
forall a b. (a -> b) -> a -> b
$ case OccName -> Set OccName -> Bool
forall a. Ord a => a -> Set a -> Bool
S.member (String -> OccName
mkVarOcc String
tn) Set OccName
in_scope of
        Bool
True  -> String
tn String -> String -> String
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show (Set OccName -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Set OccName
in_scope)
        Bool
False -> String
tn


------------------------------------------------------------------------------
-- | Like 'mkGoodName' but creates several apart names.
mkManyGoodNames
  :: (Traversable t, Monad m)
  => Set OccName
  -> t Type
  -> m (t OccName)
mkManyGoodNames :: Set OccName -> t Type -> m (t OccName)
mkManyGoodNames Set OccName
in_scope t Type
args =
  (StateT (Set OccName) m (t OccName)
 -> Set OccName -> m (t OccName))
-> Set OccName
-> StateT (Set OccName) m (t OccName)
-> m (t OccName)
forall a b c. (a -> b -> c) -> b -> a -> c
flip StateT (Set OccName) m (t OccName) -> Set OccName -> m (t OccName)
forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m a
evalStateT Set OccName
in_scope (StateT (Set OccName) m (t OccName) -> m (t OccName))
-> StateT (Set OccName) m (t OccName) -> m (t OccName)
forall a b. (a -> b) -> a -> b
$ t Type
-> (Type -> StateT (Set OccName) m OccName)
-> StateT (Set OccName) m (t OccName)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
t a -> (a -> f b) -> f (t b)
for t Type
args ((Type -> StateT (Set OccName) m OccName)
 -> StateT (Set OccName) m (t OccName))
-> (Type -> StateT (Set OccName) m OccName)
-> StateT (Set OccName) m (t OccName)
forall a b. (a -> b) -> a -> b
$ \Type
at -> do
    Set OccName
in_scope <- StateT (Set OccName) m (Set OccName)
forall s (m :: * -> *). MonadState s m => m s
get
    let n :: OccName
n = Set OccName -> Type -> OccName
mkGoodName Set OccName
in_scope Type
at
    (Set OccName -> Set OccName) -> StateT (Set OccName) m ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((Set OccName -> Set OccName) -> StateT (Set OccName) m ())
-> (Set OccName -> Set OccName) -> StateT (Set OccName) m ()
forall a b. (a -> b) -> a -> b
$ OccName -> Set OccName -> Set OccName
forall a. Ord a => a -> Set a -> Set a
S.insert OccName
n
    OccName -> StateT (Set OccName) m OccName
forall (f :: * -> *) a. Applicative f => a -> f a
pure OccName
n


------------------------------------------------------------------------------
-- | Which names are in scope?
getInScope :: Map OccName a -> [OccName]
getInScope :: Map OccName a -> [OccName]
getInScope = Map OccName a -> [OccName]
forall k a. Map k a -> [k]
M.keys