-- | An implementation of LJT proof search directly on Core terms.
module GHC.LJT where
import FastString
import Unique
import Type
import Id
import Var
import CoreSyn
import Outputable
import TyCoRep
import TyCon
import DataCon
import MkCore
import MkId
import CoreUtils
import TysWiredIn
import BasicTypes
import NameEnv
import NameSet
import Data.List
import Data.Hashable
import Control.Monad
import Data.Bifunctor
ljt :: Type -> [CoreExpr]
ljt t = [] ==> t
(==>) :: [Id] -> Type -> [CoreExpr]
-- Rule Axiom
-- (TODO: The official algorithm restricts this rule to atoms. Why?)
ante ==> goal
| Just v <- find (\v -> idType v `eqType` goal) ante
= pure $ Var v
-- Rule f⇒
ante ==> goal
| Just v <- find (\v -> isEmptyTy (idType v)) ante
= pure $ mkWildCase (Var v) (idType v) goal []
-- Rule →⇒2
ante ==> goal
| Just ((v,((tys, build, _destruct),_r)),ante') <- anyA (funLeft isProdType) ante
= let vs = map newVar tys
expr = mkLams vs (App (Var v) (build (map Var vs)))
v' = newVar (exprType expr)
in mkLetNonRec v' expr <$> (v' : ante') ==> goal
-- Rule →⇒3
ante ==> goal
| Just ((v,((tys, injs, _destruct),_r)),ante') <- anyA (funLeft isSumType) ante
= let es = [ lam ty (\vx -> App (Var v) (inj (Var vx))) | (ty,inj) <- zip tys injs ]
in letsA es $ \vs -> (vs ++ ante') ==> goal
-- Rule ∧⇒
ante ==> goal
| Just ((v,(tys, _build, destruct)),ante') <- anyA isProdType ante
= let pats = map newVar tys
in destruct (Var v) pats <$> (pats ++ ante') ==> goal
-- Rule ⇒∧
ante ==> goal
| Just (tys, build, _destruct) <- isProdType goal
= build <$> sequence [ante ==> ty | ty <- tys]
-- Rule ∨⇒
ante ==> goal
| Just ((vAorB, (tys, _injs, destruct)),ante') <- anyA isSumType ante
= let vs = map newVar tys in
destruct (Var vAorB) vs <$> sequence [ (v:ante') ==> goal | v <- vs]
-- Rule ⇒→
ante ==> FunTy t1 t2
= Lam v <$> (v : ante) ==> t2
where
v = newVar t1
-- Rule →⇒1
-- (TODO: The official algorithm restricts this rule to atoms. Why?)
ante ==> goal
| let isInAnte a = find (\v -> idType v `eqType` a) ante
, Just ((vAB, (vA,_)), ante') <- anyA (funLeft isInAnte) ante
= letA (App (Var vAB) (Var vA)) $ \vB -> (vB : ante') ==> goal
-- Rule ⇒∨
ante ==> goal
| Just (tys, injs, _destruct) <- isSumType goal
= msum [ inj <$> ante ==> ty | (ty,inj) <- zip tys injs ]
-- Rule →⇒4
ante ==> goal
| Just ((vABC, ((a,b),_)), ante') <- anyA (funLeft (funLeft Just)) ante
= do
let eBC = lam b $ \vB -> App (Var vABC) (lam a $ \_ -> Var vB)
eAB <- letA eBC $ \vBC -> (vBC : ante') ==> FunTy a b
letA (App (Var vABC) eAB) $ \vC -> (vC : ante') ==> goal
-- Nothing found :-(
_ante ==> _goal
= -- pprTrace "go" (vcat [ ppr (idType v) | v <- ante] $$ text "------" $$ ppr goal) $
mzero
-- Smart constructors
newVar :: Type -> Id
newVar ty = mkSysLocal (mkFastString "x") (mkBuiltinUnique i) ty
where i = hash (showSDocUnsafe (ppr ty))
-- We don’t mind if variables with equal types shadow each other,
-- so let’s just derive the unique from the type
lam :: Type -> (Id -> CoreExpr) -> CoreExpr
lam ty gen = Lam v $ gen v
where v = newVar ty
lamA :: Applicative f => Type -> (Id -> f CoreExpr) -> f CoreExpr
lamA ty gen = Lam v <$> gen v
where v = newVar ty
let_ :: CoreExpr -> (Id -> CoreExpr) -> CoreExpr
let_ e gen = mkLetNonRec v e $ gen v
where v = newVar (exprType e)
letA :: Applicative f => CoreExpr -> (Id -> f CoreExpr) -> f CoreExpr
letA e gen = mkLetNonRec v e <$> gen v
where v = newVar (exprType e)
letsA :: Applicative f => [CoreExpr] -> ([Id] -> f CoreExpr) -> f CoreExpr
letsA es gen = mkLets (zipWith NonRec vs es) <$> gen vs
where vs = map (newVar . exprType) es
-- Predicate on types
isProdType :: Type -> Maybe ([Type], [CoreExpr] -> CoreExpr, CoreExpr -> [Id] -> CoreExpr -> CoreExpr)
isProdType ty
| Just (tc, _, dc, repargs) <- splitDataProductType_maybe ty
, not (isRecTyCon tc)
= Just ( repargs
, \args -> mkConApp dc (map Type repargs ++ args)
, \scrut pats rhs -> mkWildCase scrut ty (exprType rhs) [(DataAlt dc, pats, rhs)]
)
| Just (tc, ty_args) <- splitTyConApp_maybe ty
, Just dc <- newTyConDataCon_maybe tc
, not (isRecTyCon tc)
, let repargs = dataConInstArgTys dc ty_args
= Just ( repargs
, \[arg] -> wrapNewTypeBody tc ty_args arg
, \scrut [pat] rhs ->
mkLetNonRec pat (unwrapNewTypeBody tc ty_args scrut) rhs
)
isProdType _ = Nothing
-- Haskell sum constructors can have multiple parameters. For our purposes, if
-- so, we wrap them in a product.
isSumType :: Type -> Maybe ([Type], [CoreExpr -> CoreExpr], CoreExpr -> [Id] -> [CoreExpr] -> CoreExpr)
isSumType ty
| Just (tc, ty_args) <- splitTyConApp_maybe ty
, Just dcs <- isDataSumTyCon_maybe tc
, not (isRecTyCon tc)
= let tys = [ mkTupleTy Boxed (dataConInstArgTys dc ty_args) | dc <- dcs ]
injs = [
let vtys = dataConInstArgTys dc ty_args
vs = map newVar vtys
in \ e -> mkSmallTupleCase vs (mkConApp dc (map Type ty_args ++ map Var vs))
(mkWildValBinder (exprType e)) e
| dc <- dcs]
destruct = \e vs alts ->
Case e (mkWildValBinder (exprType e)) (exprType (head alts))
[ let pats = map newVar (dataConInstArgTys dc ty_args) in
(DataAlt dc, pats, mkLetNonRec v (mkCoreVarTup pats) rhs)
| (dc,v,rhs) <- zip3 dcs vs alts ]
in Just (tys, injs, destruct)
isSumType _ = Nothing
-- We don’t want to look into recursive type cons.
-- Which ones are recursive? Surely those that get mentioned in their
-- arguments. Or in type cons in their arguments.
-- But that is not enough, because of higher kinded arguments. So prohibit
-- those as well.
isRecTyCon :: TyCon -> Bool
isRecTyCon tc = go emptyNameSet tc
where
go seen tc | tyConName tc `elemNameSet` seen = True
| any isHigherKind paramKinds = False
| any (go seen') mentionedTyCons = True
| otherwise = False
where mentionedTyCons = concatMap getTyCons $ concatMap dataConOrigArgTys $ tyConDataCons tc
paramKinds = map varType (tyConTyVars tc)
seen' = seen `extendNameSet` tyConName tc
isHigherKind :: Kind -> Bool
isHigherKind k = not (k `eqType` liftedTypeKind)
getTyCons :: Type -> [TyCon]
getTyCons = nameEnvElts . go
where
go (TyConApp tc tys) = unitNameEnv (tyConName tc) tc `plusNameEnv` go_s tys
go (LitTy _) = emptyNameEnv
go (TyVarTy _) = emptyNameEnv
go (AppTy a b) = go a `plusNameEnv` go b
go (FunTy a b) = go a `plusNameEnv` go b
go (ForAllTy _ ty) = go ty
go (CastTy ty _) = go ty
go (CoercionTy co) = emptyNameEnv
go_s = foldr (plusNameEnv . go) emptyNameEnv
-- Combinators to search for matching things
funLeft :: (Type -> Maybe a) -> Type -> Maybe (a,Type)
funLeft p (FunTy t1 t2) = (\x -> (x,t2)) <$> p t1
funLeft _ _ = Nothing
anyA :: (Type -> Maybe a) -> [Id] -> Maybe ((Id, a), [Id])
anyA _ [] = Nothing
anyA p (v:vs) | Just x <- p (idType v) = Just ((v,x), vs)
| otherwise = second (v:) <$> anyA p vs