content: {-# LANGUAGE ImportQualifiedPost #-}
  {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE FlexibleInstances #-}
  {-# LANGUAGE TypeFamilies #-}
  
  module GHC.Util.FreeVars (
      vars, varss, pvars,
      Vars (..), FreeVars(..) , AllVars (..)
    ) where
  
  import GHC.Types.Name.Reader
  import GHC.Types.Name.Occurrence
  import GHC.Types.Name
  import GHC.Hs
  import GHC.Types.SrcLoc
  
  import Data.Generics.Uniplate.DataOnly
  import Data.Monoid
  import Data.Semigroup
  import Data.List.Extra
  import Data.Set (Set)
  import Data.Set qualified as Set
  import Prelude
  
  ( ^+ ) :: Set OccName -> Set OccName -> Set OccName
  ( ^+ ) = Set.union
  ( ^- ) :: Set OccName -> Set OccName -> Set OccName
  ( ^- ) = Set.difference
  
  -- See [Note : Space leaks lurking here?] below.
  data Vars = Vars{bound :: Set OccName, free :: Set OccName}
  
  -- Useful for debugging.
  instance Show Vars where
    show (Vars bs fs) = "bound : " ++
      show (map occNameString (Set.toList bs)) ++
      ", free : " ++ show (map occNameString (Set.toList fs))
  
  instance Semigroup Vars where
      Vars x1 x2 <> Vars y1 y2 = Vars (x1 ^+ y1) (x2 ^+ y2)
  
  instance Monoid Vars where
      mempty = Vars Set.empty Set.empty
      mconcat vs = Vars (Set.unions $ map bound vs) (Set.unions $ map free vs)
  
  -- A type `a` is a model of `AllVars a` if exists a function
  -- `allVars` for producing a pair of the bound and free variable
  -- sets in a value of `a`.
  class AllVars a where
      -- | Return the variables, erring on the side of more free
      -- variables.
      allVars :: a -> Vars
  
  -- A type `a` is a model of `FreeVars a` if exists a function
  -- `freeVars` for producing a set of free variables of a value of
  -- `a`.
  class FreeVars a where
      -- | Return the variables, erring on the side of more free
      -- variables.
      freeVars :: a -> Set OccName
  
  -- Trivial instances.
  instance AllVars Vars  where allVars = id
  instance FreeVars (Set OccName) where freeVars = id
  -- [Note : Space leaks lurking here?]
  -- ==================================
  -- We make use of `foldr`. @cocreature suggests we want bangs on `data
  -- Vars` and replace usages of `mconcat` with `foldl`.
  instance (AllVars a) => AllVars [a] where  allVars = mconcatMap allVars
  instance (FreeVars a) => FreeVars [a] where  freeVars = Set.unions . map freeVars
  
  -- Construct a `Vars` value with no bound vars.
  freeVars_ :: (FreeVars a) => a -> Vars
  freeVars_ = Vars Set.empty . freeVars
  
  -- `inFree a b` is the set of free variables in a together with the
  -- free variables in b not bound in a.
  inFree :: (AllVars a, FreeVars b) => a -> b -> Set OccName
  inFree a b = free aa ^+ (freeVars b ^- bound aa)
      where aa = allVars a
  
  -- `inVars a b` is a value of `Vars_` with bound variables the union
  -- of the bound variables of a and b and free variables the union
  -- of the free variables of a and the free variables of b not
  -- bound by a.
  inVars :: (AllVars a, AllVars b) => a -> b -> Vars
  inVars a b =
    Vars (bound aa ^+ bound bb) (free aa ^+ (free bb ^- bound aa))
      where aa = allVars a
            bb = allVars b
  
  -- Get an `OccName` out of a reader name.
  unqualNames :: LocatedN RdrName -> [OccName]
  unqualNames (L _ (Unqual x)) = [x]
  unqualNames (L _ (Exact x)) = [nameOccName x]
  unqualNames _ = []
  
  instance FreeVars (LocatedA (HsExpr GhcPs)) where
    freeVars (L _ (HsVar _ x)) = Set.fromList $ unqualNames x -- Variable.
    freeVars (L _ (HsUnboundVar _ x)) = Set.fromList [rdrNameOcc x] -- Unbound variable; also used for "holes".
    freeVars (L _ (HsLam _ LamSingle mg)) = free (allVars mg) -- Lambda abstraction. Currently always a single match.
    freeVars (L _ (HsLam _ _ MG{mg_alts=(L _ ms)})) = free (allVars ms) -- Lambda case
    freeVars (L _ (HsCase _ of_ MG{mg_alts=(L _ ms)})) = freeVars of_ ^+ free (allVars ms) -- Case expr.
    freeVars (L _ (HsLet _ binds e)) = inFree binds e -- Let (rec).
    freeVars (L _ (HsDo _ ctxt (L _ stmts))) = snd $ foldl' alg mempty stmts -- Do block.
      where
        alg ::
          (Set OccName, Set OccName) ->
          LocatedA (StmtLR GhcPs GhcPs (LocatedA (HsExpr GhcPs))) ->
          (Set OccName, Set OccName)
        alg (accBound0, accFree0) stmt = (accBound, accFree)
          where
            accBound =
              accBound0
                ^+ ( case stmt of
                       L _ (BindStmt _ pat _) -> bound (allVars pat)
                       L _ (LetStmt _ binds) -> bound (allVars binds)
                       _ -> mempty
                   )
            accFree = accFree0 ^+ (free (allVars stmt) ^- accBound0)
    freeVars (L _ (RecordCon _ _ (HsRecFields _ flds _))) = Set.unions $ map freeVars flds -- Record construction.
    freeVars (L _ (RecordUpd _ e flds)) =
      case flds of
        RegularRecUpdFields _ fs -> Set.unions $ freeVars e : map freeVars fs
        OverloadedRecUpdFields _ ps -> Set.unions $ freeVars e : map freeVars ps
    freeVars (L _ (HsMultiIf _ grhss)) = free (allVars grhss) -- Multi-way if.
    freeVars (L _ (HsTypedBracket _ e)) = freeVars e
    freeVars (L _ (HsUntypedBracket _ (ExpBr _ e))) = freeVars e
    freeVars (L _ (HsUntypedBracket _ (VarBr _ _ v))) = Set.fromList [occName (unLoc v)]
  
    freeVars (L _ HsOverLabel{}) = mempty -- Overloaded label. The id of the in-scope fromLabel.
    freeVars (L _ HsIPVar{}) = mempty -- Implicit parameter.
    freeVars (L _ HsOverLit{}) = mempty -- Overloaded literal.
    freeVars (L _ HsLit{}) = mempty -- Simple literal.
  
    -- freeVars (e@(L _ HsAppType{})) = freeVars $ children e -- Visible type application e.g. f @ Int x y.
    -- freeVars (e@(L _ HsApp{})) = freeVars $ children e -- Application.
    -- freeVars (e@(L _ OpApp{})) = freeVars $ children e -- Operator application.
    -- freeVars (e@(L _ NegApp{})) = freeVars $ children e -- Negation operator.
    -- freeVars (e@(L _ HsPar{})) = freeVars $ children e -- Parenthesized expr.
    -- freeVars (e@(L _ SectionL{})) = freeVars $ children e -- Left section.
    -- freeVars (e@(L _ SectionR{})) = freeVars $ children e -- Right section.
    -- freeVars (e@(L _ ExplicitTuple{})) = freeVars $ children e -- Explicit tuple and sections thereof.
    -- freeVars (e@(L _ ExplicitSum{})) = freeVars $ children e -- Used for unboxed sum types.
    -- freeVars (e@(L _ HsIf{})) = freeVars $ children e -- If.
    -- freeVars (e@(L _ ExplicitList{})) = freeVars $ children e -- Syntactic list e.g. [a, b, c].
    -- freeVars (e@(L _ ExprWithTySig{})) = freeVars $ children e -- Expr with type signature.
    -- freeVars (e@(L _ ArithSeq {})) = freeVars $ children e -- Arithmetic sequence.
    -- freeVars (e@(L _ HsSCC{})) = freeVars $ children e -- Set cost center pragma (expr whose const is to be measured).
    -- freeVars (e@(L _ HsCoreAnn{})) = freeVars $ children e -- Pragma.
    -- freeVars (e@(L _ HsBracket{})) = freeVars $ children e -- Haskell bracket.
    -- freeVars (e@(L _ HsSpliceE{})) = freeVars $ children e -- Template haskell splice expr.
    -- freeVars (e@(L _ HsProc{})) = freeVars $ children e -- Proc notation for arrows.
    -- freeVars (e@(L _ HsStatic{})) = freeVars $ children e -- Static pointers extension.
    -- freeVars (e@(L _ HsArrApp{})) = freeVars $ children e -- Arrow tail or arrow application.
    -- freeVars (e@(L _ HsArrForm{})) = freeVars $ children e -- Come back to it. Arrow tail or arrow application.
    -- freeVars (e@(L _ HsTick{})) = freeVars $ children e -- Haskell program coverage (Hpc) support.
    -- freeVars (e@(L _ HsBinTick{})) = freeVars $ children e -- Haskell program coverage (Hpc) support.
    -- freeVars (e@(L _ HsTickPragma{})) = freeVars $ children e -- Haskell program coverage (Hpc) support.
    -- freeVars (e@(L _ EAsPat{})) = freeVars $ children e -- Expr as pat.
    -- freeVars (e@(L _ EViewPat{})) = freeVars $ children e -- View pattern.
    -- freeVars (e@(L _ ELazyPat{})) = freeVars $ children e -- Lazy pattern.
  
    freeVars e = freeVars $ children e
  
  instance FreeVars (HsTupArg GhcPs) where
    freeVars (Present _ args) = freeVars args
    freeVars _ = mempty
  
  instance FreeVars (LocatedA (HsFieldBind (LocatedA (FieldOcc GhcPs)) (LocatedA (HsExpr GhcPs)))) where
     freeVars o@(L _ (HsFieldBind _ x _ True)) = Set.singleton $ occName $ unLoc $ foLabel $ unLoc x -- a pun
     freeVars o@(L _ (HsFieldBind _ _ x _)) = freeVars x
  
  instance FreeVars (LocatedA (HsFieldBind (LocatedAn NoEpAnns (FieldLabelStrings GhcPs)) (LocatedA (HsExpr GhcPs)))) where
    freeVars (L _ (HsFieldBind _ _ x _)) = freeVars x
  
  instance AllVars (LocatedA (Pat GhcPs)) where
    allVars (L _ (VarPat _ (L _ x))) = Vars (Set.singleton $ rdrNameOcc x) Set.empty -- Variable pattern.
    allVars (L _ (AsPat _ n x)) = allVars (noLocA $ VarPat noExtField n :: LocatedA (Pat GhcPs)) <> allVars x -- As pattern.
    allVars (L _ (ConPat _ _ (RecCon (HsRecFields _ flds _)))) = allVars flds
    allVars (L _ (NPlusKPat _ n _ _ _ _)) = allVars (noLocA $ VarPat noExtField n :: LocatedA (Pat GhcPs)) -- n+k pattern.
    allVars (L _ (ViewPat _ e p)) = freeVars_ e <> allVars p -- View pattern.
    allVars (L _ WildPat{}) = mempty -- Wildcard pattern.
    allVars (L _ LitPat{}) = mempty -- Literal pattern.
    allVars (L _ NPat{}) = mempty -- Natural pattern.
    allVars (L _ InvisPat {}) = mempty -- since ghc-9.10.1
  
    -- allVars p@SplicePat{} = allVars $ children p -- Splice pattern (includes quasi-quotes).
    -- allVars p@SigPat{} = allVars $ children p -- Pattern with a type signature.
    -- allVars p@CoPat{} = allVars $ children p -- Coercion pattern.
    -- allVars p@LazyPat{} = allVars $ children p -- Lazy pattern.
    -- allVars p@ParPat{} = allVars $ children p -- Parenthesized pattern.
    -- allVars p@BangPat{} = allVars $ children p -- Bang pattern.
    -- allVars p@ListPat{} = allVars $ children p -- Syntactic list.
    -- allVars p@TuplePat{} = allVars $ children p -- Tuple sub patterns.
    -- allVars p@SumPat{} = allVars $ children p -- Anonymous sum pattern.
  
    allVars p = allVars $ children p
  
  instance AllVars (LocatedA (HsFieldBind (LocatedA (FieldOcc GhcPs)) (LocatedA (Pat GhcPs)))) where
     allVars (L _ (HsFieldBind _ _ x _)) = allVars x
  
  instance AllVars (LocatedA (StmtLR GhcPs GhcPs (LocatedA (HsExpr GhcPs)))) where
    allVars (L _ (LastStmt _ expr _ _)) = freeVars_ expr -- The last stmt of a ListComp, MonadComp, DoExpr,MDoExpr.
    allVars (L _ (BindStmt _ pat expr)) = allVars pat <> freeVars_ expr -- A generator e.g. x <- [1, 2, 3].
    allVars (L _ (BodyStmt _ expr _ _)) = freeVars_ expr -- A boolean guard e.g. even x.
    allVars (L _ (LetStmt _ binds)) = allVars binds -- A local declaration e.g. let y = x + 1
    allVars (L _ (TransStmt _ _ stmts _ using by _ _ fmap_)) = allVars stmts <> freeVars_ using <> maybe mempty freeVars_ by <> freeVars_ (noLocA fmap_ :: LocatedA (HsExpr GhcPs)) -- Apply a function to a list of statements in order.
    allVars (L _ (RecStmt _ stmts _ _ _ _ _)) = allVars (unLoc stmts) -- A recursive binding for a group of arrows.
    allVars (L _ ParStmt{}) = mempty -- Parallel list thing. Come back to it.
  
  instance AllVars (HsLocalBinds GhcPs) where
    allVars (HsValBinds _ (ValBinds _ binds _)) = allVars binds -- Value bindings.
    allVars (HsIPBinds _ (IPBinds _ binds)) = allVars binds -- Implicit parameter bindings.
    allVars EmptyLocalBinds{} =  mempty -- The case of no local bindings (signals the empty `let` or `where` clause).
    allVars _ = mempty -- extension points
  
  instance AllVars (LocatedA (IPBind GhcPs)) where
    allVars (L _ (IPBind _ _ e)) = freeVars_ e
  
  instance AllVars (LocatedA (HsBindLR GhcPs GhcPs)) where
    allVars (L _ FunBind{fun_id=n, fun_matches=MG{mg_alts=(L _ ms)}}) = allVars (noLocA $ VarPat noExtField n :: LocatedA (Pat GhcPs)) <> allVars ms -- Function bindings and simple variable bindings e.g. f x = e, f !x = 3, f = e, !x = e, x `f` y = e
    allVars (L _ PatBind{pat_lhs=n, pat_rhs=grhss}) = allVars n <> allVars grhss -- Ctor patterns and some other interesting cases e.g. Just x = e, (x) = e, x :: Ty = e.
  
    allVars (L _ (PatSynBind _ PSB{})) = mempty -- Come back to it.
  
  instance AllVars (MatchGroup GhcPs (LocatedA (HsExpr GhcPs))) where
    allVars (MG _ _alts@(L _ alts)) = foldMap (\m -> inVars (allVars ((unLoc . m_pats) m)) (allVars (m_grhss m))) ms
      where ms = map unLoc alts
  
  instance AllVars (LocatedA (Match GhcPs (LocatedA (HsExpr GhcPs)))) where
    allVars (L _ (Match _ FunRhs {mc_fun=name} pats grhss)) = allVars (noLocA $ VarPat noExtField name :: LocatedA (Pat GhcPs)) <> (allVars . unLoc) pats <> allVars grhss -- A pattern matching on an argument of a function binding.
    allVars (L _ (Match _ (StmtCtxt ctxt) pats grhss)) = allVars ctxt <> (allVars . unLoc) pats <> allVars grhss -- Pattern of a do-stmt, list comprehension, pattern guard etc.
    allVars (L _ (Match _ _ pats grhss)) = inVars ((allVars . unLoc) pats) (allVars grhss) -- Everything else.
  
  instance AllVars (HsStmtContext (GenLocated SrcSpanAnnN RdrName)) where
    allVars (PatGuard FunRhs{mc_fun=n}) = allVars (noLocA $ VarPat noExtField n :: LocatedA (Pat GhcPs))
    allVars ParStmtCtxt{} = mempty -- Come back to it.
    allVars TransStmtCtxt{}  = mempty -- Come back to it.
    allVars _ = mempty
  
  instance AllVars (GRHSs GhcPs (LocatedA (HsExpr GhcPs))) where
    allVars (GRHSs _ grhss binds) = inVars binds (mconcatMap allVars grhss)
  
  instance AllVars (LocatedAn NoEpAnns (GRHS GhcPs (LocatedA (HsExpr GhcPs)))) where
    allVars (L _ (GRHS _ guards expr)) = Vars (bound gs) (free gs ^+ (freeVars expr ^- bound gs)) where gs = allVars guards
  
  instance AllVars (LocatedA (HsDecl GhcPs)) where
    allVars (L l (ValD _ bind)) = allVars (L l bind :: LocatedA (HsBindLR GhcPs GhcPs))
    allVars _ = mempty
  
  
  vars :: FreeVars a => a -> [String]
  vars = Set.toList . Set.map occNameString . freeVars
  
  varss :: AllVars a => a -> [String]
  varss = Set.toList . Set.map occNameString . free . allVars
  
  pvars :: AllVars a => a -> [String]
  pvars = Set.toList . Set.map occNameString . bound . allVars
  