diff --git a/examples/simple-arith.hs b/examples/simple-arith.hs --- a/examples/simple-arith.hs +++ b/examples/simple-arith.hs @@ -1,68 +1,77 @@ {-# LANGUAGE DataKinds, TypeOperators, GADTs, TypeFamilies, ExplicitForAll, FlexibleContexts #-} {-# LANGUAGE ScopedTypeVariables, CPP #-} +#if !MIN_VERSION_singletons(2,4,1) +{-# LANGUAGE KindSignatures #-} +#endif {-# OPTIONS_GHC -fplugin GHC.TypeLits.Presburger #-} module Main where import Data.Type.Equality -import GHC.TypeLits (Nat, type (*), type (+), type (<=?), CmpNat) +import GHC.TypeLits (Nat, type (<=?), CmpNat) import Proof.Propositional (Empty(..)) import Proof.Propositional (IsTrue(Witness)) import Data.Singletons.Prelude +import Data.Singletons.Prelude.List + +#if !MIN_VERSION_singletons(2,4,1) import qualified Data.Singletons.Prelude as Sing -import Data.Void +type n <= m = (n :: Nat) Sing.:<= (m :: Nat) +infix 4 <= +type n - m = (n :: Nat) Sing.:- (m :: Nat) +infixl 6 - +type n + m = (n :: Nat) Sing.:+ (m :: Nat) +infixl 6 + +type n * m = (n :: Nat) Sing.:* (m :: Nat) +infixl 7 * +#endif type n <=! m = IsTrue (n <=? m) infix 4 <=! --- natLeqZero :: ((n <=? 0) ~ 'True) => proxy n -> n :~: 0 --- natLeqZero _ = Refl +natLen :: (Length xs <= Length ys) ~ 'True + => proxy xs -> proxy ys -> Length ys - Length xs + Length xs :~: Length ys +natLen _ _ = Refl -#if MIN_VERSION_singletons(2,4,1) natLeqZero' :: ((n <= 0) ~ 'True) => proxy n -> n :~: 0 -#else -natLeqZero' :: ((n :<= 0) ~ 'True) => proxy n -> n :~: 0 -#endif natLeqZero' _ = Refl --- (%:<=?) :: Sing n -> Sing m -> Sing (n <=? m) --- n %:<=? m = case sCompare n m of --- SLT -> STrue --- SEQ -> STrue --- SGT -> SFalse - --- natLeqZero :: IsTrue (n <=? 0) -> Sing n -> n :~: 0 --- natLeqZero Witness Zero = Refl - +plusLeq :: (n <= m) ~ 'True => proxy (n :: Nat) -> proxy m -> ((m - n) + n :~: m) +plusLeq _ _ = Refl --- hoge :: ((n + 1 <=? n) ~ 'False) => () --- hoge = () +(%:<=?) :: Sing n -> Sing m -> Sing (n <=? m) +n %:<=? m = case sCompare n m of + SLT -> STrue + SEQ -> STrue + SGT -> SFalse --- fuga :: ((n + 1 <=? 0) ~ 'False) => () --- fuga = () +hoge :: ((n + 1 <=? n) ~ 'False) => proxy n -> () +hoge _ = () --- bar :: ((2 * (n + 1)) ~ ((2 * n) + 2)) => proxy n -> () --- bar _ = () +hoge' :: (((n + 1) <= n) ~ 'False) => proxy n -> () +hoge' _ = () --- trans :: proxy n -> proxy m -> n <=! m -> (n + 1) <=! (m + 1) --- trans _ _ Witness = Witness +bar :: ((2 * (n + 1)) ~ ((2 * n) + 2)) => proxy n -> () +bar _ = () --- eqv :: proxy n -> proxy m -> (n <=? m) :~: ((n + 1) <=? (m + 1)) --- eqv _ _ = Refl +trans :: proxy n -> proxy m -> n <=! m -> (n + 1) <=! (m + 1) +trans _ _ Witness = Witness --- leqSucc :: proxy n -> proxy m -> IsTrue ((n + 1) :<= m) -> CmpNat n m :~: 'LT --- leqSucc _ _ Witness = Refl +eqv :: proxy n -> proxy m -> (n <=? m) :~: ((n + 1) <=? (m + 1)) +eqv _ _ = Refl +leqSucc :: proxy n -> proxy m -> IsTrue ((n + 1) <= m) -> CmpNat n m :~: 'LT +leqSucc _ _ Witness = Refl --- predSucc :: forall proxy n. Empty (n <=! 0) => proxy n -> IsTrue (n + 1 <=? 2 * n) --- predSucc _ = Witness +predSucc :: forall proxy n. Empty (n <=! 0) => proxy n -> IsTrue (n + 1 <=? 2 * n) +predSucc _ = Witness main :: IO () main = putStrLn "finished" --- succLEqLTSucc :: Sing m -> Compare 0 (m + 1) :~: 'LT --- succLEqLTSucc _ = Refl +succLEqLTSucc :: Sing m -> Compare 0 (m + 1) :~: 'LT +succLEqLTSucc _ = Refl --- succCompare :: Sing (n :: Nat) -> Sing m -> CmpNat n m :~: CmpNat (n + 1) (m + 1) --- succCompare _ _ = Refl +succCompare :: Sing (n :: Nat) -> Sing m -> CmpNat n m :~: CmpNat (n + 1) (m + 1) +succCompare _ _ = Refl eqToRefl :: Sing (n :: Nat) -> Sing (m :: Nat) -> CmpNat n m :~: 'EQ -> n :~: m eqToRefl _n _m Refl = Refl diff --git a/ghc-typelits-presburger.cabal b/ghc-typelits-presburger.cabal --- a/ghc-typelits-presburger.cabal +++ b/ghc-typelits-presburger.cabal @@ -1,5 +1,5 @@ name: ghc-typelits-presburger -version: 0.2.0.4 +version: 0.2.0.5 synopsis: Presburger Arithmetic Solver for GHC Type-level natural numbers. description: @ghc-typelits-presburger@ augments GHC type-system with Presburger Arithmetic Solver for Type-level natural numbers. @@ -34,8 +34,11 @@ , ghc-tcplugins-extra >= 0.2 && < 0.4 , equational-reasoning >= 0.4.0.0 , singletons + , transformers + , mtl , containers , pretty + , syb , reflection default-language: Haskell2010 diff --git a/src/GHC/Compat.hs b/src/GHC/Compat.hs --- a/src/GHC/Compat.hs +++ b/src/GHC/Compat.hs @@ -1,5 +1,8 @@ -{-# LANGUAGE CPP, PatternGuards, PatternSynonyms, ViewPatterns #-} +{-# LANGUAGE CPP, FlexibleInstances, PatternGuards, PatternSynonyms #-} +{-# LANGUAGE TypeSynonymInstances, ViewPatterns #-} +{-# OPTIONS_GHC -Wno-orphans #-} module GHC.Compat (module GHC.Compat) where +import Data.Function (on) import FamInst as GHC.Compat import FastString as GHC.Compat (fsLit) import GHC.TcPluginM.Extra as GHC.Compat (evByFiat, lookupModule, lookupName, @@ -10,7 +13,7 @@ import GhcPlugins as GHC.Compat (mkTcOcc, ppr, promotedFalseDataCon) import GhcPlugins as GHC.Compat (promotedTrueDataCon, text) import GhcPlugins as GHC.Compat (tyConAppTyCon_maybe, typeKind) -import GhcPlugins as GHC.Compat (TyCon, typeNatKind) +import GhcPlugins as GHC.Compat (typeNatKind) import Module as GHC.Compat (ModuleName, mkModuleName) import OccName as GHC.Compat (emptyOccSet, mkInstTyTcOcc) import Plugins as GHC.Compat (Plugin (..), defaultPlugin) @@ -22,6 +25,7 @@ import TcRnTypes as GHC.Compat (TcPlugin (..), ctEvPred, ctEvidence) import TcType as GHC.Compat (tcTyFamInsts) import TcTypeNats as GHC.Compat +import TyCon as GHC.Compat #if defined(__GLASGOW_HASKELL__) && __GLASGOW_HASKELL__ >= 800 import GhcPlugins (InScopeSet, Outputable, emptyUFM) import qualified PrelNames as Old @@ -30,7 +34,6 @@ emptyTCvSubst) import Type as GHC.Compat (eqType, unionTCvSubst) import qualified Type as Old -import qualified TysPrim as Old import TysWiredIn as GHC.Compat (boolTyCon) import Unify as Old (tcUnifyTy) #else @@ -41,9 +44,11 @@ import TysWiredIn as GHC.Compat (promotedBoolTyCon) import Unify as GHC.Compat (tcUnifyTy) #endif -import TcPluginM (lookupOrig) -import Type as GHC.Compat (splitTyConApp_maybe) -import Unique as GHC.Compat (getKey, getUnique) +import Data.Generics.Twins +import TcPluginM (lookupOrig) +import TyCoRep () +import Type as GHC.Compat (splitTyConApp_maybe) +import Unique as GHC.Compat (getKey, getUnique) #if defined(__GLASGOW_HASKELL__) && __GLASGOW_HASKELL__ >= 800 data TvSubst = TvSubst InScopeSet TvSubstEnv @@ -106,3 +111,11 @@ decompFunTy :: Type -> [Type] decompFunTy (FunTy t1 t2) = t1 : decompFunTy t2 decompFunTy t = [t] + +newtype TypeEq = TypeEq { runTypeEq :: Type } + +instance Eq TypeEq where + (==) = geq `on` runTypeEq + +instance Ord TypeEq where + compare = gcompare `on` runTypeEq diff --git a/src/GHC/TypeLits/Presburger.hs b/src/GHC/TypeLits/Presburger.hs --- a/src/GHC/TypeLits/Presburger.hs +++ b/src/GHC/TypeLits/Presburger.hs @@ -1,28 +1,42 @@ -{-# LANGUAGE CPP, DataKinds, FlexibleContexts, MultiWayIf #-} -{-# LANGUAGE OverloadedStrings, PatternGuards, RankNTypes, RecordWildCards #-} -{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE CPP, DataKinds, FlexibleContexts, FlexibleInstances #-} +{-# LANGUAGE LambdaCase, MultiWayIf, OverloadedStrings, PatternGuards #-} +{-# LANGUAGE RankNTypes, RecordWildCards, TypeOperators #-} +{-# LANGUAGE TypeSynonymInstances #-} {-# OPTIONS_GHC -Wno-unused-imports #-} module GHC.TypeLits.Presburger (plugin) where import GHC.Compat -import Class (classTyCon) -import Control.Monad (replicateM) -import Data.Foldable (asum) -import Data.Integer.SAT (Expr (..), Prop (..), PropSet, assert) -import Data.Integer.SAT (checkSat, noProps, toName) -import qualified Data.Integer.SAT as SAT -import Data.List (nub) -import Data.Maybe (fromMaybe, isNothing, mapMaybe) -import Data.Reflection (Given, give, given) -import Debug.Trace -import GHC.TypeLits (Nat) -import Outputable (showSDocUnsafe) -import TcPluginM (getFamInstEnvs, lookupOrig, matchFam, - newFlexiTyVar, tcLookupClass, - unsafeTcPluginTcM) -import Type (mkTyVarTy, splitTyConApp) -import TysWiredIn (promotedEQDataCon, promotedGTDataCon, - promotedLTDataCon) +import Class (Class, classTyCon) +import Control.Applicative ((<|>)) +import Control.Monad (forM_, mzero) +import Control.Monad.State.Class +import Control.Monad.Trans.Class +import Control.Monad.Trans.Maybe (MaybeT (..)) +import Control.Monad.Trans.State (StateT, runStateT) +import Data.Foldable (asum) +import Data.Integer.SAT (Expr (..), Prop (..), PropSet, + assert) +import Data.Integer.SAT (checkSat, noProps, toName) +import qualified Data.Integer.SAT as SAT +import Data.List (nub) +import qualified Data.Map.Strict as M +import qualified Data.Map.Strict as M +import Data.Maybe (catMaybes, fromJust) +import Data.Maybe (fromMaybe, isNothing, mapMaybe) +import Data.Reflection (Given, give, given) +import GHC.TypeLits (Nat) +import Outputable (showSDocUnsafe) +import TcPluginM (getFamInstEnvs, lookupOrig, + matchFam, newFlexiTyVar, newWanted, + tcLookupClass, unsafeTcPluginTcM) +import TcRnTypes +import TyCoRep (Coercion (..), KindCoercion (..)) +import Type (mkPrimEqPredRole, mkTyVarTy, + splitTyConApp) +import TysWiredIn (promotedEQDataCon, + promotedGTDataCon, + promotedLTDataCon) +import Var assert' :: Prop -> PropSet -> PropSet assert' p ps = foldr assert ps (p : varPos) @@ -117,7 +131,8 @@ withTyCons $ do let subst = emptyTvSubst -- foldr unionTvSubst emptyTvSubst $ map genSubst gs - givens = mapMaybe (\a -> (,) a <$> toPresburgerPred subst (deconsPred a)) gs + ngs <- mapM (\a -> runMachine $ (,) a <$> toPresburgerPred subst (deconsPred a)) gs + let givens = catMaybes ngs prems0 = map snd givens prems = foldr assert' noProps prems0 (solved, _) = foldr go ([], noProps) givens @@ -135,16 +150,21 @@ tcPluginTrace "wanteds" $ ppr $ map deconsPred ws tcPluginTrace "givens" $ ppr $ map (substTy subst . deconsPred) gs tcPluginTrace "deriveds" $ ppr $ map deconsPred ds - let wants = mapMaybe (\ct -> (,) ct <$> toPresburgerPred subst (substTy subst $ deconsPred ct)) $ - filter (isWanted . ctEvidence) ws - prems = foldr assert' noProps $ - mapMaybe (toPresburgerPred subst . substTy subst . deconsPred) (gs ++ ds) - solved = map fst $ filter (isProved . testIf prems . snd) wants + (prems, wants) <- do + wants <- catMaybes <$> + mapM + (\ct -> runMachine $ (,) ct <$> toPresburgerPred subst (substTy subst $ deconsPred ct)) + (filter (isWanted . ctEvidence) ws) + resls <- mapM (runMachine . toPresburgerPred subst . substTy subst . deconsPred) + (gs ++ ds) + let prems = foldr assert' noProps $ catMaybes resls + return (prems, wants) + let solved = map fst $ filter (isProved . testIf prems . snd) wants coerced = [(evByFiat "ghc-typelits-presburger" t1 t2, ct) | ct <- solved , EqPred NomEq t1 t2 <- return (classifyPredType $ deconsPred ct) ] - tcPluginTrace "prems" (text $ show $ map (toPresburgerPred subst .substTy subst . deconsPred) (gs ++ ds)) + -- tcPluginTrace "prems" (text $ show $ map (toPresburgerPred subst .substTy subst . deconsPred) (gs ++ ds)) tcPluginTrace "final goals" (text $ show $ map snd wants) case testIf prems (foldr ((:&&) . snd) PTrue wants) of Proved -> do @@ -213,7 +233,7 @@ emptyTyCon :: Given MyEnv => TyCon emptyTyCon = emptyClsTyCon given -toPresburgerPred :: Given MyEnv => TvSubst -> Type -> Maybe Prop +toPresburgerPred :: Given MyEnv => TvSubst -> Type -> Machine Prop toPresburgerPred subst (TyConApp con [t1, t2]) | con == typeNatLeqTyCon = (:<=) <$> toPresburgerExp subst t1 <*> toPresburgerExp subst t2 toPresburgerPred subst ty @@ -226,12 +246,12 @@ , con == emptyTyCon = Not <$> toPresburgerPred subst l | Just (con, [l]) <- splitTyConApp_maybe ty -- IsTrue l => , con == isTrueTyCon = toPresburgerPred subst l - | otherwise = Nothing + | otherwise = mzero boolLeqs :: Given MyEnv => [TyCon] boolLeqs = [typeNatLeqTyCon, typeLeqBoolTyCon] -toPresburgerPredTree :: Given MyEnv => TvSubst -> PredTree -> Maybe Prop +toPresburgerPredTree :: Given MyEnv => TvSubst -> PredTree -> Machine Prop toPresburgerPredTree subst (EqPred NomEq p false) -- P ~ 'False <=> Not P ~ 'True | Just promotedFalseDataCon == tyConAppTyCon_maybe (substTy subst false) = Not <$> toPresburgerPredTree subst (EqPred NomEq p (mkTyConTy promotedTrueDataCon)) @@ -285,8 +305,9 @@ ,(promotedEQDataCon, (:==)) ,(promotedGTDataCon, (:>)) ] - in lookup cmp dic <*> toPresburgerExp subst a - <*> toPresburgerExp subst b + in MaybeT (return $ lookup cmp dic) + <*> toPresburgerExp subst a + <*> toPresburgerExp subst b toPresburgerPredTree subst (EqPred NomEq t1 t2) -- x ~ CmpNat a b | Just (con, [a, b]) <- splitTyConApp_maybe (substTy subst t2) , con `elem` [singCompareCon, typeNatCmpTyCon] @@ -295,32 +316,36 @@ ,(promotedEQDataCon, (:==)) ,(promotedGTDataCon, (:>)) ] - in lookup cmp dic <*> toPresburgerExp subst a - <*> toPresburgerExp subst b + in MaybeT (return $ lookup cmp dic) + <*> toPresburgerExp subst a + <*> toPresburgerExp subst b toPresburgerPredTree subst (ClassPred con [t1, t2]) -- (n :: Nat) <= (m :: Nat) | typeNatLeqTyCon == classTyCon con , typeKind t1 `eqType` typeNatKind = (:<=) <$> toPresburgerExp subst t1 <*> toPresburgerExp subst t2 -toPresburgerPredTree _ _ = Nothing +toPresburgerPredTree _ _ = mzero -toPresburgerExp :: TvSubst -> Type -> Maybe Expr +toPresburgerExp :: TvSubst -> Type -> Machine Expr toPresburgerExp dic ty = case substTy dic ty of - TyVarTy t -> Just $ Var $ toName $ getKey $ getUnique t - TyConApp tc ts -> - let step con op - | tc == con, [tl, tr] <- ts = - op <$> toPresburgerExp dic tl <*> toPresburgerExp dic tr - | otherwise = Nothing - in case ts of - [tl, tr] | tc == typeNatMulTyCon -> - case (simpleExp tl, simpleExp tr) of - (LitTy (NumTyLit n), LitTy (NumTyLit m)) -> Just $ K $ n * m - (LitTy (NumTyLit n), x) -> (:*) <$> pure n <*> toPresburgerExp dic x - (x, LitTy (NumTyLit n)) -> (:*) <$> pure n <*> toPresburgerExp dic x - _ -> Nothing - _ -> asum [ step con op - | (con, op) <- [(typeNatAddTyCon, (:+)), (typeNatSubTyCon, (:-))]] - LitTy (NumTyLit n) -> Just (K n) - _ -> Nothing + TyVarTy t -> return $ Var $ toName $ getKey $ getUnique t + t@(TyConApp tc ts) -> body tc ts <|> Var . toName . getKey . getUnique <$> toVar t + LitTy (NumTyLit n) -> return (K n) + LitTy _ -> mzero + t -> Var . toName . getKey .getUnique <$> toVar t + where + body tc ts = + let step con op + | tc == con, [tl, tr] <- ts = + op <$> toPresburgerExp dic tl <*> toPresburgerExp dic tr + | otherwise = mzero + in case ts of + [tl, tr] | tc == typeNatMulTyCon -> + case (simpleExp tl, simpleExp tr) of + (LitTy (NumTyLit n), LitTy (NumTyLit m)) -> return $ K $ n * m + (LitTy (NumTyLit n), x) -> (:*) <$> pure n <*> toPresburgerExp dic x + (x, LitTy (NumTyLit n)) -> (:*) <$> pure n <*> toPresburgerExp dic x + _ -> mzero + _ -> asum [ step con op + | (con, op) <- [(typeNatAddTyCon, (:+)), (typeNatSubTyCon, (:-))]] -- simplTypeCmp :: Type -> Type @@ -344,3 +369,21 @@ | otherwise = Nothing simpleExp t = t +type ParseEnv = M.Map TypeEq TyVar + +type Machine = MaybeT (StateT ParseEnv TcPluginM) + +runMachine :: Machine a -> TcPluginM (Maybe a) +runMachine act = do + (ma, dic) <- runStateT (runMaybeT act) M.empty + forM_ (M.toList dic) $ \(TypeEq ty, var) -> + newWanted undefined $ mkPrimEqPredRole Nominal (mkTyVarTy var) ty + return ma + +toVar :: Type -> Machine TyVar +toVar ty = gets (M.lookup (TypeEq ty)) >>= \case + Just v -> return v + Nothing -> do + v <- lift $ lift $ newFlexiTyVar $ typeKind ty + modify $ M.insert (TypeEq ty) v + return v