--------------------------------------------------------------------------------
-- | Convert GHC Core into Administrative Normal Form (ANF) --------------------
--------------------------------------------------------------------------------

{-# LANGUAGE FlexibleInstances          #-}
{-# LANGUAGE LambdaCase                 #-}
{-# LANGUAGE NoMonomorphismRestriction  #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE OverloadedStrings          #-}
{-# LANGUAGE ViewPatterns               #-}

module Language.Haskell.Liquid.Transforms.ANF (anormalize) where

import           Debug.Trace (trace)
import           Prelude                          hiding (error)
import           Language.Haskell.Liquid.GHC.TypeRep
import           Liquid.GHC.API  as Ghc hiding ( mkTyArg
                                                                , showPpr
                                                                , DsM
                                                                , panic)
import qualified Liquid.GHC.API  as Ghc
import           Control.Monad.State.Lazy
import           System.Console.CmdArgs.Verbosity (whenLoud)
import qualified Language.Fixpoint.Types    as F

import           Language.Haskell.Liquid.UX.Config  as UX
import qualified Language.Haskell.Liquid.Misc       as Misc
import           Language.Haskell.Liquid.GHC.Misc   as GM
import           Language.Haskell.Liquid.Transforms.Rec
import           Language.Haskell.Liquid.Transforms.InlineAux
import           Language.Haskell.Liquid.Transforms.Rewrite
import           Language.Haskell.Liquid.Types.Errors

import qualified Language.Haskell.Liquid.GHC.SpanStack as Sp
import qualified Language.Haskell.Liquid.GHC.Resugar   as Rs
import           Data.Maybe                       (fromMaybe)
import           Data.List                        (sortBy, (\\))
import qualified Text.Printf as Printf
import           Data.Hashable
import Data.HashMap.Strict (HashMap)
import qualified Data.HashMap.Strict as HM

--------------------------------------------------------------------------------
-- | A-Normalize a module ------------------------------------------------------
--------------------------------------------------------------------------------
anormalize :: UX.Config -> HscEnv -> ModGuts -> IO [CoreBind]
--------------------------------------------------------------------------------
anormalize :: Config -> HscEnv -> ModGuts -> IO [CoreBind]
anormalize Config
cfg HscEnv
hscEnv ModGuts
modGuts = do
  IO () -> IO ()
whenLoud forall a b. (a -> b) -> a -> b
$ do
    String -> IO ()
putStrLn String
"***************************** GHC CoreBinds ***************************"
    String -> IO ()
putStrLn forall a b. (a -> b) -> a -> b
$ Bool -> [CoreBind] -> String
GM.showCBs Bool
untidy (ModGuts -> [CoreBind]
mg_binds ModGuts
modGuts)
    String -> IO ()
putStrLn String
"***************************** REC CoreBinds ***************************"
    String -> IO ()
putStrLn forall a b. (a -> b) -> a -> b
$ Bool -> [CoreBind] -> String
GM.showCBs Bool
untidy [CoreBind]
orig_cbs
    String -> IO ()
putStrLn String
"***************************** RWR CoreBinds ***************************"
    String -> IO ()
putStrLn forall a b. (a -> b) -> a -> b
$ Bool -> [CoreBind] -> String
GM.showCBs Bool
untidy [CoreBind]
rwr_cbs
  forall a. a -> Maybe a -> a
fromMaybe forall {a}. a
err forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> b
snd forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a.
HscEnv -> ModGuts -> DsM a -> IO (Messages DecoratedSDoc, Maybe a)
initDsWithModGuts HscEnv
hscEnv ModGuts
modGuts IOEnv (Env DsGblEnv DsLclEnv) [CoreBind]
act -- hscEnv m grEnv tEnv emptyFamInstEnv act
    where
      err :: a
err      = forall a. Maybe SrcSpan -> String -> a
panic forall a. Maybe a
Nothing String
"Oops, cannot A-Normalize GHC Core!"
      act :: IOEnv (Env DsGblEnv DsLclEnv) [CoreBind]
act      = forall (m :: * -> *) (t :: * -> *) a b.
(Monad m, Traversable t) =>
(a -> m [b]) -> t a -> m [b]
Misc.concatMapM (AnfEnv -> CoreBind -> IOEnv (Env DsGblEnv DsLclEnv) [CoreBind]
normalizeTopBind AnfEnv
γ0) [CoreBind]
rwr_cbs
      γ0 :: AnfEnv
γ0       = Config -> AnfEnv
emptyAnfEnv Config
cfg
      rwr_cbs :: [CoreBind]
rwr_cbs  = Config -> [CoreBind] -> [CoreBind]
rewriteBinds Config
cfg [CoreBind]
orig_cbs
      orig_cbs :: [CoreBind]
orig_cbs = [CoreBind] -> [CoreBind]
transformRecExpr [CoreBind]
inl_cbs
      inl_cbs :: [CoreBind]
inl_cbs  = Config -> Module -> [CoreBind] -> [CoreBind]
inlineAux Config
cfg (ModGuts -> Module
mg_module ModGuts
modGuts) forall a b. (a -> b) -> a -> b
$ ModGuts -> [CoreBind]
mg_binds ModGuts
modGuts
      untidy :: Bool
untidy   = Config -> Bool
UX.untidyCore Config
cfg

--------------------------------------------------------------------------------
-- | A-Normalize a @CoreBind@ --------------------------------------------------
--------------------------------------------------------------------------------

-- Can't make the below default for normalizeBind as it
-- fails tests/pos/lets.hs due to GHCs odd let-bindings

normalizeTopBind :: AnfEnv -> Bind CoreBndr -> Ghc.DsM [CoreBind]
normalizeTopBind :: AnfEnv -> CoreBind -> IOEnv (Env DsGblEnv DsLclEnv) [CoreBind]
normalizeTopBind AnfEnv
γ (NonRec Id
x Expr Id
e)
  = do Expr Id
e' <- forall a. DsM a -> DsM a
runDsM forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m a
evalStateT (AnfEnv -> Expr Id -> DsMW (Expr Id)
stitch AnfEnv
γ Expr Id
e) ([CoreBind] -> DsST
DsST [])
       forall (m :: * -> *) a. Monad m => a -> m a
return [CoreBind -> CoreBind
normalizeTyVars forall a b. (a -> b) -> a -> b
$ forall b. b -> Expr b -> Bind b
NonRec Id
x Expr Id
e']

normalizeTopBind AnfEnv
γ (Rec [(Id, Expr Id)]
xes)
  = do DsST
xes' <- forall a. DsM a -> DsM a
runDsM forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m s
execStateT (AnfEnv -> CoreBind -> DsMW ()
normalizeBind AnfEnv
γ (forall b. [(b, Expr b)] -> Bind b
Rec [(Id, Expr Id)]
xes)) ([CoreBind] -> DsST
DsST [])
       forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map CoreBind -> CoreBind
normalizeTyVars (DsST -> [CoreBind]
st_binds DsST
xes')

normalizeTyVars :: Bind Id -> Bind Id
normalizeTyVars :: CoreBind -> CoreBind
normalizeTyVars (NonRec Id
x Expr Id
e) = forall b. b -> Expr b -> Bind b
NonRec (Id -> Type -> Id
setVarType Id
x Type
t') forall a b. (a -> b) -> a -> b
$ Expr Id -> Expr Id
normalizeForAllTys Expr Id
e
  where
    t' :: Type
t'       = String -> [Id] -> [Id] -> Type -> Type
subst String
msg [Id]
as [Id]
as' Type
bt
    msg :: String
msg      = String
"WARNING: unable to renameVars on " forall a. [a] -> [a] -> [a]
++ forall a. Outputable a => a -> String
GM.showPpr Id
x
    as' :: [Id]
as'      = forall a b. (a, b) -> a
fst forall a b. (a -> b) -> a -> b
$ Type -> ([Id], Type)
splitForAllTyCoVars forall a b. (a -> b) -> a -> b
$ Expr Id -> Type
exprType Expr Id
e
    ([Id]
as, Type
bt) = Type -> ([Id], Type)
splitForAllTyCoVars (Id -> Type
varType Id
x)
normalizeTyVars (Rec [(Id, Expr Id)]
xes)    = forall b. [(b, Expr b)] -> Bind b
Rec [(Id, Expr Id)]
xes'
  where
    nrec :: [CoreBind]
nrec     = CoreBind -> CoreBind
normalizeTyVars forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry forall b. b -> Expr b -> Bind b
NonRec forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(Id, Expr Id)]
xes)
    xes' :: [(Id, Expr Id)]
xes'     = (\case NonRec Id
x Expr Id
e -> (Id
x, Expr Id
e); CoreBind
_ -> forall a. Maybe SrcSpan -> String -> a
impossible forall a. Maybe a
Nothing String
"This cannot happen") forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [CoreBind]
nrec

subst :: String -> [TyVar] -> [TyVar] -> Type -> Type
subst :: String -> [Id] -> [Id] -> Type -> Type
subst String
msg [Id]
as [Id]
as' Type
bt
  | forall (t :: * -> *) a. Foldable t => t a -> Int
length [Id]
as forall a. Eq a => a -> a -> Bool
== forall (t :: * -> *) a. Foldable t => t a -> Int
length [Id]
as'
  = [TyCoVarBinder] -> Type -> Type
mkForAllTys (Id -> TyCoVarBinder
mkTyArg forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Id]
as') forall a b. (a -> b) -> a -> b
$ HasCallStack => TCvSubst -> Type -> Type
substTy TCvSubst
su Type
bt
  | Bool
otherwise
  = forall a. String -> a -> a
trace String
msg forall a b. (a -> b) -> a -> b
$ [TyCoVarBinder] -> Type -> Type
mkForAllTys (Id -> TyCoVarBinder
mkTyArg forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Id]
as) Type
bt
  where su :: TCvSubst
su = [(Id, Type)] -> TCvSubst
mkTvSubstPrs forall a b. (a -> b) -> a -> b
$ forall a b. [a] -> [b] -> [(a, b)]
zip [Id]
as ([Id] -> [Type]
mkTyVarTys [Id]
as')

-- | eta-expand CoreBinds with quantified types
normalizeForAllTys :: CoreExpr -> CoreExpr
normalizeForAllTys :: Expr Id -> Expr Id
normalizeForAllTys Expr Id
e = case Expr Id
e of
  Lam Id
b Expr Id
_ | Id -> Bool
isTyVar Id
b
    -> Expr Id
e
  Expr Id
_ -> forall b. [b] -> Expr b -> Expr b
mkLams [Id]
tvs (forall b. Expr b -> [Type] -> Expr b
mkTyApps Expr Id
e (forall a b. (a -> b) -> [a] -> [b]
map Id -> Type
mkTyVarTy [Id]
tvs))
  where
  ([Id]
tvs, Type
_) = Type -> ([Id], Type)
splitForAllTyCoVars (Expr Id -> Type
exprType Expr Id
e)


newtype DsM a = DsM {forall a. DsM a -> DsM a
runDsM :: Ghc.DsM a}
   deriving (forall a b. a -> DsM b -> DsM a
forall a b. (a -> b) -> DsM a -> DsM b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: forall a b. a -> DsM b -> DsM a
$c<$ :: forall a b. a -> DsM b -> DsM a
fmap :: forall a b. (a -> b) -> DsM a -> DsM b
$cfmap :: forall a b. (a -> b) -> DsM a -> DsM b
Functor, Applicative DsM
forall a. a -> DsM a
forall a b. DsM a -> DsM b -> DsM b
forall a b. DsM a -> (a -> DsM b) -> DsM b
forall (m :: * -> *).
Applicative m
-> (forall a b. m a -> (a -> m b) -> m b)
-> (forall a b. m a -> m b -> m b)
-> (forall a. a -> m a)
-> Monad m
return :: forall a. a -> DsM a
$creturn :: forall a. a -> DsM a
>> :: forall a b. DsM a -> DsM b -> DsM b
$c>> :: forall a b. DsM a -> DsM b -> DsM b
>>= :: forall a b. DsM a -> (a -> DsM b) -> DsM b
$c>>= :: forall a b. DsM a -> (a -> DsM b) -> DsM b
Monad, Monad DsM
DsM [Unique]
DsM UniqSupply
DsM Unique
forall (m :: * -> *).
Monad m -> m UniqSupply -> m Unique -> m [Unique] -> MonadUnique m
getUniquesM :: DsM [Unique]
$cgetUniquesM :: DsM [Unique]
getUniqueM :: DsM Unique
$cgetUniqueM :: DsM Unique
getUniqueSupplyM :: DsM UniqSupply
$cgetUniqueSupplyM :: DsM UniqSupply
MonadUnique, Functor DsM
forall a. a -> DsM a
forall a b. DsM a -> DsM b -> DsM a
forall a b. DsM a -> DsM b -> DsM b
forall a b. DsM (a -> b) -> DsM a -> DsM b
forall a b c. (a -> b -> c) -> DsM a -> DsM b -> DsM c
forall (f :: * -> *).
Functor f
-> (forall a. a -> f a)
-> (forall a b. f (a -> b) -> f a -> f b)
-> (forall a b c. (a -> b -> c) -> f a -> f b -> f c)
-> (forall a b. f a -> f b -> f b)
-> (forall a b. f a -> f b -> f a)
-> Applicative f
<* :: forall a b. DsM a -> DsM b -> DsM a
$c<* :: forall a b. DsM a -> DsM b -> DsM a
*> :: forall a b. DsM a -> DsM b -> DsM b
$c*> :: forall a b. DsM a -> DsM b -> DsM b
liftA2 :: forall a b c. (a -> b -> c) -> DsM a -> DsM b -> DsM c
$cliftA2 :: forall a b c. (a -> b -> c) -> DsM a -> DsM b -> DsM c
<*> :: forall a b. DsM (a -> b) -> DsM a -> DsM b
$c<*> :: forall a b. DsM (a -> b) -> DsM a -> DsM b
pure :: forall a. a -> DsM a
$cpure :: forall a. a -> DsM a
Applicative)

newtype DsST = DsST { DsST -> [CoreBind]
st_binds :: [CoreBind] }

type DsMW = StateT DsST DsM

------------------------------------------------------------------
normalizeBind :: AnfEnv -> CoreBind -> DsMW ()
------------------------------------------------------------------
normalizeBind :: AnfEnv -> CoreBind -> DsMW ()
normalizeBind AnfEnv
γ (NonRec Id
x Expr Id
e)
  = do Expr Id
e' <- AnfEnv -> Expr Id -> DsMW (Expr Id)
normalize AnfEnv
γ Expr Id
e
       [CoreBind] -> DsMW ()
add [forall b. b -> Expr b -> Bind b
NonRec Id
x Expr Id
e']

normalizeBind AnfEnv
γ (Rec [(Id, Expr Id)]
xes)
  = do [Expr Id]
es' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (AnfEnv -> Expr Id -> DsMW (Expr Id)
stitch AnfEnv
γ) [Expr Id]
es
       [CoreBind] -> DsMW ()
add [forall b. [(b, Expr b)] -> Bind b
Rec (forall a b. [a] -> [b] -> [(a, b)]
zip [Id]
xs [Expr Id]
es')]
    where
       ([Id]
xs, [Expr Id]
es) = forall a b. [(a, b)] -> ([a], [b])
unzip [(Id, Expr Id)]
xes

--------------------------------------------------------------------
normalizeName :: AnfEnv -> CoreExpr -> DsMW CoreExpr
--------------------------------------------------------------------

-- normalizeNameDebug γ e
--   = liftM (tracePpr ("normalizeName" ++ showPpr e)) $ normalizeName γ e

normalizeName :: AnfEnv -> Expr Id -> DsMW (Expr Id)
normalizeName AnfEnv
γ e :: Expr Id
e@(Lit Literal
l)
  | Literal -> Bool
shouldNormalize Literal
l
  = AnfEnv -> Expr Id -> DsMW (Expr Id)
normalizeLiteral AnfEnv
γ Expr Id
e
  | Bool
otherwise
  = forall (m :: * -> *) a. Monad m => a -> m a
return Expr Id
e

normalizeName AnfEnv
γ (Var Id
x)
  = forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall b. Id -> Expr b
Var (AnfEnv -> Id -> Id -> Id
lookupAnfEnv AnfEnv
γ Id
x Id
x)

normalizeName AnfEnv
_ e :: Expr Id
e@(Type Type
_)
  = forall (m :: * -> *) a. Monad m => a -> m a
return Expr Id
e

normalizeName AnfEnv
γ e :: Expr Id
e@(Coercion Coercion
_)
  = do Id
x     <- forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall a b. (a -> b) -> a -> b
$ AnfEnv -> Type -> DsM Id
freshNormalVar AnfEnv
γ forall a b. (a -> b) -> a -> b
$ Expr Id -> Type
exprType Expr Id
e
       [CoreBind] -> DsMW ()
add  [forall b. b -> Expr b -> Bind b
NonRec Id
x Expr Id
e]
       forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall b. Id -> Expr b
Var Id
x

normalizeName AnfEnv
γ (Tick CoreTickish
tt Expr Id
e)
  = do Expr Id
e'    <- AnfEnv -> Expr Id -> DsMW (Expr Id)
normalizeName (AnfEnv
γ AnfEnv -> CoreTickish -> AnfEnv
`at` CoreTickish
tt) Expr Id
e
       forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall b. CoreTickish -> Expr b -> Expr b
Tick CoreTickish
tt Expr Id
e'

normalizeName AnfEnv
γ Expr Id
e
  = do Expr Id
e'   <- AnfEnv -> Expr Id -> DsMW (Expr Id)
normalize AnfEnv
γ Expr Id
e
       Id
x    <- forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall a b. (a -> b) -> a -> b
$ AnfEnv -> Type -> DsM Id
freshNormalVar AnfEnv
γ forall a b. (a -> b) -> a -> b
$ Expr Id -> Type
exprType Expr Id
e
       [CoreBind] -> DsMW ()
add [forall b. b -> Expr b -> Bind b
NonRec Id
x Expr Id
e']
       forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall b. Id -> Expr b
Var Id
x

shouldNormalize :: Literal -> Bool
shouldNormalize :: Literal -> Bool
shouldNormalize (LitNumber {})  = Bool
True
shouldNormalize (LitString {})    = Bool
True
shouldNormalize Literal
_               = Bool
False

add :: [CoreBind] -> DsMW ()
add :: [CoreBind] -> DsMW ()
add [CoreBind]
w = forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify forall a b. (a -> b) -> a -> b
$ \DsST
s -> DsST
s { st_binds :: [CoreBind]
st_binds = DsST -> [CoreBind]
st_binds DsST
s forall a. [a] -> [a] -> [a]
++ [CoreBind]
w}

--------------------------------------------------------------------------------
normalizeLiteral :: AnfEnv -> CoreExpr -> DsMW CoreExpr
--------------------------------------------------------------------------------
normalizeLiteral :: AnfEnv -> Expr Id -> DsMW (Expr Id)
normalizeLiteral AnfEnv
γ Expr Id
e =
  do Id
x <- forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall a b. (a -> b) -> a -> b
$ AnfEnv -> Type -> DsM Id
freshNormalVar AnfEnv
γ forall a b. (a -> b) -> a -> b
$ Expr Id -> Type
exprType Expr Id
e
     [CoreBind] -> DsMW ()
add [forall b. b -> Expr b -> Bind b
NonRec Id
x Expr Id
e]
     forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall b. Id -> Expr b
Var Id
x

--------------------------------------------------------------------------------
normalize :: AnfEnv -> CoreExpr -> DsMW CoreExpr
--------------------------------------------------------------------------------
normalize :: AnfEnv -> Expr Id -> DsMW (Expr Id)
normalize AnfEnv
γ Expr Id
e
  | forall t. HasConfig t => t -> Bool
UX.patternFlag AnfEnv
γ
  , Just Pattern
p <- Expr Id -> Maybe Pattern
Rs.lift Expr Id
e
  = AnfEnv -> Pattern -> DsMW (Expr Id)
normalizePattern AnfEnv
γ Pattern
p

normalize AnfEnv
γ (Lam Id
x Expr Id
e) | Id -> Bool
isTyVar Id
x
  = do Expr Id
e' <- AnfEnv -> Expr Id -> DsMW (Expr Id)
normalize AnfEnv
γ Expr Id
e
       forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall b. b -> Expr b -> Expr b
Lam Id
x Expr Id
e'

normalize AnfEnv
γ (Lam Id
x Expr Id
e)
  = do Expr Id
e' <- AnfEnv -> Expr Id -> DsMW (Expr Id)
stitch AnfEnv
γ Expr Id
e
       forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall b. b -> Expr b -> Expr b
Lam Id
x Expr Id
e'

normalize AnfEnv
γ (Let CoreBind
b Expr Id
e)
  = do AnfEnv -> CoreBind -> DsMW ()
normalizeBind AnfEnv
γ CoreBind
b
       AnfEnv -> Expr Id -> DsMW (Expr Id)
normalize AnfEnv
γ Expr Id
e
       -- Need to float bindings all the way up to the top
       -- Due to GHCs odd let-bindings (see tests/pos/lets.hs)

normalize AnfEnv
γ (Case Expr Id
e Id
x Type
t [Alt Id]
as)
  = do Expr Id
n     <- AnfEnv -> Expr Id -> DsMW (Expr Id)
normalizeName AnfEnv
γ Expr Id
e
       Id
x'    <- forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall a b. (a -> b) -> a -> b
$ AnfEnv -> Type -> DsM Id
freshNormalVar AnfEnv
γ Type
τx -- rename "wild" to avoid shadowing
       let γ' :: AnfEnv
γ' = AnfEnv -> Id -> Id -> AnfEnv
extendAnfEnv AnfEnv
γ Id
x Id
x'
       [Alt Id]
as'   <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [Alt Id]
as forall a b. (a -> b) -> a -> b
$ \(Alt AltCon
c [Id]
xs Expr Id
e') -> forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall b. AltCon -> [b] -> Expr b -> Alt b
Alt AltCon
c [Id]
xs) (AnfEnv -> Expr Id -> DsMW (Expr Id)
stitch (AltCon -> AnfEnv -> AnfEnv
incrCaseDepth AltCon
c AnfEnv
γ') Expr Id
e')
       [Alt Id]
as''  <- forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall a b. (a -> b) -> a -> b
$ AnfEnv -> Type -> [Alt Id] -> DsM [Alt Id]
expandDefaultCase AnfEnv
γ Type
τx [Alt Id]
as'
       forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall b. Expr b -> b -> Type -> [Alt b] -> Expr b
Case Expr Id
n Id
x' Type
t [Alt Id]
as''
    where τx :: Type
τx = Id -> Type
GM.expandVarType Id
x

normalize AnfEnv
γ (Var Id
x)
  = forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall b. Id -> Expr b
Var (AnfEnv -> Id -> Id -> Id
lookupAnfEnv AnfEnv
γ Id
x Id
x)

normalize AnfEnv
_ e :: Expr Id
e@(Lit Literal
_)
  = forall (m :: * -> *) a. Monad m => a -> m a
return Expr Id
e

normalize AnfEnv
_ e :: Expr Id
e@(Type Type
_)
  = forall (m :: * -> *) a. Monad m => a -> m a
return Expr Id
e

normalize AnfEnv
γ (Cast Expr Id
e Coercion
τ)
  = do Expr Id
e' <- AnfEnv -> Expr Id -> DsMW (Expr Id)
normalizeName AnfEnv
γ Expr Id
e
       forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall b. Expr b -> Coercion -> Expr b
Cast Expr Id
e' Coercion
τ

normalize AnfEnv
γ (App Expr Id
e1 e2 :: Expr Id
e2@(Type Type
_))
  = do Expr Id
e1' <- AnfEnv -> Expr Id -> DsMW (Expr Id)
normalize AnfEnv
γ Expr Id
e1
       Expr Id
e2' <- AnfEnv -> Expr Id -> DsMW (Expr Id)
normalize AnfEnv
γ Expr Id
e2
       forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall b. Expr b -> Expr b -> Expr b
App Expr Id
e1' Expr Id
e2'

normalize AnfEnv
γ (App Expr Id
e1 Expr Id
e2)
  = do Expr Id
e1' <- AnfEnv -> Expr Id -> DsMW (Expr Id)
normalize AnfEnv
γ Expr Id
e1
       Expr Id
n2  <- AnfEnv -> Expr Id -> DsMW (Expr Id)
normalizeName AnfEnv
γ Expr Id
e2
       forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall b. Expr b -> Expr b -> Expr b
App Expr Id
e1' Expr Id
n2

normalize AnfEnv
γ (Tick CoreTickish
tt Expr Id
e)
  = do Expr Id
e' <- AnfEnv -> Expr Id -> DsMW (Expr Id)
normalize (AnfEnv
γ AnfEnv -> CoreTickish -> AnfEnv
`at` CoreTickish
tt) Expr Id
e
       forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall b. CoreTickish -> Expr b -> Expr b
Tick CoreTickish
tt Expr Id
e'

normalize AnfEnv
_ (Coercion Coercion
c)
  = forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall b. Coercion -> Expr b
Coercion Coercion
c

--------------------------------------------------------------------------------
stitch :: AnfEnv -> CoreExpr -> DsMW CoreExpr
--------------------------------------------------------------------------------
stitch :: AnfEnv -> Expr Id -> DsMW (Expr Id)
stitch AnfEnv
γ Expr Id
e
  = do DsST
bs'   <- forall s (m :: * -> *). MonadState s m => m s
get
       forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify forall a b. (a -> b) -> a -> b
$ \DsST
s -> DsST
s { st_binds :: [CoreBind]
st_binds = [] }
       Expr Id
e'    <- AnfEnv -> Expr Id -> DsMW (Expr Id)
normalize AnfEnv
γ Expr Id
e
       [CoreBind]
bs    <- DsST -> [CoreBind]
st_binds forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall s (m :: * -> *). MonadState s m => m s
get
       forall s (m :: * -> *). MonadState s m => s -> m ()
put DsST
bs'
       forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ [CoreBind] -> Expr Id -> Expr Id
mkCoreLets [CoreBind]
bs Expr Id
e'

_mkCoreLets' :: [CoreBind] -> CoreExpr -> CoreExpr
_mkCoreLets' :: [CoreBind] -> Expr Id -> Expr Id
_mkCoreLets' [CoreBind]
bs Expr Id
e = [CoreBind] -> Expr Id -> Expr Id
mkCoreLets [CoreBind]
bs1 Expr Id
e1
  where
    (Expr Id
e1, [CoreBind]
bs1)    = forall a. Outputable a => String -> a -> a
GM.tracePpr String
"MKCORELETS" (Expr Id
e, [CoreBind]
bs)

--------------------------------------------------------------------------------
normalizePattern :: AnfEnv -> Rs.Pattern -> DsMW CoreExpr
--------------------------------------------------------------------------------
normalizePattern :: AnfEnv -> Pattern -> DsMW (Expr Id)
normalizePattern AnfEnv
γ p :: Pattern
p@(Rs.PatBind {}) = do
  -- don't normalize the >>= itself, we have a special typing rule for it
  Expr Id
e1'   <- AnfEnv -> Expr Id -> DsMW (Expr Id)
normalize AnfEnv
γ (Pattern -> Expr Id
Rs.patE1 Pattern
p)
  Expr Id
e2'   <- AnfEnv -> Expr Id -> DsMW (Expr Id)
stitch    AnfEnv
γ (Pattern -> Expr Id
Rs.patE2 Pattern
p)
  forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ Pattern -> Expr Id
Rs.lower Pattern
p { patE1 :: Expr Id
Rs.patE1 = Expr Id
e1', patE2 :: Expr Id
Rs.patE2 = Expr Id
e2' }

normalizePattern AnfEnv
γ p :: Pattern
p@(Rs.PatReturn {}) = do
  Expr Id
e'    <- AnfEnv -> Expr Id -> DsMW (Expr Id)
normalize AnfEnv
γ (Pattern -> Expr Id
Rs.patE Pattern
p)
  forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ Pattern -> Expr Id
Rs.lower Pattern
p { patE :: Expr Id
Rs.patE = Expr Id
e' }

normalizePattern AnfEnv
_ p :: Pattern
p@(Rs.PatProject {}) =
  forall (m :: * -> *) a. Monad m => a -> m a
return (Pattern -> Expr Id
Rs.lower Pattern
p)

normalizePattern AnfEnv
γ p :: Pattern
p@(Rs.PatSelfBind {}) = do
  AnfEnv -> Expr Id -> DsMW (Expr Id)
normalize AnfEnv
γ (Pattern -> Expr Id
Rs.patE Pattern
p)

normalizePattern AnfEnv
γ p :: Pattern
p@(Rs.PatSelfRecBind {}) = do
  Expr Id
e'    <- AnfEnv -> Expr Id -> DsMW (Expr Id)
normalize AnfEnv
γ (Pattern -> Expr Id
Rs.patE Pattern
p)
  forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ Pattern -> Expr Id
Rs.lower Pattern
p { patE :: Expr Id
Rs.patE = Expr Id
e' }


--------------------------------------------------------------------------------
expandDefault :: AnfEnv -> Bool
--------------------------------------------------------------------------------
expandDefault :: AnfEnv -> Bool
expandDefault AnfEnv
γ = AnfEnv -> Int
aeCaseDepth AnfEnv
γ forall a. Ord a => a -> a -> Bool
<= forall t. HasConfig t => t -> Int
maxCaseExpand AnfEnv
γ

--------------------------------------------------------------------------------
expandDefaultCase :: AnfEnv
                  -> Type
                  -> [CoreAlt]
                  -> DsM [CoreAlt]
--------------------------------------------------------------------------------
expandDefaultCase :: AnfEnv -> Type -> [Alt Id] -> DsM [Alt Id]
expandDefaultCase AnfEnv
γ Type
tyapp zs :: [Alt Id]
zs@(Alt AltCon
DEFAULT [Id]
_ Expr Id
_ : [Alt Id]
_) | AnfEnv -> Bool
expandDefault AnfEnv
γ
  = AnfEnv -> Type -> [Alt Id] -> DsM [Alt Id]
expandDefaultCase' AnfEnv
γ Type
tyapp [Alt Id]
zs

expandDefaultCase AnfEnv
γ tyapp :: Type
tyapp@(TyConApp TyCon
tc [Type]
_) z :: [Alt Id]
z@(Alt AltCon
DEFAULT [Id]
_ Expr Id
_:[Alt Id]
dcs)
  = case TyCon -> Maybe [DataCon]
tyConDataCons_maybe TyCon
tc of
       Just [DataCon]
ds -> do let ds' :: [DataCon]
ds' = [DataCon]
ds forall a. Eq a => [a] -> [a] -> [a]
\\ [ DataCon
d | Alt (DataAlt DataCon
d) [Id]
_ Expr Id
_ <- [Alt Id]
dcs]
                     let n :: Int
n   = forall (t :: * -> *) a. Foldable t => t a -> Int
length [DataCon]
ds'
                     if Int
n forall a. Eq a => a -> a -> Bool
== Int
1
                       then AnfEnv -> Type -> [Alt Id] -> DsM [Alt Id]
expandDefaultCase' AnfEnv
γ Type
tyapp [Alt Id]
z
                       else if forall t. HasConfig t => t -> Int
maxCaseExpand AnfEnv
γ forall a. Eq a => a -> a -> Bool
/= Int
2
                            then forall (m :: * -> *) a. Monad m => a -> m a
return [Alt Id]
z
                            else forall (m :: * -> *) a. Monad m => a -> m a
return (forall a. String -> a -> a
trace (Bool -> AnfEnv -> Int -> String
expandMessage Bool
False AnfEnv
γ Int
n) [Alt Id]
z)
       Maybe [DataCon]
Nothing -> forall (m :: * -> *) a. Monad m => a -> m a
return [Alt Id]
z --

expandDefaultCase AnfEnv
_ Type
_ [Alt Id]
z
   = forall (m :: * -> *) a. Monad m => a -> m a
return [Alt Id]
z

expandDefaultCase'
  :: AnfEnv -> Type -> [CoreAlt] -> DsM [CoreAlt]
expandDefaultCase' :: AnfEnv -> Type -> [Alt Id] -> DsM [Alt Id]
expandDefaultCase' AnfEnv
γ Type
t (Alt AltCon
DEFAULT [Id]
_ Expr Id
e : [Alt Id]
dcs)
  | Just [(DataCon, [Id], [Type])]
dtss <- Type -> [AltCon] -> Maybe [(DataCon, [Id], [Type])]
GM.defaultDataCons Type
t ((\(Alt AltCon
dc [Id]
_ Expr Id
_) -> AltCon
dc) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Alt Id]
dcs) = do
      [Alt Id]
dcs'    <- forall a. AnfEnv -> [a] -> [a]
warnCaseExpand AnfEnv
γ forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [(DataCon, [Id], [Type])]
dtss (AnfEnv -> Expr Id -> (DataCon, [Id], [Type]) -> DsM (Alt Id)
cloneCase AnfEnv
γ Expr Id
e)
      forall (m :: * -> *) a. Monad m => a -> m a
return   forall a b. (a -> b) -> a -> b
$ [Alt Id] -> [Alt Id]
sortCases ([Alt Id]
dcs' forall a. [a] -> [a] -> [a]
++ [Alt Id]
dcs)
expandDefaultCase' AnfEnv
_ Type
_ [Alt Id]
z
   = forall (m :: * -> *) a. Monad m => a -> m a
return [Alt Id]
z

cloneCase :: AnfEnv -> CoreExpr -> (DataCon, [TyVar], [Type]) -> DsM CoreAlt
cloneCase :: AnfEnv -> Expr Id -> (DataCon, [Id], [Type]) -> DsM (Alt Id)
cloneCase AnfEnv
γ Expr Id
e (DataCon
d, [Id]
as, [Type]
ts) = do
  [Id]
xs  <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (AnfEnv -> Type -> DsM Id
freshNormalVar AnfEnv
γ) [Type]
ts
  forall (m :: * -> *) a. Monad m => a -> m a
return (forall b. AltCon -> [b] -> Expr b -> Alt b
Alt (DataCon -> AltCon
DataAlt DataCon
d) ([Id]
as forall a. [a] -> [a] -> [a]
++ [Id]
xs) Expr Id
e)

sortCases :: [CoreAlt] -> [CoreAlt]
sortCases :: [Alt Id] -> [Alt Id]
sortCases = forall a. (a -> a -> Ordering) -> [a] -> [a]
sortBy forall a. Alt a -> Alt a -> Ordering
Ghc.cmpAlt

warnCaseExpand :: AnfEnv -> [a] -> [a]
warnCaseExpand :: forall a. AnfEnv -> [a] -> [a]
warnCaseExpand AnfEnv
γ [a]
xs
  | Int
10 forall a. Ord a => a -> a -> Bool
< Int
n          = forall a. String -> a -> a
trace (Bool -> AnfEnv -> Int -> String
expandMessage Bool
True AnfEnv
γ Int
n) [a]
xs
  | Bool
otherwise       = [a]
xs
  where
   n :: Int
n                = forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs

expandMessage :: Bool -> AnfEnv -> Int -> String
expandMessage :: Bool -> AnfEnv -> Int -> String
expandMessage Bool
expand AnfEnv
γ Int
n = [String] -> String
unlines [forall {t}. PrintfType t => t
msg1, forall {t}. PrintfType t => t
msg2]
  where
    msg1 :: t
msg1            = forall r. PrintfType r => String -> r
Printf.printf String
"WARNING: (%s) %s DEFAULT with %d cases at depth %d" (forall a. Outputable a => a -> String
showPpr SrcSpan
sp) String
v1 Int
n Int
d
    msg2 :: t
msg2            = forall r. PrintfType r => String -> r
Printf.printf String
"%s expansion with --max-case-expand=%d" String
v2 Int
d'
    (String
v1, String
v2, Int
d')
      | Bool
expand      = (String
"Expanding"    , String
"Disable", Int
dforall a. Num a => a -> a -> a
-Int
1) :: (String, String, Int)
      | Bool
otherwise   = (String
"Not expanding", String
"Enable" , Int
dforall a. Num a => a -> a -> a
+Int
1)
    d :: Int
d               = AnfEnv -> Int
aeCaseDepth AnfEnv
γ
    sp :: SrcSpan
sp              = SpanStack -> SrcSpan
Sp.srcSpan (AnfEnv -> SpanStack
aeSrcSpan AnfEnv
γ)

--------------------------------------------------------------------------------
-- | ANF Environments ----------------------------------------------------------
--------------------------------------------------------------------------------
freshNormalVar :: AnfEnv -> Type -> DsM Id
freshNormalVar :: AnfEnv -> Type -> DsM Id
freshNormalVar AnfEnv
γ Type
t = do
  Unique
u     <- forall (m :: * -> *). MonadUnique m => m Unique
getUniqueM
  let i :: Int
i  = Unique -> Int
getKey Unique
u
  let sp :: SrcSpan
sp = SpanStack -> SrcSpan
Sp.srcSpan (AnfEnv -> SpanStack
aeSrcSpan AnfEnv
γ)
  forall (m :: * -> *) a. Monad m => a -> m a
return (OccName -> Unique -> Type -> Type -> SrcSpan -> Id
mkUserLocal (Int -> OccName
anfOcc Int
i) Unique
u Type
Ghc.Many Type
t SrcSpan
sp)

anfOcc :: Int -> OccName
anfOcc :: Int -> OccName
anfOcc = FastString -> OccName
mkVarOccFS forall b c a. (b -> c) -> (a -> b) -> a -> c
. Symbol -> FastString
GM.symbolFastString forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Show a => Symbol -> a -> Symbol
F.intSymbol Symbol
F.anfPrefix

data AnfEnv = AnfEnv
  { AnfEnv -> HashMap StableId Id
aeVarEnv    :: HashMap StableId Id
  -- ^ A mapping between a 'StableId' (see below) and an 'Id'.
  , AnfEnv -> SpanStack
aeSrcSpan   :: Sp.SpanStack
  , AnfEnv -> Config
aeCfg       :: UX.Config
  , AnfEnv -> Int
aeCaseDepth :: !Int
  }

-- | A \"stable\" 'Id'. When transforming 'Core' into ANF notation, we need to keep around a mapping between
-- a particular 'Var' (typically an 'Id') and an 'Id'. Previously this was accomplished using a 'VarEnv',
-- a GHC data structure where keys are 'Unique's. Working with 'Unique' in GHC is not always robust enough
-- when it comes to LH. First of all, the /way/ 'Unique's are constructed might change between GHC versions,
-- and they are not stable between rebuilds/compilations. In the case of this module, in GHC 9 the test
-- BST.hs was failing because two different 'Id's, namely \"wild_X2\" and \"dOrd_X2\" were being given the
-- same 'Unique' by GHC (i.e. \"X2\") which was causing the relevant entry to be overwritten in the 'AnfEnv'
-- causing a unification error.
--
-- A 'StableId' is simply a wrapper over an 'Id' with a different 'Eq' instance that really guarantee
-- uniqueness (for our purposes, anyway).
newtype StableId = StableId Id

instance Eq StableId where
  (StableId Id
id1) == :: StableId -> StableId -> Bool
== (StableId Id
id2) =
    -- We first use the default 'Eq' instance, which works on uniques (basically, integers) and is
    -- efficient. If we get 'False' it means those 'Unique' are really different, but if we get 'True',
    -- we need to be /really/ sure that's the case by using the 'stableNameCmp' function on the 'Name's.
    -- Nothing to do when id1 == id2 as the uniques are /really/ different.
    (Id
id1 forall a. Eq a => a -> a -> Bool
== Id
id2) Bool -> Bool -> Bool
&& (Name -> Name -> Ordering
stableNameCmp (forall a. NamedThing a => a -> Name
getName Id
id1) (forall a. NamedThing a => a -> Name
getName Id
id2) forall a. Eq a => a -> a -> Bool
== Ordering
EQ) -- Avoid unique clashing.

-- For the 'Hashable' instance, we rely on the 'Unique'. This means in pratice there is a tiny chance
-- of collision, but this should only marginally affects the efficiency of the data structure.
instance Hashable StableId where
  hashWithSalt :: Int -> StableId -> Int
hashWithSalt Int
s (StableId Id
id1) = forall a. Hashable a => Int -> a -> Int
hashWithSalt Int
s (Unique -> Int
getKey forall a b. (a -> b) -> a -> b
$ forall a. Uniquable a => a -> Unique
getUnique Id
id1)

-- Shows this 'StableId' by also outputting the associated unique.
instance Show StableId where
  show :: StableId -> String
show (StableId Id
id1) = Name -> String
nameStableString (forall a. NamedThing a => a -> Name
getName Id
id1) forall a. Semigroup a => a -> a -> a
<> String
"_" forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> String
show (forall a. Uniquable a => a -> Unique
getUnique Id
id1)

instance UX.HasConfig AnfEnv where
  getConfig :: AnfEnv -> Config
getConfig = AnfEnv -> Config
aeCfg

emptyAnfEnv :: UX.Config -> AnfEnv
emptyAnfEnv :: Config -> AnfEnv
emptyAnfEnv Config
cfg = AnfEnv
  { aeVarEnv :: HashMap StableId Id
aeVarEnv    = forall a. Monoid a => a
mempty
  , aeSrcSpan :: SpanStack
aeSrcSpan   = SpanStack
Sp.empty
  , aeCfg :: Config
aeCfg       = Config
cfg
  , aeCaseDepth :: Int
aeCaseDepth = Int
1
  }

lookupAnfEnv :: AnfEnv -> Id -> Id -> Id
lookupAnfEnv :: AnfEnv -> Id -> Id -> Id
lookupAnfEnv AnfEnv
γ Id
x (Id -> StableId
StableId -> StableId
y) = forall k v. (Eq k, Hashable k) => v -> k -> HashMap k v -> v
HM.lookupDefault Id
x StableId
y (AnfEnv -> HashMap StableId Id
aeVarEnv AnfEnv
γ)

extendAnfEnv :: AnfEnv -> Id -> Id -> AnfEnv
extendAnfEnv :: AnfEnv -> Id -> Id -> AnfEnv
extendAnfEnv AnfEnv
γ (Id -> StableId
StableId -> StableId
x) Id
y = AnfEnv
γ { aeVarEnv :: HashMap StableId Id
aeVarEnv = forall k v.
(Eq k, Hashable k) =>
k -> v -> HashMap k v -> HashMap k v
HM.insert StableId
x Id
y (AnfEnv -> HashMap StableId Id
aeVarEnv AnfEnv
γ) }

incrCaseDepth :: AltCon -> AnfEnv -> AnfEnv
incrCaseDepth :: AltCon -> AnfEnv -> AnfEnv
incrCaseDepth AltCon
DEFAULT AnfEnv
γ = AnfEnv
γ { aeCaseDepth :: Int
aeCaseDepth = Int
1 forall a. Num a => a -> a -> a
+ AnfEnv -> Int
aeCaseDepth AnfEnv
γ }
incrCaseDepth AltCon
_       AnfEnv
γ = AnfEnv
γ

at :: AnfEnv -> CoreTickish -> AnfEnv
at :: AnfEnv -> CoreTickish -> AnfEnv
at AnfEnv
γ CoreTickish
tt = AnfEnv
γ { aeSrcSpan :: SpanStack
aeSrcSpan = Span -> SpanStack -> SpanStack
Sp.push (CoreTickish -> Span
Sp.Tick CoreTickish
tt) (AnfEnv -> SpanStack
aeSrcSpan AnfEnv
γ)}