content: {-
  (c) Bartosz Nitka, Facebook 2015
  
  Utilities for efficiently and deterministically computing free variables.
  
  -}
  
  {-# LANGUAGE BangPatterns #-}
  
  module Eta.Utils.FV (
          -- * Deterministic free vars computations
          FV, InterestingVarFun,
  
          -- * Running the computations
          fvVarListVarSet, fvVarList, fvVarSet, fvDVarSet,
  
          -- ** Manipulating those computations
          oneVar,
          unitFV,
          emptyFV,
          mkFVs,
          unionFV,
          unionsFV,
          delFV,
          delFVs,
          filterFV,
          mapUnionFV,
      ) where
  
  import Eta.BasicTypes.Var
  import Eta.BasicTypes.VarSet
  
  -- | Predicate on possible free variables: returns @True@ iff the variable is
  -- interesting
  type InterestingVarFun = Var -> Bool
  
  -- Note [Deterministic FV]
  -- ~~~~~~~~~~~~~~~~~~~~~~~
  -- When computing free variables, the order in which you get them affects
  -- the results of floating and specialization. If you use UniqFM to collect
  -- them and then turn that into a list, you get them in nondeterministic
  -- order as described in Note [Deterministic UniqFM] in UniqDFM.
  
  -- A naive algorithm for free variables relies on merging sets of variables.
  -- Merging costs O(n+m) for UniqFM and for UniqDFM there's an additional log
  -- factor. It's cheaper to incrementally add to a list and use a set to check
  -- for duplicates.
  type FV = InterestingVarFun
               -- Used for filtering sets as we build them
            -> VarSet
               -- Locally bound variables
            -> ([Var], VarSet)
               -- List to preserve ordering and set to check for membership,
               -- so that the list doesn't have duplicates
               -- For explanation of why using `VarSet` is not deterministic see
               -- Note [Deterministic UniqFM] in UniqDFM.
            -> ([Var], VarSet)
  
  -- Note [FV naming conventions]
  -- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~
  -- To get the performance and determinism that FV provides, FV computations
  -- need to built up from smaller FV computations and then evaluated with
  -- one of `fvVarList`, `fvDVarSet`, `fvVarListVarSet`. That means the functions
  -- returning FV need to be exported.
  --
  -- The conventions are:
  --
  -- a) non-deterministic functions:
  --   * a function that returns VarSet
  --       e.g. `tyVarsOfType`
  -- b) deterministic functions:
  --   * a worker that returns FV
  --       e.g. `tyFVsOfType`
  --   * a function that returns [Var]
  --       e.g. `tyVarsOfTypeList`
  --   * a function that returns DVarSet
  --       e.g. `tyVarsOfTypeDSet`
  --
  -- Where tyVarsOfType, tyVarsOfTypeList, tyVarsOfTypeDSet are implemented
  -- in terms of the worker evaluated with fvVarSet, fvVarList, fvDVarSet
  -- respectively.
  
  -- | Run a free variable computation, returning a list of distinct free
  -- variables in deterministic order and a non-deterministic set containing
  -- those variables.
  fvVarListVarSet :: FV ->  ([Var], VarSet)
  fvVarListVarSet fv = fv (const True) emptyVarSet ([], emptyVarSet)
  
  -- | Run a free variable computation, returning a list of distinct free
  -- variables in deterministic order.
  fvVarList :: FV -> [Var]
  fvVarList = fst . fvVarListVarSet
  
  -- | Run a free variable computation, returning a deterministic set of free
  -- variables. Note that this is just a wrapper around the version that
  -- returns a deterministic list. If you need a list you should use
  -- `fvVarList`.
  fvDVarSet :: FV -> DVarSet
  fvDVarSet = mkDVarSet . fst . fvVarListVarSet
  
  -- | Run a free variable computation, returning a non-deterministic set of
  -- free variables. Don't use if the set will be later converted to a list
  -- and the order of that list will impact the generated code.
  fvVarSet :: FV -> VarSet
  fvVarSet = snd . fvVarListVarSet
  
  -- Note [FV eta expansion]
  -- ~~~~~~~~~~~~~~~~~~~~~~~
  -- Let's consider an eta-reduced implementation of freeVarsOf using FV:
  --
  -- freeVarsOf (App a b) = freeVarsOf a `unionFV` freeVarsOf b
  --
  -- If GHC doesn't eta-expand it, after inlining unionFV we end up with
  --
  -- freeVarsOf = \x ->
  --   case x of
  --     App a b -> \fv_cand in_scope acc ->
  --       freeVarsOf a fv_cand in_scope $! freeVarsOf b fv_cand in_scope $! acc
  --
  -- which has to create a thunk, resulting in more allocations.
  --
  -- On the other hand if it is eta-expanded:
  --
  -- freeVarsOf (App a b) fv_cand in_scope acc =
  --   (freeVarsOf a `unionFV` freeVarsOf b) fv_cand in_scope acc
  --
  -- after inlining unionFV we have:
  --
  -- freeVarsOf = \x fv_cand in_scope acc ->
  --   case x of
  --     App a b ->
  --       freeVarsOf a fv_cand in_scope $! freeVarsOf b fv_cand in_scope $! acc
  --
  -- which saves allocations.
  --
  -- GHC when presented with knowledge about all the call sites, correctly
  -- eta-expands in this case. Unfortunately due to the fact that freeVarsOf gets
  -- exported to be composed with other functions, GHC doesn't have that
  -- information and has to be more conservative here.
  --
  -- Hence functions that get exported and return FV need to be manually
  -- eta-expanded. See also #11146.
  
  -- | Add a variable - when free, to the returned free variables.
  -- Ignores duplicates and respects the filtering function.
  unitFV :: Id -> FV
  unitFV var fv_cand in_scope acc@(have, haveSet)
    | var `elemVarSet` in_scope = acc
    | var `elemVarSet` haveSet = acc
    | fv_cand var = (var:have, extendVarSet haveSet var)
    | otherwise = acc
  {-# INLINE unitFV #-}
  
  -- | Return no free variables.
  emptyFV :: FV
  emptyFV _ _ acc = acc
  {-# INLINE emptyFV #-}
  
  -- | Union two free variable computations.
  unionFV :: FV -> FV -> FV
  unionFV fv1 fv2 fv_cand in_scope acc =
    fv1 fv_cand in_scope $! fv2 fv_cand in_scope $! acc
  {-# INLINE unionFV #-}
  
  -- | Mark the variable as not free by putting it in scope.
  delFV :: Var -> FV -> FV
  delFV var fv fv_cand !in_scope acc =
    fv fv_cand (extendVarSet in_scope var) acc
  {-# INLINE delFV #-}
  
  -- | Mark many free variables as not free.
  delFVs :: VarSet -> FV -> FV
  delFVs vars fv fv_cand !in_scope acc =
    fv fv_cand (in_scope `unionVarSet` vars) acc
  {-# INLINE delFVs #-}
  
  -- | Filter a free variable computation.
  filterFV :: InterestingVarFun -> FV -> FV
  filterFV fv_cand2 fv fv_cand1 in_scope acc =
    fv (\v -> fv_cand1 v && fv_cand2 v) in_scope acc
  {-# INLINE filterFV #-}
  
  -- | Map a free variable computation over a list and union the results.
  mapUnionFV :: (a -> FV) -> [a] -> FV
  mapUnionFV _f [] _fv_cand _in_scope acc = acc
  mapUnionFV f (a:as) fv_cand in_scope acc =
    mapUnionFV f as fv_cand in_scope $! f a fv_cand in_scope $! acc
  {-# INLINABLE mapUnionFV #-}
  
  -- | Union many free variable computations.
  unionsFV :: [FV] -> FV
  unionsFV fvs fv_cand in_scope acc = mapUnionFV id fvs fv_cand in_scope acc
  {-# INLINE unionsFV #-}
  
  -- | Add multiple variables - when free, to the returned free variables.
  -- Ignores duplicates and respects the filtering function.
  mkFVs :: [Var] -> FV
  mkFVs vars fv_cand in_scope acc =
    mapUnionFV unitFV vars fv_cand in_scope acc
  {-# INLINE mkFVs #-}
  
  {-# INLINE oneVar #-}
  oneVar :: Id -> FV
  oneVar var fv_cand in_scope acc@(have, haveSet)
    = {- ASSERT( isId var ) probably not going to work -} fvs
    where
    fvs | var `elemVarSet` in_scope = acc
        | var `elemVarSet` haveSet = acc
        | fv_cand var = (var:have, extendVarSet haveSet var)
        | otherwise = acc
  