-- | An implementation of LJT proof search directly on Core terms.
module GHC.LJT where
import GHC.Plugins
import GHC.Core.TyCo.Rep
import GHC.Types.Id.Make
import GHC.Types.Unique
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) (unrestricted (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 _af _mult 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 VisArg Many 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) Many 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
, let repargs = map scaledThing repargs'
, not (isRecTyCon tc)
= Just ( repargs
, \args -> mkConApp dc (map Type repargs ++ args)
, \scrut pats rhs -> mkWildCase scrut (unrestricted ty) (exprType rhs) [(DataAlt dc, pats, rhs)]
)
| Just (tc, ty_args) <- splitTyConApp_maybe ty
, Just dc <- newTyConDataCon_maybe tc
, not (isRecTyCon tc)
, let repargs = map scaledThing $ 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 (map scaledThing (dataConInstArgTys dc ty_args)) | dc <- dcs ]
injs = [
let vtys = dataConInstArgTys dc ty_args
vs = map (newVar . scaledThing) vtys
in \ e -> mkSmallTupleCase vs (mkConApp dc (map Type ty_args ++ map Var vs))
(mkWildValBinder Many (exprType e)) e
| dc <- dcs]
destruct = \e vs alts ->
Case e (mkWildValBinder Many (exprType e)) (exprType (head alts))
[ let pats = map (newVar . scaledThing) (dataConInstArgTys dc ty_args) in
(DataAlt dc, pats, mkLetNonRec v (mkCoreTup (map Var 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 $
map scaledThing $
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
-- A copy from MkId.hs, no longer exported there :-(
wrapNewTypeBody :: TyCon -> [Type] -> CoreExpr -> CoreExpr
wrapNewTypeBody tycon args result_expr
= wrapFamInstBody tycon args $
mkCast result_expr (mkSymCo co)
where
co = mkUnbranchedAxInstCo Representational (newTyConCo tycon) args []
-- Combinators to search for matching things
funLeft :: (Type -> Maybe a) -> Type -> Maybe (a,Type)
funLeft p (FunTy _af _mult 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