{-# LANGUAGE CPP #-}

{- | This module provides a variant of `ghc-typelits-presburger`,
 which can be also solve symbols added in this package, such as
 @Min@, @Max@, @<@, @>@, and @>=@.
-}
module Data.Type.Natural.Presburger.MinMaxSolver (plugin) where

import Control.Monad (mzero)
import GHC.TypeLits.Presburger.Compat
import GHC.TypeLits.Presburger.Types

#if MIN_VERSION_ghc(9,0,0)
import GHC.Plugins
  ( Plugin,
    fsLit,
    mkModuleName,
    mkTcOcc,
    splitTyConApp_maybe,
  )
import GHC.Tc.Plugin
#else
import GhcPlugins
  ( Plugin,
    fsLit,
    mkModuleName,
    mkTcOcc,
    splitTyConApp_maybe,
  )
import TcPluginM
#endif

plugin :: Plugin
plugin :: Plugin
plugin =
  TcPluginM Translation -> Plugin
pluginWith (TcPluginM Translation -> Plugin)
-> TcPluginM Translation -> Plugin
forall a b. (a -> b) -> a -> b
$
    Translation -> Translation -> Translation
forall a. Semigroup a => a -> a -> a
(<>) (Translation -> Translation -> Translation)
-> TcPluginM Translation -> TcPluginM (Translation -> Translation)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> TcPluginM Translation
defaultTranslation TcPluginM (Translation -> Translation)
-> TcPluginM Translation -> TcPluginM Translation
forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> TcPluginM Translation
genTypeNatsTranslation

genTypeNatsTranslation :: TcPluginM Translation
genTypeNatsTranslation :: TcPluginM Translation
genTypeNatsTranslation = do
  Module
orderMod <- ModuleName -> FastString -> TcPluginM Module
lookupModule (String -> ModuleName
mkModuleName String
"Data.Type.Natural.Lemma.Order") (String -> FastString
fsLit String
"type-natural")
  TyCon
singNatLt <- Name -> TcPluginM TyCon
tcLookupTyCon (Name -> TcPluginM TyCon) -> TcPluginM Name -> TcPluginM TyCon
forall (m :: Type -> Type) a b. Monad m => (a -> m b) -> m a -> m b
=<< Module -> OccName -> TcPluginM Name
lookupOrig Module
orderMod (String -> OccName
mkTcOcc String
"<?")
  TyCon
singNatGeq <- Name -> TcPluginM TyCon
tcLookupTyCon (Name -> TcPluginM TyCon) -> TcPluginM Name -> TcPluginM TyCon
forall (m :: Type -> Type) a b. Monad m => (a -> m b) -> m a -> m b
=<< Module -> OccName -> TcPluginM Name
lookupOrig Module
orderMod (String -> OccName
mkTcOcc String
">=?")
  TyCon
singNatGt <- Name -> TcPluginM TyCon
tcLookupTyCon (Name -> TcPluginM TyCon) -> TcPluginM Name -> TcPluginM TyCon
forall (m :: Type -> Type) a b. Monad m => (a -> m b) -> m a -> m b
=<< Module -> OccName -> TcPluginM Name
lookupOrig Module
orderMod (String -> OccName
mkTcOcc String
">?")

  TyCon
singNatLtProp <- Name -> TcPluginM TyCon
tcLookupTyCon (Name -> TcPluginM TyCon) -> TcPluginM Name -> TcPluginM TyCon
forall (m :: Type -> Type) a b. Monad m => (a -> m b) -> m a -> m b
=<< Module -> OccName -> TcPluginM Name
lookupOrig Module
orderMod (String -> OccName
mkTcOcc String
"<")
  TyCon
singNatGeqProp <- Name -> TcPluginM TyCon
tcLookupTyCon (Name -> TcPluginM TyCon) -> TcPluginM Name -> TcPluginM TyCon
forall (m :: Type -> Type) a b. Monad m => (a -> m b) -> m a -> m b
=<< Module -> OccName -> TcPluginM Name
lookupOrig Module
orderMod (String -> OccName
mkTcOcc String
">=")
  TyCon
singNatGtProp <- Name -> TcPluginM TyCon
tcLookupTyCon (Name -> TcPluginM TyCon) -> TcPluginM Name -> TcPluginM TyCon
forall (m :: Type -> Type) a b. Monad m => (a -> m b) -> m a -> m b
=<< Module -> OccName -> TcPluginM Name
lookupOrig Module
orderMod (String -> OccName
mkTcOcc String
">")

  TyCon
singMin <- Name -> TcPluginM TyCon
tcLookupTyCon (Name -> TcPluginM TyCon) -> TcPluginM Name -> TcPluginM TyCon
forall (m :: Type -> Type) a b. Monad m => (a -> m b) -> m a -> m b
=<< Module -> OccName -> TcPluginM Name
lookupOrig Module
orderMod (String -> OccName
mkTcOcc String
"Min")
  TyCon
singMax <- Name -> TcPluginM TyCon
tcLookupTyCon (Name -> TcPluginM TyCon) -> TcPluginM Name -> TcPluginM TyCon
forall (m :: Type -> Type) a b. Monad m => (a -> m b) -> m a -> m b
=<< Module -> OccName -> TcPluginM Name
lookupOrig Module
orderMod (String -> OccName
mkTcOcc String
"Max")
  TyCon
caseMinAux <- Name -> TcPluginM TyCon
tcLookupTyCon (Name -> TcPluginM TyCon) -> TcPluginM Name -> TcPluginM TyCon
forall (m :: Type -> Type) a b. Monad m => (a -> m b) -> m a -> m b
=<< Module -> OccName -> TcPluginM Name
lookupOrig Module
orderMod (String -> OccName
mkTcOcc String
"MinAux")
  TyCon
caseMaxAux <- Name -> TcPluginM TyCon
tcLookupTyCon (Name -> TcPluginM TyCon) -> TcPluginM Name -> TcPluginM TyCon
forall (m :: Type -> Type) a b. Monad m => (a -> m b) -> m a -> m b
=<< Module -> OccName -> TcPluginM Name
lookupOrig Module
orderMod (String -> OccName
mkTcOcc String
"MaxAux")
  Translation -> TcPluginM Translation
forall (m :: Type -> Type) a. Monad m => a -> m a
return
    Translation
forall a. Monoid a => a
mempty
      { natGeqBool :: [TyCon]
natGeqBool = [TyCon
singNatGeq]
      , natLtBool :: [TyCon]
natLtBool = [TyCon
singNatLt]
      , natGtBool :: [TyCon]
natGtBool = [TyCon
singNatGt]
      , natMin :: [TyCon]
natMin = [TyCon
singMin]
      , natMax :: [TyCon]
natMax = [TyCon
singMax]
      , parsePred :: (Type -> Machine Expr) -> Type -> Machine Prop
parsePred = \Type -> Machine Expr
toE Type
ty ->
          case HasDebugCallStack => Type -> Maybe (TyCon, [Type])
Type -> Maybe (TyCon, [Type])
splitTyConApp_maybe Type
ty of
            Just (TyCon
con, [Type
l, Type
r])
              | TyCon
con TyCon -> TyCon -> Bool
forall a. Eq a => a -> a -> Bool
== TyCon
singNatLtProp -> Expr -> Expr -> Prop
(:<) (Expr -> Expr -> Prop)
-> Machine Expr
-> MaybeT (StateT ParseEnv TcPluginM) (Expr -> Prop)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> Type -> Machine Expr
toE Type
l MaybeT (StateT ParseEnv TcPluginM) (Expr -> Prop)
-> Machine Expr -> Machine Prop
forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> Type -> Machine Expr
toE Type
r
              | TyCon
con TyCon -> TyCon -> Bool
forall a. Eq a => a -> a -> Bool
== TyCon
singNatGtProp -> Expr -> Expr -> Prop
(:>) (Expr -> Expr -> Prop)
-> Machine Expr
-> MaybeT (StateT ParseEnv TcPluginM) (Expr -> Prop)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> Type -> Machine Expr
toE Type
l MaybeT (StateT ParseEnv TcPluginM) (Expr -> Prop)
-> Machine Expr -> Machine Prop
forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> Type -> Machine Expr
toE Type
r
              | TyCon
con TyCon -> TyCon -> Bool
forall a. Eq a => a -> a -> Bool
== TyCon
singNatGeqProp -> Expr -> Expr -> Prop
(:>=) (Expr -> Expr -> Prop)
-> Machine Expr
-> MaybeT (StateT ParseEnv TcPluginM) (Expr -> Prop)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> Type -> Machine Expr
toE Type
l MaybeT (StateT ParseEnv TcPluginM) (Expr -> Prop)
-> Machine Expr -> Machine Prop
forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> Type -> Machine Expr
toE Type
r
            Maybe (TyCon, [Type])
_ -> Machine Prop
forall (m :: Type -> Type) a. MonadPlus m => m a
mzero
      , parseExpr :: (Type -> Machine Expr) -> Type -> Machine Expr
parseExpr = \Type -> Machine Expr
toE Type
ty ->
          case HasDebugCallStack => Type -> Maybe (TyCon, [Type])
Type -> Maybe (TyCon, [Type])
splitTyConApp_maybe Type
ty of
            Just (TyCon
con, [Type
_, Type
n, Type
m])
              | TyCon
con TyCon -> TyCon -> Bool
forall a. Eq a => a -> a -> Bool
== TyCon
caseMinAux ->
                Expr -> Expr -> Expr
Min (Expr -> Expr -> Expr)
-> Machine Expr
-> MaybeT (StateT ParseEnv TcPluginM) (Expr -> Expr)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> Type -> Machine Expr
toE Type
n MaybeT (StateT ParseEnv TcPluginM) (Expr -> Expr)
-> Machine Expr -> Machine Expr
forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> Type -> Machine Expr
toE Type
m
              | TyCon
con TyCon -> TyCon -> Bool
forall a. Eq a => a -> a -> Bool
== TyCon
caseMaxAux ->
                Expr -> Expr -> Expr
Max (Expr -> Expr -> Expr)
-> Machine Expr
-> MaybeT (StateT ParseEnv TcPluginM) (Expr -> Expr)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> Type -> Machine Expr
toE Type
n MaybeT (StateT ParseEnv TcPluginM) (Expr -> Expr)
-> Machine Expr -> Machine Expr
forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> Type -> Machine Expr
toE Type
m
            Maybe (TyCon, [Type])
_ -> Machine Expr
forall (m :: Type -> Type) a. MonadPlus m => m a
mzero
      }