{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TypeSynonymInstances #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE DeriveDataTypeable #-}
module HieAst ( mkHieFile ) where

import GhcPrelude

import Avail                      ( Avails )
import Bag                        ( Bag, bagToList )
import BasicTypes
import BooleanFormula
import Class                      ( FunDep )
import CoreUtils                  ( exprType )
import ConLike                    ( conLikeName )
import Config                     ( cProjectVersion )
import Desugar                    ( deSugarExpr )
import FieldLabel
import HsSyn
import HscTypes
import Module                     ( ModuleName, ml_hs_file )
import MonadUtils                 ( concatMapM, liftIO )
import Name                       ( Name, nameSrcSpan, setNameLoc )
import NameEnv                    ( NameEnv, emptyNameEnv, extendNameEnv, lookupNameEnv )
import SrcLoc
import TcHsSyn                    ( hsLitType, hsPatType )
import Type                       ( mkVisFunTys, Type )
import TysWiredIn                 ( mkListTy, mkSumTy )
import Var                        ( Id, Var, setVarName, varName, varType )
import TcRnTypes
import MkIface                    ( mkIfaceExports )

import HieTypes
import HieUtils

import qualified Data.Array as A
import qualified Data.ByteString as BS
import qualified Data.ByteString.Char8 as BSC
import qualified Data.Map as M
import qualified Data.Set as S
import Data.Data                  ( Data, Typeable )
import Data.List                  ( foldl1' )
import Data.Maybe                 ( listToMaybe )
import Control.Monad.Trans.Reader
import Control.Monad.Trans.Class  ( lift )

-- These synonyms match those defined in main/GHC.hs
type RenamedSource     = ( HsGroup GhcRn, [LImportDecl GhcRn]
                         , Maybe [(LIE GhcRn, Avails)]
                         , Maybe LHsDocString )
type TypecheckedSource = LHsBinds GhcTc

{- Note [Name Remapping]
The Typechecker introduces new names for mono names in AbsBinds.
We don't care about the distinction between mono and poly bindings,
so we replace all occurrences of the mono name with the poly name.
-}
newtype HieState = HieState
  { name_remapping :: NameEnv Id
  }

initState :: HieState
initState = HieState emptyNameEnv

class ModifyState a where -- See Note [Name Remapping]
  addSubstitution :: a -> a -> HieState -> HieState

instance ModifyState Name where
  addSubstitution _ _ hs = hs

instance ModifyState Id where
  addSubstitution mono poly hs =
    hs{name_remapping = extendNameEnv (name_remapping hs) (varName mono) poly}

modifyState :: ModifyState (IdP p) => [ABExport p] -> HieState -> HieState
modifyState = foldr go id
  where
    go ABE{abe_poly=poly,abe_mono=mono} f = addSubstitution mono poly . f
    go _ f = f

type HieM = ReaderT HieState Hsc

-- | Construct an 'HieFile' from the outputs of the typechecker.
mkHieFile :: ModSummary
          -> TcGblEnv
          -> RenamedSource -> Hsc HieFile
mkHieFile ms ts rs = do
  let tc_binds = tcg_binds ts
  (asts', arr) <- getCompressedAsts tc_binds rs
  let Just src_file = ml_hs_file $ ms_location ms
  src <- liftIO $ BS.readFile src_file
  return $ HieFile
      { hie_version = curHieVersion
      , hie_ghc_version = BSC.pack cProjectVersion
      , hie_hs_file = src_file
      , hie_module = ms_mod ms
      , hie_types = arr
      , hie_asts = asts'
      -- mkIfaceExports sorts the AvailInfos for stability
      , hie_exports = mkIfaceExports (tcg_exports ts)
      , hie_hs_src = src
      }

getCompressedAsts :: TypecheckedSource -> RenamedSource
  -> Hsc (HieASTs TypeIndex, A.Array TypeIndex HieTypeFlat)
getCompressedAsts ts rs = do
  asts <- enrichHie ts rs
  return $ compressTypes asts

enrichHie :: TypecheckedSource -> RenamedSource -> Hsc (HieASTs Type)
enrichHie ts (hsGrp, imports, exports, _) = flip runReaderT initState $ do
    tasts <- toHie $ fmap (BC RegularBind ModuleScope) ts
    rasts <- processGrp hsGrp
    imps <- toHie $ filter (not . ideclImplicit . unLoc) imports
    exps <- toHie $ fmap (map $ IEC Export . fst) exports
    let spanFile children = case children of
          [] -> mkRealSrcSpan (mkRealSrcLoc "" 1 1) (mkRealSrcLoc "" 1 1)
          _ -> mkRealSrcSpan (realSrcSpanStart $ nodeSpan $ head children)
                             (realSrcSpanEnd   $ nodeSpan $ last children)

        modulify xs =
          Node (simpleNodeInfo "Module" "Module") (spanFile xs) xs

        asts = HieASTs
          $ resolveTyVarScopes
          $ M.map (modulify . mergeSortAsts)
          $ M.fromListWith (++)
          $ map (\x -> (srcSpanFile (nodeSpan x),[x])) flat_asts

        flat_asts = concat
          [ tasts
          , rasts
          , imps
          , exps
          ]
    return asts
  where
    processGrp grp = concatM
      [ toHie $ fmap (RS ModuleScope ) hs_valds grp
      , toHie $ hs_splcds grp
      , toHie $ hs_tyclds grp
      , toHie $ hs_derivds grp
      , toHie $ hs_fixds grp
      , toHie $ hs_defds grp
      , toHie $ hs_fords grp
      , toHie $ hs_warnds grp
      , toHie $ hs_annds grp
      , toHie $ hs_ruleds grp
      ]

getRealSpan :: SrcSpan -> Maybe Span
getRealSpan (RealSrcSpan sp) = Just sp
getRealSpan _ = Nothing

grhss_span :: GRHSs p body -> SrcSpan
grhss_span (GRHSs _ xs bs) = foldl' combineSrcSpans (getLoc bs) (map getLoc xs)
grhss_span (XGRHSs _) = error "XGRHS has no span"

bindingsOnly :: [Context Name] -> [HieAST a]
bindingsOnly [] = []
bindingsOnly (C c n : xs) = case nameSrcSpan n of
  RealSrcSpan span -> Node nodeinfo span [] : bindingsOnly xs
    where nodeinfo = NodeInfo S.empty [] (M.singleton (Right n) info)
          info = mempty{identInfo = S.singleton c}
  _ -> bindingsOnly xs

concatM :: Monad m => [m [a]] -> m [a]
concatM xs = concat <$> sequence xs

{- Note [Capturing Scopes and other non local information]
toHie is a local tranformation, but scopes of bindings cannot be known locally,
hence we have to push the relevant info down into the binding nodes.
We use the following types (*Context and *Scoped) to wrap things and
carry the required info
(Maybe Span) always carries the span of the entire binding, including rhs
-}
data Context a = C ContextInfo a -- Used for names and bindings

data RContext a = RC RecFieldContext a
data RFContext a = RFC RecFieldContext (Maybe Span) a
-- ^ context for record fields

data IEContext a = IEC IEType a
-- ^ context for imports/exports

data BindContext a = BC BindType Scope a
-- ^ context for imports/exports

data PatSynFieldContext a = PSC (Maybe Span) a
-- ^ context for pattern synonym fields.

data SigContext a = SC SigInfo a
-- ^ context for type signatures

data SigInfo = SI SigType (Maybe Span)

data SigType = BindSig | ClassSig | InstSig

data RScoped a = RS Scope a
-- ^ Scope spans over everything to the right of a, (mostly) not
-- including a itself
-- (Includes a in a few special cases like recursive do bindings) or
-- let/where bindings

-- | Pattern scope
data PScoped a = PS (Maybe Span)
                    Scope       -- ^ use site of the pattern
                    Scope       -- ^ pattern to the right of a, not including a
                    a
  deriving (Typeable, Data) -- Pattern Scope

{- Note [TyVar Scopes]
Due to -XScopedTypeVariables, type variables can be in scope quite far from
their original binding. We resolve the scope of these type variables
in a separate pass
-}
data TScoped a = TS TyVarScope a -- TyVarScope

data TVScoped a = TVS TyVarScope Scope a -- TyVarScope
-- ^ First scope remains constant
-- Second scope is used to build up the scope of a tyvar over
-- things to its right, ala RScoped

-- | Each element scopes over the elements to the right
listScopes :: Scope -> [Located a] -> [RScoped (Located a)]
listScopes _ [] = []
listScopes rhsScope [pat] = [RS rhsScope pat]
listScopes rhsScope (pat : pats) = RS sc pat : pats'
  where
    pats'@((RS scope p):_) = listScopes rhsScope pats
    sc = combineScopes scope $ mkScope $ getLoc p

-- | 'listScopes' specialised to 'PScoped' things
patScopes
  :: Maybe Span
  -> Scope
  -> Scope
  -> [LPat (GhcPass p)]
  -> [PScoped (LPat (GhcPass p))]
patScopes rsp useScope patScope xs =
  map (\(RS sc a) -> PS rsp useScope sc (unLoc a)) $
    listScopes patScope (map dL xs)

-- | 'listScopes' specialised to 'TVScoped' things
tvScopes
  :: TyVarScope
  -> Scope
  -> [LHsTyVarBndr a]
  -> [TVScoped (LHsTyVarBndr a)]
tvScopes tvScope rhsScope xs =
  map (\(RS sc a)-> TVS tvScope sc a) $ listScopes rhsScope xs

{- Note [Scoping Rules for SigPat]
Explicitly quantified variables in pattern type signatures are not
brought into scope in the rhs, but implicitly quantified variables
are (HsWC and HsIB).
This is unlike other signatures, where explicitly quantified variables
are brought into the RHS Scope
For example
foo :: forall a. ...;
foo = ... -- a is in scope here

bar (x :: forall a. a -> a) = ... -- a is not in scope here
--   ^ a is in scope here (pattern body)

bax (x :: a) = ... -- a is in scope here
Because of HsWC and HsIB pass on their scope to their children
we must wrap the LHsType in pattern signatures in a
Shielded explictly, so that the HsWC/HsIB scope is not passed
on the the LHsType
-}

data Shielded a = SH Scope a -- Ignores its TScope, uses its own scope instead

type family ProtectedSig a where
  ProtectedSig GhcRn = HsWildCardBndrs GhcRn (HsImplicitBndrs
                                                GhcRn
                                                (Shielded (LHsType GhcRn)))
  ProtectedSig GhcTc = NoExt

class ProtectSig a where
  protectSig :: Scope -> LHsSigWcType (NoGhcTc a) -> ProtectedSig a

instance (HasLoc a) => HasLoc (Shielded a) where
  loc (SH _ a) = loc a

instance (ToHie (TScoped a)) => ToHie (TScoped (Shielded a)) where
  toHie (TS _ (SH sc a)) = toHie (TS (ResolvedScopes [sc]) a)

instance ProtectSig GhcTc where
  protectSig _ _ = NoExt

instance ProtectSig GhcRn where
  protectSig sc (HsWC a (HsIB b sig)) =
    HsWC a (HsIB b (SH sc sig))
  protectSig _ _ = error "protectSig not given HsWC (HsIB)"

class HasLoc a where
  -- ^ defined so that HsImplicitBndrs and HsWildCardBndrs can
  -- know what their implicit bindings are scoping over
  loc :: a -> SrcSpan

instance HasLoc thing => HasLoc (TScoped thing) where
  loc (TS _ a) = loc a

instance HasLoc thing => HasLoc (PScoped thing) where
  loc (PS _ _ _ a) = loc a

instance HasLoc (LHsQTyVars GhcRn) where
  loc (HsQTvs _ vs) = loc vs
  loc _ = noSrcSpan

instance HasLoc thing => HasLoc (HsImplicitBndrs a thing) where
  loc (HsIB _ a) = loc a
  loc _ = noSrcSpan

instance HasLoc thing => HasLoc (HsWildCardBndrs a thing) where
  loc (HsWC _ a) = loc a
  loc _ = noSrcSpan

instance HasLoc (Located a) where
  loc (L l _) = l

instance HasLoc a => HasLoc [a] where
  loc [] = noSrcSpan
  loc xs = foldl1' combineSrcSpans $ map loc xs

instance (HasLoc a, HasLoc b) => HasLoc (FamEqn s a b) where
  loc (FamEqn _ a Nothing b _ c) = foldl1' combineSrcSpans [loc a, loc b, loc c]
  loc (FamEqn _ a (Just tvs) b _ c) = foldl1' combineSrcSpans
                                              [loc a, loc tvs, loc b, loc c]
  loc _ = noSrcSpan
instance (HasLoc tm, HasLoc ty) => HasLoc (HsArg tm ty) where
  loc (HsValArg tm) = loc tm
  loc (HsTypeArg _ ty) = loc ty
  loc (HsArgPar sp)  = sp

instance HasLoc (HsDataDefn GhcRn) where
  loc def@(HsDataDefn{}) = loc $ dd_cons def
    -- Only used for data family instances, so we only need rhs
    -- Most probably the rest will be unhelpful anyway
  loc _ = noSrcSpan

instance HasLoc (Pat (GhcPass a)) where
  loc (dL -> L l _) = l

-- | The main worker class
class ToHie a where
  toHie :: a -> HieM [HieAST Type]

-- | Used to collect type info
class Data a => HasType a where
  getTypeNode :: a -> HieM [HieAST Type]

instance (ToHie a) => ToHie [a] where
  toHie = concatMapM toHie

instance (ToHie a) => ToHie (Bag a) where
  toHie = toHie . bagToList

instance (ToHie a) => ToHie (Maybe a) where
  toHie = maybe (pure []) toHie

instance ToHie (Context (Located NoExt)) where
  toHie _ = pure []

instance ToHie (TScoped NoExt) where
  toHie _ = pure []

instance ToHie (IEContext (Located ModuleName)) where
  toHie (IEC c (L (RealSrcSpan span) mname)) =
      pure $ [Node (NodeInfo S.empty [] idents) span []]
    where details = mempty{identInfo = S.singleton (IEThing c)}
          idents = M.singleton (Left mname) details
  toHie _ = pure []

instance ToHie (Context (Located Var)) where
  toHie c = case c of
      C context (L (RealSrcSpan span) name')
        -> do
        m <- asks name_remapping
        let name = case lookupNameEnv m (varName name') of
              Just var -> var
              Nothing-> name'
        pure
          [Node
            (NodeInfo S.empty [] $
              M.singleton (Right $ varName name)
                          (IdentifierDetails (Just $ varType name')
                                             (S.singleton context)))
            span
            []]
      _ -> pure []

instance ToHie (Context (Located Name)) where
  toHie c = case c of
      C context (L (RealSrcSpan span) name') -> do
        m <- asks name_remapping
        let name = case lookupNameEnv m name' of
              Just var -> varName var
              Nothing -> name'
        pure
          [Node
            (NodeInfo S.empty [] $
              M.singleton (Right name)
                          (IdentifierDetails Nothing
                                             (S.singleton context)))
            span
            []]
      _ -> pure []

-- | Dummy instances - never called
instance ToHie (TScoped (LHsSigWcType GhcTc)) where
  toHie _ = pure []
instance ToHie (TScoped (LHsWcType GhcTc)) where
  toHie _ = pure []
instance ToHie (SigContext (LSig GhcTc)) where
  toHie _ = pure []
instance ToHie (TScoped Type) where
  toHie _ = pure []

instance HasType (LHsBind GhcRn) where
  getTypeNode (L spn bind) = makeNode bind spn

instance HasType (LHsBind GhcTc) where
  getTypeNode (L spn bind) = case bind of
      FunBind{fun_id = name} -> makeTypeNode bind spn (varType $ unLoc name)
      _ -> makeNode bind spn

instance HasType (LPat GhcRn) where
  getTypeNode (dL -> L spn pat) = makeNode pat spn

instance HasType (LPat GhcTc) where
  getTypeNode (dL -> L spn opat) = makeTypeNode opat spn (hsPatType opat)

instance HasType (LHsExpr GhcRn) where
  getTypeNode (L spn e) = makeNode e spn

-- | This instance tries to construct 'HieAST' nodes which include the type of
-- the expression. It is not yet possible to do this efficiently for all
-- expression forms, so we skip filling in the type for those inputs.
--
-- 'HsApp', for example, doesn't have any type information available directly on
-- the node. Our next recourse would be to desugar it into a 'CoreExpr' then
-- query the type of that. Yet both the desugaring call and the type query both
-- involve recursive calls to the function and argument! This is particularly
-- problematic when you realize that the HIE traversal will eventually visit
-- those nodes too and ask for their types again.
--
-- Since the above is quite costly, we just skip cases where computing the
-- expression's type is going to be expensive.
--
-- See #16233
instance HasType (LHsExpr GhcTc) where
  getTypeNode e@(L spn e') = lift $
    -- Some expression forms have their type immediately available
    let tyOpt = case e' of
          HsLit _ l -> Just (hsLitType l)
          HsOverLit _ o -> Just (overLitType o)

          HsLam     _ (MG { mg_ext = groupTy }) -> Just (matchGroupType groupTy)
          HsLamCase _ (MG { mg_ext = groupTy }) -> Just (matchGroupType groupTy)
          HsCase _  _ (MG { mg_ext = groupTy }) -> Just (mg_res_ty groupTy)

          ExplicitList  ty _ _   -> Just (mkListTy ty)
          ExplicitSum   ty _ _ _ -> Just (mkSumTy ty)
          HsDo          ty _ _   -> Just ty
          HsMultiIf     ty _     -> Just ty

          _ -> Nothing

    in
    case tyOpt of
      _ | skipDesugaring e' -> fallback
        | otherwise -> do
            hs_env <- Hsc $ \e w -> return (e,w)
            (_,mbe) <- liftIO $ deSugarExpr hs_env e
            maybe fallback (makeTypeNode e' spn . exprType) mbe
    where
      fallback = makeNode e' spn

      matchGroupType :: MatchGroupTc -> Type
      matchGroupType (MatchGroupTc args res) = mkVisFunTys args res

      -- | Skip desugaring of these expressions for performance reasons.
      --
      -- See impact on Haddock output (esp. missing type annotations or links)
      -- before marking more things here as 'False'. See impact on Haddock
      -- performance before marking more things as 'True'.
      skipDesugaring :: HsExpr a -> Bool
      skipDesugaring e = case e of
        HsVar{}        -> False
        HsUnboundVar{} -> False
        HsConLikeOut{} -> False
        HsRecFld{}     -> False
        HsOverLabel{}  -> False
        HsIPVar{}      -> False
        HsWrap{}       -> False
        _              -> True

instance ( ToHie (Context (Located (IdP a)))
         , ToHie (MatchGroup a (LHsExpr a))
         , ToHie (PScoped (LPat a))
         , ToHie (GRHSs a (LHsExpr a))
         , ToHie (LHsExpr a)
         , ToHie (Located (PatSynBind a a))
         , HasType (LHsBind a)
         , ModifyState (IdP a)
         , Data (HsBind a)
         ) => ToHie (BindContext (LHsBind a)) where
  toHie (BC context scope b@(L span bind)) =
    concatM $ getTypeNode b : case bind of
      FunBind{fun_id = name, fun_matches = matches} ->
        [ toHie $ C (ValBind context scope $ getRealSpan span) name
        , toHie matches
        ]
      PatBind{pat_lhs = lhs, pat_rhs = rhs} ->
        [ toHie $ PS (getRealSpan span) scope NoScope lhs
        , toHie rhs
        ]
      VarBind{var_rhs = expr} ->
        [ toHie expr
        ]
      AbsBinds{abs_exports = xs, abs_binds = binds} ->
        [ local (modifyState xs) $ -- Note [Name Remapping]
            toHie $ fmap (BC context scope) binds
        ]
      PatSynBind _ psb ->
        [ toHie $ L span psb -- PatSynBinds only occur at the top level
        ]
      XHsBindsLR _ -> []

instance ( ToHie (LMatch a body)
         ) => ToHie (MatchGroup a body) where
  toHie mg = concatM $ case mg of
    MG{ mg_alts = (L span alts) , mg_origin = FromSource } ->
      [ pure $ locOnly span
      , toHie alts
      ]
    MG{} -> []
    XMatchGroup _ -> []

instance ( ToHie (Context (Located (IdP a)))
         , ToHie (PScoped (LPat a))
         , ToHie (HsPatSynDir a)
         ) => ToHie (Located (PatSynBind a a)) where
    toHie (L sp psb) = concatM $ case psb of
      PSB{psb_id=var, psb_args=dets, psb_def=pat, psb_dir=dir} ->
        [ toHie $ C (Decl PatSynDec $ getRealSpan sp) var
        , toHie $ toBind dets
        , toHie $ PS Nothing lhsScope NoScope pat
        , toHie dir
        ]
        where
          lhsScope = combineScopes varScope detScope
          varScope = mkLScope var
          detScope = case dets of
            (PrefixCon args) -> foldr combineScopes NoScope $ map mkLScope args
            (InfixCon a b) -> combineScopes (mkLScope a) (mkLScope b)
            (RecCon r) -> foldr go NoScope r
          go (RecordPatSynField a b) c = combineScopes c
            $ combineScopes (mkLScope a) (mkLScope b)
          detSpan = case detScope of
            LocalScope a -> Just a
            _ -> Nothing
          toBind (PrefixCon args) = PrefixCon $ map (C Use) args
          toBind (InfixCon a b) = InfixCon (C Use a) (C Use b)
          toBind (RecCon r) = RecCon $ map (PSC detSpan) r
      XPatSynBind _ -> []

instance ( ToHie (MatchGroup a (LHsExpr a))
         ) => ToHie (HsPatSynDir a) where
  toHie dir = case dir of
    ExplicitBidirectional mg -> toHie mg
    _ -> pure []

instance ( a ~ GhcPass p
         , ToHie body
         , ToHie (HsMatchContext (NameOrRdrName (IdP a)))
         , ToHie (PScoped (LPat a))
         , ToHie (GRHSs a body)
         , Data (Match a body)
         ) => ToHie (LMatch (GhcPass p) body) where
  toHie (L span m ) = concatM $ makeNode m span : case m of
    Match{m_ctxt=mctx, m_pats = pats, m_grhss =  grhss } ->
      [ toHie mctx
      , let rhsScope = mkScope $ grhss_span grhss
          in toHie $ patScopes Nothing rhsScope NoScope pats
      , toHie grhss
      ]
    XMatch _ -> []

instance ( ToHie (Context (Located a))
         ) => ToHie (HsMatchContext a) where
  toHie (FunRhs{mc_fun=name}) = toHie $ C MatchBind name
  toHie (StmtCtxt a) = toHie a
  toHie _ = pure []

instance ( ToHie (HsMatchContext a)
         ) => ToHie (HsStmtContext a) where
  toHie (PatGuard a) = toHie a
  toHie (ParStmtCtxt a) = toHie a
  toHie (TransStmtCtxt a) = toHie a
  toHie _ = pure []

instance ( a ~ GhcPass p
         , ToHie (Context (Located (IdP a)))
         , ToHie (RContext (HsRecFields a (PScoped (LPat a))))
         , ToHie (LHsExpr a)
         , ToHie (TScoped (LHsSigWcType a))
         , ProtectSig a
         , ToHie (TScoped (ProtectedSig a))
         , HasType (LPat a)
         , Data (HsSplice a)
         ) => ToHie (PScoped (LPat (GhcPass p))) where
  toHie (PS rsp scope pscope lpat@(dL -> L ospan opat)) =
    concatM $ getTypeNode lpat : case opat of
      WildPat _ ->
        []
      VarPat _ lname ->
        [ toHie $ C (PatternBind scope pscope rsp) lname
        ]
      LazyPat _ p ->
        [ toHie $ PS rsp scope pscope p
        ]
      AsPat _ lname pat ->
        [ toHie $ C (PatternBind scope
                                 (combineScopes (mkLScope (dL pat)) pscope)
                                 rsp)
                    lname
        , toHie $ PS rsp scope pscope pat
        ]
      ParPat _ pat ->
        [ toHie $ PS rsp scope pscope pat
        ]
      BangPat _ pat ->
        [ toHie $ PS rsp scope pscope pat
        ]
      ListPat _ pats ->
        [ toHie $ patScopes rsp scope pscope pats
        ]
      TuplePat _ pats _ ->
        [ toHie $ patScopes rsp scope pscope pats
        ]
      SumPat _ pat _ _ ->
        [ toHie $ PS rsp scope pscope pat
        ]
      ConPatIn c dets ->
        [ toHie $ C Use c
        , toHie $ contextify dets
        ]
      ConPatOut {pat_con = con, pat_args = dets}->
        [ toHie $ C Use $ fmap conLikeName con
        , toHie $ contextify dets
        ]
      ViewPat _ expr pat ->
        [ toHie expr
        , toHie $ PS rsp scope pscope pat
        ]
      SplicePat _ sp ->
        [ toHie $ L ospan sp
        ]
      LitPat _ _ ->
        []
      NPat _ _ _ _ ->
        []
      NPlusKPat _ n _ _ _ _ ->
        [ toHie $ C (PatternBind scope pscope rsp) n
        ]
      SigPat _ pat sig ->
        [ toHie $ PS rsp scope pscope pat
        , let cscope = mkLScope (dL pat) in
            toHie $ TS (ResolvedScopes [cscope, scope, pscope])
                       (protectSig @a cscope sig)
              -- See Note [Scoping Rules for SigPat]
        ]
      CoPat _ _ _ _ ->
        []
      XPat _ -> []
    where
      contextify (PrefixCon args) = PrefixCon $ patScopes rsp scope pscope args
      contextify (InfixCon a b) = InfixCon a' b'
        where [a', b'] = patScopes rsp scope pscope [a,b]
      contextify (RecCon r) = RecCon $ RC RecFieldMatch $ contextify_rec r
      contextify_rec (HsRecFields fds a) = HsRecFields (map go scoped_fds) a
        where
          go (RS fscope (L spn (HsRecField lbl pat pun))) =
            L spn $ HsRecField lbl (PS rsp scope fscope pat) pun
          scoped_fds = listScopes pscope fds

instance ( ToHie body
         , ToHie (LGRHS a body)
         , ToHie (RScoped (LHsLocalBinds a))
         ) => ToHie (GRHSs a body) where
  toHie grhs = concatM $ case grhs of
    GRHSs _ grhss binds ->
     [ toHie grhss
     , toHie $ RS (mkScope $ grhss_span grhs) binds
     ]
    XGRHSs _ -> []

instance ( ToHie (Located body)
         , ToHie (RScoped (GuardLStmt a))
         , Data (GRHS a (Located body))
         ) => ToHie (LGRHS a (Located body)) where
  toHie (L span g) = concatM $ makeNode g span : case g of
    GRHS _ guards body ->
      [ toHie $ listScopes (mkLScope body) guards
      , toHie body
      ]
    XGRHS _ -> []

instance ( a ~ GhcPass p
         , ToHie (Context (Located (IdP a)))
         , HasType (LHsExpr a)
         , ToHie (PScoped (LPat a))
         , ToHie (MatchGroup a (LHsExpr a))
         , ToHie (LGRHS a (LHsExpr a))
         , ToHie (RContext (HsRecordBinds a))
         , ToHie (RFContext (Located (AmbiguousFieldOcc a)))
         , ToHie (ArithSeqInfo a)
         , ToHie (LHsCmdTop a)
         , ToHie (RScoped (GuardLStmt a))
         , ToHie (RScoped (LHsLocalBinds a))
         , ToHie (TScoped (LHsWcType (NoGhcTc a)))
         , ToHie (TScoped (LHsSigWcType (NoGhcTc a)))
         , Data (HsExpr a)
         , Data (HsSplice a)
         , Data (HsTupArg a)
         , Data (AmbiguousFieldOcc a)
         ) => ToHie (LHsExpr (GhcPass p)) where
  toHie e@(L mspan oexpr) = concatM $ getTypeNode e : case oexpr of
      HsVar _ (L _ var) ->
        [ toHie $ C Use (L mspan var)
             -- Patch up var location since typechecker removes it
        ]
      HsUnboundVar _ _ ->
        []
      HsConLikeOut _ con ->
        [ toHie $ C Use $ L mspan $ conLikeName con
        ]
      HsRecFld _ fld ->
        [ toHie $ RFC RecFieldOcc Nothing (L mspan fld)
        ]
      HsOverLabel _ _ _ -> []
      HsIPVar _ _ -> []
      HsOverLit _ _ -> []
      HsLit _ _ -> []
      HsLam _ mg ->
        [ toHie mg
        ]
      HsLamCase _ mg ->
        [ toHie mg
        ]
      HsApp _ a b ->
        [ toHie a
        , toHie b
        ]
      HsAppType _ expr sig ->
        [ toHie expr
        , toHie $ TS (ResolvedScopes []) sig
        ]
      OpApp _ a b c ->
        [ toHie a
        , toHie b
        , toHie c
        ]
      NegApp _ a _ ->
        [ toHie a
        ]
      HsPar _ a ->
        [ toHie a
        ]
      SectionL _ a b ->
        [ toHie a
        , toHie b
        ]
      SectionR _ a b ->
        [ toHie a
        , toHie b
        ]
      ExplicitTuple _ args _ ->
        [ toHie args
        ]
      ExplicitSum _ _ _ expr ->
        [ toHie expr
        ]
      HsCase _ expr matches ->
        [ toHie expr
        , toHie matches
        ]
      HsIf _ _ a b c ->
        [ toHie a
        , toHie b
        , toHie c
        ]
      HsMultiIf _ grhss ->
        [ toHie grhss
        ]
      HsLet _ binds expr ->
        [ toHie $ RS (mkLScope expr) binds
        , toHie expr
        ]
      HsDo _ _ (L ispan stmts) ->
        [ pure $ locOnly ispan
        , toHie $ listScopes NoScope stmts
        ]
      ExplicitList _ _ exprs ->
        [ toHie exprs
        ]
      RecordCon {rcon_con_name = name, rcon_flds = binds}->
        [ toHie $ C Use name
        , toHie $ RC RecFieldAssign $ binds
        ]
      RecordUpd {rupd_expr = expr, rupd_flds = upds}->
        [ toHie expr
        , toHie $ map (RC RecFieldAssign) upds
        ]
      ExprWithTySig _ expr sig ->
        [ toHie expr
        , toHie $ TS (ResolvedScopes [mkLScope expr]) sig
        ]
      ArithSeq _ _ info ->
        [ toHie info
        ]
      HsSCC _ _ _ expr ->
        [ toHie expr
        ]
      HsCoreAnn _ _ _ expr ->
        [ toHie expr
        ]
      HsProc _ pat cmdtop ->
        [ toHie $ PS Nothing (mkLScope cmdtop) NoScope pat
        , toHie cmdtop
        ]
      HsStatic _ expr ->
        [ toHie expr
        ]
      HsTick _ _ expr ->
        [ toHie expr
        ]
      HsBinTick _ _ _ expr ->
        [ toHie expr
        ]
      HsTickPragma _ _ _ _ expr ->
        [ toHie expr
        ]
      HsWrap _ _ a ->
        [ toHie $ L mspan a
        ]
      HsBracket _ b ->
        [ toHie b
        ]
      HsRnBracketOut _ b p ->
        [ toHie b
        , toHie p
        ]
      HsTcBracketOut _ b p ->
        [ toHie b
        , toHie p
        ]
      HsSpliceE _ x ->
        [ toHie $ L mspan x
        ]
      EWildPat _ -> []
      EAsPat _ a b ->
        [ toHie $ C Use a
        , toHie b
        ]
      EViewPat _ a b ->
        [ toHie a
        , toHie b
        ]
      ELazyPat _ a ->
        [ toHie a
        ]
      XExpr _ -> []

instance ( a ~ GhcPass p
         , ToHie (LHsExpr a)
         , Data (HsTupArg a)
         ) => ToHie (LHsTupArg (GhcPass p)) where
  toHie (L span arg) = concatM $ makeNode arg span : case arg of
    Present _ expr ->
      [ toHie expr
      ]
    Missing _ -> []
    XTupArg _ -> []

instance ( a ~ GhcPass p
         , ToHie (PScoped (LPat a))
         , ToHie (LHsExpr a)
         , ToHie (SigContext (LSig a))
         , ToHie (RScoped (LHsLocalBinds a))
         , ToHie (RScoped (ApplicativeArg a))
         , ToHie (Located body)
         , Data (StmtLR a a (Located body))
         , Data (StmtLR a a (Located (HsExpr a)))
         ) => ToHie (RScoped (LStmt (GhcPass p) (Located body))) where
  toHie (RS scope (L span stmt)) = concatM $ makeNode stmt span : case stmt of
      LastStmt _ body _ _ ->
        [ toHie body
        ]
      BindStmt _ pat body _ _ ->
        [ toHie $ PS (getRealSpan $ getLoc body) scope NoScope pat
        , toHie body
        ]
      ApplicativeStmt _ stmts _ ->
        [ concatMapM (toHie . RS scope . snd) stmts
        ]
      BodyStmt _ body _ _ ->
        [ toHie body
        ]
      LetStmt _ binds ->
        [ toHie $ RS scope binds
        ]
      ParStmt _ parstmts _ _ ->
        [ concatMapM (\(ParStmtBlock _ stmts _ _) ->
                          toHie $ listScopes NoScope stmts)
                     parstmts
        ]
      TransStmt {trS_stmts = stmts, trS_using = using, trS_by = by} ->
        [ toHie $ listScopes scope stmts
        , toHie using
        , toHie by
        ]
      RecStmt {recS_stmts = stmts} ->
        [ toHie $ map (RS $ combineScopes scope (mkScope span)) stmts
        ]
      XStmtLR _ -> []

instance ( ToHie (LHsExpr a)
         , ToHie (PScoped (LPat a))
         , ToHie (BindContext (LHsBind a))
         , ToHie (SigContext (LSig a))
         , ToHie (RScoped (HsValBindsLR a a))
         , Data (HsLocalBinds a)
         ) => ToHie (RScoped (LHsLocalBinds a)) where
  toHie (RS scope (L sp binds)) = concatM $ makeNode binds sp : case binds of
      EmptyLocalBinds _ -> []
      HsIPBinds _ _ -> []
      HsValBinds _ valBinds ->
        [ toHie $ RS (combineScopes scope $ mkScope sp)
                      valBinds
        ]
      XHsLocalBindsLR _ -> []

instance ( ToHie (BindContext (LHsBind a))
         , ToHie (SigContext (LSig a))
         , ToHie (RScoped (XXValBindsLR a a))
         ) => ToHie (RScoped (HsValBindsLR a a)) where
  toHie (RS sc v) = concatM $ case v of
    ValBinds _ binds sigs ->
      [ toHie $ fmap (BC RegularBind sc) binds
      , toHie $ fmap (SC (SI BindSig Nothing)) sigs
      ]
    XValBindsLR x -> [ toHie $ RS sc x ]

instance ToHie (RScoped (NHsValBindsLR GhcTc)) where
  toHie (RS sc (NValBinds binds sigs)) = concatM $
    [ toHie (concatMap (map (BC RegularBind sc) . bagToList . snd) binds)
    , toHie $ fmap (SC (SI BindSig Nothing)) sigs
    ]
instance ToHie (RScoped (NHsValBindsLR GhcRn)) where
  toHie (RS sc (NValBinds binds sigs)) = concatM $
    [ toHie (concatMap (map (BC RegularBind sc) . bagToList . snd) binds)
    , toHie $ fmap (SC (SI BindSig Nothing)) sigs
    ]

instance ( ToHie (RContext (LHsRecField a arg))
         ) => ToHie (RContext (HsRecFields a arg)) where
  toHie (RC c (HsRecFields fields _)) = toHie $ map (RC c) fields

instance ( ToHie (RFContext (Located label))
         , ToHie arg
         , HasLoc arg
         , Data label
         , Data arg
         ) => ToHie (RContext (LHsRecField' label arg)) where
  toHie (RC c (L span recfld)) = concatM $ makeNode recfld span : case recfld of
    HsRecField label expr _ ->
      [ toHie $ RFC c (getRealSpan $ loc expr) label
      , toHie expr
      ]

removeDefSrcSpan :: Name -> Name
removeDefSrcSpan n = setNameLoc n noSrcSpan

instance ToHie (RFContext (LFieldOcc GhcRn)) where
  toHie (RFC c rhs (L nspan f)) = concatM $ case f of
    FieldOcc name _ ->
      [ toHie $ C (RecField c rhs) (L nspan $ removeDefSrcSpan name)
      ]
    XFieldOcc _ -> []

instance ToHie (RFContext (LFieldOcc GhcTc)) where
  toHie (RFC c rhs (L nspan f)) = concatM $ case f of
    FieldOcc var _ ->
      let var' = setVarName var (removeDefSrcSpan $ varName var)
      in [ toHie $ C (RecField c rhs) (L nspan var')
         ]
    XFieldOcc _ -> []

instance ToHie (RFContext (Located (AmbiguousFieldOcc GhcRn))) where
  toHie (RFC c rhs (L nspan afo)) = concatM $ case afo of
    Unambiguous name _ ->
      [ toHie $ C (RecField c rhs) $ L nspan $ removeDefSrcSpan name
      ]
    Ambiguous _name _ ->
      [ ]
    XAmbiguousFieldOcc _ -> []

instance ToHie (RFContext (Located (AmbiguousFieldOcc GhcTc))) where
  toHie (RFC c rhs (L nspan afo)) = concatM $ case afo of
    Unambiguous var _ ->
      let var' = setVarName var (removeDefSrcSpan $ varName var)
      in [ toHie $ C (RecField c rhs) (L nspan var')
         ]
    Ambiguous var _ ->
      let var' = setVarName var (removeDefSrcSpan $ varName var)
      in [ toHie $ C (RecField c rhs) (L nspan var')
         ]
    XAmbiguousFieldOcc _ -> []

instance ( a ~ GhcPass p
         , ToHie (PScoped (LPat a))
         , ToHie (BindContext (LHsBind a))
         , ToHie (LHsExpr a)
         , ToHie (SigContext (LSig a))
         , ToHie (RScoped (HsValBindsLR a a))
         , Data (StmtLR a a (Located (HsExpr a)))
         , Data (HsLocalBinds a)
         ) => ToHie (RScoped (ApplicativeArg (GhcPass p))) where
  toHie (RS sc (ApplicativeArgOne _ pat expr _)) = concatM
    [ toHie $ PS Nothing sc NoScope pat
    , toHie expr
    ]
  toHie (RS sc (ApplicativeArgMany _ stmts _ pat)) = concatM
    [ toHie $ listScopes NoScope stmts
    , toHie $ PS Nothing sc NoScope pat
    ]
  toHie (RS _ (XApplicativeArg _)) = pure []

instance (ToHie arg, ToHie rec) => ToHie (HsConDetails arg rec) where
  toHie (PrefixCon args) = toHie args
  toHie (RecCon rec) = toHie rec
  toHie (InfixCon a b) = concatM [ toHie a, toHie b]

instance ( ToHie (LHsCmd a)
         , Data  (HsCmdTop a)
         ) => ToHie (LHsCmdTop a) where
  toHie (L span top) = concatM $ makeNode top span : case top of
    HsCmdTop _ cmd ->
      [ toHie cmd
      ]
    XCmdTop _ -> []

instance ( a ~ GhcPass p
         , ToHie (PScoped (LPat a))
         , ToHie (BindContext (LHsBind a))
         , ToHie (LHsExpr a)
         , ToHie (MatchGroup a (LHsCmd a))
         , ToHie (SigContext (LSig a))
         , ToHie (RScoped (HsValBindsLR a a))
         , Data (HsCmd a)
         , Data (HsCmdTop a)
         , Data (StmtLR a a (Located (HsCmd a)))
         , Data (HsLocalBinds a)
         , Data (StmtLR a a (Located (HsExpr a)))
         ) => ToHie (LHsCmd (GhcPass p)) where
  toHie (L span cmd) = concatM $ makeNode cmd span : case cmd of
      HsCmdArrApp _ a b _ _ ->
        [ toHie a
        , toHie b
        ]
      HsCmdArrForm _ a _ _ cmdtops ->
        [ toHie a
        , toHie cmdtops
        ]
      HsCmdApp _ a b ->
        [ toHie a
        , toHie b
        ]
      HsCmdLam _ mg ->
        [ toHie mg
        ]
      HsCmdPar _ a ->
        [ toHie a
        ]
      HsCmdCase _ expr alts ->
        [ toHie expr
        , toHie alts
        ]
      HsCmdIf _ _ a b c ->
        [ toHie a
        , toHie b
        , toHie c
        ]
      HsCmdLet _ binds cmd' ->
        [ toHie $ RS (mkLScope cmd') binds
        , toHie cmd'
        ]
      HsCmdDo _ (L ispan stmts) ->
        [ pure $ locOnly ispan
        , toHie $ listScopes NoScope stmts
        ]
      HsCmdWrap _ _ _ -> []
      XCmd _ -> []

instance ToHie (TyClGroup GhcRn) where
  toHie (TyClGroup _ classes roles instances) = concatM
    [ toHie classes
    , toHie roles
    , toHie instances
    ]
  toHie (XTyClGroup _) = pure []

instance ToHie (LTyClDecl GhcRn) where
  toHie (L span decl) = concatM $ makeNode decl span : case decl of
      FamDecl {tcdFam = fdecl} ->
        [ toHie (L span fdecl)
        ]
      SynDecl {tcdLName = name, tcdTyVars = vars, tcdRhs = typ} ->
        [ toHie $ C (Decl SynDec $ getRealSpan span) name
        , toHie $ TS (ResolvedScopes [mkScope $ getLoc typ]) vars
        , toHie typ
        ]
      DataDecl {tcdLName = name, tcdTyVars = vars, tcdDataDefn = defn} ->
        [ toHie $ C (Decl DataDec $ getRealSpan span) name
        , toHie $ TS (ResolvedScopes [quant_scope, rhs_scope]) vars
        , toHie defn
        ]
        where
          quant_scope = mkLScope $ dd_ctxt defn
          rhs_scope = sig_sc `combineScopes` con_sc `combineScopes` deriv_sc
          sig_sc = maybe NoScope mkLScope $ dd_kindSig defn
          con_sc = foldr combineScopes NoScope $ map mkLScope $ dd_cons defn
          deriv_sc = mkLScope $ dd_derivs defn
      ClassDecl { tcdCtxt = context
                , tcdLName = name
                , tcdTyVars = vars
                , tcdFDs = deps
                , tcdSigs = sigs
                , tcdMeths = meths
                , tcdATs = typs
                , tcdATDefs = deftyps
                } ->
        [ toHie $ C (Decl ClassDec $ getRealSpan span) name
        , toHie context
        , toHie $ TS (ResolvedScopes [context_scope, rhs_scope]) vars
        , toHie deps
        , toHie $ map (SC $ SI ClassSig $ getRealSpan span) sigs
        , toHie $ fmap (BC InstanceBind ModuleScope) meths
        , toHie typs
        , concatMapM (pure . locOnly . getLoc) deftyps
        , toHie $ map (go . unLoc) deftyps
        ]
        where
          context_scope = mkLScope context
          rhs_scope = foldl1' combineScopes $ map mkScope
            [ loc deps, loc sigs, loc (bagToList meths), loc typs, loc deftyps]

          go :: TyFamDefltEqn GhcRn
             -> FamEqn GhcRn (TScoped (LHsQTyVars GhcRn)) (LHsType GhcRn)
          go (FamEqn a var bndrs pat b rhs) =
             FamEqn a var bndrs (TS (ResolvedScopes [mkLScope rhs]) pat) b rhs
          go (XFamEqn NoExt) = XFamEqn NoExt
      XTyClDecl _ -> []

instance ToHie (LFamilyDecl GhcRn) where
  toHie (L span decl) = concatM $ makeNode decl span : case decl of
      FamilyDecl _ info name vars _ sig inj ->
        [ toHie $ C (Decl FamDec $ getRealSpan span) name
        , toHie $ TS (ResolvedScopes [rhsSpan]) vars
        , toHie info
        , toHie $ RS injSpan sig
        , toHie inj
        ]
        where
          rhsSpan = sigSpan `combineScopes` injSpan
          sigSpan = mkScope $ getLoc sig
          injSpan = maybe NoScope (mkScope . getLoc) inj
      XFamilyDecl _ -> []

instance ToHie (FamilyInfo GhcRn) where
  toHie (ClosedTypeFamily (Just eqns)) = concatM $
    [ concatMapM (pure . locOnly . getLoc) eqns
    , toHie $ map go eqns
    ]
    where
      go (L l ib) = TS (ResolvedScopes [mkScope l]) ib
  toHie _ = pure []

instance ToHie (RScoped (LFamilyResultSig GhcRn)) where
  toHie (RS sc (L span sig)) = concatM $ makeNode sig span : case sig of
      NoSig _ ->
        []
      KindSig _ k ->
        [ toHie k
        ]
      TyVarSig _ bndr ->
        [ toHie $ TVS (ResolvedScopes [sc]) NoScope bndr
        ]
      XFamilyResultSig _ -> []

instance ToHie (Located (FunDep (Located Name))) where
  toHie (L span fd@(lhs, rhs)) = concatM $
    [ makeNode fd span
    , toHie $ map (C Use) lhs
    , toHie $ map (C Use) rhs
    ]

instance (ToHie pats, ToHie rhs, HasLoc pats, HasLoc rhs)
    => ToHie (TScoped (FamEqn GhcRn pats rhs)) where
  toHie (TS _ f) = toHie f

instance ( ToHie pats
         , ToHie rhs
         , HasLoc pats
         , HasLoc rhs
         ) => ToHie (FamEqn GhcRn pats rhs) where
  toHie fe@(FamEqn _ var tybndrs pats _ rhs) = concatM $
    [ toHie $ C (Decl InstDec $ getRealSpan $ loc fe) var
    , toHie $ fmap (tvScopes (ResolvedScopes []) scope) tybndrs
    , toHie pats
    , toHie rhs
    ]
    where scope = combineScopes patsScope rhsScope
          patsScope = mkScope (loc pats)
          rhsScope = mkScope (loc rhs)
  toHie (XFamEqn _) = pure []

instance ToHie (LInjectivityAnn GhcRn) where
  toHie (L span ann) = concatM $ makeNode ann span : case ann of
      InjectivityAnn lhs rhs ->
        [ toHie $ C Use lhs
        , toHie $ map (C Use) rhs
        ]

instance ToHie (HsDataDefn GhcRn) where
  toHie (HsDataDefn _ _ ctx _ mkind cons derivs) = concatM
    [ toHie ctx
    , toHie mkind
    , toHie cons
    , toHie derivs
    ]
  toHie (XHsDataDefn _) = pure []

instance ToHie (HsDeriving GhcRn) where
  toHie (L span clauses) = concatM
    [ pure $ locOnly span
    , toHie clauses
    ]

instance ToHie (LHsDerivingClause GhcRn) where
  toHie (L span cl) = concatM $ makeNode cl span : case cl of
      HsDerivingClause _ strat (L ispan tys) ->
        [ toHie strat
        , pure $ locOnly ispan
        , toHie $ map (TS (ResolvedScopes [])) tys
        ]
      XHsDerivingClause _ -> []

instance ToHie (Located (DerivStrategy GhcRn)) where
  toHie (L span strat) = concatM $ makeNode strat span : case strat of
      StockStrategy -> []
      AnyclassStrategy -> []
      NewtypeStrategy -> []
      ViaStrategy s -> [ toHie $ TS (ResolvedScopes []) s ]

instance ToHie (Located OverlapMode) where
  toHie (L span _) = pure $ locOnly span

instance ToHie (LConDecl GhcRn) where
  toHie (L span decl) = concatM $ makeNode decl span : case decl of
      ConDeclGADT { con_names = names, con_qvars = qvars
                  , con_mb_cxt = ctx, con_args = args, con_res_ty = typ } ->
        [ toHie $ map (C (Decl ConDec $ getRealSpan span)) names
        , toHie $ TS (ResolvedScopes [ctxScope, rhsScope]) qvars
        , toHie ctx
        , toHie args
        , toHie typ
        ]
        where
          rhsScope = combineScopes argsScope tyScope
          ctxScope = maybe NoScope mkLScope ctx
          argsScope = condecl_scope args
          tyScope = mkLScope typ
      ConDeclH98 { con_name = name, con_ex_tvs = qvars
                 , con_mb_cxt = ctx, con_args = dets } ->
        [ toHie $ C (Decl ConDec $ getRealSpan span) name
        , toHie $ tvScopes (ResolvedScopes []) rhsScope qvars
        , toHie ctx
        , toHie dets
        ]
        where
          rhsScope = combineScopes ctxScope argsScope
          ctxScope = maybe NoScope mkLScope ctx
          argsScope = condecl_scope dets
      XConDecl _ -> []
    where condecl_scope args = case args of
            PrefixCon xs -> foldr combineScopes NoScope $ map mkLScope xs
            InfixCon a b -> combineScopes (mkLScope a) (mkLScope b)
            RecCon x -> mkLScope x

instance ToHie (Located [LConDeclField GhcRn]) where
  toHie (L span decls) = concatM $
    [ pure $ locOnly span
    , toHie decls
    ]

instance ( HasLoc thing
         , ToHie (TScoped thing)
         ) => ToHie (TScoped (HsImplicitBndrs GhcRn thing)) where
  toHie (TS sc (HsIB ibrn a)) = concatM $
      [ pure $ bindingsOnly $ map (C $ TyVarBind (mkScope span) sc) ibrn
      , toHie $ TS sc a
      ]
    where span = loc a
  toHie (TS _ (XHsImplicitBndrs _)) = pure []

instance ( HasLoc thing
         , ToHie (TScoped thing)
         ) => ToHie (TScoped (HsWildCardBndrs GhcRn thing)) where
  toHie (TS sc (HsWC names a)) = concatM $
      [ pure $ bindingsOnly $ map (C $ TyVarBind (mkScope span) sc) names
      , toHie $ TS sc a
      ]
    where span = loc a
  toHie (TS _ (XHsWildCardBndrs _)) = pure []

instance ToHie (SigContext (LSig GhcRn)) where
  toHie (SC (SI styp msp) (L sp sig)) = concatM $ makeNode sig sp : case sig of
      TypeSig _ names typ ->
        [ toHie $ map (C TyDecl) names
        , toHie $ TS (UnresolvedScope (map unLoc names) Nothing) typ
        ]
      PatSynSig _ names typ ->
        [ toHie $ map (C TyDecl) names
        , toHie $ TS (UnresolvedScope (map unLoc names) Nothing) typ
        ]
      ClassOpSig _ _ names typ ->
        [ case styp of
            ClassSig -> toHie $ map (C $ ClassTyDecl $ getRealSpan sp) names
            _  -> toHie $ map (C $ TyDecl) names
        , toHie $ TS (UnresolvedScope (map unLoc names) msp) typ
        ]
      IdSig _ _ -> []
      FixSig _ fsig ->
        [ toHie $ L sp fsig
        ]
      InlineSig _ name _ ->
        [ toHie $ (C Use) name
        ]
      SpecSig _ name typs _ ->
        [ toHie $ (C Use) name
        , toHie $ map (TS (ResolvedScopes [])) typs
        ]
      SpecInstSig _ _ typ ->
        [ toHie $ TS (ResolvedScopes []) typ
        ]
      MinimalSig _ _ form ->
        [ toHie form
        ]
      SCCFunSig _ _ name mtxt ->
        [ toHie $ (C Use) name
        , pure $ maybe [] (locOnly . getLoc) mtxt
        ]
      CompleteMatchSig _ _ (L ispan names) typ ->
        [ pure $ locOnly ispan
        , toHie $ map (C Use) names
        , toHie $ fmap (C Use) typ
        ]
      XSig _ -> []

instance ToHie (LHsType GhcRn) where
  toHie x = toHie $ TS (ResolvedScopes []) x

instance ToHie (TScoped (LHsType GhcRn)) where
  toHie (TS tsc (L span t)) = concatM $ makeNode t span : case t of
      HsForAllTy _ _ bndrs body ->
        [ toHie $ tvScopes tsc (mkScope $ getLoc body) bndrs
        , toHie body
        ]
      HsQualTy _ ctx body ->
        [ toHie ctx
        , toHie body
        ]
      HsTyVar _ _ var ->
        [ toHie $ C Use var
        ]
      HsAppTy _ a b ->
        [ toHie a
        , toHie b
        ]
      HsAppKindTy _ ty ki ->
        [ toHie ty
        , toHie $ TS (ResolvedScopes []) ki
        ]
      HsFunTy _ a b ->
        [ toHie a
        , toHie b
        ]
      HsListTy _ a ->
        [ toHie a
        ]
      HsTupleTy _ _ tys ->
        [ toHie tys
        ]
      HsSumTy _ tys ->
        [ toHie tys
        ]
      HsOpTy _ a op b ->
        [ toHie a
        , toHie $ C Use op
        , toHie b
        ]
      HsParTy _ a ->
        [ toHie a
        ]
      HsIParamTy _ ip ty ->
        [ toHie ip
        , toHie ty
        ]
      HsKindSig _ a b ->
        [ toHie a
        , toHie b
        ]
      HsSpliceTy _ a ->
        [ toHie $ L span a
        ]
      HsDocTy _ a _ ->
        [ toHie a
        ]
      HsBangTy _ _ ty ->
        [ toHie ty
        ]
      HsRecTy _ fields ->
        [ toHie fields
        ]
      HsExplicitListTy _ _ tys ->
        [ toHie tys
        ]
      HsExplicitTupleTy _ tys ->
        [ toHie tys
        ]
      HsTyLit _ _ -> []
      HsWildCardTy _ -> []
      HsStarTy _ _ -> []
      XHsType _ -> []

instance (ToHie tm, ToHie ty) => ToHie (HsArg tm ty) where
  toHie (HsValArg tm) = toHie tm
  toHie (HsTypeArg _ ty) = toHie ty
  toHie (HsArgPar sp) = pure $ locOnly sp

instance ToHie (TVScoped (LHsTyVarBndr GhcRn)) where
  toHie (TVS tsc sc (L span bndr)) = concatM $ makeNode bndr span : case bndr of
      UserTyVar _ var ->
        [ toHie $ C (TyVarBind sc tsc) var
        ]
      KindedTyVar _ var kind ->
        [ toHie $ C (TyVarBind sc tsc) var
        , toHie kind
        ]
      XTyVarBndr _ -> []

instance ToHie (TScoped (LHsQTyVars GhcRn)) where
  toHie (TS sc (HsQTvs (HsQTvsRn implicits _) vars)) = concatM $
    [ pure $ bindingsOnly bindings
    , toHie $ tvScopes sc NoScope vars
    ]
    where
      varLoc = loc vars
      bindings = map (C $ TyVarBind (mkScope varLoc) sc) implicits
  toHie (TS _ (XLHsQTyVars _)) = pure []

instance ToHie (LHsContext GhcRn) where
  toHie (L span tys) = concatM $
      [ pure $ locOnly span
      , toHie tys
      ]

instance ToHie (LConDeclField GhcRn) where
  toHie (L span field) = concatM $ makeNode field span : case field of
      ConDeclField _ fields typ _ ->
        [ toHie $ map (RFC RecFieldDecl (getRealSpan $ loc typ)) fields
        , toHie typ
        ]
      XConDeclField _ -> []

instance ToHie (LHsExpr a) => ToHie (ArithSeqInfo a) where
  toHie (From expr) = toHie expr
  toHie (FromThen a b) = concatM $
    [ toHie a
    , toHie b
    ]
  toHie (FromTo a b) = concatM $
    [ toHie a
    , toHie b
    ]
  toHie (FromThenTo a b c) = concatM $
    [ toHie a
    , toHie b
    , toHie c
    ]

instance ToHie (LSpliceDecl GhcRn) where
  toHie (L span decl) = concatM $ makeNode decl span : case decl of
      SpliceDecl _ splice _ ->
        [ toHie splice
        ]
      XSpliceDecl _ -> []

instance ToHie (HsBracket a) where
  toHie _ = pure []

instance ToHie PendingRnSplice where
  toHie _ = pure []

instance ToHie PendingTcSplice where
  toHie _ = pure []

instance ToHie (LBooleanFormula (Located Name)) where
  toHie (L span form) = concatM $ makeNode form span : case form of
      Var a ->
        [ toHie $ C Use a
        ]
      And forms ->
        [ toHie forms
        ]
      Or forms ->
        [ toHie forms
        ]
      Parens f ->
        [ toHie f
        ]

instance ToHie (Located HsIPName) where
  toHie (L span e) = makeNode e span

instance ( ToHie (LHsExpr a)
         , Data (HsSplice a)
         ) => ToHie (Located (HsSplice a)) where
  toHie (L span sp) = concatM $ makeNode sp span : case sp of
      HsTypedSplice _ _ _ expr ->
        [ toHie expr
        ]
      HsUntypedSplice _ _ _ expr ->
        [ toHie expr
        ]
      HsQuasiQuote _ _ _ ispan _ ->
        [ pure $ locOnly ispan
        ]
      HsSpliced _ _ _ ->
        []
      HsSplicedT _ ->
        []
      XSplice _ -> []

instance ToHie (LRoleAnnotDecl GhcRn) where
  toHie (L span annot) = concatM $ makeNode annot span : case annot of
      RoleAnnotDecl _ var roles ->
        [ toHie $ C Use var
        , concatMapM (pure . locOnly . getLoc) roles
        ]
      XRoleAnnotDecl _ -> []

instance ToHie (LInstDecl GhcRn) where
  toHie (L span decl) = concatM $ makeNode decl span : case decl of
      ClsInstD _ d ->
        [ toHie $ L span d
        ]
      DataFamInstD _ d ->
        [ toHie $ L span d
        ]
      TyFamInstD _ d ->
        [ toHie $ L span d
        ]
      XInstDecl _ -> []

instance ToHie (LClsInstDecl GhcRn) where
  toHie (L span decl) = concatM
    [ toHie $ TS (ResolvedScopes [mkScope span]) $ cid_poly_ty decl
    , toHie $ fmap (BC InstanceBind ModuleScope) $ cid_binds decl
    , toHie $ map (SC $ SI InstSig $ getRealSpan span) $ cid_sigs decl
    , pure $ concatMap (locOnly . getLoc) $ cid_tyfam_insts decl
    , toHie $ cid_tyfam_insts decl
    , pure $ concatMap (locOnly . getLoc) $ cid_datafam_insts decl
    , toHie $ cid_datafam_insts decl
    , toHie $ cid_overlap_mode decl
    ]

instance ToHie (LDataFamInstDecl GhcRn) where
  toHie (L sp (DataFamInstDecl d)) = toHie $ TS (ResolvedScopes [mkScope sp]) d

instance ToHie (LTyFamInstDecl GhcRn) where
  toHie (L sp (TyFamInstDecl d)) = toHie $ TS (ResolvedScopes [mkScope sp]) d

instance ToHie (Context a)
         => ToHie (PatSynFieldContext (RecordPatSynField a)) where
  toHie (PSC sp (RecordPatSynField a b)) = concatM $
    [ toHie $ C (RecField RecFieldDecl sp) a
    , toHie $ C Use b
    ]

instance ToHie (LDerivDecl GhcRn) where
  toHie (L span decl) = concatM $ makeNode decl span : case decl of
      DerivDecl _ typ strat overlap ->
        [ toHie $ TS (ResolvedScopes []) typ
        , toHie strat
        , toHie overlap
        ]
      XDerivDecl _ -> []

instance ToHie (LFixitySig GhcRn) where
  toHie (L span sig) = concatM $ makeNode sig span : case sig of
      FixitySig _ vars _ ->
        [ toHie $ map (C Use) vars
        ]
      XFixitySig _ -> []

instance ToHie (LDefaultDecl GhcRn) where
  toHie (L span decl) = concatM $ makeNode decl span : case decl of
      DefaultDecl _ typs ->
        [ toHie typs
        ]
      XDefaultDecl _ -> []

instance ToHie (LForeignDecl GhcRn) where
  toHie (L span decl) = concatM $ makeNode decl span : case decl of
      ForeignImport {fd_name = name, fd_sig_ty = sig, fd_fi = fi} ->
        [ toHie $ C (ValBind RegularBind ModuleScope $ getRealSpan span) name
        , toHie $ TS (ResolvedScopes []) sig
        , toHie fi
        ]
      ForeignExport {fd_name = name, fd_sig_ty = sig, fd_fe = fe} ->
        [ toHie $ C Use name
        , toHie $ TS (ResolvedScopes []) sig
        , toHie fe
        ]
      XForeignDecl _ -> []

instance ToHie ForeignImport where
  toHie (CImport (L a _) (L b _) _ _ (L c _)) = pure $ concat $
    [ locOnly a
    , locOnly b
    , locOnly c
    ]

instance ToHie ForeignExport where
  toHie (CExport (L a _) (L b _)) = pure $ concat $
    [ locOnly a
    , locOnly b
    ]

instance ToHie (LWarnDecls GhcRn) where
  toHie (L span decl) = concatM $ makeNode decl span : case decl of
      Warnings _ _ warnings ->
        [ toHie warnings
        ]
      XWarnDecls _ -> []

instance ToHie (LWarnDecl GhcRn) where
  toHie (L span decl) = concatM $ makeNode decl span : case decl of
      Warning _ vars _ ->
        [ toHie $ map (C Use) vars
        ]
      XWarnDecl _ -> []

instance ToHie (LAnnDecl GhcRn) where
  toHie (L span decl) = concatM $ makeNode decl span : case decl of
      HsAnnotation _ _ prov expr ->
        [ toHie prov
        , toHie expr
        ]
      XAnnDecl _ -> []

instance ToHie (Context (Located a)) => ToHie (AnnProvenance a) where
  toHie (ValueAnnProvenance a) = toHie $ C Use a
  toHie (TypeAnnProvenance a) = toHie $ C Use a
  toHie ModuleAnnProvenance = pure []

instance ToHie (LRuleDecls GhcRn) where
  toHie (L span decl) = concatM $ makeNode decl span : case decl of
      HsRules _ _ rules ->
        [ toHie rules
        ]
      XRuleDecls _ -> []

instance ToHie (LRuleDecl GhcRn) where
  toHie (L _ (XRuleDecl _)) = pure []
  toHie (L span r@(HsRule _ rname _ tybndrs bndrs exprA exprB)) = concatM
        [ makeNode r span
        , pure $ locOnly $ getLoc rname
        , toHie $ fmap (tvScopes (ResolvedScopes []) scope) tybndrs
        , toHie $ map (RS $ mkScope span) bndrs
        , toHie exprA
        , toHie exprB
        ]
    where scope = bndrs_sc `combineScopes` exprA_sc `combineScopes` exprB_sc
          bndrs_sc = maybe NoScope mkLScope (listToMaybe bndrs)
          exprA_sc = mkLScope exprA
          exprB_sc = mkLScope exprB

instance ToHie (RScoped (LRuleBndr GhcRn)) where
  toHie (RS sc (L span bndr)) = concatM $ makeNode bndr span : case bndr of
      RuleBndr _ var ->
        [ toHie $ C (ValBind RegularBind sc Nothing) var
        ]
      RuleBndrSig _ var typ ->
        [ toHie $ C (ValBind RegularBind sc Nothing) var
        , toHie $ TS (ResolvedScopes [sc]) typ
        ]
      XRuleBndr _ -> []

instance ToHie (LImportDecl GhcRn) where
  toHie (L span decl) = concatM $ makeNode decl span : case decl of
      ImportDecl { ideclName = name, ideclAs = as, ideclHiding = hidden } ->
        [ toHie $ IEC Import name
        , toHie $ fmap (IEC ImportAs) as
        , maybe (pure []) goIE hidden
        ]
      XImportDecl _ -> []
    where
      goIE (hiding, (L sp liens)) = concatM $
        [ pure $ locOnly sp
        , toHie $ map (IEC c) liens
        ]
        where
         c = if hiding then ImportHiding else Import

instance ToHie (IEContext (LIE GhcRn)) where
  toHie (IEC c (L span ie)) = concatM $ makeNode ie span : case ie of
      IEVar _ n ->
        [ toHie $ IEC c n
        ]
      IEThingAbs _ n ->
        [ toHie $ IEC c n
        ]
      IEThingAll _ n ->
        [ toHie $ IEC c n
        ]
      IEThingWith _ n _ ns flds ->
        [ toHie $ IEC c n
        , toHie $ map (IEC c) ns
        , toHie $ map (IEC c) flds
        ]
      IEModuleContents _ n ->
        [ toHie $ IEC c n
        ]
      IEGroup _ _ _ -> []
      IEDoc _ _ -> []
      IEDocNamed _ _ -> []
      XIE _ -> []

instance ToHie (IEContext (LIEWrappedName Name)) where
  toHie (IEC c (L span iewn)) = concatM $ makeNode iewn span : case iewn of
      IEName n ->
        [ toHie $ C (IEThing c) n
        ]
      IEPattern p ->
        [ toHie $ C (IEThing c) p
        ]
      IEType n ->
        [ toHie $ C (IEThing c) n
        ]

instance ToHie (IEContext (Located (FieldLbl Name))) where
  toHie (IEC c (L span lbl)) = concatM $ makeNode lbl span : case lbl of
      FieldLabel _ _ n ->
        [ toHie $ C (IEThing c) $ L span n
        ]


{-# LANGUAGE ScopedTypeVariables #-}
module HieBin ( readHieFile, writeHieFile, HieName(..), toHieName ) where

import GhcPrelude

import Binary
import BinIface                   ( getDictFastString )
import FastMutInt
import FastString                 ( FastString )
import Module                     ( Module )
import Name
import NameCache
import Outputable
import PrelInfo
import SrcLoc
import UniqSupply                 ( takeUniqFromSupply )
import Unique
import UniqFM

import qualified Data.Array as A
import Data.IORef
import Data.List                  ( mapAccumR )
import Data.Word                  ( Word32 )
import Control.Monad              ( replicateM )
import System.Directory           ( createDirectoryIfMissing )
import System.FilePath            ( takeDirectory )

-- | `Name`'s get converted into `HieName`'s before being written into @.hie@
-- files. See 'toHieName' and 'fromHieName' for logic on how to convert between
-- these two types.
data HieName
  = ExternalName !Module !OccName !SrcSpan
  | LocalName !OccName !SrcSpan
  | KnownKeyName !Unique
  deriving (Eq)

instance Ord HieName where
  compare (ExternalName a b c) (ExternalName d e f) = compare (a,b,c) (d,e,f)
  compare (LocalName a b) (LocalName c d) = compare (a,b) (c,d)
  compare (KnownKeyName a) (KnownKeyName b) = nonDetCmpUnique a b
    -- Not actually non determinstic as it is a KnownKey
  compare ExternalName{} _ = LT
  compare LocalName{} ExternalName{} = GT
  compare LocalName{} _ = LT
  compare KnownKeyName{} _ = GT

instance Outputable HieName where
  ppr (ExternalName m n sp) = text "ExternalName" <+> ppr m <+> ppr n <+> ppr sp
  ppr (LocalName n sp) = text "LocalName" <+> ppr n <+> ppr sp
  ppr (KnownKeyName u) = text "KnownKeyName" <+> ppr u

data HieSymbolTable = HieSymbolTable
  { hie_symtab_next :: !FastMutInt
  , hie_symtab_map  :: !(IORef (UniqFM (Int, HieName)))
  }

data HieDictionary = HieDictionary
  { hie_dict_next :: !FastMutInt -- The next index to use
  , hie_dict_map  :: !(IORef (UniqFM (Int,FastString))) -- indexed by FastString
  }

initBinMemSize :: Int
initBinMemSize = 1024*1024

writeHieFile :: Binary a => FilePath -> a -> IO ()
writeHieFile hie_file_path hiefile = do
  bh0 <- openBinMem initBinMemSize

  -- remember where the dictionary pointer will go
  dict_p_p <- tellBin bh0
  put_ bh0 dict_p_p

  -- remember where the symbol table pointer will go
  symtab_p_p <- tellBin bh0
  put_ bh0 symtab_p_p

  -- Make some intial state
  symtab_next <- newFastMutInt
  writeFastMutInt symtab_next 0
  symtab_map <- newIORef emptyUFM
  let hie_symtab = HieSymbolTable {
                      hie_symtab_next = symtab_next,
                      hie_symtab_map  = symtab_map }
  dict_next_ref <- newFastMutInt
  writeFastMutInt dict_next_ref 0
  dict_map_ref <- newIORef emptyUFM
  let hie_dict = HieDictionary {
                      hie_dict_next = dict_next_ref,
                      hie_dict_map  = dict_map_ref }

  -- put the main thing
  let bh = setUserData bh0 $ newWriteState (putName hie_symtab)
                                           (putName hie_symtab)
                                           (putFastString hie_dict)
  put_ bh hiefile

  -- write the symtab pointer at the front of the file
  symtab_p <- tellBin bh
  putAt bh symtab_p_p symtab_p
  seekBin bh symtab_p

  -- write the symbol table itself
  symtab_next' <- readFastMutInt symtab_next
  symtab_map'  <- readIORef symtab_map
  putSymbolTable bh symtab_next' symtab_map'

  -- write the dictionary pointer at the fornt of the file
  dict_p <- tellBin bh
  putAt bh dict_p_p dict_p
  seekBin bh dict_p

  -- write the dictionary itself
  dict_next <- readFastMutInt dict_next_ref
  dict_map  <- readIORef dict_map_ref
  putDictionary bh dict_next dict_map

  -- and send the result to the file
  createDirectoryIfMissing True (takeDirectory hie_file_path)
  writeBinMem bh hie_file_path
  return ()

readHieFile :: Binary a => NameCache -> FilePath -> IO (a, NameCache)
readHieFile nc file = do
  bh0 <- readBinMem file

  dict  <- get_dictionary bh0

  -- read the symbol table so we are capable of reading the actual data
  (bh1, nc') <- do
      let bh1 = setUserData bh0 $ newReadState (error "getSymtabName")
                                               (getDictFastString dict)
      (nc', symtab) <- get_symbol_table bh1
      let bh1' = setUserData bh1
               $ newReadState (getSymTabName symtab)
                              (getDictFastString dict)
      return (bh1', nc')

  -- load the actual data
  hiefile <- get bh1
  return (hiefile, nc')
  where
    get_dictionary bin_handle = do
      dict_p <- get bin_handle
      data_p <- tellBin bin_handle
      seekBin bin_handle dict_p
      dict <- getDictionary bin_handle
      seekBin bin_handle data_p
      return dict

    get_symbol_table bh1 = do
      symtab_p <- get bh1
      data_p'  <- tellBin bh1
      seekBin bh1 symtab_p
      (nc', symtab) <- getSymbolTable bh1 nc
      seekBin bh1 data_p'
      return (nc', symtab)

putFastString :: HieDictionary -> BinHandle -> FastString -> IO ()
putFastString HieDictionary { hie_dict_next = j_r,
                              hie_dict_map  = out_r}  bh f
  = do
    out <- readIORef out_r
    let unique = getUnique f
    case lookupUFM out unique of
        Just (j, _)  -> put_ bh (fromIntegral j :: Word32)
        Nothing -> do
           j <- readFastMutInt j_r
           put_ bh (fromIntegral j :: Word32)
           writeFastMutInt j_r (j + 1)
           writeIORef out_r $! addToUFM out unique (j, f)

putSymbolTable :: BinHandle -> Int -> UniqFM (Int,HieName) -> IO ()
putSymbolTable bh next_off symtab = do
  put_ bh next_off
  let names = A.elems (A.array (0,next_off-1) (nonDetEltsUFM symtab))
  mapM_ (putHieName bh) names

getSymbolTable :: BinHandle -> NameCache -> IO (NameCache, SymbolTable)
getSymbolTable bh namecache = do
  sz <- get bh
  od_names <- replicateM sz (getHieName bh)
  let arr = A.listArray (0,sz-1) names
      (namecache', names) = mapAccumR fromHieName namecache od_names
  return (namecache', arr)

getSymTabName :: SymbolTable -> BinHandle -> IO Name
getSymTabName st bh = do
  i :: Word32 <- get bh
  return $ st A.! (fromIntegral i)

putName :: HieSymbolTable -> BinHandle -> Name -> IO ()
putName (HieSymbolTable next ref) bh name = do
  symmap <- readIORef ref
  case lookupUFM symmap name of
    Just (off, ExternalName mod occ (UnhelpfulSpan _))
      | isGoodSrcSpan (nameSrcSpan name) -> do
      let hieName = ExternalName mod occ (nameSrcSpan name)
      writeIORef ref $! addToUFM symmap name (off, hieName)
      put_ bh (fromIntegral off :: Word32)
    Just (off, LocalName _occ span)
      | notLocal (toHieName name) || nameSrcSpan name /= span -> do
      writeIORef ref $! addToUFM symmap name (off, toHieName name)
      put_ bh (fromIntegral off :: Word32)
    Just (off, _) -> put_ bh (fromIntegral off :: Word32)
    Nothing -> do
        off <- readFastMutInt next
        writeFastMutInt next (off+1)
        writeIORef ref $! addToUFM symmap name (off, toHieName name)
        put_ bh (fromIntegral off :: Word32)

  where
    notLocal :: HieName -> Bool
    notLocal LocalName{} = False
    notLocal _ = True

-- ** Converting to and from `HieName`'s

toHieName :: Name -> HieName
toHieName name
  | isKnownKeyName name = KnownKeyName (nameUnique name)
  | isExternalName name = ExternalName (nameModule name)
                                       (nameOccName name)
                                       (nameSrcSpan name)
  | otherwise = LocalName (nameOccName name) (nameSrcSpan name)

fromHieName :: NameCache -> HieName -> (NameCache, Name)
fromHieName nc (ExternalName mod occ span) =
    let cache = nsNames nc
    in case lookupOrigNameCache cache mod occ of
         Just name -> (nc, name)
         Nothing ->
           let (uniq, us) = takeUniqFromSupply (nsUniqs nc)
               name       = mkExternalName uniq mod occ span
               new_cache  = extendNameCache cache mod occ name
           in ( nc{ nsUniqs = us, nsNames = new_cache }, name )
fromHieName nc (LocalName occ span) =
    let (uniq, us) = takeUniqFromSupply (nsUniqs nc)
        name       = mkInternalName uniq occ span
    in ( nc{ nsUniqs = us }, name )
fromHieName nc (KnownKeyName u) = case lookupKnownKeyName u of
    Nothing -> pprPanic "fromHieName:unknown known-key unique"
                        (ppr (unpkUnique u))
    Just n -> (nc, n)

-- ** Reading and writing `HieName`'s

putHieName :: BinHandle -> HieName -> IO ()
putHieName bh (ExternalName mod occ span) = do
  putByte bh 0
  put_ bh (mod, occ, span)
putHieName bh (LocalName occName span) = do
  putByte bh 1
  put_ bh (occName, span)
putHieName bh (KnownKeyName uniq) = do
  putByte bh 2
  put_ bh $ unpkUnique uniq

getHieName :: BinHandle -> IO HieName
getHieName bh = do
  t <- getByte bh
  case t of
    0 -> do
      (modu, occ, span) <- get bh
      return $ ExternalName modu occ span
    1 -> do
      (occ, span) <- get bh
      return $ LocalName occ span
    2 -> do
      (c,i) <- get bh
      return $ KnownKeyName $ mkUnique c i
    _ -> panic "HieBin.getHieName: invalid tag"


{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE FlexibleContexts #-}
module HieDebug where

import GhcPrelude

import SrcLoc
import Module
import FastString
import Outputable

import HieTypes
import HieBin
import HieUtils

import qualified Data.Map as M
import qualified Data.Set as S
import Data.Function    ( on )
import Data.List        ( sortOn )
import Data.Foldable    ( toList )

ppHies :: Outputable a => (HieASTs a) -> SDoc
ppHies (HieASTs asts) = M.foldrWithKey go "" asts
  where
    go k a rest = vcat $
      [ "File: " <> ppr k
      , ppHie a
      , rest
      ]

ppHie :: Outputable a => HieAST a -> SDoc
ppHie = go 0
  where
    go n (Node inf sp children) = hang header n rest
      where
        rest = vcat $ map (go (n+2)) children
        header = hsep
          [ "Node"
          , ppr sp
          , ppInfo inf
          ]

ppInfo :: Outputable a => NodeInfo a -> SDoc
ppInfo ni = hsep
  [ ppr $ toList $ nodeAnnotations ni
  , ppr $ nodeType ni
  , ppr $ M.toList $ nodeIdentifiers ni
  ]

type Diff a = a -> a -> [SDoc]

diffFile :: Diff HieFile
diffFile = diffAsts eqDiff `on` (getAsts . hie_asts)

diffAsts :: (Outputable a, Eq a) => Diff a -> Diff (M.Map FastString (HieAST a))
diffAsts f = diffList (diffAst f) `on` M.elems

diffAst :: (Outputable a, Eq a) => Diff a -> Diff (HieAST a)
diffAst diffType (Node info1 span1 xs1) (Node info2 span2 xs2) =
    infoDiff ++ spanDiff ++ diffList (diffAst diffType) xs1 xs2
  where
    spanDiff
      | span1 /= span2 = [hsep ["Spans", ppr span1, "and", ppr span2, "differ"]]
      | otherwise = []
    infoDiff
      = (diffList eqDiff `on` (S.toAscList . nodeAnnotations)) info1 info2
     ++ (diffList diffType `on` nodeType) info1 info2
     ++ (diffIdents `on` nodeIdentifiers) info1 info2
    diffIdents a b = (diffList diffIdent `on` normalizeIdents) a b
    diffIdent (a,b) (c,d) = diffName a c
                         ++ eqDiff b d
    diffName (Right a) (Right b) = case (a,b) of
      (ExternalName m o _, ExternalName m' o' _) -> eqDiff (m,o) (m',o')
      (LocalName o _, ExternalName _ o' _) -> eqDiff o o'
      _ -> eqDiff a b
    diffName a b = eqDiff a b

type DiffIdent = Either ModuleName HieName

normalizeIdents :: NodeIdentifiers a -> [(DiffIdent,IdentifierDetails a)]
normalizeIdents = sortOn fst . map (first toHieName) . M.toList
  where
    first f (a,b) = (fmap f a, b)

diffList :: Diff a -> Diff [a]
diffList f xs ys
  | length xs == length ys = concat $ zipWith f xs ys
  | otherwise = ["length of lists doesn't match"]

eqDiff :: (Outputable a, Eq a) => Diff a
eqDiff a b
  | a == b = []
  | otherwise = [hsep [ppr a, "and", ppr b, "do not match"]]

validAst :: HieAST a -> Either SDoc ()
validAst (Node _ span children) = do
  checkContainment children
  checkSorted children
  mapM_ validAst children
  where
    checkSorted [] = return ()
    checkSorted [_] = return ()
    checkSorted (x:y:xs)
      | nodeSpan x `leftOf` nodeSpan y = checkSorted (y:xs)
      | otherwise = Left $ hsep
          [ ppr $ nodeSpan x
          , "is not to the left of"
          , ppr $ nodeSpan y
          ]
    checkContainment [] = return ()
    checkContainment (x:xs)
      | span `containsSpan` (nodeSpan x) = checkContainment xs
      | otherwise = Left $ hsep
          [ ppr $ span
          , "does not contain"
          , ppr $ nodeSpan x
          ]

-- | Look for any identifiers which occur outside of their supposed scopes.
-- Returns a list of error messages.
validateScopes :: M.Map FastString (HieAST a) -> [SDoc]
validateScopes asts = M.foldrWithKey (\k a b -> valid k a ++ b) [] refMap
  where
    refMap = generateReferencesMap asts
    valid (Left _) _ = []
    valid (Right n) refs = concatMap inScope refs
      where
        mapRef = foldMap getScopeFromContext . identInfo . snd
        scopes = case foldMap mapRef refs of
          Just xs -> xs
          Nothing -> []
        inScope (sp, dets)
          |  definedInAsts asts n
          && any isOccurrence (identInfo dets)
            = case scopes of
              [] -> []
              _ -> if any (`scopeContainsSpan` sp) scopes
                   then []
                   else return $ hsep $
                     [ "Name", ppr n, "at position", ppr sp
                     , "doesn't occur in calculated scope", ppr scopes]
          | otherwise = []


{-# LANGUAGE DeriveTraversable #-}
{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE TypeSynonymInstances #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE ScopedTypeVariables #-}
module HieTypes where

import GhcPrelude

import Binary
import FastString                 ( FastString )
import IfaceType
import Module                     ( ModuleName, Module )
import Name                       ( Name )
import Outputable hiding ( (<>) )
import SrcLoc                     ( RealSrcSpan )
import Avail

import qualified Data.Array as A
import qualified Data.Map as M
import qualified Data.Set as S
import Data.ByteString            ( ByteString )
import Data.Data                  ( Typeable, Data )
import Data.Semigroup             ( Semigroup(..) )
import Data.Word                  ( Word8 )
import Control.Applicative        ( (<|>) )

type Span = RealSrcSpan

-- | Current version of @.hie@ files
curHieVersion :: Word8
curHieVersion = 0

{- |
GHC builds up a wealth of information about Haskell source as it compiles it.
@.hie@ files are a way of persisting some of this information to disk so that
external tools that need to work with haskell source don't need to parse,
typecheck, and rename all over again. These files contain:

  * a simplified AST

       * nodes are annotated with source positions and types
       * identifiers are annotated with scope information

  * the raw bytes of the initial Haskell source

Besides saving compilation cycles, @.hie@ files also offer a more stable
interface than the GHC API.
-}
data HieFile = HieFile
    { hie_version :: Word8
    -- ^ version of the HIE format

    , hie_ghc_version :: ByteString
    -- ^ Version of GHC that produced this file

    , hie_hs_file :: FilePath
    -- ^ Initial Haskell source file path

    , hie_module :: Module
    -- ^ The module this HIE file is for

    , hie_types :: A.Array TypeIndex HieTypeFlat
    -- ^ Types referenced in the 'hie_asts'.
    --
    -- See Note [Efficient serialization of redundant type info]

    , hie_asts :: HieASTs TypeIndex
    -- ^ Type-annotated abstract syntax trees

    , hie_exports :: [AvailInfo]
    -- ^ The names that this module exports

    , hie_hs_src :: ByteString
    -- ^ Raw bytes of the initial Haskell source
    }

instance Binary HieFile where
  put_ bh hf = do
    put_ bh $ hie_version hf
    put_ bh $ hie_ghc_version hf
    put_ bh $ hie_hs_file hf
    put_ bh $ hie_module hf
    put_ bh $ hie_types hf
    put_ bh $ hie_asts hf
    put_ bh $ hie_exports hf
    put_ bh $ hie_hs_src hf

  get bh = HieFile
    <$> get bh
    <*> get bh
    <*> get bh
    <*> get bh
    <*> get bh
    <*> get bh
    <*> get bh
    <*> get bh

{-
Note [Efficient serialization of redundant type info]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

The type information in .hie files is highly repetitive and redundant. For
example, consider the expression

    const True 'a'

There is a lot of shared structure between the types of subterms:

  * const True 'a' ::                 Bool
  * const True     ::         Char -> Bool
  * const          :: Bool -> Char -> Bool

Since all 3 of these types need to be stored in the .hie file, it is worth
making an effort to deduplicate this shared structure. The trick is to define
a new data type that is a flattened version of 'Type':

    data HieType a = HAppTy a a  -- data Type = AppTy Type Type
                   | HFunTy a a  --           | FunTy Type Type
                   | ...

    type TypeIndex = Int

Types in the final AST are stored in an 'A.Array TypeIndex (HieType TypeIndex)',
where the 'TypeIndex's in the 'HieType' are references to other elements of the
array. Types recovered from GHC are deduplicated and stored in this compressed
form with sharing of subtrees.
-}

type TypeIndex = Int

-- | A flattened version of 'Type'.
--
-- See Note [Efficient serialization of redundant type info]
data HieType a
  = HTyVarTy Name
  | HAppTy a (HieArgs a)
  | HTyConApp IfaceTyCon (HieArgs a)
  | HForAllTy ((Name, a),ArgFlag) a
  | HFunTy  a a
  | HQualTy a a           -- ^ type with constraint: @t1 => t2@ (see 'IfaceDFunTy')
  | HLitTy IfaceTyLit
  | HCastTy a
  | HCoercionTy
    deriving (Functor, Foldable, Traversable, Eq)

type HieTypeFlat = HieType TypeIndex

-- | Roughly isomorphic to the original core 'Type'.
newtype HieTypeFix = Roll (HieType (HieTypeFix))

instance Binary (HieType TypeIndex) where
  put_ bh (HTyVarTy n) = do
    putByte bh 0
    put_ bh n
  put_ bh (HAppTy a b) = do
    putByte bh 1
    put_ bh a
    put_ bh b
  put_ bh (HTyConApp n xs) = do
    putByte bh 2
    put_ bh n
    put_ bh xs
  put_ bh (HForAllTy bndr a) = do
    putByte bh 3
    put_ bh bndr
    put_ bh a
  put_ bh (HFunTy a b) = do
    putByte bh 4
    put_ bh a
    put_ bh b
  put_ bh (HQualTy a b) = do
    putByte bh 5
    put_ bh a
    put_ bh b
  put_ bh (HLitTy l) = do
    putByte bh 6
    put_ bh l
  put_ bh (HCastTy a) = do
    putByte bh 7
    put_ bh a
  put_ bh (HCoercionTy) = putByte bh 8

  get bh = do
    (t :: Word8) <- get bh
    case t of
      0 -> HTyVarTy <$> get bh
      1 -> HAppTy <$> get bh <*> get bh
      2 -> HTyConApp <$> get bh <*> get bh
      3 -> HForAllTy <$> get bh <*> get bh
      4 -> HFunTy <$> get bh <*> get bh
      5 -> HQualTy <$> get bh <*> get bh
      6 -> HLitTy <$> get bh
      7 -> HCastTy <$> get bh
      8 -> return HCoercionTy
      _ -> panic "Binary (HieArgs Int): invalid tag"

-- | A list of type arguments along with their respective visibilities (ie. is
-- this an argument that would return 'True' for 'isVisibleArgFlag'?).
newtype HieArgs a = HieArgs [(Bool,a)]
  deriving (Functor, Foldable, Traversable, Eq)

instance Binary (HieArgs TypeIndex) where
  put_ bh (HieArgs xs) = put_ bh xs
  get bh = HieArgs <$> get bh

-- | Mapping from filepaths (represented using 'FastString') to the
-- corresponding AST
newtype HieASTs a = HieASTs { getAsts :: (M.Map FastString (HieAST a)) }
  deriving (Functor, Foldable, Traversable)

instance Binary (HieASTs TypeIndex) where
  put_ bh asts = put_ bh $ M.toAscList $ getAsts asts
  get bh = HieASTs <$> fmap M.fromDistinctAscList (get bh)

data HieAST a =
  Node
    { nodeInfo :: NodeInfo a
    , nodeSpan :: Span
    , nodeChildren :: [HieAST a]
    } deriving (Functor, Foldable, Traversable)

instance Binary (HieAST TypeIndex) where
  put_ bh ast = do
    put_ bh $ nodeInfo ast
    put_ bh $ nodeSpan ast
    put_ bh $ nodeChildren ast

  get bh = Node
    <$> get bh
    <*> get bh
    <*> get bh

-- | The information stored in one AST node.
--
-- The type parameter exists to provide flexibility in representation of types
-- (see Note [Efficient serialization of redundant type info]).
data NodeInfo a = NodeInfo
    { nodeAnnotations :: S.Set (FastString,FastString)
    -- ^ (name of the AST node constructor, name of the AST node Type)

    , nodeType :: [a]
    -- ^ The Haskell types of this node, if any.

    , nodeIdentifiers :: NodeIdentifiers a
    -- ^ All the identifiers and their details
    } deriving (Functor, Foldable, Traversable)

instance Binary (NodeInfo TypeIndex) where
  put_ bh ni = do
    put_ bh $ S.toAscList $ nodeAnnotations ni
    put_ bh $ nodeType ni
    put_ bh $ M.toList $ nodeIdentifiers ni
  get bh = NodeInfo
    <$> fmap (S.fromDistinctAscList) (get bh)
    <*> get bh
    <*> fmap (M.fromList) (get bh)

type Identifier = Either ModuleName Name

type NodeIdentifiers a = M.Map Identifier (IdentifierDetails a)

-- | Information associated with every identifier
--
-- We need to include types with identifiers because sometimes multiple
-- identifiers occur in the same span(Overloaded Record Fields and so on)
data IdentifierDetails a = IdentifierDetails
  { identType :: Maybe a
  , identInfo :: S.Set ContextInfo
  } deriving (Eq, Functor, Foldable, Traversable)

instance Outputable a => Outputable (IdentifierDetails a) where
  ppr x = text "IdentifierDetails" <+> ppr (identType x) <+> ppr (identInfo x)

instance Semigroup (IdentifierDetails a) where
  d1 <> d2 = IdentifierDetails (identType d1 <|> identType d2)
                               (S.union (identInfo d1) (identInfo d2))

instance Monoid (IdentifierDetails a) where
  mempty = IdentifierDetails Nothing S.empty

instance Binary (IdentifierDetails TypeIndex) where
  put_ bh dets = do
    put_ bh $ identType dets
    put_ bh $ S.toAscList $ identInfo dets
  get bh =  IdentifierDetails
    <$> get bh
    <*> fmap (S.fromDistinctAscList) (get bh)

-- | Different contexts under which identifiers exist
data ContextInfo
  = Use                -- ^ regular variable
  | MatchBind
  | IEThing IEType     -- ^ import/export
  | TyDecl

  -- | Value binding
  | ValBind
      BindType     -- ^ whether or not the binding is in an instance
      Scope        -- ^ scope over which the value is bound
      (Maybe Span) -- ^ span of entire binding

  -- | Pattern binding
  --
  -- This case is tricky because the bound identifier can be used in two
  -- distinct scopes. Consider the following example (with @-XViewPatterns@)
  --
  -- @
  -- do (b, a, (a -> True)) <- bar
  --    foo a
  -- @
  --
  -- The identifier @a@ has two scopes: in the view pattern @(a -> True)@ and
  -- in the rest of the @do@-block in @foo a@.
  | PatternBind
      Scope        -- ^ scope /in the pattern/ (the variable bound can be used
                   -- further in the pattern)
      Scope        -- ^ rest of the scope outside the pattern
      (Maybe Span) -- ^ span of entire binding

  | ClassTyDecl (Maybe Span)

  -- | Declaration
  | Decl
      DeclType     -- ^ type of declaration
      (Maybe Span) -- ^ span of entire binding

  -- | Type variable
  | TyVarBind Scope TyVarScope

  -- | Record field
  | RecField RecFieldContext (Maybe Span)
    deriving (Eq, Ord, Show)

instance Outputable ContextInfo where
  ppr = text . show

instance Binary ContextInfo where
  put_ bh Use = putByte bh 0
  put_ bh (IEThing t) = do
    putByte bh 1
    put_ bh t
  put_ bh TyDecl = putByte bh 2
  put_ bh (ValBind bt sc msp) = do
    putByte bh 3
    put_ bh bt
    put_ bh sc
    put_ bh msp
  put_ bh (PatternBind a b c) = do
    putByte bh 4
    put_ bh a
    put_ bh b
    put_ bh c
  put_ bh (ClassTyDecl sp) = do
    putByte bh 5
    put_ bh sp
  put_ bh (Decl a b) = do
    putByte bh 6
    put_ bh a
    put_ bh b
  put_ bh (TyVarBind a b) = do
    putByte bh 7
    put_ bh a
    put_ bh b
  put_ bh (RecField a b) = do
    putByte bh 8
    put_ bh a
    put_ bh b
  put_ bh MatchBind = putByte bh 9

  get bh = do
    (t :: Word8) <- get bh
    case t of
      0 -> return Use
      1 -> IEThing <$> get bh
      2 -> return TyDecl
      3 -> ValBind <$> get bh <*> get bh <*> get bh
      4 -> PatternBind <$> get bh <*> get bh <*> get bh
      5 -> ClassTyDecl <$> get bh
      6 -> Decl <$> get bh <*> get bh
      7 -> TyVarBind <$> get bh <*> get bh
      8 -> RecField <$> get bh <*> get bh
      9 -> return MatchBind
      _ -> panic "Binary ContextInfo: invalid tag"

-- | Types of imports and exports
data IEType
  = Import
  | ImportAs
  | ImportHiding
  | Export
    deriving (Eq, Enum, Ord, Show)

instance Binary IEType where
  put_ bh b = putByte bh (fromIntegral (fromEnum b))
  get bh = do x <- getByte bh; pure $! (toEnum (fromIntegral x))

data RecFieldContext
  = RecFieldDecl
  | RecFieldAssign
  | RecFieldMatch
  | RecFieldOcc
    deriving (Eq, Enum, Ord, Show)

instance Binary RecFieldContext where
  put_ bh b = putByte bh (fromIntegral (fromEnum b))
  get bh = do x <- getByte bh; pure $! (toEnum (fromIntegral x))

data BindType
  = RegularBind
  | InstanceBind
    deriving (Eq, Ord, Show, Enum)

instance Binary BindType where
  put_ bh b = putByte bh (fromIntegral (fromEnum b))
  get bh = do x <- getByte bh; pure $! (toEnum (fromIntegral x))

data DeclType
  = FamDec     -- ^ type or data family
  | SynDec     -- ^ type synonym
  | DataDec    -- ^ data declaration
  | ConDec     -- ^ constructor declaration
  | PatSynDec  -- ^ pattern synonym
  | ClassDec   -- ^ class declaration
  | InstDec    -- ^ instance declaration
    deriving (Eq, Ord, Show, Enum)

instance Binary DeclType where
  put_ bh b = putByte bh (fromIntegral (fromEnum b))
  get bh = do x <- getByte bh; pure $! (toEnum (fromIntegral x))

data Scope
  = NoScope
  | LocalScope Span
  | ModuleScope
    deriving (Eq, Ord, Show, Typeable, Data)

instance Outputable Scope where
  ppr NoScope = text "NoScope"
  ppr (LocalScope sp) = text "LocalScope" <+> ppr sp
  ppr ModuleScope = text "ModuleScope"

instance Binary Scope where
  put_ bh NoScope = putByte bh 0
  put_ bh (LocalScope span) = do
    putByte bh 1
    put_ bh span
  put_ bh ModuleScope = putByte bh 2

  get bh = do
    (t :: Word8) <- get bh
    case t of
      0 -> return NoScope
      1 -> LocalScope <$> get bh
      2 -> return ModuleScope
      _ -> panic "Binary Scope: invalid tag"

-- | Scope of a type variable.
--
-- This warrants a data type apart from 'Scope' because of complexities
-- introduced by features like @-XScopedTypeVariables@ and @-XInstanceSigs@. For
-- example, consider:
--
-- @
-- foo, bar, baz :: forall a. a -> a
-- @
--
-- Here @a@ is in scope in all the definitions of @foo@, @bar@, and @baz@, so we
-- need a list of scopes to keep track of this. Furthermore, this list cannot be
-- computed until we resolve the binding sites of @foo@, @bar@, and @baz@.
--
-- Consequently, @a@ starts with an @'UnresolvedScope' [foo, bar, baz] Nothing@
-- which later gets resolved into a 'ResolvedScopes'.
data TyVarScope
  = ResolvedScopes [Scope]

  -- | Unresolved scopes should never show up in the final @.hie@ file
  | UnresolvedScope
        [Name]        -- ^ names of the definitions over which the scope spans
        (Maybe Span)  -- ^ the location of the instance/class declaration for
                      -- the case where the type variable is declared in a
                      -- method type signature
    deriving (Eq, Ord)

instance Show TyVarScope where
  show (ResolvedScopes sc) = show sc
  show _ = error "UnresolvedScope"

instance Binary TyVarScope where
  put_ bh (ResolvedScopes xs) = do
    putByte bh 0
    put_ bh xs
  put_ bh (UnresolvedScope ns span) = do
    putByte bh 1
    put_ bh ns
    put_ bh span

  get bh = do
    (t :: Word8) <- get bh
    case t of
      0 -> ResolvedScopes <$> get bh
      1 -> UnresolvedScope <$> get bh <*> get bh
      _ -> panic "Binary TyVarScope: invalid tag"


{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE FlexibleInstances #-}
module HieUtils where

import GhcPrelude

import CoreMap
import DynFlags                   ( DynFlags )
import FastString                 ( FastString, mkFastString )
import IfaceType
import Name hiding (varName)
import Outputable                 ( renderWithStyle, ppr, defaultUserStyle )
import SrcLoc
import ToIface
import TyCon
import TyCoRep
import Type
import Var
import VarEnv

import HieTypes

import qualified Data.Map as M
import qualified Data.Set as S
import qualified Data.IntMap.Strict as IM
import qualified Data.Array as A
import Data.Data                  ( typeOf, typeRepTyCon, Data(toConstr) )
import Data.Maybe                 ( maybeToList )
import Data.Monoid
import Data.Traversable           ( for )
import Control.Monad.Trans.State.Strict hiding (get)

generateReferencesMap
  :: Foldable f
  => f (HieAST a)
  -> M.Map Identifier [(Span, IdentifierDetails a)]
generateReferencesMap = foldr (\ast m -> M.unionWith (++) (go ast) m) M.empty
  where
    go ast = M.unionsWith (++) (this : map go (nodeChildren ast))
      where
        this = fmap (pure . (nodeSpan ast,)) $ nodeIdentifiers $ nodeInfo ast

renderHieType :: DynFlags -> HieTypeFix -> String
renderHieType df ht = renderWithStyle df (ppr $ hieTypeToIface ht) sty
  where sty = defaultUserStyle df

resolveVisibility :: Type -> [Type] -> [(Bool,Type)]
resolveVisibility kind ty_args
  = go (mkEmptyTCvSubst in_scope) kind ty_args
  where
    in_scope = mkInScopeSet (tyCoVarsOfTypes ty_args)

    go _   _                   []     = []
    go env ty                  ts
      | Just ty' <- coreView ty
      = go env ty' ts
    go env (ForAllTy (Bndr tv vis) res) (t:ts)
      | isVisibleArgFlag vis = (True , t) : ts'
      | otherwise            = (False, t) : ts'
      where
        ts' = go (extendTvSubst env tv t) res ts

    go env (FunTy { ft_res = res }) (t:ts) -- No type-class args in tycon apps
      = (True,t) : (go env res ts)

    go env (TyVarTy tv) ts
      | Just ki <- lookupTyVar env tv = go env ki ts
    go env kind (t:ts) = (True, t) : (go env kind ts) -- Ill-kinded

foldType :: (HieType a -> a) -> HieTypeFix -> a
foldType f (Roll t) = f $ fmap (foldType f) t

hieTypeToIface :: HieTypeFix -> IfaceType
hieTypeToIface = foldType go
  where
    go (HTyVarTy n) = IfaceTyVar $ occNameFS $ getOccName n
    go (HAppTy a b) = IfaceAppTy a (hieToIfaceArgs b)
    go (HLitTy l) = IfaceLitTy l
    go (HForAllTy ((n,k),af) t) = let b = (occNameFS $ getOccName n, k)
                                  in IfaceForAllTy (Bndr (IfaceTvBndr b) af) t
    go (HFunTy a b)     = IfaceFunTy VisArg   a    b
    go (HQualTy pred b) = IfaceFunTy InvisArg pred b
    go (HCastTy a) = a
    go HCoercionTy = IfaceTyVar "<coercion type>"
    go (HTyConApp a xs) = IfaceTyConApp a (hieToIfaceArgs xs)

    -- This isn't fully faithful - we can't produce the 'Inferred' case
    hieToIfaceArgs :: HieArgs IfaceType -> IfaceAppArgs
    hieToIfaceArgs (HieArgs xs) = go' xs
      where
        go' [] = IA_Nil
        go' ((True ,x):xs) = IA_Arg x Required $ go' xs
        go' ((False,x):xs) = IA_Arg x Specified $ go' xs

data HieTypeState
  = HTS
    { tyMap      :: !(TypeMap TypeIndex)
    , htyTable   :: !(IM.IntMap HieTypeFlat)
    , freshIndex :: !TypeIndex
    }

initialHTS :: HieTypeState
initialHTS = HTS emptyTypeMap IM.empty 0

freshTypeIndex :: State HieTypeState TypeIndex
freshTypeIndex = do
  index <- gets freshIndex
  modify' $ \hts -> hts { freshIndex = index+1 }
  return index

compressTypes
  :: HieASTs Type
  -> (HieASTs TypeIndex, A.Array TypeIndex HieTypeFlat)
compressTypes asts = (a, arr)
  where
    (a, (HTS _ m i)) = flip runState initialHTS $
      for asts $ \typ -> do
        i <- getTypeIndex typ
        return i
    arr = A.array (0,i-1) (IM.toList m)

recoverFullType :: TypeIndex -> A.Array TypeIndex HieTypeFlat -> HieTypeFix
recoverFullType i m = go i
  where
    go i = Roll $ fmap go (m A.! i)

getTypeIndex :: Type -> State HieTypeState TypeIndex
getTypeIndex t
  | otherwise = do
      tm <- gets tyMap
      case lookupTypeMap tm t of
        Just i -> return i
        Nothing -> do
          ht <- go t
          extendHTS t ht
  where
    extendHTS t ht = do
      i <- freshTypeIndex
      modify' $ \(HTS tm tt fi) ->
        HTS (extendTypeMap tm t i) (IM.insert i ht tt) fi
      return i

    go (TyVarTy v) = return $ HTyVarTy $ varName v
    go ty@(AppTy _ _) = do
      let (head,args) = splitAppTys ty
          visArgs = HieArgs $ resolveVisibility (typeKind head) args
      ai <- getTypeIndex head
      argsi <- mapM getTypeIndex visArgs
      return $ HAppTy ai argsi
    go (TyConApp f xs) = do
      let visArgs = HieArgs $ resolveVisibility (tyConKind f) xs
      is <- mapM getTypeIndex visArgs
      return $ HTyConApp (toIfaceTyCon f) is
    go (ForAllTy (Bndr v a) t) = do
      k <- getTypeIndex (varType v)
      i <- getTypeIndex t
      return $ HForAllTy ((varName v,k),a) i
    go (FunTy { ft_af = af, ft_arg = a, ft_res = b }) = do
      ai <- getTypeIndex a
      bi <- getTypeIndex b
      return $ case af of
                 InvisArg -> HQualTy ai bi
                 VisArg   -> HFunTy ai bi
    go (LitTy a) = return $ HLitTy $ toIfaceTyLit a
    go (CastTy t _) = do
      i <- getTypeIndex t
      return $ HCastTy i
    go (CoercionTy _) = return HCoercionTy

resolveTyVarScopes :: M.Map FastString (HieAST a) -> M.Map FastString (HieAST a)
resolveTyVarScopes asts = M.map go asts
  where
    go ast = resolveTyVarScopeLocal ast asts

resolveTyVarScopeLocal :: HieAST a -> M.Map FastString (HieAST a) -> HieAST a
resolveTyVarScopeLocal ast asts = go ast
  where
    resolveNameScope dets = dets{identInfo =
      S.map resolveScope (identInfo dets)}
    resolveScope (TyVarBind sc (UnresolvedScope names Nothing)) =
      TyVarBind sc $ ResolvedScopes
        [ LocalScope binding
        | name <- names
        , Just binding <- [getNameBinding name asts]
        ]
    resolveScope (TyVarBind sc (UnresolvedScope names (Just sp))) =
      TyVarBind sc $ ResolvedScopes
        [ LocalScope binding
        | name <- names
        , Just binding <- [getNameBindingInClass name sp asts]
        ]
    resolveScope scope = scope
    go (Node info span children) = Node info' span $ map go children
      where
        info' = info { nodeIdentifiers = idents }
        idents = M.map resolveNameScope $ nodeIdentifiers info

getNameBinding :: Name -> M.Map FastString (HieAST a) -> Maybe Span
getNameBinding n asts = do
  (_,msp) <- getNameScopeAndBinding n asts
  msp

getNameScope :: Name -> M.Map FastString (HieAST a) -> Maybe [Scope]
getNameScope n asts = do
  (scopes,_) <- getNameScopeAndBinding n asts
  return scopes

getNameBindingInClass
  :: Name
  -> Span
  -> M.Map FastString (HieAST a)
  -> Maybe Span
getNameBindingInClass n sp asts = do
  ast <- M.lookup (srcSpanFile sp) asts
  getFirst $ foldMap First $ do
    child <- flattenAst ast
    dets <- maybeToList
      $ M.lookup (Right n) $ nodeIdentifiers $ nodeInfo child
    let binding = foldMap (First . getBindSiteFromContext) (identInfo dets)
    return (getFirst binding)

getNameScopeAndBinding
  :: Name
  -> M.Map FastString (HieAST a)
  -> Maybe ([Scope], Maybe Span)
getNameScopeAndBinding n asts = case nameSrcSpan n of
  RealSrcSpan sp -> do -- @Maybe
    ast <- M.lookup (srcSpanFile sp) asts
    defNode <- selectLargestContainedBy sp ast
    getFirst $ foldMap First $ do -- @[]
      node <- flattenAst defNode
      dets <- maybeToList
        $ M.lookup (Right n) $ nodeIdentifiers $ nodeInfo node
      scopes <- maybeToList $ foldMap getScopeFromContext (identInfo dets)
      let binding = foldMap (First . getBindSiteFromContext) (identInfo dets)
      return $ Just (scopes, getFirst binding)
  _ -> Nothing

getScopeFromContext :: ContextInfo -> Maybe [Scope]
getScopeFromContext (ValBind _ sc _) = Just [sc]
getScopeFromContext (PatternBind a b _) = Just [a, b]
getScopeFromContext (ClassTyDecl _) = Just [ModuleScope]
getScopeFromContext (Decl _ _) = Just [ModuleScope]
getScopeFromContext (TyVarBind a (ResolvedScopes xs)) = Just $ a:xs
getScopeFromContext (TyVarBind a _) = Just [a]
getScopeFromContext _ = Nothing

getBindSiteFromContext :: ContextInfo -> Maybe Span
getBindSiteFromContext (ValBind _ _ sp) = sp
getBindSiteFromContext (PatternBind _ _ sp) = sp
getBindSiteFromContext _ = Nothing

flattenAst :: HieAST a -> [HieAST a]
flattenAst n =
  n : concatMap flattenAst (nodeChildren n)

smallestContainingSatisfying
  :: Span
  -> (HieAST a -> Bool)
  -> HieAST a
  -> Maybe (HieAST a)
smallestContainingSatisfying sp cond node
  | nodeSpan node `containsSpan` sp = getFirst $ mconcat
      [ foldMap (First . smallestContainingSatisfying sp cond) $
          nodeChildren node
      , First $ if cond node then Just node else Nothing
      ]
  | sp `containsSpan` nodeSpan node = Nothing
  | otherwise = Nothing

selectLargestContainedBy :: Span -> HieAST a -> Maybe (HieAST a)
selectLargestContainedBy sp node
  | sp `containsSpan` nodeSpan node = Just node
  | nodeSpan node `containsSpan` sp =
      getFirst $ foldMap (First . selectLargestContainedBy sp) $
        nodeChildren node
  | otherwise = Nothing

selectSmallestContaining :: Span -> HieAST a -> Maybe (HieAST a)
selectSmallestContaining sp node
  | nodeSpan node `containsSpan` sp = getFirst $ mconcat
      [ foldMap (First . selectSmallestContaining sp) $ nodeChildren node
      , First (Just node)
      ]
  | sp `containsSpan` nodeSpan node = Nothing
  | otherwise = Nothing

definedInAsts :: M.Map FastString (HieAST a) -> Name -> Bool
definedInAsts asts n = case nameSrcSpan n of
  RealSrcSpan sp -> srcSpanFile sp `elem` M.keys asts
  _ -> False

isOccurrence :: ContextInfo -> Bool
isOccurrence Use = True
isOccurrence _ = False

scopeContainsSpan :: Scope -> Span -> Bool
scopeContainsSpan NoScope _ = False
scopeContainsSpan ModuleScope _ = True
scopeContainsSpan (LocalScope a) b = a `containsSpan` b

-- | One must contain the other. Leaf nodes cannot contain anything
combineAst :: HieAST Type -> HieAST Type -> HieAST Type
combineAst a@(Node aInf aSpn xs) b@(Node bInf bSpn ys)
  | aSpn == bSpn = Node (aInf `combineNodeInfo` bInf) aSpn (mergeAsts xs ys)
  | aSpn `containsSpan` bSpn = combineAst b a
combineAst a (Node xs span children) = Node xs span (insertAst a children)

-- | Insert an AST in a sorted list of disjoint Asts
insertAst :: HieAST Type -> [HieAST Type] -> [HieAST Type]
insertAst x = mergeAsts [x]

-- | Merge two nodes together.
--
-- Precondition and postcondition: elements in 'nodeType' are ordered.
combineNodeInfo :: NodeInfo Type -> NodeInfo Type -> NodeInfo Type
(NodeInfo as ai ad) `combineNodeInfo` (NodeInfo bs bi bd) =
  NodeInfo (S.union as bs) (mergeSorted ai bi) (M.unionWith (<>) ad bd)
  where
    mergeSorted :: [Type] -> [Type] -> [Type]
    mergeSorted la@(a:as) lb@(b:bs) = case nonDetCmpType a b of
                                        LT -> a : mergeSorted as lb
                                        EQ -> a : mergeSorted as bs
                                        GT -> b : mergeSorted la bs
    mergeSorted as [] = as
    mergeSorted [] bs = bs

{- | Merge two sorted, disjoint lists of ASTs, combining when necessary.

In the absence of position-altering pragmas (ex: @# line "file.hs" 3@),
different nodes in an AST tree should either have disjoint spans (in
which case you can say for sure which one comes first) or one span
should be completely contained in the other (in which case the contained
span corresponds to some child node).

However, since Haskell does have position-altering pragmas it /is/
possible for spans to be overlapping. Here is an example of a source file
in which @foozball@ and @quuuuuux@ have overlapping spans:

@
module Baz where

# line 3 "Baz.hs"
foozball :: Int
foozball = 0

# line 3 "Baz.hs"
bar, quuuuuux :: Int
bar = 1
quuuuuux = 2
@

In these cases, we just do our best to produce sensible `HieAST`'s. The blame
should be laid at the feet of whoever wrote the line pragmas in the first place
(usually the C preprocessor...).
-}
mergeAsts :: [HieAST Type] -> [HieAST Type] -> [HieAST Type]
mergeAsts xs [] = xs
mergeAsts [] ys = ys
mergeAsts xs@(a:as) ys@(b:bs)
  | span_a `containsSpan`   span_b = mergeAsts (combineAst a b : as) bs
  | span_b `containsSpan`   span_a = mergeAsts as (combineAst a b : bs)
  | span_a `rightOf`        span_b = b : mergeAsts xs bs
  | span_a `leftOf`         span_b = a : mergeAsts as ys

  -- These cases are to work around ASTs that are not fully disjoint
  | span_a `startsRightOf`  span_b = b : mergeAsts as ys
  | otherwise                      = a : mergeAsts as ys
  where
    span_a = nodeSpan a
    span_b = nodeSpan b

rightOf :: Span -> Span -> Bool
rightOf s1 s2
  = (srcSpanStartLine s1, srcSpanStartCol s1)
       >= (srcSpanEndLine s2, srcSpanEndCol s2)
    && (srcSpanFile s1 == srcSpanFile s2)

leftOf :: Span -> Span -> Bool
leftOf s1 s2
  = (srcSpanEndLine s1, srcSpanEndCol s1)
       <= (srcSpanStartLine s2, srcSpanStartCol s2)
    && (srcSpanFile s1 == srcSpanFile s2)

startsRightOf :: Span -> Span -> Bool
startsRightOf s1 s2
  = (srcSpanStartLine s1, srcSpanStartCol s1)
       >= (srcSpanStartLine s2, srcSpanStartCol s2)

-- | combines and sorts ASTs using a merge sort
mergeSortAsts :: [HieAST Type] -> [HieAST Type]
mergeSortAsts = go . map pure
  where
    go [] = []
    go [xs] = xs
    go xss = go (mergePairs xss)
    mergePairs [] = []
    mergePairs [xs] = [xs]
    mergePairs (xs:ys:xss) = mergeAsts xs ys : mergePairs xss

simpleNodeInfo :: FastString -> FastString -> NodeInfo a
simpleNodeInfo cons typ = NodeInfo (S.singleton (cons, typ)) [] M.empty

locOnly :: SrcSpan -> [HieAST a]
locOnly (RealSrcSpan span) =
  [Node e span []]
    where e = NodeInfo S.empty [] M.empty
locOnly _ = []

mkScope :: SrcSpan -> Scope
mkScope (RealSrcSpan sp) = LocalScope sp
mkScope _ = NoScope

mkLScope :: Located a -> Scope
mkLScope = mkScope . getLoc

combineScopes :: Scope -> Scope -> Scope
combineScopes ModuleScope _ = ModuleScope
combineScopes _ ModuleScope = ModuleScope
combineScopes NoScope x = x
combineScopes x NoScope = x
combineScopes (LocalScope a) (LocalScope b) =
  mkScope $ combineSrcSpans (RealSrcSpan a) (RealSrcSpan b)

{-# INLINEABLE makeNode #-}
makeNode
  :: (Applicative m, Data a)
  => a                       -- ^ helps fill in 'nodeAnnotations' (with 'Data')
  -> SrcSpan                 -- ^ return an empty list if this is unhelpful
  -> m [HieAST b]
makeNode x spn = pure $ case spn of
  RealSrcSpan span -> [Node (simpleNodeInfo cons typ) span []]
  _ -> []
  where
    cons = mkFastString . show . toConstr $ x
    typ = mkFastString . show . typeRepTyCon . typeOf $ x

{-# INLINEABLE makeTypeNode #-}
makeTypeNode
  :: (Applicative m, Data a)
  => a                       -- ^ helps fill in 'nodeAnnotations' (with 'Data')
  -> SrcSpan                 -- ^ return an empty list if this is unhelpful
  -> Type                    -- ^ type to associate with the node
  -> m [HieAST Type]
makeTypeNode x spn etyp = pure $ case spn of
  RealSrcSpan span ->
    [Node (NodeInfo (S.singleton (cons,typ)) [etyp] M.empty) span []]
  _ -> []
  where
    cons = mkFastString . show . toConstr $ x
    typ = mkFastString . show . typeRepTyCon . typeOf $ x