Built-in type families: CharToNat, NatToChar (#19535)

Co-authored-by: Daniel Rogozin <daniel.rogozin@serokell.io>
Co-authored-by: Rinat Stryungis <rinat.stryungis@serokell.io>

2 files changed, 123 insertions, 31 deletions

diff --git a/compiler/GHC/Builtin/Names.hs b/compiler/GHC/Builtin/Names.hs
index 4779ca8de2..ec8d5a0afc 100644
--- a/compiler/GHC/Builtin/Names.hs
+++ b/compiler/GHC/Builtin/Names.hs
@@ -2000,41 +2000,10 @@ uFloatTyConKey  = mkPreludeTyConUnique 161
 uIntTyConKey    = mkPreludeTyConUnique 162
 uWordTyConKey   = mkPreludeTyConUnique 163
 
--- Type-level naturals
-typeSymbolKindConNameKey, typeCharKindConNameKey,
-  typeNatAddTyFamNameKey, typeNatMulTyFamNameKey, typeNatExpTyFamNameKey,
-  typeNatSubTyFamNameKey
-  , typeSymbolCmpTyFamNameKey, typeNatCmpTyFamNameKey, typeCharCmpTyFamNameKey
-  , typeLeqCharTyFamNameKey
-  , typeNatDivTyFamNameKey
-  , typeNatModTyFamNameKey
-  , typeNatLogTyFamNameKey
-  , typeConsSymbolTyFamNameKey, typeUnconsSymbolTyFamNameKey
-  :: Unique
-typeSymbolKindConNameKey  = mkPreludeTyConUnique 165
-typeCharKindConNameKey    = mkPreludeTyConUnique 166
-typeNatAddTyFamNameKey    = mkPreludeTyConUnique 167
-typeNatMulTyFamNameKey    = mkPreludeTyConUnique 168
-typeNatExpTyFamNameKey    = mkPreludeTyConUnique 169
-typeNatSubTyFamNameKey    = mkPreludeTyConUnique 171
-typeSymbolCmpTyFamNameKey = mkPreludeTyConUnique 172
-typeNatCmpTyFamNameKey    = mkPreludeTyConUnique 173
-typeCharCmpTyFamNameKey   = mkPreludeTyConUnique 174
-typeLeqCharTyFamNameKey   = mkPreludeTyConUnique 175
-typeNatDivTyFamNameKey  = mkPreludeTyConUnique 176
-typeNatModTyFamNameKey  = mkPreludeTyConUnique 177
-typeNatLogTyFamNameKey  = mkPreludeTyConUnique 178
-typeConsSymbolTyFamNameKey = mkPreludeTyConUnique 179
-typeUnconsSymbolTyFamNameKey = mkPreludeTyConUnique 180
-
 -- Custom user type-errors
 errorMessageTypeErrorFamKey :: Unique
 errorMessageTypeErrorFamKey =  mkPreludeTyConUnique 181
 
-
-
-ntTyConKey:: Unique
-ntTyConKey = mkPreludeTyConUnique 182
 coercibleTyConKey :: Unique
 coercibleTyConKey = mkPreludeTyConUnique 183
 
@@ -2093,6 +2062,38 @@ multMulTyConKey = mkPreludeTyConUnique 199
 
 #include "primop-vector-uniques.hs-incl"
 
+------------- Type-level Symbol, Nat, Char ----------
+--      USES TyConUniques 400-499
+-----------------------------------------------------
+typeSymbolKindConNameKey, typeCharKindConNameKey,
+  typeNatAddTyFamNameKey, typeNatMulTyFamNameKey, typeNatExpTyFamNameKey,
+  typeNatSubTyFamNameKey
+  , typeSymbolCmpTyFamNameKey, typeNatCmpTyFamNameKey, typeCharCmpTyFamNameKey
+  , typeLeqCharTyFamNameKey
+  , typeNatDivTyFamNameKey
+  , typeNatModTyFamNameKey
+  , typeNatLogTyFamNameKey
+  , typeConsSymbolTyFamNameKey, typeUnconsSymbolTyFamNameKey
+  , typeCharToNatTyFamNameKey, typeNatToCharTyFamNameKey
+  :: Unique
+typeSymbolKindConNameKey  = mkPreludeTyConUnique 400
+typeCharKindConNameKey    = mkPreludeTyConUnique 401
+typeNatAddTyFamNameKey    = mkPreludeTyConUnique 402
+typeNatMulTyFamNameKey    = mkPreludeTyConUnique 403
+typeNatExpTyFamNameKey    = mkPreludeTyConUnique 404
+typeNatSubTyFamNameKey    = mkPreludeTyConUnique 405
+typeSymbolCmpTyFamNameKey = mkPreludeTyConUnique 406
+typeNatCmpTyFamNameKey    = mkPreludeTyConUnique 407
+typeCharCmpTyFamNameKey   = mkPreludeTyConUnique 408
+typeLeqCharTyFamNameKey   = mkPreludeTyConUnique 409
+typeNatDivTyFamNameKey  = mkPreludeTyConUnique 410
+typeNatModTyFamNameKey  = mkPreludeTyConUnique 411
+typeNatLogTyFamNameKey  = mkPreludeTyConUnique 412
+typeConsSymbolTyFamNameKey = mkPreludeTyConUnique 413
+typeUnconsSymbolTyFamNameKey = mkPreludeTyConUnique 414
+typeCharToNatTyFamNameKey = mkPreludeTyConUnique 415
+typeNatToCharTyFamNameKey = mkPreludeTyConUnique 416
+
 {-
 ************************************************************************
 *                                                                      *
diff --git a/compiler/GHC/Builtin/Types/Literals.hs b/compiler/GHC/Builtin/Types/Literals.hs
index 6b21a68bd3..19dcba7116 100644
--- a/compiler/GHC/Builtin/Types/Literals.hs
+++ b/compiler/GHC/Builtin/Types/Literals.hs
@@ -21,6 +21,8 @@ module GHC.Builtin.Types.Literals
   , typeCharCmpTyCon
   , typeConsSymbolTyCon
   , typeUnconsSymbolTyCon
+  , typeCharToNatTyCon
+  , typeNatToCharTyCon
   ) where
 
 import GHC.Prelude
@@ -55,10 +57,13 @@ import GHC.Builtin.Names
                   , typeCharCmpTyFamNameKey
                   , typeConsSymbolTyFamNameKey
                   , typeUnconsSymbolTyFamNameKey
+                  , typeCharToNatTyFamNameKey
+                  , typeNatToCharTyFamNameKey
                   )
 import GHC.Data.FastString
 import Control.Monad ( guard )
 import Data.List  ( isPrefixOf, isSuffixOf )
+import qualified Data.Char as Char
 
 {-
 Note [Type-level literals]
@@ -155,6 +160,8 @@ typeNatTyCons =
   , typeCharCmpTyCon
   , typeConsSymbolTyCon
   , typeUnconsSymbolTyCon
+  , typeCharToNatTyCon
+  , typeNatToCharTyCon
   ]
 
 typeNatAddTyCon :: TyCon
@@ -321,6 +328,43 @@ typeUnconsSymbolTyCon =
       , sfInteractInert = interactInertUnconsSymbol
       }
 
+typeCharToNatTyCon :: TyCon
+typeCharToNatTyCon =
+  mkFamilyTyCon name
+    (mkTemplateAnonTyConBinders [ charTy ])
+    naturalTy
+    Nothing
+    (BuiltInSynFamTyCon ops)
+    Nothing
+    (Injective [True])
+  where
+  name = mkWiredInTyConName UserSyntax gHC_TYPELITS (fsLit "CharToNat")
+                  typeCharToNatTyFamNameKey typeCharToNatTyCon
+  ops = BuiltInSynFamily
+      { sfMatchFam      = matchFamCharToNat
+      , sfInteractTop   = interactTopCharToNat
+      , sfInteractInert = \_ _ _ _ -> []
+      }
+
+
+typeNatToCharTyCon :: TyCon
+typeNatToCharTyCon =
+  mkFamilyTyCon name
+    (mkTemplateAnonTyConBinders [ naturalTy ])
+    charTy
+    Nothing
+    (BuiltInSynFamTyCon ops)
+    Nothing
+    (Injective [True])
+  where
+  name = mkWiredInTyConName UserSyntax gHC_TYPELITS (fsLit "NatToChar")
+                  typeNatToCharTyFamNameKey typeNatToCharTyCon
+  ops = BuiltInSynFamily
+      { sfMatchFam      = matchFamNatToChar
+      , sfInteractTop   = interactTopNatToChar
+      , sfInteractInert = \_ _ _ _ -> []
+      }
+
 -- Make a unary built-in constructor of kind: Nat -> Nat
 mkTypeNatFunTyCon1 :: Name -> BuiltInSynFamily -> TyCon
 mkTypeNatFunTyCon1 op tcb =
@@ -369,6 +413,8 @@ axAddDef
   , axAppendSymbolDef
   , axConsSymbolDef
   , axUnconsSymbolDef
+  , axCharToNatDef
+  , axNatToCharDef
   , axAdd0L
   , axAdd0R
   , axMul0L
@@ -436,6 +482,14 @@ axUnconsSymbolDef =
     \str -> Just $
       mkPromotedMaybeTy charSymbolPairKind (fmap reifyCharSymbolPairTy (unconsFS str))
 
+axCharToNatDef =
+  mkUnAxiom "CharToNatDef" typeCharToNatTyCon isCharLitTy $
+    \c -> Just $ num (charToInteger c)
+
+axNatToCharDef =
+  mkUnAxiom "NatToCharDef" typeNatToCharTyCon isNumLitTy $
+    \n -> fmap mkCharLitTy (integerToChar n)
+
 axSubDef = mkBinAxiom "SubDef" typeNatSubTyCon isNumLitTy isNumLitTy $
               \x y -> fmap num (minus x y)
 
@@ -740,6 +794,29 @@ matchFamUnconsSymbol [s]
   mbX = isStrLitTy s
 matchFamUnconsSymbol _ = Nothing
 
+matchFamCharToNat :: [Type] -> Maybe (CoAxiomRule, [Type], Type)
+matchFamCharToNat [c]
+  | Just c' <- isCharLitTy c, n <- charToInteger c'
+  = Just (axCharToNatDef, [c], mkNumLitTy n)
+  | otherwise = Nothing
+matchFamCharToNat _ = Nothing
+
+matchFamNatToChar :: [Type] -> Maybe (CoAxiomRule, [Type], Type)
+matchFamNatToChar [n]
+  | Just n' <- isNumLitTy n, Just c <- integerToChar n'
+  = Just (axNatToCharDef, [n], mkCharLitTy c)
+  | otherwise = Nothing
+matchFamNatToChar _ = Nothing
+
+charToInteger :: Char -> Integer
+charToInteger c = fromIntegral (Char.ord c)
+
+integerToChar :: Integer -> Maybe Char
+integerToChar n | inBounds = Just (Char.chr (fromInteger n))
+  where inBounds = n >= charToInteger minBound &&
+                   n <= charToInteger maxBound
+integerToChar _ = Nothing
+
 {-------------------------------------------------------------------------------
 Interact with axioms
 -------------------------------------------------------------------------------}
@@ -887,6 +964,20 @@ interactTopUnconsSymbol [s] r
 
 interactTopUnconsSymbol _ _ = []
 
+interactTopCharToNat :: [Xi] -> Xi -> [Pair Type]
+interactTopCharToNat [s] r
+  -- (CharToNat c ~ 122) => (c ~ 'z')
+  | Just n <- isNumLitTy r
+  , Just c <- integerToChar n
+  = [ s === mkCharLitTy c ]
+interactTopCharToNat _ _ = []
+
+interactTopNatToChar :: [Xi] -> Xi -> [Pair Type]
+interactTopNatToChar [s] r
+  -- (NatToChar n ~ 'z') => (n ~ 122)
+  | Just c <- isCharLitTy r
+  = [ s === mkNumLitTy (charToInteger c) ]
+interactTopNatToChar _ _ = []
 
 {-------------------------------------------------------------------------------
 Interaction with inerts