--------------------------------------------------------------------------------
-- | 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 ( get, mkTyArg
                                                                , showPpr
                                                                , DsM
                                                                , panic)
import qualified Liquid.GHC.API  as Ghc
import           Control.Monad (forM)
import           Control.Monad.State.Lazy
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.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
  [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. HasCallStack => 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) (ModGuts -> [CoreBind]
mg_binds ModGuts
modGuts)
      γ0 :: AnfEnv
γ0       = Config -> AnfEnv
emptyAnfEnv 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 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 [])
       return [normalizeTyVars $ NonRec x e']

normalizeTopBind AnfEnv
γ (Rec [(Id, Expr Id)]
xes)
  = do 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 [])
       return $ map normalizeTyVars (st_binds 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
$ HasDebugCallStack => 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
$ HasDebugCallStack => 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 (HasDebugCallStack => 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 e' <- AnfEnv -> Expr Id -> StateT DsST DsM (Expr Id)
normalize AnfEnv
γ Expr Id
e
       add [NonRec x e']

normalizeBind AnfEnv
γ (Rec [(Id, Expr Id)]
xes)
  = do 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
       add [Rec (zip xs 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 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
$ HasDebugCallStack => Expr Id -> Type
Expr Id -> Type
exprType Expr Id
e
       add  [NonRec x e]
       return $ Var x

normalizeName AnfEnv
γ (Tick CoreTickish
tt Expr Id
e)
  = do e'    <- AnfEnv -> Expr Id -> StateT DsST DsM (Expr Id)
normalizeName (AnfEnv
γ AnfEnv -> CoreTickish -> AnfEnv
`at` CoreTickish
tt) Expr Id
e
       return $ Tick tt e'

normalizeName AnfEnv
γ Expr Id
e
  = do e'   <- AnfEnv -> Expr Id -> StateT DsST DsM (Expr Id)
normalize AnfEnv
γ Expr Id
e
       x    <- lift $ freshNormalVar γ $ exprType e
       add [NonRec x e']
       return $ Var 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 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
$ HasDebugCallStack => Expr Id -> Type
Expr Id -> Type
exprType Expr Id
e
     add [NonRec x e]
     return $ Var 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 e' <- AnfEnv -> Expr Id -> StateT DsST DsM (Expr Id)
normalize AnfEnv
γ Expr Id
e
       return $ Lam x e'

normalize AnfEnv
γ (Lam Id
x Expr Id
e)
  = do e' <- AnfEnv -> Expr Id -> StateT DsST DsM (Expr Id)
stitch AnfEnv
γ Expr Id
e
       return $ Lam x 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 n     <- AnfEnv -> Expr Id -> StateT DsST DsM (Expr Id)
normalizeName AnfEnv
γ Expr Id
e
       x'    <- lift $ freshNormalVar γ τx -- rename "wild" to avoid shadowing
       let γ' = AnfEnv -> Id -> Id -> AnfEnv
extendAnfEnv AnfEnv
γ Id
x Id
x'
       as'   <- forM as $ \(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')
       as''  <- lift $ expandDefaultCase γ τx as'
       return $ Case n x' t 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 e' <- AnfEnv -> Expr Id -> StateT DsST DsM (Expr Id)
normalizeName AnfEnv
γ Expr Id
e
       return $ Cast e' τ

normalize AnfEnv
γ (App Expr Id
e1 e2 :: Expr Id
e2@(Type Type
_))
  = do e1' <- AnfEnv -> Expr Id -> StateT DsST DsM (Expr Id)
normalize AnfEnv
γ Expr Id
e1
       e2' <- normalize γ e2
       return $ App e1' e2'

normalize AnfEnv
γ (App Expr Id
e1 Expr Id
e2)
  = do e1' <- AnfEnv -> Expr Id -> StateT DsST DsM (Expr Id)
normalize AnfEnv
γ Expr Id
e1
       n2  <- normalizeName γ e2
       return $ App e1' n2

normalize AnfEnv
γ (Tick CoreTickish
tt Expr Id
e)
  = do e' <- AnfEnv -> Expr Id -> StateT DsST DsM (Expr Id)
normalize (AnfEnv
γ AnfEnv -> CoreTickish -> AnfEnv
`at` CoreTickish
tt) Expr Id
e
       return $ Tick tt 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 bs'   <- StateT DsST DsM DsST
forall s (m :: * -> *). MonadState s m => m s
get
       modify $ \DsST
s -> DsST
s { st_binds = [] }
       e'    <- normalize γ e
       bs    <- st_binds <$> get
       put bs'
       -- See Note [Shape of normalized terms]
       let (tvs, e'') = collectTyBinders e'
       return $ mkLams tvs (mkCoreLets bs 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
  e1'   <- AnfEnv -> Expr Id -> StateT DsST DsM (Expr Id)
normalize AnfEnv
γ (Pattern -> Expr Id
Rs.patE1 Pattern
p)
  e2'   <- stitch    γ (Rs.patE2 p)
  return $ Rs.lower p { Rs.patE1 = e1', Rs.patE2 = e2' }

normalizePattern AnfEnv
γ p :: Pattern
p@(Rs.PatReturn {}) = do
  e'    <- AnfEnv -> Expr Id -> StateT DsST DsM (Expr Id)
normalize AnfEnv
γ (Pattern -> Expr Id
Rs.patE Pattern
p)
  return $ Rs.lower 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
  e'    <- AnfEnv -> Expr Id -> StateT DsST DsM (Expr Id)
normalize AnfEnv
γ (Pattern -> Expr Id
Rs.patE Pattern
p)
  return $ Rs.lower 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
      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)
      return   $ sortCases (dcs' ++ 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
  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
  return (Alt (DataAlt d) (as ++ xs) 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
  u     <- DsM Unique
forall (m :: * -> *). MonadUnique m => m Unique
getUniqueM
  let i  = Unique -> UniqueId
getKey Unique
u
  let sp = SpanStack -> SrcSpan
Sp.srcSpan (AnfEnv -> SpanStack
aeSrcSpan AnfEnv
γ)
  return (mkUserLocalOrCoVar (anfOcc i) u Ghc.ManyTy t sp)

anfOcc :: UniqueId -> OccName
anfOcc :: UniqueId -> OccName
anfOcc = FastString -> OccName
mkVarOccFS (FastString -> OccName)
-> (UniqueId -> FastString) -> UniqueId -> OccName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Symbol -> FastString
GM.symbolFastString (Symbol -> FastString)
-> (UniqueId -> Symbol) -> UniqueId -> FastString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Symbol -> UniqueId -> 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 -> UniqueId -> Int
forall a. Hashable a => Int -> a -> Int
hashWithSalt Int
s (Unique -> UniqueId
getKey (Unique -> UniqueId) -> Unique -> UniqueId
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 γ)}