--------------------------------------------------------------------------------
-- | 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 (forM)
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 (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
    String -> IO ()
putStrLn String
"***************************** GHC CoreBinds ***************************"
    String -> IO ()
putStrLn (String -> IO ()) -> String -> IO ()
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 (String -> IO ()) -> String -> IO ()
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 (String -> IO ()) -> String -> IO ()
forall a b. (a -> b) -> a -> b
$ Bool -> [CoreBind] -> String
GM.showCBs Bool
untidy [CoreBind]
rwr_cbs
  [CoreBind] -> Maybe [CoreBind] -> [CoreBind]
forall a. a -> Maybe a -> a
fromMaybe [CoreBind]
forall {a}. a
err (Maybe [CoreBind] -> [CoreBind])
-> ((Messages DsMessage, Maybe [CoreBind]) -> Maybe [CoreBind])
-> (Messages DsMessage, Maybe [CoreBind])
-> [CoreBind]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Messages DsMessage, Maybe [CoreBind]) -> Maybe [CoreBind]
forall a b. (a, b) -> b
snd ((Messages DsMessage, Maybe [CoreBind]) -> [CoreBind])
-> IO (Messages DsMessage, Maybe [CoreBind]) -> IO [CoreBind]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> HscEnv
-> ModGuts
-> DsM [CoreBind]
-> IO (Messages DsMessage, Maybe [CoreBind])
forall a.
HscEnv -> ModGuts -> DsM a -> IO (Messages DsMessage, Maybe a)
initDsWithModGuts HscEnv
hscEnv ModGuts
modGuts DsM [CoreBind]
act -- hscEnv m grEnv tEnv emptyFamInstEnv act
    where
      err :: a
err      = Maybe SrcSpan -> String -> a
forall a. Maybe SrcSpan -> String -> a
panic Maybe SrcSpan
forall a. Maybe a
Nothing String
"Oops, cannot A-Normalize GHC Core!"
      act :: DsM [CoreBind]
act      = (CoreBind -> DsM [CoreBind]) -> [CoreBind] -> DsM [CoreBind]
forall (m :: * -> *) (t :: * -> *) a b.
(Monad m, Traversable t) =>
(a -> m [b]) -> t a -> m [b]
Misc.concatMapM (AnfEnv -> CoreBind -> DsM [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) ([CoreBind] -> [CoreBind]) -> [CoreBind] -> [CoreBind]
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 -> DsM [CoreBind]
normalizeTopBind AnfEnv
γ (NonRec Id
x Expr Id
e)
  = do Expr Id
e' <- DsM (Expr Id) -> DsM (Expr Id)
forall a. DsM a -> DsM a
runDsM (DsM (Expr Id) -> DsM (Expr Id)) -> DsM (Expr Id) -> DsM (Expr Id)
forall a b. (a -> b) -> a -> b
$ StateT DsST DsM (Expr Id) -> DsST -> DsM (Expr Id)
forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m a
evalStateT (AnfEnv -> Expr Id -> StateT DsST DsM (Expr Id)
stitch AnfEnv
γ Expr Id
e) ([CoreBind] -> DsST
DsST [])
       [CoreBind] -> DsM [CoreBind]
forall a. a -> IOEnv (Env DsGblEnv DsLclEnv) a
forall (m :: * -> *) a. Monad m => a -> m a
return [CoreBind -> CoreBind
normalizeTyVars (CoreBind -> CoreBind) -> CoreBind -> CoreBind
forall a b. (a -> b) -> a -> b
$ Id -> Expr Id -> CoreBind
forall b. b -> Expr b -> Bind b
NonRec Id
x Expr Id
e']

normalizeTopBind AnfEnv
γ (Rec [(Id, Expr Id)]
xes)
  = do DsST
xes' <- DsM DsST -> DsM DsST
forall a. DsM a -> DsM a
runDsM (DsM DsST -> DsM DsST) -> DsM DsST -> DsM DsST
forall a b. (a -> b) -> a -> b
$ StateT DsST DsM () -> DsST -> DsM DsST
forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m s
execStateT (AnfEnv -> CoreBind -> StateT DsST DsM ()
normalizeBind AnfEnv
γ ([(Id, Expr Id)] -> CoreBind
forall b. [(b, Expr b)] -> Bind b
Rec [(Id, Expr Id)]
xes)) ([CoreBind] -> DsST
DsST [])
       [CoreBind] -> DsM [CoreBind]
forall a. a -> IOEnv (Env DsGblEnv DsLclEnv) a
forall (m :: * -> *) a. Monad m => a -> m a
return ([CoreBind] -> DsM [CoreBind]) -> [CoreBind] -> DsM [CoreBind]
forall a b. (a -> b) -> a -> b
$ (CoreBind -> CoreBind) -> [CoreBind] -> [CoreBind]
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) = Id -> Expr Id -> CoreBind
forall b. b -> Expr b -> Bind b
NonRec (Id -> Type -> Id
setVarType Id
x Type
t') (Expr Id -> CoreBind) -> Expr Id -> CoreBind
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 " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Id -> String
forall a. Outputable a => a -> String
GM.showPpr Id
x
    as' :: [Id]
as'      = ([Id], Type) -> [Id]
forall a b. (a, b) -> a
fst (([Id], Type) -> [Id]) -> ([Id], Type) -> [Id]
forall a b. (a -> b) -> a -> b
$ Type -> ([Id], Type)
splitForAllTyCoVars (Type -> ([Id], Type)) -> Type -> ([Id], Type)
forall a b. (a -> b) -> a -> b
$ (() :: Constraint) => Expr Id -> Type
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)    = [(Id, Expr Id)] -> CoreBind
forall b. [(b, Expr b)] -> Bind b
Rec [(Id, Expr Id)]
xes'
  where
    nrec :: [CoreBind]
nrec     = CoreBind -> CoreBind
normalizeTyVars (CoreBind -> CoreBind) -> [CoreBind] -> [CoreBind]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((Id -> Expr Id -> CoreBind) -> (Id, Expr Id) -> CoreBind
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Id -> Expr Id -> CoreBind
forall b. b -> Expr b -> Bind b
NonRec ((Id, Expr Id) -> CoreBind) -> [(Id, Expr Id)] -> [CoreBind]
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
_ -> Maybe SrcSpan -> String -> (Id, Expr Id)
forall a. Maybe SrcSpan -> String -> a
impossible Maybe SrcSpan
forall a. Maybe a
Nothing String
"This cannot happen") (CoreBind -> (Id, Expr Id)) -> [CoreBind] -> [(Id, Expr Id)]
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
  | [Id] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Id]
as Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [Id] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Id]
as'
  = [ForAllTyBinder] -> Type -> Type
mkForAllTys (Id -> ForAllTyBinder
mkTyArg (Id -> ForAllTyBinder) -> [Id] -> [ForAllTyBinder]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Id]
as') (Type -> Type) -> Type -> Type
forall a b. (a -> b) -> a -> b
$ (() :: Constraint) => Subst -> Type -> Type
Subst -> Type -> Type
substTy Subst
su Type
bt
  | Bool
otherwise
  = String -> Type -> Type
forall a. String -> a -> a
trace String
msg (Type -> Type) -> Type -> Type
forall a b. (a -> b) -> a -> b
$ [ForAllTyBinder] -> Type -> Type
mkForAllTys (Id -> ForAllTyBinder
mkTyArg (Id -> ForAllTyBinder) -> [Id] -> [ForAllTyBinder]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Id]
as) Type
bt
  where su :: Subst
su = [(Id, Type)] -> Subst
mkTvSubstPrs ([(Id, Type)] -> Subst) -> [(Id, Type)] -> Subst
forall a b. (a -> b) -> a -> b
$ [Id] -> [Type] -> [(Id, Type)]
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
_ -> [Id] -> Expr Id -> Expr Id
forall b. [b] -> Expr b -> Expr b
mkLams [Id]
tvs (Expr Id -> [Type] -> Expr Id
forall b. Expr b -> [Type] -> Expr b
mkTyApps Expr Id
e ((Id -> Type) -> [Id] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Id -> Type
mkTyVarTy [Id]
tvs))
  where
  ([Id]
tvs, Type
_) = Type -> ([Id], Type)
splitForAllTyCoVars ((() :: Constraint) => Expr Id -> Type
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 -> b) -> DsM a -> DsM b)
-> (forall a b. a -> DsM b -> DsM a) -> Functor DsM
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
$cfmap :: forall a b. (a -> b) -> DsM a -> DsM b
fmap :: forall a b. (a -> b) -> DsM a -> DsM b
$c<$ :: forall a b. a -> DsM b -> DsM a
<$ :: forall a b. a -> DsM b -> DsM a
Functor, Applicative DsM
Applicative DsM =>
(forall a b. DsM a -> (a -> DsM b) -> DsM b)
-> (forall a b. DsM a -> DsM b -> DsM b)
-> (forall a. a -> DsM a)
-> Monad 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
$c>>= :: forall a b. DsM a -> (a -> DsM b) -> DsM b
>>= :: forall a b. DsM a -> (a -> DsM b) -> DsM b
$c>> :: forall a b. DsM a -> DsM b -> DsM b
>> :: forall a b. DsM a -> DsM b -> DsM b
$creturn :: forall a. a -> DsM a
return :: forall a. a -> DsM a
Monad, Monad DsM
DsM [Unique]
DsM Unique
DsM UniqSupply
Monad DsM =>
DsM UniqSupply -> DsM Unique -> DsM [Unique] -> MonadUnique DsM
forall (m :: * -> *).
Monad m =>
m UniqSupply -> m Unique -> m [Unique] -> MonadUnique m
$cgetUniqueSupplyM :: DsM UniqSupply
getUniqueSupplyM :: DsM UniqSupply
$cgetUniqueM :: DsM Unique
getUniqueM :: DsM Unique
$cgetUniquesM :: DsM [Unique]
getUniquesM :: DsM [Unique]
MonadUnique, Functor DsM
Functor DsM =>
(forall a. a -> DsM a)
-> (forall a b. DsM (a -> b) -> DsM a -> DsM b)
-> (forall a b c. (a -> b -> c) -> DsM a -> DsM b -> DsM c)
-> (forall a b. DsM a -> DsM b -> DsM b)
-> (forall a b. DsM a -> DsM b -> DsM a)
-> Applicative 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
$cpure :: forall a. a -> DsM a
pure :: forall a. a -> DsM a
$c<*> :: forall a b. DsM (a -> b) -> DsM a -> DsM b
<*> :: forall a b. DsM (a -> b) -> DsM a -> DsM b
$cliftA2 :: forall a b c. (a -> b -> c) -> DsM a -> DsM b -> DsM c
liftA2 :: forall a b c. (a -> b -> c) -> DsM a -> DsM b -> DsM c
$c*> :: forall a b. DsM a -> DsM b -> DsM b
*> :: forall a b. DsM a -> DsM b -> DsM b
$c<* :: forall a b. DsM a -> DsM b -> DsM a
<* :: forall a b. DsM a -> DsM b -> DsM a
Applicative)

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

type DsMW = StateT DsST DsM

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

normalizeBind AnfEnv
γ (Rec [(Id, Expr Id)]
xes)
  = do [Expr Id]
es' <- (Expr Id -> StateT DsST DsM (Expr Id))
-> [Expr Id] -> StateT DsST DsM [Expr Id]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (AnfEnv -> Expr Id -> StateT DsST DsM (Expr Id)
stitch AnfEnv
γ) [Expr Id]
es
       [CoreBind] -> StateT DsST DsM ()
add [[(Id, Expr Id)] -> CoreBind
forall b. [(b, Expr b)] -> Bind b
Rec ([Id] -> [Expr Id] -> [(Id, Expr Id)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Id]
xs [Expr Id]
es')]
    where
       ([Id]
xs, [Expr Id]
es) = [(Id, Expr Id)] -> ([Id], [Expr Id])
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 -> StateT DsST DsM (Expr Id)
normalizeName AnfEnv
γ e :: Expr Id
e@(Lit Literal
l)
  | Literal -> Bool
shouldNormalize Literal
l
  = AnfEnv -> Expr Id -> StateT DsST DsM (Expr Id)
normalizeLiteral AnfEnv
γ Expr Id
e
  | Bool
otherwise
  = Expr Id -> StateT DsST DsM (Expr Id)
forall a. a -> StateT DsST DsM a
forall (m :: * -> *) a. Monad m => a -> m a
return Expr Id
e

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

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

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

normalizeName AnfEnv
γ (Tick CoreTickish
tt Expr Id
e)
  = do Expr Id
e'    <- AnfEnv -> Expr Id -> StateT DsST DsM (Expr Id)
normalizeName (AnfEnv
γ AnfEnv -> CoreTickish -> AnfEnv
`at` CoreTickish
tt) Expr Id
e
       Expr Id -> StateT DsST DsM (Expr Id)
forall a. a -> StateT DsST DsM a
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr Id -> StateT DsST DsM (Expr Id))
-> Expr Id -> StateT DsST DsM (Expr Id)
forall a b. (a -> b) -> a -> b
$ CoreTickish -> Expr Id -> Expr Id
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 -> StateT DsST DsM (Expr Id)
normalize AnfEnv
γ Expr Id
e
       Id
x    <- DsM Id -> StateT DsST DsM Id
forall (m :: * -> *) a. Monad m => m a -> StateT DsST m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (DsM Id -> StateT DsST DsM Id) -> DsM Id -> StateT DsST DsM Id
forall a b. (a -> b) -> a -> b
$ AnfEnv -> Type -> DsM Id
freshNormalVar AnfEnv
γ (Type -> DsM Id) -> Type -> DsM Id
forall a b. (a -> b) -> a -> b
$ (() :: Constraint) => Expr Id -> Type
Expr Id -> Type
exprType Expr Id
e
       [CoreBind] -> StateT DsST DsM ()
add [Id -> Expr Id -> CoreBind
forall b. b -> Expr b -> Bind b
NonRec Id
x Expr Id
e']
       Expr Id -> StateT DsST DsM (Expr Id)
forall a. a -> StateT DsST DsM a
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr Id -> StateT DsST DsM (Expr Id))
-> Expr Id -> StateT DsST DsM (Expr Id)
forall a b. (a -> b) -> a -> b
$ Id -> Expr Id
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] -> StateT DsST DsM ()
add [CoreBind]
w = (DsST -> DsST) -> StateT DsST DsM ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((DsST -> DsST) -> StateT DsST DsM ())
-> (DsST -> DsST) -> StateT DsST DsM ()
forall a b. (a -> b) -> a -> b
$ \DsST
s -> DsST
s { st_binds = st_binds s ++ w}

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

--------------------------------------------------------------------------------
normalize :: AnfEnv -> CoreExpr -> DsMW CoreExpr
--------------------------------------------------------------------------------
normalize :: AnfEnv -> Expr Id -> StateT DsST DsM (Expr Id)
normalize AnfEnv
γ Expr Id
e
  | AnfEnv -> Bool
forall t. HasConfig t => t -> Bool
UX.patternFlag AnfEnv
γ
  , Just Pattern
p <- Expr Id -> Maybe Pattern
Rs.lift Expr Id
e
  = AnfEnv -> Pattern -> StateT DsST DsM (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 -> StateT DsST DsM (Expr Id)
normalize AnfEnv
γ Expr Id
e
       Expr Id -> StateT DsST DsM (Expr Id)
forall a. a -> StateT DsST DsM a
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr Id -> StateT DsST DsM (Expr Id))
-> Expr Id -> StateT DsST DsM (Expr Id)
forall a b. (a -> b) -> a -> b
$ Id -> Expr Id -> Expr Id
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 -> StateT DsST DsM (Expr Id)
stitch AnfEnv
γ Expr Id
e
       Expr Id -> StateT DsST DsM (Expr Id)
forall a. a -> StateT DsST DsM a
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr Id -> StateT DsST DsM (Expr Id))
-> Expr Id -> StateT DsST DsM (Expr Id)
forall a b. (a -> b) -> a -> b
$ Id -> Expr Id -> Expr Id
forall b. b -> Expr b -> Expr b
Lam Id
x Expr Id
e'

normalize AnfEnv
γ (Let CoreBind
b Expr Id
e)
  = do AnfEnv -> CoreBind -> StateT DsST DsM ()
normalizeBind AnfEnv
γ CoreBind
b
       AnfEnv -> Expr Id -> StateT DsST DsM (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 -> StateT DsST DsM (Expr Id)
normalizeName AnfEnv
γ Expr Id
e
       Id
x'    <- DsM Id -> StateT DsST DsM Id
forall (m :: * -> *) a. Monad m => m a -> StateT DsST m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (DsM Id -> StateT DsST DsM Id) -> DsM Id -> StateT DsST DsM Id
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'   <- [Alt Id]
-> (Alt Id -> StateT DsST DsM (Alt Id)) -> StateT DsST DsM [Alt Id]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [Alt Id]
as ((Alt Id -> StateT DsST DsM (Alt Id)) -> StateT DsST DsM [Alt Id])
-> (Alt Id -> StateT DsST DsM (Alt Id)) -> StateT DsST DsM [Alt Id]
forall a b. (a -> b) -> a -> b
$ \(Alt AltCon
c [Id]
xs Expr Id
e') -> (Expr Id -> Alt Id)
-> StateT DsST DsM (Expr Id) -> StateT DsST DsM (Alt Id)
forall a b. (a -> b) -> StateT DsST DsM a -> StateT DsST DsM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (AltCon -> [Id] -> Expr Id -> Alt Id
forall b. AltCon -> [b] -> Expr b -> Alt b
Alt AltCon
c [Id]
xs) (AnfEnv -> Expr Id -> StateT DsST DsM (Expr Id)
stitch (AltCon -> AnfEnv -> AnfEnv
incrCaseDepth AltCon
c AnfEnv
γ') Expr Id
e')
       [Alt Id]
as''  <- DsM [Alt Id] -> StateT DsST DsM [Alt Id]
forall (m :: * -> *) a. Monad m => m a -> StateT DsST m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (DsM [Alt Id] -> StateT DsST DsM [Alt Id])
-> DsM [Alt Id] -> StateT DsST DsM [Alt Id]
forall a b. (a -> b) -> a -> b
$ AnfEnv -> Type -> [Alt Id] -> DsM [Alt Id]
expandDefaultCase AnfEnv
γ Type
τx [Alt Id]
as'
       Expr Id -> StateT DsST DsM (Expr Id)
forall a. a -> StateT DsST DsM a
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr Id -> StateT DsST DsM (Expr Id))
-> Expr Id -> StateT DsST DsM (Expr Id)
forall a b. (a -> b) -> a -> b
$ Expr Id -> Id -> Type -> [Alt Id] -> Expr Id
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)
  = Expr Id -> StateT DsST DsM (Expr Id)
forall a. a -> StateT DsST DsM a
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr Id -> StateT DsST DsM (Expr Id))
-> Expr Id -> StateT DsST DsM (Expr Id)
forall a b. (a -> b) -> a -> b
$ Id -> Expr Id
forall b. Id -> Expr b
Var (AnfEnv -> Id -> Id -> Id
lookupAnfEnv AnfEnv
γ Id
x Id
x)

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

normalize AnfEnv
_ e :: Expr Id
e@(Type Type
_)
  = Expr Id -> StateT DsST DsM (Expr Id)
forall a. a -> StateT DsST DsM a
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 -> StateT DsST DsM (Expr Id)
normalizeName AnfEnv
γ Expr Id
e
       Expr Id -> StateT DsST DsM (Expr Id)
forall a. a -> StateT DsST DsM a
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr Id -> StateT DsST DsM (Expr Id))
-> Expr Id -> StateT DsST DsM (Expr Id)
forall a b. (a -> b) -> a -> b
$ Expr Id -> Coercion -> Expr Id
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 -> StateT DsST DsM (Expr Id)
normalize AnfEnv
γ Expr Id
e1
       Expr Id
e2' <- AnfEnv -> Expr Id -> StateT DsST DsM (Expr Id)
normalize AnfEnv
γ Expr Id
e2
       Expr Id -> StateT DsST DsM (Expr Id)
forall a. a -> StateT DsST DsM a
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr Id -> StateT DsST DsM (Expr Id))
-> Expr Id -> StateT DsST DsM (Expr Id)
forall a b. (a -> b) -> a -> b
$ Expr Id -> Expr Id -> Expr Id
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 -> StateT DsST DsM (Expr Id)
normalize AnfEnv
γ Expr Id
e1
       Expr Id
n2  <- AnfEnv -> Expr Id -> StateT DsST DsM (Expr Id)
normalizeName AnfEnv
γ Expr Id
e2
       Expr Id -> StateT DsST DsM (Expr Id)
forall a. a -> StateT DsST DsM a
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr Id -> StateT DsST DsM (Expr Id))
-> Expr Id -> StateT DsST DsM (Expr Id)
forall a b. (a -> b) -> a -> b
$ Expr Id -> Expr Id -> Expr Id
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 -> StateT DsST DsM (Expr Id)
normalize (AnfEnv
γ AnfEnv -> CoreTickish -> AnfEnv
`at` CoreTickish
tt) Expr Id
e
       Expr Id -> StateT DsST DsM (Expr Id)
forall a. a -> StateT DsST DsM a
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr Id -> StateT DsST DsM (Expr Id))
-> Expr Id -> StateT DsST DsM (Expr Id)
forall a b. (a -> b) -> a -> b
$ CoreTickish -> Expr Id -> Expr Id
forall b. CoreTickish -> Expr b -> Expr b
Tick CoreTickish
tt Expr Id
e'

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

--------------------------------------------------------------------------------
stitch :: AnfEnv -> CoreExpr -> DsMW CoreExpr
--------------------------------------------------------------------------------
stitch :: AnfEnv -> Expr Id -> StateT DsST DsM (Expr Id)
stitch AnfEnv
γ Expr Id
e
  = do DsST
bs'   <- StateT DsST DsM DsST
forall s (m :: * -> *). MonadState s m => m s
get
       (DsST -> DsST) -> StateT DsST DsM ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((DsST -> DsST) -> StateT DsST DsM ())
-> (DsST -> DsST) -> StateT DsST DsM ()
forall a b. (a -> b) -> a -> b
$ \DsST
s -> DsST
s { st_binds = [] }
       Expr Id
e'    <- AnfEnv -> Expr Id -> StateT DsST DsM (Expr Id)
normalize AnfEnv
γ Expr Id
e
       [CoreBind]
bs    <- DsST -> [CoreBind]
st_binds (DsST -> [CoreBind])
-> StateT DsST DsM DsST -> StateT DsST DsM [CoreBind]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> StateT DsST DsM DsST
forall s (m :: * -> *). MonadState s m => m s
get
       DsST -> StateT DsST DsM ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put DsST
bs'
       -- See Note [Shape of normalized terms]
       let ([Id]
tvs, Expr Id
e'') = Expr Id -> ([Id], Expr Id)
collectTyBinders Expr Id
e'
       Expr Id -> StateT DsST DsM (Expr Id)
forall a. a -> StateT DsST DsM a
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr Id -> StateT DsST DsM (Expr Id))
-> Expr Id -> StateT DsST DsM (Expr Id)
forall a b. (a -> b) -> a -> b
$ [Id] -> Expr Id -> Expr Id
forall b. [b] -> Expr b -> Expr b
mkLams [Id]
tvs ([CoreBind] -> Expr Id -> Expr Id
mkCoreLets [CoreBind]
bs Expr Id
e'')

-- Note [Shape of normalized terms]
--
-- The termination checker in
-- Language.Haskell.Liquid.Constraint.Termination.collectArguments expects the
-- type binders to come before lets:
--
-- > \ (@a) -> let ... in \ b c d -> ...
--
-- Therefore, stitch makes sure to insert new lets after the type binders
--
-- > \ (@a) -> let lqanf... = ... in let ... in \ b c d -> ...
--

_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)    = String -> (Expr Id, [CoreBind]) -> (Expr Id, [CoreBind])
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 -> StateT DsST DsM (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 -> StateT DsST DsM (Expr Id)
normalize AnfEnv
γ (Pattern -> Expr Id
Rs.patE1 Pattern
p)
  Expr Id
e2'   <- AnfEnv -> Expr Id -> StateT DsST DsM (Expr Id)
stitch    AnfEnv
γ (Pattern -> Expr Id
Rs.patE2 Pattern
p)
  Expr Id -> StateT DsST DsM (Expr Id)
forall a. a -> StateT DsST DsM a
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr Id -> StateT DsST DsM (Expr Id))
-> Expr Id -> StateT DsST DsM (Expr Id)
forall a b. (a -> b) -> a -> b
$ Pattern -> Expr Id
Rs.lower Pattern
p { Rs.patE1 = e1', Rs.patE2 = e2' }

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

normalizePattern AnfEnv
_ p :: Pattern
p@(Rs.PatProject {}) =
  Expr Id -> StateT DsST DsM (Expr Id)
forall a. a -> StateT DsST DsM a
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 -> StateT DsST DsM (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 -> StateT DsST DsM (Expr Id)
normalize AnfEnv
γ (Pattern -> Expr Id
Rs.patE Pattern
p)
  Expr Id -> StateT DsST DsM (Expr Id)
forall a. a -> StateT DsST DsM a
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr Id -> StateT DsST DsM (Expr Id))
-> Expr Id -> StateT DsST DsM (Expr Id)
forall a b. (a -> b) -> a -> b
$ Pattern -> Expr Id
Rs.lower Pattern
p { Rs.patE = e' }


--------------------------------------------------------------------------------
expandDefault :: AnfEnv -> Bool
--------------------------------------------------------------------------------
expandDefault :: AnfEnv -> Bool
expandDefault AnfEnv
γ = AnfEnv -> Int
aeCaseDepth AnfEnv
γ Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= AnfEnv -> Int
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 [DataCon] -> [DataCon] -> [DataCon]
forall a. Eq a => [a] -> [a] -> [a]
\\ [ DataCon
d | Alt (DataAlt DataCon
d) [Id]
_ Expr Id
_ <- [Alt Id]
dcs]
                     let n :: Int
n   = [DataCon] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [DataCon]
ds'
                     if Int
n Int -> Int -> Bool
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 AnfEnv -> Int
forall t. HasConfig t => t -> Int
maxCaseExpand AnfEnv
γ Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
2
                            then [Alt Id] -> DsM [Alt Id]
forall a. a -> DsM a
forall (m :: * -> *) a. Monad m => a -> m a
return [Alt Id]
z
                            else [Alt Id] -> DsM [Alt Id]
forall a. a -> DsM a
forall (m :: * -> *) a. Monad m => a -> m a
return (String -> [Alt Id] -> [Alt Id]
forall a. String -> a -> a
trace (Bool -> AnfEnv -> Int -> String
expandMessage Bool
False AnfEnv
γ Int
n) [Alt Id]
z)
       Maybe [DataCon]
Nothing -> [Alt Id] -> DsM [Alt Id]
forall a. a -> DsM a
forall (m :: * -> *) a. Monad m => a -> m a
return [Alt Id]
z --

expandDefaultCase AnfEnv
_ Type
_ [Alt Id]
z
   = [Alt Id] -> DsM [Alt Id]
forall a. a -> DsM a
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) (Alt Id -> AltCon) -> [Alt Id] -> [AltCon]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Alt Id]
dcs) = do
      [Alt Id]
dcs'    <- AnfEnv -> [Alt Id] -> [Alt Id]
forall a. AnfEnv -> [a] -> [a]
warnCaseExpand AnfEnv
γ ([Alt Id] -> [Alt Id]) -> DsM [Alt Id] -> DsM [Alt Id]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(DataCon, [Id], [Type])]
-> ((DataCon, [Id], [Type]) -> DsM (Alt Id)) -> DsM [Alt Id]
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)
      [Alt Id] -> DsM [Alt Id]
forall a. a -> DsM a
forall (m :: * -> *) a. Monad m => a -> m a
return   ([Alt Id] -> DsM [Alt Id]) -> [Alt Id] -> DsM [Alt Id]
forall a b. (a -> b) -> a -> b
$ [Alt Id] -> [Alt Id]
sortCases ([Alt Id]
dcs' [Alt Id] -> [Alt Id] -> [Alt Id]
forall a. [a] -> [a] -> [a]
++ [Alt Id]
dcs)
expandDefaultCase' AnfEnv
_ Type
_ [Alt Id]
z
   = [Alt Id] -> DsM [Alt Id]
forall a. a -> DsM a
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  <- (Type -> DsM Id) -> [Type] -> DsM [Id]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (AnfEnv -> Type -> DsM Id
freshNormalVar AnfEnv
γ) [Type]
ts
  Alt Id -> DsM (Alt Id)
forall a. a -> DsM a
forall (m :: * -> *) a. Monad m => a -> m a
return (AltCon -> [Id] -> Expr Id -> Alt Id
forall b. AltCon -> [b] -> Expr b -> Alt b
Alt (DataCon -> AltCon
DataAlt DataCon
d) ([Id]
as [Id] -> [Id] -> [Id]
forall a. [a] -> [a] -> [a]
++ [Id]
xs) Expr Id
e)

sortCases :: [CoreAlt] -> [CoreAlt]
sortCases :: [Alt Id] -> [Alt Id]
sortCases = (Alt Id -> Alt Id -> Ordering) -> [Alt Id] -> [Alt Id]
forall a. (a -> a -> Ordering) -> [a] -> [a]
sortBy Alt Id -> Alt Id -> Ordering
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 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n          = String -> [a] -> [a]
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                = [a] -> Int
forall a. [a] -> Int
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 [String
forall {t}. PrintfType t => t
msg1, String
forall {t}. PrintfType t => t
msg2]
  where
    msg1 :: t
msg1            = String -> String -> String -> Int -> Int -> t
forall r. PrintfType r => String -> r
Printf.printf String
"WARNING: (%s) %s DEFAULT with %d cases at depth %d" (SrcSpan -> String
forall a. Outputable a => a -> String
showPpr SrcSpan
sp) String
v1 Int
n Int
d
    msg2 :: t
msg2            = String -> String -> Int -> t
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
dInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) :: (String, String, Int)
      | Bool
otherwise   = (String
"Not expanding", String
"Enable" , Int
dInt -> Int -> Int
forall 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     <- DsM Unique
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
γ)
  Id -> DsM Id
forall a. a -> DsM a
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.ManyTy Type
t SrcSpan
sp)

anfOcc :: Int -> OccName
anfOcc :: Int -> OccName
anfOcc = FastString -> OccName
mkVarOccFS (FastString -> OccName) -> (Int -> FastString) -> Int -> OccName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Symbol -> FastString
GM.symbolFastString (Symbol -> FastString) -> (Int -> Symbol) -> Int -> FastString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Symbol -> Int -> Symbol
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 Id -> Id -> Bool
forall a. Eq a => a -> a -> Bool
== Id
id2) Bool -> Bool -> Bool
&& (Name -> Name -> Ordering
stableNameCmp (Id -> Name
forall a. NamedThing a => a -> Name
getName Id
id1) (Id -> Name
forall a. NamedThing a => a -> Name
getName Id
id2) Ordering -> Ordering -> Bool
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) = Int -> Int -> Int
forall a. Hashable a => Int -> a -> Int
hashWithSalt Int
s (Unique -> Int
getKey (Unique -> Int) -> Unique -> Int
forall a b. (a -> b) -> a -> b
$ Id -> Unique
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 (Id -> Name
forall a. NamedThing a => a -> Name
getName Id
id1) String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
"_" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Unique -> String
forall a. Show a => a -> String
show (Id -> Unique
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    = HashMap StableId Id
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) = Id -> StableId -> HashMap StableId Id -> Id
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 = HM.insert x y (aeVarEnv γ) }

incrCaseDepth :: AltCon -> AnfEnv -> AnfEnv
incrCaseDepth :: AltCon -> AnfEnv -> AnfEnv
incrCaseDepth AltCon
DEFAULT AnfEnv
γ = AnfEnv
γ { aeCaseDepth = 1 + aeCaseDepth γ }
incrCaseDepth AltCon
_       AnfEnv
γ = AnfEnv
γ

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