{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE RecordWildCards #-}
module Haddock.Interface.Specialize
( specializeInstHead
) where
import Haddock.Syb
import Haddock.Types
import GHC
import Name
import FastString
import Control.Monad
import Control.Monad.Trans.Reader
import Control.Monad.Trans.State
import Data.Data
import qualified Data.List as List
import Data.Maybe
import Data.Map (Map)
import qualified Data.Map as Map
import Data.Set (Set)
import qualified Data.Set as Set
specialize :: forall name a. (Ord name, DataId name, NamedThing name)
=> Data a
=> [(name, HsType name)] -> a -> a
specialize specs = go
where
go :: forall x. Data x => x -> x
go = everywhereButType @name $ mkT $ sugar . specialize_ty_var
specialize_ty_var (HsTyVar _ (L _ name'))
| Just t <- Map.lookup name' spec_map = t
specialize_ty_var typ = typ
spec_map = Map.fromList [ (n, go t) | (n, t) <- specs]
specializeTyVarBndrs :: (Ord name, DataId name, NamedThing name)
=> Data a
=> LHsQTyVars name -> [HsType name]
-> a -> a
specializeTyVarBndrs bndrs typs =
specialize $ zip bndrs' typs
where
bndrs' = map (bname . unLoc) . hsq_explicit $ bndrs
bname (UserTyVar (L _ name)) = name
bname (KindedTyVar (L _ name) _) = name
specializePseudoFamilyDecl :: (Ord name, DataId name, NamedThing name)
=> LHsQTyVars name -> [HsType name]
-> PseudoFamilyDecl name
-> PseudoFamilyDecl name
specializePseudoFamilyDecl bndrs typs decl =
decl {pfdTyVars = map (specializeTyVarBndrs bndrs typs) (pfdTyVars decl)}
specializeSig :: forall name . (Ord name, DataId name, SetName name, NamedThing name)
=> LHsQTyVars name -> [HsType name]
-> Sig name
-> Sig name
specializeSig bndrs typs (TypeSig lnames typ) =
TypeSig lnames (typ {hswc_body = (hswc_body typ) {hsib_body = noLoc typ'}})
where
true_type :: HsType name
true_type = unLoc (hsSigWcType typ)
typ' :: HsType name
typ' = rename fv $ specializeTyVarBndrs bndrs typs true_type
fv = foldr Set.union Set.empty . map freeVariables $ typs
specializeSig _ _ sig = sig
specializeInstHead :: (Ord name, DataId name, SetName name, NamedThing name)
=> InstHead name -> InstHead name
specializeInstHead ihd@InstHead { ihdInstType = clsi@ClassInst { .. }, .. } =
ihd { ihdInstType = instType' }
where
instType' = clsi
{ clsiSigs = map specializeSig' clsiSigs
, clsiAssocTys = map specializeFamilyDecl' clsiAssocTys
}
specializeSig' = specializeSig clsiTyVars ihdTypes
specializeFamilyDecl' = specializePseudoFamilyDecl clsiTyVars ihdTypes
specializeInstHead ihd = ihd
sugar :: forall name. (NamedThing name, DataId name)
=> HsType name -> HsType name
sugar = sugarOperators . sugarTuples . sugarLists
sugarLists :: NamedThing name => HsType name -> HsType name
sugarLists (HsAppTy (L _ (HsTyVar _ (L _ name))) ltyp)
| isBuiltInSyntax name' && strName == "[]" = HsListTy ltyp
where
name' = getName name
strName = occNameString . nameOccName $ name'
sugarLists typ = typ
sugarTuples :: NamedThing name => HsType name -> HsType name
sugarTuples typ =
aux [] typ
where
aux apps (HsAppTy (L _ ftyp) atyp) = aux (atyp:apps) ftyp
aux apps (HsParTy (L _ typ')) = aux apps typ'
aux apps (HsTyVar _ (L _ name))
| isBuiltInSyntax name' && suitable = HsTupleTy HsBoxedTuple apps
where
name' = getName name
strName = occNameString . nameOccName $ name'
suitable = case parseTupleArity strName of
Just arity -> arity == length apps
Nothing -> False
aux _ _ = typ
sugarOperators :: NamedThing name => HsType name -> HsType name
sugarOperators (HsAppTy (L _ (HsAppTy (L _ (HsTyVar _ (L l name))) la)) lb)
| isSymOcc $ getOccName name' = mkHsOpTy la (L l name) lb
| isBuiltInSyntax name' && getOccString name == "(->)" = HsFunTy la lb
where
name' = getName name
sugarOperators typ = typ
parseTupleArity :: String -> Maybe Int
parseTupleArity ('(':commas) = do
n <- parseCommas commas
guard $ n /= 0
return $ n + 1
where
parseCommas (',':rest) = (+ 1) <$> parseCommas rest
parseCommas ")" = Just 0
parseCommas _ = Nothing
parseTupleArity _ = Nothing
type NameRep = FastString
getNameRep :: NamedThing name => name -> NameRep
getNameRep = occNameFS . getOccName
nameRepString :: NameRep -> String
nameRepString = unpackFS
stringNameRep :: String -> NameRep
stringNameRep = mkFastString
setInternalNameRep :: SetName name => NameRep -> name -> name
setInternalNameRep = setInternalOccName . mkVarOccFS
setInternalOccName :: SetName name => OccName -> name -> name
setInternalOccName occ name =
setName nname' name
where
nname = getName name
nname' = mkInternalName (nameUnique nname) occ (nameSrcSpan nname)
freeVariables :: forall name. (NamedThing name, DataId name)
=> HsType name -> Set NameRep
freeVariables =
everythingWithState Set.empty Set.union query
where
query term ctx = case cast term :: Maybe (HsType name) of
Just (HsForAllTy bndrs _) ->
(Set.empty, Set.union ctx (bndrsNames bndrs))
Just (HsTyVar _ (L _ name))
| getName name `Set.member` ctx -> (Set.empty, ctx)
| otherwise -> (Set.singleton $ getNameRep name, ctx)
_ -> (Set.empty, ctx)
bndrsNames = Set.fromList . map (getName . tyVarName . unLoc)
rename :: SetName name => Set NameRep -> HsType name -> HsType name
rename fv typ = runReader (renameType typ) $ RenameEnv
{ rneFV = fv
, rneCtx = Map.empty
}
type Rename name = Reader (RenameEnv name)
type Rebind name = State (RenameEnv name)
data RenameEnv name = RenameEnv
{ rneFV :: Set NameRep
, rneCtx :: Map Name name
}
renameType :: SetName name => HsType name -> Rename name (HsType name)
renameType (HsForAllTy bndrs lt) = rebind bndrs $ \bndrs' ->
HsForAllTy
<$> pure bndrs'
<*> renameLType lt
renameType (HsQualTy lctxt lt) =
HsQualTy
<$> located renameContext lctxt
<*> renameLType lt
renameType (HsTyVar ip name) = HsTyVar ip <$> located renameName name
renameType (HsAppTy lf la) = HsAppTy <$> renameLType lf <*> renameLType la
renameType (HsFunTy la lr) = HsFunTy <$> renameLType la <*> renameLType lr
renameType (HsListTy lt) = HsListTy <$> renameLType lt
renameType (HsPArrTy lt) = HsPArrTy <$> renameLType lt
renameType (HsTupleTy srt lt) = HsTupleTy srt <$> mapM renameLType lt
renameType (HsSumTy lt) = HsSumTy <$> mapM renameLType lt
renameType (HsOpTy la lop lb) =
HsOpTy <$> renameLType la <*> located renameName lop <*> renameLType lb
renameType (HsParTy lt) = HsParTy <$> renameLType lt
renameType (HsIParamTy ip lt) = HsIParamTy ip <$> renameLType lt
renameType (HsEqTy la lb) = HsEqTy <$> renameLType la <*> renameLType lb
renameType (HsKindSig lt lk) = HsKindSig <$> renameLType lt <*> pure lk
renameType t@(HsSpliceTy _ _) = pure t
renameType (HsDocTy lt doc) = HsDocTy <$> renameLType lt <*> pure doc
renameType (HsBangTy bang lt) = HsBangTy bang <$> renameLType lt
renameType t@(HsRecTy _) = pure t
renameType t@(HsCoreTy _) = pure t
renameType (HsExplicitListTy ip ph ltys) =
HsExplicitListTy ip ph <$> renameLTypes ltys
renameType (HsExplicitTupleTy phs ltys) =
HsExplicitTupleTy phs <$> renameLTypes ltys
renameType t@(HsTyLit _) = pure t
renameType (HsWildCardTy wc) = pure (HsWildCardTy wc)
renameType (HsAppsTy _) = error "HsAppsTy: Only used before renaming"
renameLType :: SetName name => LHsType name -> Rename name (LHsType name)
renameLType = located renameType
renameLTypes :: SetName name => [LHsType name] -> Rename name [LHsType name]
renameLTypes = mapM renameLType
renameContext :: SetName name => HsContext name -> Rename name (HsContext name)
renameContext = renameLTypes
renameName :: SetName name => name -> Rename name name
renameName name = do
RenameEnv { rneCtx = ctx } <- ask
pure $ fromMaybe name (Map.lookup (getName name) ctx)
rebind :: SetName name
=> [LHsTyVarBndr name] -> ([LHsTyVarBndr name] -> Rename name a)
-> Rename name a
rebind lbndrs action = do
(lbndrs', env') <- runState (rebindLTyVarBndrs lbndrs) <$> ask
local (const env') (action lbndrs')
rebindLTyVarBndrs :: SetName name
=> [LHsTyVarBndr name] -> Rebind name [LHsTyVarBndr name]
rebindLTyVarBndrs lbndrs = mapM (located rebindTyVarBndr) lbndrs
rebindTyVarBndr :: SetName name
=> HsTyVarBndr name -> Rebind name (HsTyVarBndr name)
rebindTyVarBndr (UserTyVar (L l name)) =
UserTyVar . L l <$> rebindName name
rebindTyVarBndr (KindedTyVar name kinds) =
KindedTyVar <$> located rebindName name <*> pure kinds
rebindName :: SetName name => name -> Rebind name name
rebindName name = do
RenameEnv { .. } <- get
taken <- takenNames
case Map.lookup (getName name) rneCtx of
Just name' -> pure name'
Nothing | getNameRep name `Set.member` taken -> freshName name
Nothing -> reuseName name
freshName :: SetName name => name -> Rebind name name
freshName name = do
env@RenameEnv { .. } <- get
taken <- takenNames
let name' = setInternalNameRep (findFreshName taken rep) name
put $ env { rneCtx = Map.insert nname name' rneCtx }
return name'
where
nname = getName name
rep = getNameRep nname
reuseName :: SetName name => name -> Rebind name name
reuseName name = do
env@RenameEnv { .. } <- get
put $ env { rneCtx = Map.insert (getName name) name rneCtx }
return name
takenNames :: NamedThing name => Rebind name (Set NameRep)
takenNames = do
RenameEnv { .. } <- get
return $ Set.union rneFV (ctxElems rneCtx)
where
ctxElems = Set.fromList . map getNameRep . Map.elems
findFreshName :: Set NameRep -> NameRep -> NameRep
findFreshName taken =
fromJust . List.find isFresh . alternativeNames
where
isFresh = not . flip Set.member taken
alternativeNames :: NameRep -> [NameRep]
alternativeNames name
| [_] <- nameRepString name = letterNames ++ alternativeNames' name
where
letterNames = map (stringNameRep . pure) ['a'..'z']
alternativeNames name = alternativeNames' name
alternativeNames' :: NameRep -> [NameRep]
alternativeNames' name =
[ stringNameRep $ str ++ show i | i :: Int <- [0..] ]
where
str = nameRepString name
located :: Functor f => (a -> f b) -> Located a -> f (Located b)
located f (L loc e) = L loc <$> f e
tyVarName :: HsTyVarBndr name -> name
tyVarName (UserTyVar name) = unLoc name
tyVarName (KindedTyVar (L _ name) _) = name