{-
(c) Bartosz Nitka, Facebook 2015

-}

{-# LANGUAGE BangPatterns #-}

-- | Utilities for efficiently and deterministically computing free variables.
module GHC.Utils.FV (
        -- * Deterministic free vars computations
        FV, InterestingVarFun,

        -- * Running the computations
        fvVarList, fvVarSet, fvDVarSet,

        -- ** Manipulating those computations
        unitFV,
        emptyFV,
        mkFVs,
        unionFV,
        unionsFV,
        delFV,
        delFVs,
        filterFV,
        mapUnionFV,
    ) where

import GHC.Prelude

import GHC.Types.Var
import GHC.Types.Var.Set

-- | 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 GHC.Types.Unique.DFM.

-- 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
        -> VarAcc           -- Accumulator
        -> VarAcc

type VarAcc = ([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 GHC.Types.Unique.DFM.

-- 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` 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.
fvVarAcc :: FV ->  ([Var], VarSet)
fvVarAcc :: FV -> ([Var], VarSet)
fvVarAcc FV
fv = FV
fv (Bool -> Var -> Bool
forall a b. a -> b -> a
const Bool
True) VarSet
emptyVarSet ([], VarSet
emptyVarSet)

-- | Run a free variable computation, returning a list of distinct free
-- variables in deterministic order.
fvVarList :: FV -> [Var]
fvVarList :: FV -> [Var]
fvVarList = ([Var], VarSet) -> [Var]
forall a b. (a, b) -> a
fst (([Var], VarSet) -> [Var])
-> (FV -> ([Var], VarSet)) -> FV -> [Var]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. FV -> ([Var], VarSet)
fvVarAcc

-- | 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 :: FV -> DVarSet
fvDVarSet = [Var] -> DVarSet
mkDVarSet ([Var] -> DVarSet) -> (FV -> [Var]) -> FV -> DVarSet
forall b c a. (b -> c) -> (a -> b) -> a -> c
. FV -> [Var]
fvVarList

-- | 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 :: FV -> VarSet
fvVarSet = ([Var], VarSet) -> VarSet
forall a b. (a, b) -> b
snd (([Var], VarSet) -> VarSet)
-> (FV -> ([Var], VarSet)) -> FV -> VarSet
forall b c a. (b -> c) -> (a -> b) -> a -> c
. FV -> ([Var], VarSet)
fvVarAcc

-- 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
unitFV Var
var Var -> Bool
fv_cand VarSet
in_scope acc :: ([Var], VarSet)
acc@([Var]
have, VarSet
haveSet)
  | Var
var Var -> VarSet -> Bool
`elemVarSet` VarSet
in_scope = ([Var], VarSet)
acc
  | Var
var Var -> VarSet -> Bool
`elemVarSet` VarSet
haveSet = ([Var], VarSet)
acc
  | Var -> Bool
fv_cand Var
var = (Var
varVar -> [Var] -> [Var]
forall a. a -> [a] -> [a]
:[Var]
have, VarSet -> Var -> VarSet
extendVarSet VarSet
haveSet Var
var)
  | Bool
otherwise = ([Var], VarSet)
acc
{-# INLINE unitFV #-}

-- | Return no free variables.
emptyFV :: FV
emptyFV :: FV
emptyFV Var -> Bool
_ VarSet
_ ([Var], VarSet)
acc = ([Var], VarSet)
acc
{-# INLINE emptyFV #-}

-- | Union two free variable computations.
unionFV :: FV -> FV -> FV
unionFV :: FV -> FV -> FV
unionFV FV
fv1 FV
fv2 Var -> Bool
fv_cand VarSet
in_scope ([Var], VarSet)
acc =
  FV
fv1 Var -> Bool
fv_cand VarSet
in_scope (([Var], VarSet) -> ([Var], VarSet))
-> ([Var], VarSet) -> ([Var], VarSet)
forall a b. (a -> b) -> a -> b
$! FV
fv2 Var -> Bool
fv_cand VarSet
in_scope (([Var], VarSet) -> ([Var], VarSet))
-> ([Var], VarSet) -> ([Var], VarSet)
forall a b. (a -> b) -> a -> b
$! ([Var], VarSet)
acc
{-# INLINE unionFV #-}

-- | Mark the variable as not free by putting it in scope.
delFV :: Var -> FV -> FV
delFV :: Var -> FV -> FV
delFV Var
var FV
fv Var -> Bool
fv_cand !VarSet
in_scope ([Var], VarSet)
acc =
  FV
fv Var -> Bool
fv_cand (VarSet -> Var -> VarSet
extendVarSet VarSet
in_scope Var
var) ([Var], VarSet)
acc
{-# INLINE delFV #-}

-- | Mark many free variables as not free.
delFVs :: VarSet -> FV -> FV
delFVs :: VarSet -> FV -> FV
delFVs VarSet
vars FV
fv Var -> Bool
fv_cand !VarSet
in_scope ([Var], VarSet)
acc =
  FV
fv Var -> Bool
fv_cand (VarSet
in_scope VarSet -> VarSet -> VarSet
`unionVarSet` VarSet
vars) ([Var], VarSet)
acc
{-# INLINE delFVs #-}

-- | Filter a free variable computation.
filterFV :: InterestingVarFun -> FV -> FV
filterFV :: (Var -> Bool) -> FV -> FV
filterFV Var -> Bool
fv_cand2 FV
fv Var -> Bool
fv_cand1 VarSet
in_scope ([Var], VarSet)
acc =
  FV
fv (\Var
v -> Var -> Bool
fv_cand1 Var
v Bool -> Bool -> Bool
&& Var -> Bool
fv_cand2 Var
v) VarSet
in_scope ([Var], VarSet)
acc
{-# INLINE filterFV #-}

-- | Map a free variable computation over a list and union the results.
mapUnionFV :: (a -> FV) -> [a] -> FV
mapUnionFV :: forall a. (a -> FV) -> [a] -> FV
mapUnionFV a -> FV
_f [] Var -> Bool
_fv_cand VarSet
_in_scope ([Var], VarSet)
acc = ([Var], VarSet)
acc
mapUnionFV a -> FV
f (a
a:[a]
as) Var -> Bool
fv_cand VarSet
in_scope ([Var], VarSet)
acc =
  (a -> FV) -> [a] -> FV
forall a. (a -> FV) -> [a] -> FV
mapUnionFV a -> FV
f [a]
as Var -> Bool
fv_cand VarSet
in_scope (([Var], VarSet) -> ([Var], VarSet))
-> ([Var], VarSet) -> ([Var], VarSet)
forall a b. (a -> b) -> a -> b
$! a -> FV
f a
a Var -> Bool
fv_cand VarSet
in_scope (([Var], VarSet) -> ([Var], VarSet))
-> ([Var], VarSet) -> ([Var], VarSet)
forall a b. (a -> b) -> a -> b
$! ([Var], VarSet)
acc
{-# INLINABLE mapUnionFV #-}

-- | Union many free variable computations.
unionsFV :: [FV] -> FV
unionsFV :: [FV] -> FV
unionsFV [FV]
fvs Var -> Bool
fv_cand VarSet
in_scope ([Var], VarSet)
acc = (FV -> FV) -> [FV] -> FV
forall a. (a -> FV) -> [a] -> FV
mapUnionFV FV -> FV
forall a. a -> a
id [FV]
fvs Var -> Bool
fv_cand VarSet
in_scope ([Var], VarSet)
acc
{-# INLINE unionsFV #-}

-- | Add multiple variables - when free, to the returned free variables.
-- Ignores duplicates and respects the filtering function.
mkFVs :: [Var] -> FV
mkFVs :: [Var] -> FV
mkFVs [Var]
vars Var -> Bool
fv_cand VarSet
in_scope ([Var], VarSet)
acc =
  (Var -> FV) -> [Var] -> FV
forall a. (a -> FV) -> [a] -> FV
mapUnionFV Var -> FV
unitFV [Var]
vars Var -> Bool
fv_cand VarSet
in_scope ([Var], VarSet)
acc
{-# INLINE mkFVs #-}