{-
(c) Bartosz Nitka, Facebook 2015

Utilities for efficiently and deterministically computing free variables.

-}

{-# LANGUAGE BangPatterns #-}

module FV (
        -- * Deterministic free vars computations
        FV, InterestingVarFun,

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

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

import GhcPrelude

import Var
import 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 -> ([Var], VarSet)
fvVarListVarSet fv :: 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)
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 :: FV -> DVarSet
fvDVarSet = [Var] -> DVarSet
mkDVarSet ([Var] -> DVarSet) -> (FV -> [Var]) -> FV -> DVarSet
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([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)
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 :: 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)
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
unitFV var :: Var
var fv_cand :: Var -> Bool
fv_cand in_scope :: VarSet
in_scope acc :: ([Var], VarSet)
acc@(have :: [Var]
have, haveSet :: 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 _ _ acc :: ([Var], VarSet)
acc = ([Var], VarSet)
acc
{-# INLINE emptyFV #-}

-- | Union two free variable computations.
unionFV :: FV -> FV -> FV
unionFV :: FV -> FV -> FV
unionFV fv1 :: FV
fv1 fv2 :: FV
fv2 fv_cand :: Var -> Bool
fv_cand in_scope :: VarSet
in_scope acc :: ([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
var fv :: FV
fv fv_cand :: Var -> Bool
fv_cand !VarSet
in_scope acc :: ([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 vars :: VarSet
vars fv :: FV
fv fv_cand :: Var -> Bool
fv_cand !VarSet
in_scope acc :: ([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 fv_cand2 :: Var -> Bool
fv_cand2 fv :: FV
fv fv_cand1 :: Var -> Bool
fv_cand1 in_scope :: VarSet
in_scope acc :: ([Var], VarSet)
acc =
  FV
fv (\v :: 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 :: (a -> FV) -> [a] -> FV
mapUnionFV _f :: a -> FV
_f [] _fv_cand :: Var -> Bool
_fv_cand _in_scope :: VarSet
_in_scope acc :: ([Var], VarSet)
acc = ([Var], VarSet)
acc
mapUnionFV f :: a -> FV
f (a :: a
a:as :: [a]
as) fv_cand :: Var -> Bool
fv_cand in_scope :: VarSet
in_scope acc :: ([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 fvs :: [FV]
fvs fv_cand :: Var -> Bool
fv_cand in_scope :: VarSet
in_scope acc :: ([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 vars :: [Var]
vars fv_cand :: Var -> Bool
fv_cand in_scope :: VarSet
in_scope acc :: ([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 #-}