diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,4 @@ +# Changelog for the [`ghc-typelits-natnormalise`](http://hackage.haskell.org/package/ghc-typelits-natnormalise) package + +## 0.1 *March 30th 2015* +* Initial release diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,26 @@ +Copyright (c) 2015, University of Twente +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are +met: + +1. Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + +2. Redistributions in binary form must reproduce the above copyright + notice, this list of conditions and the following disclaimer in the + documentation and/or other materials provided with the + distribution. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR +A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT +OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT +LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, +DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY +THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/README.md b/README.md new file mode 100644 --- /dev/null +++ b/README.md @@ -0,0 +1,35 @@ +# ghc-tynat-normalise + +A type checker plugin for GHC that can solve _equalities_ +of types of kind `Nat`, where these types are either: + +* Type-level naturals +* Type variables +* Applications of the arithmetic expressions `(+,-,*,^)`. + +It solves these equalities by normalising them to _sort-of_ +`SOP` (Sum-of-Products) form, and then perform a +simple syntactic equality. + +For example, this solver can prove the equality between: + +``` +(x + 2)^(y + 2) +``` + +and + +``` +4*x*(2 + x)^y + 4*(2 + x)^y + (2 + x)^y*x^2 +``` + +Because the latter is actually the `SOP` normal form +of the former. + +To use the plugin, add + +``` +{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-} +``` + +To the header of your file. diff --git a/Setup.hs b/Setup.hs new file mode 100644 --- /dev/null +++ b/Setup.hs @@ -0,0 +1,2 @@ +import Distribution.Simple +main = defaultMain diff --git a/ghc-typelits-natnormalise.cabal b/ghc-typelits-natnormalise.cabal new file mode 100644 --- /dev/null +++ b/ghc-typelits-natnormalise.cabal @@ -0,0 +1,75 @@ +name: ghc-typelits-natnormalise +version: 0.1 +synopsis: GHC typechecker plugin for types of kind GHC.TypeLits.Nat +description: + A type checker plugin for GHC that can solve /equalities/ of types of kind + 'GHC.TypeLits.Nat', where these types are either: + . + * Type-level naturals + * Type variables + * Applications of the arithmetic expressions @(+,-,*,^)@. + . + It solves these equalities by normalising them to /sort-of/ + 'GHC.TypeLits.Normalise.SOP.SOP' (Sum-of-Products) form, and then perform a + simple syntactic equality. + . + For example, this solver can prove the equality between: + . + @ + (x + 2)^(y + 2) + @ + . + and + . + @ + 4*x*(2 + x)^y + 4*(2 + x)^y + (2 + x)^y*x^2 + @ + . + Because the latter is actually the 'GHC.TypeLits.Normalise.SOP.SOP' normal + form of the former. + . + To use the plugin, add the + . + @ + OPTIONS_GHC -fplugin GHC.TypeLits.Normalise + @ + . + Pragma to the header of your file. +homepage: http://www.clash-lang.org/ +bug-reports: http://github.com/christiaanb/ghc-typelits-natnormalise/issues +license: BSD2 +license-file: LICENSE +author: Christiaan Baaij +maintainer: christiaan.baaij@gmail.com +copyright: Copyright © 2015 University of Twente +category: Type System +build-type: Simple +extra-source-files: README.md + CHANGELOG.md +cabal-version: >=1.10 + +source-repository head + type: git + location: https://github.com/christiaanb/ghc-typelits-natnormalise.git + +library + exposed-modules: GHC.TypeLits.Normalise, + GHC.TypeLits.Normalise.SOP, + GHC.TypeLits.Normalise.Unify + Other-Modules: GHC.Type.Instances + build-depends: base >=4.8 && <5, + ghc >=7.10 && <7.12 + hs-source-dirs: src + default-language: Haskell2010 + ghc-options: -Wall + +test-suite test-ghc-tynat-normalise + type: exitcode-stdio-1.0 + main-is: Tests.hs + build-depends: base >=4.8 && <4.9, + ghc-typelits-natnormalise >= 0.1, + tasty >= 0.10, + tasty-hunit >= 0.9 + hs-source-dirs: tests + default-language: Haskell2010 + ghc-options: -O0 -dcore-lint diff --git a/src/GHC/Type/Instances.hs b/src/GHC/Type/Instances.hs new file mode 100644 --- /dev/null +++ b/src/GHC/Type/Instances.hs @@ -0,0 +1,14 @@ +{-# OPTIONS_GHC -fno-warn-orphans #-} +{-| +Copyright : (C) 2015, University of Twente +License : BSD2 (see the file LICENSE) +Maintainer : Christiaan Baaij + +Ord instance for 'Type' +-} +module GHC.Type.Instances where + +import Type (Type,cmpType) + +instance Ord Type where + compare = cmpType diff --git a/src/GHC/TypeLits/Normalise.hs b/src/GHC/TypeLits/Normalise.hs new file mode 100644 --- /dev/null +++ b/src/GHC/TypeLits/Normalise.hs @@ -0,0 +1,168 @@ +{-# LANGUAGE TupleSections #-} + +{-# OPTIONS_HADDOCK show-extensions #-} + +{-| +Copyright : (C) 2015, University of Twente +License : BSD2 (see the file LICENSE) +Maintainer : Christiaan Baaij + +A type checker plugin for GHC that can solve /equalities/ of types of kind +'GHC.TypeLits.Nat', where these types are either: + +* Type-level naturals +* Type variables +* Applications of the arithmetic expressions @(+,-,*,^)@. + +It solves these equalities by normalising them to /sort-of/ +'GHC.TypeLits.Normalise.SOP.SOP' (Sum-of-Products) form, and then perform a +simple syntactic equality. + +For example, this solver can prove the equality between: + +@ +(x + 2)^(y + 2) +@ + +and + +@ +4*x*(2 + x)^y + 4*(2 + x)^y + (2 + x)^y*x^2 +@ + +Because the latter is actually the 'GHC.TypeLits.Normalise.SOP.SOP' normal form +of the former. + +To use the plugin, add + +@ +{\-\# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise \#-\} +@ + +To the header of your file. +-} +module GHC.TypeLits.Normalise + ( plugin ) +where + +-- external +import Data.Maybe (catMaybes, mapMaybe) + +-- GHC API +import Coercion (Role (Nominal), mkUnivCo) +import FastString (fsLit) +import Outputable (Outputable (..), (<+>), ($$), text) +import Plugins (Plugin (..), defaultPlugin) +import TcEvidence (EvTerm (EvCoercion), TcCoercion (..)) +import TcPluginM (TcPluginM, tcPluginTrace, unsafeTcPluginTcM, zonkCt) +import qualified TcMType +import TcRnTypes (Ct, CtEvidence (..), CtOrigin, TcPlugin(..), + TcPluginResult(..), ctEvidence, ctEvPred, + ctLoc, ctLocOrigin, isGiven, isWanted, mkNonCanonical) +import TcType (mkEqPred, typeKind) +import Type (EqRel (NomEq), Kind, PredTree (EqPred), PredType, Type, + TyVar, classifyPredType, mkTyVarTy) +import TysWiredIn (typeNatKind) + +-- internal +import GHC.TypeLits.Normalise.Unify + +-- | To use the plugin, add +-- +-- @ +-- {\-\# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise \#-\} +-- @ +-- +-- To the header of your file. +plugin :: Plugin +plugin = defaultPlugin { tcPlugin = const $ Just normalisePlugin } + +normalisePlugin :: TcPlugin +normalisePlugin = + TcPlugin { tcPluginInit = return () + , tcPluginSolve = decideEqualSOP + , tcPluginStop = const (return ()) + } + +decideEqualSOP :: () -> [Ct] -> [Ct] -> [Ct] -> TcPluginM TcPluginResult +decideEqualSOP _ _givens _deriveds [] = return (TcPluginOk [] []) +decideEqualSOP _ givens _deriveds wanteds = do + let unit_wanteds = mapMaybe toNatEquality wanteds + case unit_wanteds of + [] -> return (TcPluginOk [] []) + _ -> do + unit_givens <- mapMaybe toNatEquality <$> mapM zonkCt givens + sr <- simplifyNats (unit_givens ++ unit_wanteds) + tcPluginTrace "normalised" (ppr sr) + case sr of + Simplified subst evs -> + TcPluginOk (filter (isWanted . ctEvidence . snd) evs) <$> + mapM substItemToCt (filter (isWanted . ctEvidence . siNote) subst) + Impossible eq -> return (TcPluginContradiction [fromNatEquality eq]) + +substItemToCt :: SubstItem TyVar Type Ct -> TcPluginM Ct +substItemToCt si + | isGiven (ctEvidence ct) = return $ mkNonCanonical + $ CtGiven predicate + (evByFiat "units" (ty1, ty2)) loc + | otherwise = newSimpleWanted (ctLocOrigin loc) predicate + where + predicate = mkEqPred ty1 ty2 + ty1 = mkTyVarTy (siVar si) + ty2 = reifySOP (siSOP si) + ct = siNote si + loc = ctLoc ct + +type NatEquality = (Ct,CoreSOP,CoreSOP) + +fromNatEquality :: NatEquality -> Ct +fromNatEquality (ct, _, _) = ct + +data SimplifyResult + = Simplified CoreSubst [(EvTerm,Ct)] + | Impossible NatEquality + +instance Outputable SimplifyResult where + ppr (Simplified subst evs) = text "Simplified" $$ ppr subst $$ ppr evs + ppr (Impossible eq) = text "Impossible" <+> ppr eq + +simplifyNats :: [NatEquality] -> TcPluginM SimplifyResult +simplifyNats eqs = tcPluginTrace "simplifyNats" (ppr eqs) >> simples [] [] [] eqs + where + simples :: CoreSubst -> [Maybe (EvTerm, Ct)] -> [NatEquality] + -> [NatEquality] -> TcPluginM SimplifyResult + simples subst evs _xs [] = return (Simplified subst (catMaybes evs)) + simples subst evs xs (eq@(ct,u,v):eqs') = do + ur <- unifyNats ct (substsSOP subst u) (substsSOP subst v) + tcPluginTrace "unifyNats result" (ppr ur) + case ur of + Win -> simples subst (((,) <$> evMagic ct <*> pure ct):evs) [] + (xs ++ eqs') + Lose -> return (Impossible eq) + Draw [] -> simples subst evs (eq:xs) eqs' + Draw subst' -> simples (substsSubst subst' subst ++ subst') evs [eq] + (xs ++ eqs') + +-- Extract the Nat equality constraints +toNatEquality :: Ct -> Maybe NatEquality +toNatEquality ct = case classifyPredType $ ctEvPred $ ctEvidence ct of + EqPred NomEq t1 t2 + | isNatKind (typeKind t1) || isNatKind (typeKind t1) + -> Just (ct,normaliseNat t1,normaliseNat t2) + _ -> Nothing + where + isNatKind :: Kind -> Bool + isNatKind = (== typeNatKind) + +-- Utils +newSimpleWanted :: CtOrigin -> PredType -> TcPluginM Ct +newSimpleWanted orig = unsafeTcPluginTcM . TcMType.newSimpleWanted orig + +evMagic :: Ct -> Maybe EvTerm +evMagic ct = case classifyPredType $ ctEvPred $ ctEvidence ct of + EqPred NomEq t1 t2 -> Just (evByFiat "tylits_magic" (t1, t2)) + _ -> Nothing + +evByFiat :: String -> (Type, Type) -> EvTerm +evByFiat name (t1,t2) = EvCoercion $ TcCoercion + $ mkUnivCo (fsLit name) Nominal t1 t2 diff --git a/src/GHC/TypeLits/Normalise/SOP.hs b/src/GHC/TypeLits/Normalise/SOP.hs new file mode 100644 --- /dev/null +++ b/src/GHC/TypeLits/Normalise/SOP.hs @@ -0,0 +1,285 @@ +{-| +Copyright : (C) 2015, University of Twente +License : BSD2 (see the file LICENSE) +Maintainer : Christiaan Baaij + += SOP: Sum-of-Products, sorta + +The arithmetic operation for 'GHC.TypeLits.Nat' are, addition +(@'GHC.TypeLits.+'@), subtraction (@'GHC.TypeLits.-'@), multiplication +(@'GHC.TypeLits.*'@), and exponentiation (@'GHC.TypeLits.^'@). This means we +cannot write expressions in a canonical SOP normal form. We can get rid of +subtraction by working with integers, and translating @a - b@ to @a + (-1)*b@. +Exponentation cannot be getten rid of that way. So we define the following +grammar for our canonical SOP-like normal form of arithmetic expressions: + +@ +SOP ::= Product \'+\' SOP | Product +Product ::= Symbol \'*\' Product | Symbol +Symbol ::= Integer + | Var + | Var \'^\' Product + | SOP \'^\' ProductE + +ProductE ::= SymbolE \'*\' ProductE | SymbolE +SymbolE ::= Var + | Var \'^\' Product + | SOP \'^\' ProductE +@ + +So a valid SOP terms are: + +@ +x*y + y^2 +(x+y)^(k*z) +@ + +, but, + +@ +(x*y)^2 +@ + +is not, and should be: + +@ +x^2 * y^2 +@ + +Exponents are thus not allowed to have products, so for example, the expression: + +@ +(x + 2)^(y + 2) +@ + +in valid SOP form is: + +@ +4*x*(2 + x)^y + 4*(2 + x)^y + (2 + x)^y*x^2 +@ + +Also, exponents can only be integer values when the base is a variable. Although +not enforced by the grammar, the exponentials are flatted as far as possible in +SOP form. So: + +@ +(x^y)^z +@ + +is flattened to: + +@ +x^(y*z) +@ +-} +module GHC.TypeLits.Normalise.SOP + ( -- * SOP types + Symbol (..) + , Product (..) + , SOP (..) + -- * Simplification + , reduceExp + , mergeS + , mergeP + , mergeSOPAdd + , mergeSOPMul + , normaliseExp + ) +where + +-- External +import Data.Either (partitionEithers) +import Data.List (sort) + +-- GHC API +import Outputable (Outputable (..), (<+>), text, hcat, integer, punctuate) + +data Symbol v c + = I Integer -- ^ Integer constant + | C c -- ^ Non-integer constant + | E (SOP v c) (Product v c) -- ^ Exponentiation + | V v -- ^ Variable + deriving (Eq,Ord) + +newtype Product v c = P { unP :: [Symbol v c] } + deriving (Eq,Ord) + +newtype SOP v c = S { unS :: [Product v c] } + deriving (Eq,Ord) + +instance (Outputable v, Outputable c) => Outputable (SOP v c) where + ppr = hcat . punctuate (text " + ") . map ppr . unS + +instance (Outputable v, Outputable c) => Outputable (Product v c) where + ppr = hcat . punctuate (text " * ") . map ppr . unP + +instance (Outputable v, Outputable c) => Outputable (Symbol v c) where + ppr (I i) = integer i + ppr (C c) = ppr c + ppr (V s) = ppr s + ppr (E b e) = case (pprSimple b, pprSimple (S [e])) of + (bS,eS) -> bS <+> text "^" <+> eS + where + pprSimple (S [P [I i]]) = integer i + pprSimple (S [P [V v]]) = ppr v + pprSimple sop = text "(" <+> ppr sop <+> text ")" + +mergeWith :: (a -> a -> Either a a) -> [a] -> [a] +mergeWith _ [] = [] +mergeWith op (f:fs) = case partitionEithers $ map (`op` f) fs of + ([],_) -> f : mergeWith op fs + (updated,untouched) -> mergeWith op (updated ++ untouched) + +-- | reduce exponentials +-- +-- Performs the following rewrites: +-- +-- @ +-- x^0 ==> 1 +-- 0^x ==> 0 +-- 2^3 ==> 8 +-- (k ^ i) ^ j ==> k ^ (i * j) +-- @ +reduceExp :: (Ord v, Ord c) => Symbol v c -> Symbol v c +reduceExp (E _ (P [(I 0)])) = I 1 -- x^0 ==> 1 +reduceExp (E (S [P [I 0]]) _ ) = I 0 -- 0^x ==> 0 +reduceExp (E (S [P [(I i)]]) (P [(I j)])) = I (i ^ j) -- 2^3 ==> 8 + +-- (k ^ i) ^ j ==> k ^ (i * j) +reduceExp (E (S [P [(E k i)]]) j) = case normaliseExp k (S [e]) of + (S [P [s]]) -> s + _ -> E k e + where + e = P . sort . map reduceExp $ mergeWith mergeS (unP i ++ unP j) + +reduceExp s = s + +-- | Merge two symbols of a Product term +-- +-- Performs the following rewrites: +-- +-- @ +-- 8 * 7 ==> 56 +-- 1 * x ==> x +-- x * 1 ==> x +-- 0 * x ==> 0 +-- x * 0 ==> 0 +-- x * x^4 ==> x^5 +-- x^4 * x ==> x^5 +-- y*y ==> y^2 +-- @ +mergeS :: (Ord v, Ord c) => Symbol v c -> Symbol v c + -> Either (Symbol v c) (Symbol v c) +mergeS (I i) (I j) = Left (I (i * j)) -- 8 * 7 ==> 56 +mergeS (I 1) r = Left r -- 1 * x ==> x +mergeS l (I 1) = Left l -- x * 1 ==> x +mergeS (I 0) _ = Left (I 0) -- 0 * x ==> 0 +mergeS _ (I 0) = Left (I 0) -- x * 0 ==> 0 + +-- x * x^4 ==> x^5 +mergeS s (E (S [P [s']]) (P [I i])) + | s == s' + = Left (E (S [P [s']]) (P [I (i + 1)])) + +-- x^4 * x ==> x^5 +mergeS (E (S [P [s']]) (P [I i])) s + | s == s' + = Left (E (S [P [s']]) (P [I (i + 1)])) + +-- y*y ==> y^2 +mergeS l r + | l == r + = case normaliseExp (S [P [l]]) (S [P [I 2]]) of + (S [P [e]]) -> Left e + _ -> Right l + +mergeS l _ = Right l + +-- | Merge two products of a SOP term +-- +-- Performs the following rewrites: +-- +-- @ +-- 2xy + 3xy ==> 5xy +-- 2xy + xy ==> 3xy +-- xy + 2xy ==> 3xy +-- xy + xy ==> 2xy +-- @ +mergeP :: (Eq v, Eq c) => Product v c -> Product v c + -> Either (Product v c) (Product v c) +-- 2xy + 3xy ==> 5xy +mergeP (P ((I i):is)) (P ((I j):js)) + | is == js = Left . P $ (I (i + j)) : is +-- 2xy + xy ==> 3xy +mergeP (P ((I i):is)) (P js) + | is == js = Left . P $ (I (i + 1)) : is +-- xy + 2xy ==> 3xy +mergeP (P is) (P ((I j):js)) + | is == js = Left . P $ (I (j + 1)) : is +-- xy + xy ==> 2xy +mergeP (P is) (P js) + | is == js = Left . P $ (I 2) : is + | otherwise = Right $ P is + +-- | Expand or Simplify 'complex' exponentials +-- +-- Performs the following rewrites: +-- +-- @ +-- b^1 ==> b +-- 2^(y^2) ==> 4^y +-- (x + 2)^2 ==> x^2 + 4xy + 4 +-- (x + 2)^(2x) ==> (x^2 + 4xy + 4)^x +-- (x + 2)^(y + 2) ==> 4x(2 + x)^y + 4(2 + x)^y + (2 + x)^yx^2 +-- @ +normaliseExp :: (Ord v, Ord c) => SOP v c -> SOP v c -> SOP v c +-- b^1 ==> b +normaliseExp b (S [P [I 1]]) = b + +-- x^(2xy) ==> x^(2xy) +normaliseExp b@(S [P [V _]]) (S [e]) = S [P [E b e]] + +-- 2^(y^2) ==> 4^y +normaliseExp b@(S [P [_]]) (S [e@(P [_])]) = S [P [reduceExp (E b e)]] + +-- (x + 2)^2 ==> x^2 + 4xy + 4 +normaliseExp b (S [P [(I i)]]) = + foldr1 mergeSOPMul (replicate (fromInteger i) b) + +-- (x + 2)^(2x) ==> (x^2 + 4xy + 4)^x +normaliseExp b (S [P (e@(I _):es)]) = + normaliseExp (normaliseExp b (S [P [e]])) (S [P es]) + +-- (x + 2)^(xy) ==> (x+2)^(xy) +normaliseExp b (S [e]) = S [P [reduceExp (E b e)]] + +-- (x + 2)^(y + 2) ==> 4x(2 + x)^y + 4(2 + x)^y + (2 + x)^yx^2 +normaliseExp b (S e) = foldr1 mergeSOPMul (map (normaliseExp b . S . (:[])) e) + +zeroP :: Product v c -> Bool +zeroP (P ((I 0):_)) = True +zeroP _ = False + +-- | Simplifies SOP terms using +-- +-- * 'mergeS' +-- * 'mergeP' +-- * 'reduceExp' +simplifySOP :: (Ord v, Ord c) => SOP v c -> SOP v c +simplifySOP + = S + . sort . filter (not . zeroP) + . mergeWith mergeP + . map (P . sort . map reduceExp . mergeWith mergeS . unP) + . unS + +-- | Merge two SOP terms by additions +mergeSOPAdd :: (Ord v, Ord c) => SOP v c -> SOP v c -> SOP v c +mergeSOPAdd (S sop1) (S sop2) = simplifySOP $ S (sop1 ++ sop2) + +-- | Merge two SOP terms by multiplication +mergeSOPMul :: (Ord v, Ord c) => SOP v c -> SOP v c -> SOP v c +mergeSOPMul (S sop1) (S sop2) + = simplifySOP + . S + $ concatMap (zipWith (\p1 p2 -> P (unP p1 ++ unP p2)) sop1 . repeat) sop2 diff --git a/src/GHC/TypeLits/Normalise/Unify.hs b/src/GHC/TypeLits/Normalise/Unify.hs new file mode 100644 --- /dev/null +++ b/src/GHC/TypeLits/Normalise/Unify.hs @@ -0,0 +1,236 @@ +{-# OPTIONS_GHC -fno-warn-unused-imports #-} + +{-| +Copyright : (C) 2015, University of Twente +License : BSD2 (see the file LICENSE) +Maintainer : Christiaan Baaij +-} +module GHC.TypeLits.Normalise.Unify + ( -- * 'Nat' expressions \<-\> 'SOP' terms + CoreSOP + , normaliseNat + , reifySOP + -- * Substitution on 'SOP' terms + , SubstItem (..) + , TySubst + , CoreSubst + , substsSOP + , substsSubst + -- * Find unifiers + , UnifyResult (..) + , unifyNats + , unifiers + -- * Free variables in 'SOP' terms + , fvSOP + ) +where + +-- External +import Data.Function (on) +import Data.List ((\\), intersect) + +-- GHC API +import Outputable (Outputable (..), (<+>), ($$), text) +import TcPluginM (TcPluginM, tcPluginTrace) +import TcRnMonad (Ct, ctEvidence, isGiven) +import TcTypeNats (typeNatAddTyCon, typeNatExpTyCon, typeNatMulTyCon, + typeNatSubTyCon) +import Type (TyVar, mkNumLitTy, mkTyConApp, mkTyVarTy, tcView) +import TypeRep (Type (..), TyLit (..)) +import UniqSet (UniqSet, unionManyUniqSets, emptyUniqSet, unionUniqSets, + unitUniqSet) + +-- Internal +import GHC.Type.Instances () -- Ord instance for Type +import GHC.TypeLits.Normalise.SOP + +-- Used for haddock +import GHC.TypeLits (Nat) + +-- | 'SOP' with 'TyVar' variables +type CoreSOP = SOP TyVar Type +type CoreProduct = Product TyVar Type +type CoreSymbol = Symbol TyVar Type + +-- | Convert a type of /kind/ 'GHC.TypeLits.Nat' to an 'SOP' term, but +-- only when the type is constructed out of: +-- +-- * literals +-- * type variables +-- * Applications of the arithmetic operators @(+,-,*,^)@ +normaliseNat :: Type -> CoreSOP +normaliseNat ty | Just ty1 <- tcView ty = normaliseNat ty1 +normaliseNat (TyVarTy v) = S [P [V v]] +normaliseNat (LitTy (NumTyLit i)) = S [P [I i]] +normaliseNat (TyConApp tc [x,y]) + | tc == typeNatAddTyCon = mergeSOPAdd (normaliseNat x) (normaliseNat y) + | tc == typeNatSubTyCon = mergeSOPAdd (normaliseNat x) + (mergeSOPMul (S [P [I (-1)]]) + (normaliseNat y)) + | tc == typeNatMulTyCon = mergeSOPMul (normaliseNat x) (normaliseNat y) + | tc == typeNatExpTyCon = normaliseExp (normaliseNat x) (normaliseNat y) +normaliseNat t = S [P [C t]] + +-- | Convert a 'SOP' term back to a type of /kind/ 'GHC.TypeLits.Nat' +reifySOP :: CoreSOP -> Type +reifySOP = combineP . map negateP . unS + where + negateP :: CoreProduct -> Either CoreProduct CoreProduct + negateP (P ((I i):ps)) | i < 0 = Left (P ps) + negateP ps = Right ps + + combineP :: [Either CoreProduct CoreProduct] -> Type + combineP [] = mkNumLitTy 0 + combineP [p] = either (\p' -> mkTyConApp typeNatSubTyCon + [mkNumLitTy 0, reifyProduct p']) + reifyProduct p + combineP (p:ps) = let es = combineP ps + in either (\x -> mkTyConApp typeNatSubTyCon + [es, reifyProduct x]) + (\x -> mkTyConApp typeNatAddTyCon + [reifyProduct x, es]) + p + +reifyProduct :: CoreProduct -> Type +reifyProduct = foldr1 (\t1 t2 -> mkTyConApp typeNatMulTyCon [t1,t2]) + . map reifySymbol . unP + +reifySymbol :: CoreSymbol -> Type +reifySymbol (I i) = mkNumLitTy i +reifySymbol (C c) = c +reifySymbol (V v) = mkTyVarTy v +reifySymbol (E s p) = mkTyConApp typeNatExpTyCon [reifySOP s,reifyProduct p] + +-- | A substitution is essentially a list of (variable, 'SOP') pairs, +-- but we keep the original 'Ct' that lead to the substitution being +-- made, for use when turning the substitution back into constraints. +type CoreSubst = TySubst TyVar Type Ct +type TySubst v c n = [SubstItem v c n] + +data SubstItem v c n = SubstItem { siVar :: v + , siSOP :: SOP v c + , siNote :: n + } + +instance (Outputable v, Outputable c) => Outputable (SubstItem v c n) where + ppr si = ppr (siVar si) <+> text " := " <+> ppr (siSOP si) + +-- | Apply a substitution to a single normalised 'SOP' term +substsSOP :: (Ord v, Ord c) => TySubst v c n -> SOP v c -> SOP v c +substsSOP [] u = u +substsSOP (si:s) u = substsSOP s (substSOP (siVar si) (siSOP si) u) + +substSOP :: (Ord v, Ord c) => v -> SOP v c -> SOP v c -> SOP v c +substSOP tv e = foldr1 mergeSOPAdd . map (substProduct tv e) . unS + +substProduct :: (Ord v, Ord c) => v -> SOP v c -> Product v c -> SOP v c +substProduct tv e = foldr1 mergeSOPMul . map (substSymbol tv e) . unP + +substSymbol :: (Ord v, Ord c) => v -> SOP v c -> Symbol v c -> SOP v c +substSymbol _ _ s@(I _) = S [P [s]] +substSymbol _ _ s@(C _) = S [P [s]] +substSymbol tv e (V tv') + | tv == tv' = e + | otherwise = S [P [V tv']] +substSymbol tv e (E s p) = normaliseExp (substSOP tv e s) (substProduct tv e p) + +-- | Apply a substitution to a substitution +substsSubst :: (Ord v, Ord c) => TySubst v c n -> TySubst v c n -> TySubst v c n +substsSubst s = map (\si -> si {siSOP = substsSOP s (siSOP si)}) + +-- | Result of comparing two 'SOP' terms, returning a potential substitution +-- list under which the two terms are equal. +data UnifyResult + = Win -- ^ Two terms are equal + | Lose -- ^ Two terms are /not/ equal + | Draw CoreSubst -- ^ Two terms are only equal if the given substitution holds + +instance Outputable UnifyResult where + ppr Win = text "Win" + ppr (Draw subst) = text "Draw" <+> ppr subst + ppr Lose = text "Lose" + +-- | Given two 'SOP's @u@ and @v@, when their free variables ('fvSOP') are the +-- same, then we 'Win' if @u@ and @v@ are equal, and 'Lose' otherwise. +-- +-- If @u@ and @v@ do not have the same free variables, we result in a 'Draw', +-- ware @u@ and @v@ are only equal when the returned 'CoreSubst' holds. +unifyNats :: Ct -> CoreSOP -> CoreSOP -> TcPluginM UnifyResult +unifyNats ct u v = do + tcPluginTrace "unifyNats" (ppr ct $$ ppr u $$ ppr v) + return (unifyNats' ct u v) + +unifyNats' :: Ct -> CoreSOP -> CoreSOP -> UnifyResult +unifyNats' ct u v + | eqFV u v = if u == v then Win else Lose + | otherwise = Draw (unifiers ct u v) + +-- | Find unifiers for two SOP terms +-- +-- Can find the following unifiers: +-- +-- @ +-- t ~ a + b ==> [t := a + b] +-- a + b ~ t ==> [t := a + b] +-- (a + c) ~ (b + c) ==> \[a := b\] +-- (2*a) ~ (2*b) ==> [a := b] +-- @ +-- +-- However, given a wanted: +-- +-- @ +-- [W] t ~ a + b +-- @ +-- +-- this function returns @[]@, or otherwise we \"solve\" the contstraint by +-- finding a unifier equal to the constraint. +-- +-- However, given a wanted: +-- +-- @ +-- [W] (a + c) ~ (b + c) +-- @ +-- +-- we do return the unifier: +-- +-- @ +-- [a := b] +-- @ +unifiers :: Ct -> CoreSOP -> CoreSOP -> CoreSubst +unifiers ct (S [P [V x]]) s + | isGiven (ctEvidence ct) = [SubstItem x s ct] + | otherwise = [] +unifiers ct s (S [P [V x]]) + | isGiven (ctEvidence ct) = [SubstItem x s ct] + | otherwise = [] +unifiers ct u v = unifiers' ct u v + +unifiers' :: Ct -> CoreSOP -> CoreSOP -> CoreSubst +unifiers' ct (S [P [V x]]) (S []) = [SubstItem x (S [P [I 0]]) ct] +unifiers' ct (S []) (S [P [V x]]) = [SubstItem x (S [P [I 0]]) ct] +unifiers' ct (S [P [V x]]) s = [SubstItem x s ct] +unifiers' ct s (S [P [V x]]) = [SubstItem x s ct] +unifiers' ct (S [P (p:ps1)]) (S [P (p':ps2)]) + | p == p' = unifiers' ct (S [P ps1]) (S [P ps2]) + | otherwise = [] +unifiers' ct (S ps1) (S ps2) + | null psx = [] + | otherwise = unifiers' ct (S (ps1 \\ psx)) (S (ps2 \\ psx)) + where + psx = intersect ps1 ps2 + +-- | Find the 'TyVar' in a 'CoreSOP' +fvSOP :: CoreSOP -> UniqSet TyVar +fvSOP = unionManyUniqSets . map fvProduct . unS + +fvProduct :: CoreProduct -> UniqSet TyVar +fvProduct = unionManyUniqSets . map fvSymbol . unP + +fvSymbol :: CoreSymbol -> UniqSet TyVar +fvSymbol (I _) = emptyUniqSet +fvSymbol (C _) = emptyUniqSet +fvSymbol (V v) = unitUniqSet v +fvSymbol (E s p) = fvSOP s `unionUniqSets` fvProduct p + +eqFV :: CoreSOP -> CoreSOP -> Bool +eqFV = (==) `on` fvSOP diff --git a/tests/Tests.hs b/tests/Tests.hs new file mode 100644 --- /dev/null +++ b/tests/Tests.hs @@ -0,0 +1,259 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE KindSignatures #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE NoImplicitPrelude #-} + +{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-} + +import Data.Proxy +import GHC.TypeLits +import Unsafe.Coerce +import Prelude hiding (head,tail,init,(++),splitAt,concat,drop) +import qualified Prelude as P + +import Test.Tasty +import Test.Tasty.HUnit + +data Vec :: Nat -> * -> * where + Nil :: Vec 0 a + (:>) :: a -> Vec n a -> Vec (n + 1) a + +instance Show a => Show (Vec n a) where + show vs = "<" P.++ punc vs P.++ ">" + where + punc :: Show a => Vec m a -> String + punc Nil = "" + punc (x :> Nil) = show x + punc (x :> xs) = show x P.++ "," P.++ punc xs + +infixr 5 :> + +data SNat (n :: Nat) = KnownNat n => SNat (Proxy n) + +instance Show (SNat n) where + show (SNat p) = 'd' : show (natVal p) + +{-# INLINE snat #-} +-- | Create a singleton literal for a type-level natural number +snat :: KnownNat n => SNat n +snat = SNat Proxy + +{-# INLINE withSNat #-} +-- | Supply a function with a singleton natural 'n' according to the context +withSNat :: KnownNat n => (SNat n -> a) -> a +withSNat f = f (SNat Proxy) + +{-# INLINE snatToInteger #-} +snatToInteger :: SNat n -> Integer +snatToInteger (SNat p) = natVal p + +data UNat :: Nat -> * where + UZero :: UNat 0 + USucc :: UNat n -> UNat (n + 1) + +-- | Convert a singleton natural number to its unary representation +-- +-- __NB__: Not synthesisable +toUNat :: SNat n -> UNat n +toUNat (SNat p) = fromI (natVal p) + where + fromI :: Integer -> UNat m + fromI 0 = unsafeCoerce UZero + fromI n = unsafeCoerce (USucc (fromI (n - 1))) + +-- | Add two singleton natural numbers +-- +-- __NB__: Not synthesisable +addUNat :: UNat n -> UNat m -> UNat (n + m) +addUNat UZero y = y +addUNat x UZero = x +addUNat (USucc x) y = USucc (addUNat x y) + +-- | Multiply two singleton natural numbers +-- +-- __NB__: Not synthesisable +multUNat :: UNat n -> UNat m -> UNat (n * m) +multUNat UZero _ = UZero +multUNat _ UZero = UZero +multUNat (USucc x) y = addUNat y (multUNat x y) + +-- | Exponential of two singleton natural numbers +-- +-- __NB__: Not synthesisable +powUNat :: UNat n -> UNat m -> UNat (n ^ m) +powUNat _ UZero = USucc UZero +powUNat x (USucc y) = multUNat x (powUNat x y) + +-- | Extract the first element of a vector +-- +-- >>> head (1:>2:>3:>Nil) +-- 1 +head :: Vec (n + 1) a -> a +head (x :> _) = x + +-- | Extract the elements after the head of a vector +-- +-- >>> tail (1:>2:>3:>Nil) +-- <2,3> +tail :: Vec (n + 1) a -> Vec n a +tail (_ :> xs) = xs + +-- | Extract all the elements of a vector except the last element +-- +-- >>> init (1:>2:>3:>Nil) +-- <1,2> +init :: Vec (n + 1) a -> Vec n a +init (_ :> Nil) = Nil +init (x :> y :> ys) = x :> init (y :> ys) + +infixr 5 ++ +-- | Append two vectors +-- +-- >>> (1:>2:>3:>Nil) ++ (7:>8:>Nil) +-- <1,2,3,7,8> +(++) :: Vec n a -> Vec m a -> Vec (n + m) a +Nil ++ ys = ys +(x :> xs) ++ ys = x :> xs ++ ys + +-- | Split a vector into two vectors at the given point +-- +-- >>> splitAt (snat :: SNat 3) (1:>2:>3:>7:>8:>Nil) +-- (<1,2,3>, <7,8>) +-- >>> splitAt d3 (1:>2:>3:>7:>8:>Nil) +-- (<1,2,3>, <7,8>) +splitAt :: SNat m -> Vec (m + n) a -> (Vec m a, Vec n a) +splitAt n xs = splitAtU (toUNat n) xs + +splitAtU :: UNat m -> Vec (m + n) a -> (Vec m a, Vec n a) +splitAtU UZero ys = (Nil,ys) +splitAtU (USucc s) (y :> ys) = let (as,bs) = splitAtU s ys + in (y :> as, bs) + +{-# INLINE splitAtI #-} +-- | Split a vector into two vectors where the length of the two is determined +-- by the context +-- +-- >>> splitAtI (1:>2:>3:>7:>8:>Nil) :: (Vec 2 Int, Vec 3 Int) +-- (<1,2>,<3,7,8>) +splitAtI :: KnownNat m => Vec (m + n) a -> (Vec m a, Vec n a) +splitAtI = withSNat splitAt + +-- | Shift in elements to the head of a vector, bumping out elements at the +-- tail. The result is a tuple containing: +-- +-- * The new vector +-- * The shifted out elements +-- +-- >>> shiftInAt0 (1 :> 2 :> 3 :> 4 :> Nil) ((-1) :> 0 :> Nil) +-- (<-1,0,1,2,>,<3,4>) +-- >>> shiftInAt0 (1 :> Nil) ((-1) :> 0 :> Nil) +-- (<-1>,<0,1>) +shiftInAt0 :: KnownNat n + => Vec n a -- ^ The old vector + -> Vec m a -- ^ The elements to shift in at the head + -> (Vec n a, Vec m a) -- ^ (The new vector, shifted out elements) +shiftInAt0 xs ys = splitAtI zs + where + zs = ys ++ xs + +-- | Shift in element to the tail of a vector, bumping out elements at the head. +-- The result is a tuple containing: +-- +-- * The new vector +-- * The shifted out elements +-- +-- >>> shiftInAtN (1 :> 2 :> 3 :> 4 :> Nil) (5 :> 6 :> Nil) +-- (<3,4,5,6>,<1,2>) +-- >>> shiftInAtN (1 :> Nil) (2 :> 3 :> Nil) +-- (<3>,<1,2>) +shiftInAtN :: KnownNat m + => Vec n a -- ^ The old vector + -> Vec m a -- ^ The elements to shift in at the tail + -> (Vec n a,Vec m a) -- ^ (The new vector, shifted out elements) +shiftInAtN xs ys = (zsR, zsL) + where + zs = xs ++ ys + (zsL,zsR) = splitAtI zs + +-- | Concatenate a vector of vectors +-- +-- >>> concat ((1:>2:>3:>Nil) :> (4:>5:>6:>Nil) :> (7:>8:>9:>Nil) :> (10:>11:>12:>Nil) :> Nil) +-- <1,2,3,4,5,6,7,8,9,10,11,12> +concat :: Vec n (Vec m a) -> Vec (n * m) a +concat Nil = Nil +concat (x :> xs) = x ++ concat xs + +-- | Split a vector of (n * m) elements into a vector of vectors with length m, +-- where m is given +-- +-- >>> unconcat d4 (1:>2:>3:>4:>5:>6:>7:>8:>9:>10:>11:>12:>Nil) +-- <<1,2,3,4>,<5,6,7,8>,<9,10,11,12>> +unconcat :: KnownNat n => SNat m -> Vec (n * m) a -> Vec n (Vec m a) +unconcat n xs = unconcatU (withSNat toUNat) (toUNat n) xs + +unconcatU :: UNat n -> UNat m -> Vec (n * m) a -> Vec n (Vec m a) +unconcatU UZero _ _ = Nil +unconcatU (USucc n') m ys = let (as,bs) = splitAtU m ys + in as :> unconcatU n' m bs + +-- | Merge two vectors, alternating their elements, i.e., +-- +-- >>> merge (1 :> 2 :> 3 :> 4 :> Nil) (5 :> 6 :> 7 :> 8 :> Nil) +-- <1,5,2,6,3,7,4,8> +merge :: Vec n a -> Vec n a -> Vec (n + n) a +merge Nil Nil = Nil +merge (x :> xs) (y :> ys) = x :> y :> merge xs ys + +-- | 'drop' @n xs@ returns the suffix of @xs@ after the first @n@ elements +-- +-- >>> drop (snat :: SNat 3) (1:>2:>3:>4:>5:>Nil) +-- <4,5> +-- >>> drop d3 (1:>2:>3:>4:>5:>Nil) +-- <4,5> +-- >>> drop d0 (1:>2:>Nil) +-- <1,2> +drop :: SNat m -> Vec (m + n) a -> Vec n a +drop n = snd . splitAt n + +-- | 'at' @n xs@ returns @n@'th element of @xs@ +-- +-- __NB__: vector elements have an __ASCENDING__ subscript starting from 0 and +-- ending at 'maxIndex'. +-- +-- >>> at (snat :: SNat 1) (1:>2:>3:>4:>5:>Nil) +-- 2 +-- >>> at d1 (1:>2:>3:>4:>5:>Nil) +-- 2 +at :: SNat m -> Vec (m + (n + 1)) a -> a +at n xs = head $ snd $ splitAt n xs + +main :: IO () +main = defaultMain tests + +tests :: TestTree +tests = testGroup "ghc-typelits-natnormalise" + [ testGroup "Basic functionality" + [ testCase "show (head (1:>2:>3:>Nil))" $ + show (head (1:>2:>3:>Nil)) @?= + "1" + , testCase "show (tail (1:>2:>3:>Nil))" $ + show (tail (1:>2:>3:>Nil)) @?= + "<2,3>" + , testCase "show (init (1:>2:>3:>Nil))" $ + show (init (1:>2:>3:>Nil)) @?= + "<1,2>" + , testCase "show ((1:>2:>3:>Nil) ++ (7:>8:>Nil))" $ + show ((1:>2:>3:>Nil) ++ (7:>8:>Nil)) @?= + "<1,2,3,7,8>" + , testCase "show (splitAt (snat :: SNat 3) (1:>2:>3:>7:>8:>Nil))" $ + show (splitAt (snat :: SNat 3) (1:>2:>3:>7:>8:>Nil)) @?= + "(<1,2,3>,<7,8>)" + , testCase "show (concat ((1:>2:>3:>Nil) :> (4:>5:>6:>Nil) :> (7:>8:>9:>Nil) :> (10:>11:>12:>Nil) :> Nil))" $ + show (concat ((1:>2:>3:>Nil) :> (4:>5:>6:>Nil) :> (7:>8:>9:>Nil) :> (10:>11:>12:>Nil) :> Nil)) @?= + "<1,2,3,4,5,6,7,8,9,10,11,12>" + , testCase "show (unconcat (snat :: SNat 4) (1:>2:>3:>4:>5:>6:>7:>8:>9:>10:>11:>12:>Nil))" $ + show (unconcat (snat :: SNat 4) (1:>2:>3:>4:>5:>6:>7:>8:>9:>10:>11:>12:>Nil)) @?= + "<<1,2,3,4>,<5,6,7,8>,<9,10,11,12>>" + ] + ]