Implement Div, Mod, and Log for type-level nats.

Reviewers: austin, hvr, bgamari

Reviewed By: bgamari

Subscribers: RyanGlScott, dfeuer, adamgundry, rwbarton, thomie

Differential Revision: https://phabricator.haskell.org/D4002

7 files changed, 206 insertions, 15 deletions

diff --git a/compiler/prelude/PrelNames.hs b/compiler/prelude/PrelNames.hs
index 0159037606..a36f77b0e9 100644
--- a/compiler/prelude/PrelNames.hs
+++ b/compiler/prelude/PrelNames.hs
@@ -1816,6 +1816,9 @@ typeNatKindConNameKey, typeSymbolKindConNameKey,
   typeNatAddTyFamNameKey, typeNatMulTyFamNameKey, typeNatExpTyFamNameKey,
   typeNatLeqTyFamNameKey, typeNatSubTyFamNameKey
   , typeSymbolCmpTyFamNameKey, typeNatCmpTyFamNameKey
+  , typeNatDivTyFamNameKey
+  , typeNatModTyFamNameKey
+  , typeNatLogTyFamNameKey
   :: Unique
 typeNatKindConNameKey     = mkPreludeTyConUnique 164
 typeSymbolKindConNameKey  = mkPreludeTyConUnique 165
@@ -1826,48 +1829,51 @@ typeNatLeqTyFamNameKey    = mkPreludeTyConUnique 169
 typeNatSubTyFamNameKey    = mkPreludeTyConUnique 170
 typeSymbolCmpTyFamNameKey = mkPreludeTyConUnique 171
 typeNatCmpTyFamNameKey    = mkPreludeTyConUnique 172
+typeNatDivTyFamNameKey  = mkPreludeTyConUnique 173
+typeNatModTyFamNameKey  = mkPreludeTyConUnique 174
+typeNatLogTyFamNameKey  = mkPreludeTyConUnique 175
 
 -- Custom user type-errors
 errorMessageTypeErrorFamKey :: Unique
-errorMessageTypeErrorFamKey =  mkPreludeTyConUnique 173
+errorMessageTypeErrorFamKey =  mkPreludeTyConUnique 176
 
 
 
 ntTyConKey:: Unique
-ntTyConKey = mkPreludeTyConUnique 174
+ntTyConKey = mkPreludeTyConUnique 177
 coercibleTyConKey :: Unique
-coercibleTyConKey = mkPreludeTyConUnique 175
+coercibleTyConKey = mkPreludeTyConUnique 178
 
 proxyPrimTyConKey :: Unique
-proxyPrimTyConKey = mkPreludeTyConUnique 176
+proxyPrimTyConKey = mkPreludeTyConUnique 179
 
 specTyConKey :: Unique
-specTyConKey = mkPreludeTyConUnique 177
+specTyConKey = mkPreludeTyConUnique 180
 
 anyTyConKey :: Unique
-anyTyConKey = mkPreludeTyConUnique 178
+anyTyConKey = mkPreludeTyConUnique 181
 
-smallArrayPrimTyConKey        = mkPreludeTyConUnique  179
-smallMutableArrayPrimTyConKey = mkPreludeTyConUnique  180
+smallArrayPrimTyConKey        = mkPreludeTyConUnique  182
+smallMutableArrayPrimTyConKey = mkPreludeTyConUnique  183
 
 staticPtrTyConKey  :: Unique
-staticPtrTyConKey  = mkPreludeTyConUnique 181
+staticPtrTyConKey  = mkPreludeTyConUnique 184
 
 staticPtrInfoTyConKey :: Unique
-staticPtrInfoTyConKey = mkPreludeTyConUnique 182
+staticPtrInfoTyConKey = mkPreludeTyConUnique 185
 
 callStackTyConKey :: Unique
-callStackTyConKey = mkPreludeTyConUnique 183
+callStackTyConKey = mkPreludeTyConUnique 186
 
 -- Typeables
 typeRepTyConKey, someTypeRepTyConKey, someTypeRepDataConKey :: Unique
-typeRepTyConKey       = mkPreludeTyConUnique 184
-someTypeRepTyConKey   = mkPreludeTyConUnique 185
-someTypeRepDataConKey = mkPreludeTyConUnique 186
+typeRepTyConKey       = mkPreludeTyConUnique 187
+someTypeRepTyConKey   = mkPreludeTyConUnique 188
+someTypeRepDataConKey = mkPreludeTyConUnique 189
 
 
 typeSymbolAppendFamNameKey :: Unique
-typeSymbolAppendFamNameKey = mkPreludeTyConUnique 187
+typeSymbolAppendFamNameKey = mkPreludeTyConUnique 190
 
 ---------------- Template Haskell -------------------
 --      THNames.hs: USES TyConUniques 200-299
diff --git a/compiler/typecheck/TcTypeNats.hs b/compiler/typecheck/TcTypeNats.hs
index a3b45b3890..051f9985e6 100644
--- a/compiler/typecheck/TcTypeNats.hs
+++ b/compiler/typecheck/TcTypeNats.hs
@@ -35,6 +35,9 @@ import PrelNames  ( gHC_TYPELITS
                   , typeNatExpTyFamNameKey
                   , typeNatLeqTyFamNameKey
                   , typeNatSubTyFamNameKey
+                  , typeNatDivTyFamNameKey
+                  , typeNatModTyFamNameKey
+                  , typeNatLogTyFamNameKey
                   , typeNatCmpTyFamNameKey
                   , typeSymbolCmpTyFamNameKey
                   , typeSymbolAppendFamNameKey
@@ -44,6 +47,7 @@ import FastString ( FastString
                   )
 import qualified Data.Map as Map
 import Data.Maybe ( isJust )
+import Control.Monad ( guard )
 import Data.List  ( isPrefixOf, isSuffixOf )
 
 {-------------------------------------------------------------------------------
@@ -57,6 +61,9 @@ typeNatTyCons =
   , typeNatExpTyCon
   , typeNatLeqTyCon
   , typeNatSubTyCon
+  , typeNatDivTyCon
+  , typeNatModTyCon
+  , typeNatLogTyCon
   , typeNatCmpTyCon
   , typeSymbolCmpTyCon
   , typeSymbolAppendTyCon
@@ -95,6 +102,32 @@ typeNatMulTyCon = mkTypeNatFunTyCon2 name
   name = mkWiredInTyConName UserSyntax gHC_TYPENATS (fsLit "*")
             typeNatMulTyFamNameKey typeNatMulTyCon
 
+typeNatDivTyCon :: TyCon
+typeNatDivTyCon = mkTypeNatFunTyCon2 name
+  BuiltInSynFamily
+    { sfMatchFam      = matchFamDiv
+    , sfInteractTop   = interactTopDiv
+    , sfInteractInert = interactInertDiv
+    }
+  where
+  name = mkWiredInTyConName UserSyntax gHC_TYPENATS (fsLit "Div")
+            typeNatDivTyFamNameKey typeNatDivTyCon
+
+typeNatModTyCon :: TyCon
+typeNatModTyCon = mkTypeNatFunTyCon2 name
+  BuiltInSynFamily
+    { sfMatchFam      = matchFamMod
+    , sfInteractTop   = interactTopMod
+    , sfInteractInert = interactInertMod
+    }
+  where
+  name = mkWiredInTyConName UserSyntax gHC_TYPENATS (fsLit "Mod")
+            typeNatModTyFamNameKey typeNatModTyCon
+
+
+
+
+
 typeNatExpTyCon :: TyCon
 typeNatExpTyCon = mkTypeNatFunTyCon2 name
   BuiltInSynFamily
@@ -106,6 +139,19 @@ typeNatExpTyCon = mkTypeNatFunTyCon2 name
   name = mkWiredInTyConName UserSyntax gHC_TYPENATS (fsLit "^")
                 typeNatExpTyFamNameKey typeNatExpTyCon
 
+typeNatLogTyCon :: TyCon
+typeNatLogTyCon = mkTypeNatFunTyCon1 name
+  BuiltInSynFamily
+    { sfMatchFam      = matchFamLog
+    , sfInteractTop   = interactTopLog
+    , sfInteractInert = interactInertLog
+    }
+  where
+  name = mkWiredInTyConName UserSyntax gHC_TYPENATS (fsLit "Log2")
+            typeNatLogTyFamNameKey typeNatLogTyCon
+
+
+
 typeNatLeqTyCon :: TyCon
 typeNatLeqTyCon =
   mkFamilyTyCon name
@@ -176,6 +222,17 @@ typeSymbolAppendTyCon = mkTypeSymbolFunTyCon2 name
 
 
 
+-- Make a unary built-in constructor of kind: Nat -> Nat
+mkTypeNatFunTyCon1 :: Name -> BuiltInSynFamily -> TyCon
+mkTypeNatFunTyCon1 op tcb =
+  mkFamilyTyCon op
+    (mkTemplateAnonTyConBinders [ typeNatKind ])
+    typeNatKind
+    Nothing
+    (BuiltInSynFamTyCon tcb)
+    Nothing
+    NotInjective
+
 
 -- Make a binary built-in constructor of kind: Nat -> Nat -> Nat
 mkTypeNatFunTyCon2 :: Name -> BuiltInSynFamily -> TyCon
@@ -230,6 +287,11 @@ axAddDef
   , axSub0R
   , axAppendSymbol0R
   , axAppendSymbol0L
+  , axDivDef
+  , axDiv1
+  , axModDef
+  , axMod1
+  , axLogDef
   :: CoAxiomRule
 
 axAddDef = mkBinAxiom "AddDef" typeNatAddTyCon $
@@ -274,6 +336,18 @@ axAppendSymbolDef = CoAxiomRule
 axSubDef = mkBinAxiom "SubDef" typeNatSubTyCon $
               \x y -> fmap num (minus x y)
 
+axDivDef = mkBinAxiom "DivDef" typeNatDivTyCon $
+              \x y -> do guard (y /= 0)
+                         return (num (div x y))
+
+axModDef = mkBinAxiom "ModDef" typeNatModTyCon $
+              \x y -> do guard (y /= 0)
+                         return (num (mod x y))
+
+axLogDef = mkUnAxiom "LogDef" typeNatLogTyCon $
+              \x -> do (a,_) <- genLog x 2
+                       return (num a)
+
 axAdd0L     = mkAxiom1 "Add0L"    $ \(Pair s t) -> (num 0 .+. s) === t
 axAdd0R     = mkAxiom1 "Add0R"    $ \(Pair s t) -> (s .+. num 0) === t
 axSub0R     = mkAxiom1 "Sub0R"    $ \(Pair s t) -> (s .-. num 0) === t
@@ -281,6 +355,9 @@ axMul0L     = mkAxiom1 "Mul0L"    $ \(Pair s _) -> (num 0 .*. s) === num 0
 axMul0R     = mkAxiom1 "Mul0R"    $ \(Pair s _) -> (s .*. num 0) === num 0
 axMul1L     = mkAxiom1 "Mul1L"    $ \(Pair s t) -> (num 1 .*. s) === t
 axMul1R     = mkAxiom1 "Mul1R"    $ \(Pair s t) -> (s .*. num 1) === t
+axDiv1      = mkAxiom1 "Div1"     $ \(Pair s t) -> (tDiv s (num 1) === t)
+axMod1      = mkAxiom1 "Mod1"     $ \(Pair s _) -> (tMod s (num 1) === num 0)
+                                    -- XXX: Shouldn't we check that _ is 0?
 axExp1L     = mkAxiom1 "Exp1L"    $ \(Pair s _) -> (num 1 .^. s) === num 1
 axExp0R     = mkAxiom1 "Exp0R"    $ \(Pair s _) -> (s .^. num 0) === num 1
 axExp1R     = mkAxiom1 "Exp1R"    $ \(Pair s t) -> (s .^. num 1) === t
@@ -320,6 +397,11 @@ typeNatCoAxiomRules = Map.fromList $ map (\x -> (coaxrName x, x))
   , axSubDef
   , axAppendSymbol0R
   , axAppendSymbol0L
+  , axDivDef
+  , axDiv1
+  , axModDef
+  , axMod1
+  , axLogDef
   ]
 
 
@@ -337,6 +419,12 @@ s .-. t = mkTyConApp typeNatSubTyCon [s,t]
 (.*.) :: Type -> Type -> Type
 s .*. t = mkTyConApp typeNatMulTyCon [s,t]
 
+tDiv :: Type -> Type -> Type
+tDiv s t = mkTyConApp typeNatDivTyCon [s,t]
+
+tMod :: Type -> Type -> Type
+tMod s t = mkTyConApp typeNatModTyCon [s,t]
+
 (.^.) :: Type -> Type -> Type
 s .^. t = mkTyConApp typeNatExpTyCon [s,t]
 
@@ -395,6 +483,19 @@ known p x = case isNumLitTy x of
               Nothing -> False
 
 
+mkUnAxiom :: String -> TyCon -> (Integer -> Maybe Type) -> CoAxiomRule
+mkUnAxiom str tc f =
+  CoAxiomRule
+    { coaxrName      = fsLit str
+    , coaxrAsmpRoles = [Nominal]
+    , coaxrRole      = Nominal
+    , coaxrProves    = \cs ->
+        do [Pair s1 s2] <- return cs
+           s2' <- isNumLitTy s2
+           z   <- f s2'
+           return (mkTyConApp tc [s1] === z)
+    }
+
 
 
 -- For the definitional axioms
@@ -461,6 +562,24 @@ matchFamMul [s,t]
         mbY = isNumLitTy t
 matchFamMul _ = Nothing
 
+matchFamDiv :: [Type] -> Maybe (CoAxiomRule, [Type], Type)
+matchFamDiv [s,t]
+  | Just 1 <- mbY = Just (axDiv1, [s], s)
+  | Just x <- mbX, Just y <- mbY, y /= 0 = Just (axDivDef, [s,t], num (div x y))
+  where mbX = isNumLitTy s
+        mbY = isNumLitTy t
+matchFamDiv _ = Nothing
+
+matchFamMod :: [Type] -> Maybe (CoAxiomRule, [Type], Type)
+matchFamMod [s,t]
+  | Just 1 <- mbY = Just (axMod1, [s], num 0)
+  | Just x <- mbX, Just y <- mbY, y /= 0 = Just (axModDef, [s,t], num (mod x y))
+  where mbX = isNumLitTy s
+        mbY = isNumLitTy t
+matchFamMod _ = Nothing
+
+
+
 matchFamExp :: [Type] -> Maybe (CoAxiomRule, [Type], Type)
 matchFamExp [s,t]
   | Just 0 <- mbY = Just (axExp0R, [s], num 1)
@@ -472,6 +591,13 @@ matchFamExp [s,t]
         mbY = isNumLitTy t
 matchFamExp _ = Nothing
 
+matchFamLog :: [Type] -> Maybe (CoAxiomRule, [Type], Type)
+matchFamLog [s]
+  | Just x <- mbX, Just (n,_) <- genLog x 2 = Just (axLogDef, [s], num n)
+  where mbX = isNumLitTy s
+matchFamLog _ = Nothing
+
+
 matchFamLeq :: [Type] -> Maybe (CoAxiomRule, [Type], Type)
 matchFamLeq [s,t]
   | Just 0 <- mbX = Just (axLeq0L, [t], bool True)
@@ -579,6 +705,12 @@ interactTopMul [s,t] r
   mbZ = isNumLitTy r
 interactTopMul _ _ = []
 
+interactTopDiv :: [Xi] -> Xi -> [Pair Type]
+interactTopDiv _ _ = []   -- I can't think of anything...
+
+interactTopMod :: [Xi] -> Xi -> [Pair Type]
+interactTopMod _ _ = []   -- I can't think of anything...
+
 interactTopExp :: [Xi] -> Xi -> [Pair Type]
 interactTopExp [s,t] r
   | Just 0 <- mbZ = [ s === num 0 ]                                       -- (s ^ t ~ 0) => (s ~ 0)
@@ -590,6 +722,11 @@ interactTopExp [s,t] r
   mbZ = isNumLitTy r
 interactTopExp _ _ = []
 
+interactTopLog :: [Xi] -> Xi -> [Pair Type]
+interactTopLog _ _ = []   -- I can't think of anything...
+
+
+
 interactTopLeq :: [Xi] -> Xi -> [Pair Type]
 interactTopLeq [s,t] r
   | Just 0 <- mbY, Just True <- mbZ = [ s === num 0 ]                     -- (s <= 0) => (s ~ 0)
@@ -655,6 +792,12 @@ interactInertMul [x1,y1] z1 [x2,y2] z2
 
 interactInertMul _ _ _ _ = []
 
+interactInertDiv :: [Xi] -> Xi -> [Xi] -> Xi -> [Pair Type]
+interactInertDiv _ _ _ _ = []
+
+interactInertMod :: [Xi] -> Xi -> [Xi] -> Xi -> [Pair Type]
+interactInertMod _ _ _ _ = []
+
 interactInertExp :: [Xi] -> Xi -> [Xi] -> Xi -> [Pair Type]
 interactInertExp [x1,y1] z1 [x2,y2] z2
   | sameZ && known (> 1) x1 && tcEqType x1 x2 = [ y1 === y2 ]
@@ -663,6 +806,9 @@ interactInertExp [x1,y1] z1 [x2,y2] z2
 
 interactInertExp _ _ _ _ = []
 
+interactInertLog :: [Xi] -> Xi -> [Xi] -> Xi -> [Pair Type]
+interactInertLog _ _ _ _ = []
+
 
 interactInertLeq :: [Xi] -> Xi -> [Xi] -> Xi -> [Pair Type]
 interactInertLeq [x1,y1] z1 [x2,y2] z2
diff --git a/libraries/base/GHC/TypeLits.hs b/libraries/base/GHC/TypeLits.hs
index 0ea866a5bb..ba0ce974b2 100644
--- a/libraries/base/GHC/TypeLits.hs
+++ b/libraries/base/GHC/TypeLits.hs
@@ -34,6 +34,7 @@ module GHC.TypeLits
 
     -- * Functions on type literals
   , type (N.<=), type (N.<=?), type (N.+), type (N.*), type (N.^), type (N.-)
+  , type N.Div, type N.Mod, type N.Log2
   , AppendSymbol
   , N.CmpNat, CmpSymbol
 
diff --git a/libraries/base/GHC/TypeNats.hs b/libraries/base/GHC/TypeNats.hs
index 7d9ca0d8e4..e3322a2a25 100644
--- a/libraries/base/GHC/TypeNats.hs
+++ b/libraries/base/GHC/TypeNats.hs
@@ -33,6 +33,7 @@ module GHC.TypeNats
     -- * Functions on type literals
   , type (<=), type (<=?), type (+), type (*), type (^), type (-)
   , CmpNat
+  , Div, Mod, Log2
 
   ) where
 
@@ -129,6 +130,18 @@ type family (m :: Nat) ^ (n :: Nat) :: Nat
 -- @since 4.7.0.0
 type family (m :: Nat) - (n :: Nat) :: Nat
 
+-- | Division (round down) of natural numbers.
+-- @Div x 0@ is undefined (i.e., it cannot be reduced).
+type family Div (m :: Nat) (n :: Nat) :: Nat
+
+-- | Modulus of natural numbers.
+-- @Mod x 0@ is undefined (i.e., it cannot be reduced).
+type family Mod (m :: Nat) (n :: Nat) :: Nat
+
+-- | Log base 2 (round down) of natural numbers.
+-- @Log 0@ is undefined (i.e., it cannot be reduced).
+type family Log2 (m :: Nat) :: Nat
+
 --------------------------------------------------------------------------------
 
 -- | We either get evidence that this function was instantiated with the
diff --git a/libraries/base/changelog.md b/libraries/base/changelog.md
index ed8ab13580..7778cebc14 100644
--- a/libraries/base/changelog.md
+++ b/libraries/base/changelog.md
@@ -3,6 +3,9 @@
 ## 4.11.0.0 *TBA*
   * Bundled with GHC 8.4.1
 
+  * Add `Div`, `Mod`, and `Log2` functions on type-level naturals 
+    in `GHC.TypeLits`.
+
   * Add `Alternative` instance for `ZipList` (#13520)
 
   * Add instances `Num`, `Functor`, `Applicative`, `Monad`, `Semigroup`
diff --git a/testsuite/tests/ghci/scripts/T9181.stdout b/testsuite/tests/ghci/scripts/T9181.stdout
index 1ae9bd54d0..5e520fb0a2 100644
--- a/testsuite/tests/ghci/scripts/T9181.stdout
+++ b/testsuite/tests/ghci/scripts/T9181.stdout
@@ -53,9 +53,16 @@ type family (GHC.TypeNats.<=?) (a :: GHC.Types.Nat)
 type family GHC.TypeNats.CmpNat (a :: GHC.Types.Nat)
                                 (b :: GHC.Types.Nat)
             :: Ordering
+type family GHC.TypeNats.Div (a :: GHC.Types.Nat)
+                             (b :: GHC.Types.Nat)
+            :: GHC.Types.Nat
 class GHC.TypeNats.KnownNat (n :: GHC.Types.Nat) where
   GHC.TypeNats.natSing :: GHC.TypeNats.SNat n
   {-# MINIMAL natSing #-}
+type family GHC.TypeNats.Log2 (a :: GHC.Types.Nat) :: GHC.Types.Nat
+type family GHC.TypeNats.Mod (a :: GHC.Types.Nat)
+                             (b :: GHC.Types.Nat)
+            :: GHC.Types.Nat
 data GHC.Types.Nat
 data GHC.TypeNats.SomeNat
   = forall (n :: GHC.Types.Nat).
diff --git a/testsuite/tests/typecheck/should_compile/TcTypeNatSimple.hs b/testsuite/tests/typecheck/should_compile/TcTypeNatSimple.hs
index c692c3f725..911a43e507 100644
--- a/testsuite/tests/typecheck/should_compile/TcTypeNatSimple.hs
+++ b/testsuite/tests/typecheck/should_compile/TcTypeNatSimple.hs
@@ -59,6 +59,21 @@ e17 = id
 e18 :: Proxy (a - 0) -> Proxy a
 e18 = id
 
+e19 :: Proxy (Div 10 3) -> Proxy 3
+e19 = id
+
+e20 :: Proxy (Div x 1) -> Proxy x
+e20 = id
+
+e21 :: Proxy (Mod 10 3) -> Proxy 1
+e21 = id
+
+e22 :: Proxy (Mod x 1) -> Proxy 0
+e22 = id
+
+e23 :: Proxy (Log2 10) -> Proxy 3
+e23 = id
+
 --------------------------------------------------------------------------------
 -- Test interactions with inerts