{-# LANGUAGE PatternGuards #-}
-- | Eliminate applications of the '?' operator from Core after ANF.
--
-- After ANF, @e ? s@ becomes:
--
-- @
-- let a1 = e
-- let a2 = s
-- ... (?) \@A \@B a1 a2
-- @
--
-- This pass replaces @(?) \@A \@B a1 a2@ with just @a1@, so that no KVars
-- are introduced in the signature of @?@. The binding @a2 = s@ remains in
-- scope, so the lemma's postcondition still enters the environment during
-- constraint generation.

module Language.Haskell.Liquid.Transforms.QuestionMark (eliminateQuestionMark) where

import Liquid.GHC.API as Ghc

-- | Look up the '?' operator in the 'GlobalRdrEnv'. If found, eliminate all
-- its applications from the Core program. If not found (module doesn't import
-- ProofCombinators), return the bindings unchanged.
eliminateQuestionMark :: GlobalRdrEnv -> [CoreBind] -> [CoreBind]
eliminateQuestionMark :: GlobalRdrEnv -> [CoreBind] -> [CoreBind]
eliminateQuestionMark GlobalRdrEnv
rdrEnv [CoreBind]
cbs =
  case GlobalRdrEnv -> Maybe Name
lookupQuestionMark GlobalRdrEnv
rdrEnv of
    Maybe Name
Nothing   -> [CoreBind]
cbs
    Just Name
name -> (CoreBind -> CoreBind) -> [CoreBind] -> [CoreBind]
forall a b. (a -> b) -> [a] -> [b]
map (Name -> CoreBind -> CoreBind
goBind Name
name) [CoreBind]
cbs

-- | Find the 'Name' of '?' from @Language.Haskell.Liquid.ProofCombinators@
-- in the renamer environment.
lookupQuestionMark :: GlobalRdrEnv -> Maybe Name
lookupQuestionMark :: GlobalRdrEnv -> Maybe Name
lookupQuestionMark GlobalRdrEnv
rdrEnv =
  case GlobalRdrEnv -> LookupGRE GREInfo -> [GlobalRdrEltX GREInfo]
forall info.
GlobalRdrEnvX info -> LookupGRE info -> [GlobalRdrEltX info]
lookupGRE GlobalRdrEnv
rdrEnv (RdrName -> WhichGREs GREInfo -> LookupGRE GREInfo
forall info. RdrName -> WhichGREs info -> LookupGRE info
LookupRdrName RdrName
rdrName WhichGREs GREInfo
forall info. WhichGREs info
SameNameSpace) of
    [GlobalRdrEltX GREInfo
gre] -> Name -> Maybe Name
forall a. a -> Maybe a
Just (GlobalRdrEltX GREInfo -> Name
forall info. GlobalRdrEltX info -> Name
greName GlobalRdrEltX GREInfo
gre)
    [GlobalRdrEltX GREInfo]
_     -> Maybe Name
forall a. Maybe a
Nothing
  where
    rdrName :: RdrName
rdrName = ModuleName -> OccName -> RdrName
mkRdrQual (String -> ModuleName
mkModuleName String
"Language.Haskell.Liquid.ProofCombinators")
                        (String -> OccName
mkVarOcc String
"?")

goBind :: Name -> CoreBind -> CoreBind
goBind :: Name -> CoreBind -> CoreBind
goBind Name
n (NonRec CoreBndr
x Expr CoreBndr
e) = CoreBndr -> Expr CoreBndr -> CoreBind
forall b. b -> Expr b -> Bind b
NonRec CoreBndr
x (Name -> Expr CoreBndr -> Expr CoreBndr
goExpr Name
n Expr CoreBndr
e)
goBind Name
n (Rec [(CoreBndr, Expr CoreBndr)]
xes)    = [(CoreBndr, Expr CoreBndr)] -> CoreBind
forall b. [(b, Expr b)] -> Bind b
Rec [(CoreBndr
x, Name -> Expr CoreBndr -> Expr CoreBndr
goExpr Name
n Expr CoreBndr
e) | (CoreBndr
x, Expr CoreBndr
e) <- [(CoreBndr, Expr CoreBndr)]
xes]

goExpr :: Name -> CoreExpr -> CoreExpr
goExpr :: Name -> Expr CoreBndr -> Expr CoreBndr
goExpr Name
n Expr CoreBndr
e
  | Just Expr CoreBndr
firstArg <- Name -> Expr CoreBndr -> Maybe (Expr CoreBndr)
isQuestionMarkApp Name
n Expr CoreBndr
e = Name -> Expr CoreBndr -> Expr CoreBndr
goExpr Name
n Expr CoreBndr
firstArg
goExpr Name
n (Lam CoreBndr
x Expr CoreBndr
e)         = CoreBndr -> Expr CoreBndr -> Expr CoreBndr
forall b. b -> Expr b -> Expr b
Lam CoreBndr
x (Name -> Expr CoreBndr -> Expr CoreBndr
goExpr Name
n Expr CoreBndr
e)
goExpr Name
n (Let CoreBind
b Expr CoreBndr
e)          = CoreBind -> Expr CoreBndr -> Expr CoreBndr
forall b. Bind b -> Expr b -> Expr b
Let (Name -> CoreBind -> CoreBind
goBind Name
n CoreBind
b) (Name -> Expr CoreBndr -> Expr CoreBndr
goExpr Name
n Expr CoreBndr
e)
goExpr Name
n (Case Expr CoreBndr
s CoreBndr
x Type
t [Alt CoreBndr]
alts) = Expr CoreBndr
-> CoreBndr -> Type -> [Alt CoreBndr] -> Expr CoreBndr
forall b. Expr b -> b -> Type -> [Alt b] -> Expr b
Case (Name -> Expr CoreBndr -> Expr CoreBndr
goExpr Name
n Expr CoreBndr
s) CoreBndr
x Type
t [Name -> Alt CoreBndr -> Alt CoreBndr
goAlt Name
n Alt CoreBndr
a | Alt CoreBndr
a <- [Alt CoreBndr]
alts]
goExpr Name
n (Cast Expr CoreBndr
e CoercionR
co)        = Expr CoreBndr -> CoercionR -> Expr CoreBndr
forall b. Expr b -> CoercionR -> Expr b
Cast (Name -> Expr CoreBndr -> Expr CoreBndr
goExpr Name
n Expr CoreBndr
e) CoercionR
co
goExpr Name
n (Tick CoreTickish
t Expr CoreBndr
e)         = CoreTickish -> Expr CoreBndr -> Expr CoreBndr
forall b. CoreTickish -> Expr b -> Expr b
Tick CoreTickish
t (Name -> Expr CoreBndr -> Expr CoreBndr
goExpr Name
n Expr CoreBndr
e)
goExpr Name
n (App Expr CoreBndr
f Expr CoreBndr
a)          = Expr CoreBndr -> Expr CoreBndr -> Expr CoreBndr
forall b. Expr b -> Expr b -> Expr b
App (Name -> Expr CoreBndr -> Expr CoreBndr
goExpr Name
n Expr CoreBndr
f) (Name -> Expr CoreBndr -> Expr CoreBndr
goExpr Name
n Expr CoreBndr
a)
goExpr Name
_ Expr CoreBndr
e                  = Expr CoreBndr
e -- Var, Lit, Type, Coercion

goAlt :: Name -> CoreAlt -> CoreAlt
goAlt :: Name -> Alt CoreBndr -> Alt CoreBndr
goAlt Name
n (Alt AltCon
con [CoreBndr]
bs Expr CoreBndr
e) = AltCon -> [CoreBndr] -> Expr CoreBndr -> Alt CoreBndr
forall b. AltCon -> [b] -> Expr b -> Alt b
Alt AltCon
con [CoreBndr]
bs (Name -> Expr CoreBndr -> Expr CoreBndr
goExpr Name
n Expr CoreBndr
e)

-- | Detect a fully-saturated application of '?': four arguments total
-- (two types + two values), possibly wrapped in ticks.
-- Returns @Just arg1@ (the first value argument).
isQuestionMarkApp :: Name -> CoreExpr -> Maybe CoreExpr
isQuestionMarkApp :: Name -> Expr CoreBndr -> Maybe (Expr CoreBndr)
isQuestionMarkApp Name
name Expr CoreBndr
expr =
  case (CoreTickish -> Bool)
-> Expr CoreBndr -> (Expr CoreBndr, [Expr CoreBndr], [CoreTickish])
forall b.
(CoreTickish -> Bool)
-> Expr b -> (Expr b, [Expr b], [CoreTickish])
collectArgsTicks (Bool -> CoreTickish -> Bool
forall a b. a -> b -> a
const Bool
True) Expr CoreBndr
expr of
    (Var CoreBndr
v, [Expr CoreBndr]
args, [CoreTickish]
_ticks)
      | CoreBndr -> Name
varName CoreBndr
v Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
== Name
name
      , [Expr CoreBndr
_tA, Expr CoreBndr
_tB, Expr CoreBndr
a, Expr CoreBndr
_b] <- [Expr CoreBndr]
args
      -> Expr CoreBndr -> Maybe (Expr CoreBndr)
forall a. a -> Maybe a
Just Expr CoreBndr
a
    (Expr CoreBndr, [Expr CoreBndr], [CoreTickish])
_ -> Maybe (Expr CoreBndr)
forall a. Maybe a
Nothing