SchemeVar newtype

src/Expander.hs

@@ -405,7 +405,7 @@
       [ ( "open-syntax"
         , Scheme [] $ tFun [tSyntax] (Prims.primitiveDatatype "Syntax-Contents" [tSyntax])
         , ValueClosure $ HO $
           \(ValueSyntax stx) ->
             case syntaxE stx of
               Id name ->
                 primitiveCtor "identifier-contents" [ValueString name]
@@ -539,15 +539,17 @@
       | (name, fun) <- [("<", (<)), ("<=", (<=)), (">", (>)), (">=", (>=)), ("=", (==)), ("/=", (/=))]
       ] ++
       [ ("pure-IO"
-        , Scheme [KStar, KStar] $ tFun [tSchemaVar 0 []] (tIO (tSchemaVar 0 []))
+        , let a = tSchemaVar firstSchemeVar []
+          in Scheme [KStar] $ tFun [a] (tIO a)
         , ValueClosure $ HO $ \v -> ValueIOAction (pure v)
         )
       , ("bind-IO"
-        , Scheme [KStar, KStar] $
-          tFun [ tIO (tSchemaVar 0 [])
-               , tFun [tSchemaVar 0 []] (tIO (tSchemaVar 1 []))
-               ]
-               (tIO (tSchemaVar 1 []))
+        , let a:b:_ = [tSchemaVar i [] | i <- [firstSchemeVar..]]
+          in Scheme [KStar, KStar] $
+             tFun [ tIO a
+                  , tFun [a] (tIO b)
+                  ]
+                  (tIO b)
         , ValueClosure $ HO $ \(ValueIOAction mx) -> do
             ValueClosure $ HO $ \(ValueClosure f) -> do
               ValueIOAction $ do
@@ -587,22 +589,25 @@
       , ("Unit", [], [("unit", [])])
       , ("Bool", [], [("true", []), ("false", [])])
       , ("Problem", [], [("module", []), ("declaration", []), ("type", []), ("expression", [tType]), ("pattern", []), ("type-pattern", [])])
-      , ("Maybe", [KStar], [("nothing", []), ("just", [tSchemaVar 0 []])])
-      , ("List"
-        , [KStar]
-        , [ ("nil", [])
-          , ("::", [tSchemaVar 0 [], Prims.primitiveDatatype "List" [tSchemaVar 0 []]])
-          ]
-        )
+      , let a = tSchemaVar firstSchemeVar []
+        in ("Maybe", [KStar], [("nothing", []), ("just", [a])])
+      , let a = tSchemaVar firstSchemeVar []
+        in ("List"
+           , [KStar]
+           , [ ("nil", [])
+             , ("::", [a, Prims.primitiveDatatype "List" [a]])
+             ]
+           )
       , ("Syntax-Contents"
         , [KStar]
-        , [ ("list-contents", [Prims.primitiveDatatype "List" [tSchemaVar 0 []]])
-          , ("integer-contents", [tInteger])
-          , ("string-contents", [tString])
-          , ("identifier-contents", [tString])
-          -- if you add a constructor here, remember to also add a
-          -- corresponding pattern in close-syntax!
-          ]
+        , let a = tSchemaVar firstSchemeVar []
+          in [ ("list-contents", [Prims.primitiveDatatype "List" [a]])
+             , ("integer-contents", [tInteger])
+             , ("string-contents", [tString])
+             , ("identifier-contents", [tString])
+             -- if you add a constructor here, remember to also add a
+             -- corresponding pattern in close-syntax!
+             ]
         )
       ]
 
@@ -1220,8 +1225,9 @@
           case prob of
             ExprDest t dest -> do
               argTys <- traverse makeTypeMeta argKinds
+              let tyStore = S.fromList $ zip [firstSchemeVar..] argTys
               unify dest t $ tDatatype dt argTys
-              args' <- for args \a -> inst dest (Scheme argKinds a) argTys
+              args' <- for args \a -> inst dest (Scheme argKinds a) tyStore
               Stx _ _ (foundName, foundArgs) <- mustBeCons stx
               _ <- mustBeIdent foundName
               argDests <-
@@ -1234,8 +1240,9 @@
             PatternDest patTy dest -> do
               Stx _ loc (_cname, subPats) <- mustBeCons stx
               tyArgs <- traverse makeTypeMeta argKinds
+              let tyStore = S.fromList $ zip [firstSchemeVar..] tyArgs
               argTypes <- for args \ a ->
-                            inst loc (Scheme argKinds a) tyArgs
+                            inst loc (Scheme argKinds a) tyStore
               unify loc (tDatatype dt tyArgs) patTy
               if length subPats /= length argTypes
                 then throwError $ WrongArgCount stx ctor (length argTypes) (length subPats)

src/Expander/Monad.hs

@@ -241,7 +241,7 @@ data EValue
   | EPrimPatternMacro (Either (Ty, PatternPtr) TypePatternPtr -> Syntax -> Expand ())
   | EPrimUniversalMacro (MacroDest -> Syntax -> Expand ())
   | EVarMacro !Var -- ^ For bound variables (the Unique is the binding site of the var)
-  | ETypeVar !Kind !Natural
+  | ETypeVar !Kind !SchemeVar
   -- ^ For bound type variables (user-written Skolem variables or in datatype definitions)
   | EUserMacro !MacroVar -- ^ For user-written macros
   | EIncompleteMacro !MacroVar !Ident !SplitCorePtr -- ^ Macros that are themselves not yet ready to go

src/Expander/Primitives.hs

@@ -68,7 +68,6 @@
 import Data.Text (Text)
 import qualified Data.Text as T
 import Data.Traversable
-import Numeric.Natural
 
 import Binding
 import Core
@@ -130,7 +129,7 @@
   Stx scs loc (_ :: Syntax, more) <- mustBeCons stx
   Stx _ _ (nameAndArgs, ctorSpecs) <- mustBeCons (Syntax (Stx scs loc (List more)))
   Stx _ _ (name, args) <- mustBeCons nameAndArgs
-  typeArgs <- for (zip [0..] args) $ \(i, a) ->
+  typeArgs <- for (zip [firstSchemeVar..] args) $ \(i, a) ->
     prepareTypeVar i a
   sc <- freshScope $ T.pack $ "For datatype at " ++ shortShow (stxLoc stx)
   let typeScopes = map (view _1) typeArgs ++ [sc]
@@ -712,7 +711,7 @@
   forkExpanderTask $ AwaitingTypePattern dest exprTy rhsDest rhsStx
   return (dest, rhsDest)
 
-prepareTypeVar :: Natural -> Syntax -> Expand (Scope, Ident, Kind)
+prepareTypeVar :: SchemeVar -> Syntax -> Expand (Scope, Ident, Kind)
 prepareTypeVar i varStx = do
   (sc, α, b) <- varPrepHelper varStx
   k <- KMetaVar <$> liftIO newKindVar
@@ -748,23 +747,23 @@
 unaryIntegerPrim :: (Integer -> Integer) -> Value
 unaryIntegerPrim f =
   ValueClosure $ HO $
   \(ValueInteger i) ->
     ValueInteger (f i)
 
 binaryIntegerPrim :: (Integer -> Integer -> Integer) -> Value
 binaryIntegerPrim f =
   ValueClosure $ HO $
   \(ValueInteger i1) ->
     ValueClosure $ HO $
     \(ValueInteger i2) ->
       ValueInteger (f i1 i2)
 
 binaryIntegerPred :: (Integer -> Integer -> Bool) -> Value
 binaryIntegerPred f =
   ValueClosure $ HO $
   \(ValueInteger i1) ->
     ValueClosure $ HO $
     \(ValueInteger i2) ->
       if f i1 i2
         then primitiveCtor "true" []
         else primitiveCtor "false" []
@@ -773,9 +772,9 @@
 binaryStringPred :: (Text -> Text -> Bool) -> Value
 binaryStringPred f =
   ValueClosure $ HO $
   \(ValueString str1) ->
     ValueClosure $ HO $
     \(ValueString str2) ->
       if f str1 str2
         then primitiveCtor "true" []
         else primitiveCtor "false" []

src/Expander/TC.hs

@@ -14,6 +14,7 @@ import Control.Monad
 import Control.Monad.Except
 import Control.Monad.State
 import Data.Foldable
+import Data.Traversable (for)
 import Numeric.Natural
 
 import Expander.Monad
@@ -107,12 +108,14 @@ freshMeta kind = do
   return ptr
 
 
-inst :: UnificationErrorBlame blame => blame -> Scheme Ty -> [Ty] -> Expand Ty
+inst :: UnificationErrorBlame blame => blame -> Scheme Ty -> Store SchemeVar Ty -> Expand Ty
 inst blame (Scheme argKinds ty) ts
   | length ts /= length argKinds =
     throwError $ InternalError "Mismatch in number of type vars"
   | otherwise = do
-      traverse_ (uncurry $ checkKind blame) (zip argKinds ts)
+      let tys :: [Ty]
+          tys = fmap snd $ St.toAscList ts
+      traverse_ (uncurry $ checkKind blame) (zip argKinds tys)
       instTy ty
   where
     instTy :: Ty -> Expand Ty
@@ -131,13 +134,17 @@ inst blame (Scheme argKinds ty) ts
       pure $ TyF ctor' (tArgsPrefix ++ tArgsSuffix)
 
     instCtor :: TypeConstructor -> TyF Ty
-    instCtor (TSchemaVar i) = unTy $ ts !! fromIntegral i
+    instCtor (TSchemaVar v) = unTy $ ts St.! v
     instCtor ctor           = TyF ctor []
 
 
 specialize :: UnificationErrorBlame blame => blame -> Scheme Ty -> Expand Ty
 specialize blame sch@(Scheme argKinds _) = do
-  freshVars <- traverse makeTypeMeta argKinds
+  pairs <- for (zip [firstSchemeVar..] argKinds) $ \(v, k) -> do
+    meta <- makeTypeMeta k
+    pure (v, meta)
+  let freshVars :: Store SchemeVar Ty
+      freshVars = St.fromList pairs
   inst blame sch freshVars
 
 varType :: Var -> Expand (Maybe (Scheme Ty))
@@ -162,35 +169,35 @@ notFreeInCtx var = do
 generalizeType :: Ty -> Expand (Scheme Ty)
 generalizeType ty = do
   canGeneralize <- filterM notFreeInCtx =<< metas ty
-  (body, (_, _, argKinds)) <- flip runStateT (0, mempty, []) $ do
+  (body, (_, _, argKinds)) <- flip runStateT (firstSchemeVar, mempty, []) $ do
     genTyVars canGeneralize ty
   pure $ Scheme argKinds body
 
   where
     genTyVars ::
       [MetaPtr] -> Ty ->
-      StateT (Natural, Store MetaPtr Natural, [Kind]) Expand Ty
+      StateT (SchemeVar, Store MetaPtr SchemeVar, [Kind]) Expand Ty
     genTyVars vars t = do
       (Ty t') <- lift $ normType t
       Ty <$> genVarsTyF vars t'
 
     genVarsTyF ::
       [MetaPtr] -> TyF Ty ->
-      StateT (Natural, Store MetaPtr Natural, [Kind]) Expand (TyF Ty)
+      StateT (SchemeVar, Store MetaPtr SchemeVar, [Kind]) Expand (TyF Ty)
     genVarsTyF vars (TyF ctor args) =
       TyF <$> genVarsCtor vars ctor
           <*> traverse (genTyVars vars) args
 
     genVarsCtor ::
       [MetaPtr] -> TypeConstructor ->
-      StateT (Natural, Store MetaPtr Natural, [Kind]) Expand TypeConstructor
+      StateT (SchemeVar, Store MetaPtr SchemeVar, [Kind]) Expand TypeConstructor
     genVarsCtor vars (TMetaVar v)
       | v `elem` vars = do
           (i, indices, argKinds) <- get
           case St.lookup v indices of
             Nothing -> do
               k <- lift $ typeVarKind v
-              put (i + 1, St.insert v i indices, argKinds ++ [k])
+              put (succ i, St.insert v i indices, argKinds ++ [k])
               pure $ TSchemaVar i
             Just j -> pure $ TSchemaVar j
       | otherwise = pure $ TMetaVar v

src/Pretty.hs

@@ -311,14 +311,14 @@ instance Pretty VarInfo (Scheme Ty) where
     text "∀" <>
     (align $ group $
      vsep [ group $
-            vsep (zipWith ppArgKind typeVarNames argKinds) <> text "."
+            vsep (zipWith ppArgKind schemeVarNames argKinds) <> text "."
           , pp env t
           ])
     where
       ppArgKind varName kind = parens (text varName <+> text ":" <+> pp env kind)
 
-typeVarNames :: [Text]
-typeVarNames =
+schemeVarNames :: [Text]
+schemeVarNames =
   greek ++
   [ base <> T.pack (show i)
   | (base, i) <- greekNum
@@ -332,6 +332,8 @@ typeVarNames =
                , base <- greek
                ]
 
+schemeVarName :: SchemeVar -> Text
+schemeVarName (SchemeVar n) = schemeVarNames !! fromIntegral n
 
 instance Pretty VarInfo TypeConstructor where
   pp _   TSyntax        = text "Syntax"
@@ -343,7 +345,7 @@ instance Pretty VarInfo TypeConstructor where
   pp _   TIO            = text "IO"
   pp _   TType          = text "Type"
   pp env (TDatatype t)  = pp env t
-  pp _   (TSchemaVar n) = text $ typeVarNames !! fromIntegral n
+  pp _   (TSchemaVar v) = text $ schemeVarName v
   pp _   (TMetaVar v)   = text "META" <> viaShow v -- TODO
 
 instance Pretty VarInfo a => Pretty VarInfo (TyF a) where
@@ -390,7 +392,7 @@ instance (Pretty VarInfo s, Pretty VarInfo t, PrettyBinder VarInfo a, Pretty Var
     (hang 2 $ group $
      vsep ( text "data" <+> text x <+>
             hsep [ parens (text α <+> ":" <+> pp env k)
-                 | α <- typeVarNames
+                 | α <- schemeVarNames
                  | k <- argKinds
                  ] <+>
             text "="

src/Type.hs

@@ -36,6 +36,13 @@ newMetaPtr = MetaPtr <$> newUnique
 instance Show MetaPtr where
   show (MetaPtr i) = "(MetaPtr " ++ show (hashUnique i) ++ ")"
 
+newtype SchemeVar = SchemeVar Natural
+  deriving newtype (Enum, Eq, HasKey, Ord, Show)
+  deriving stock   Data
+
+firstSchemeVar :: SchemeVar
+firstSchemeVar = SchemeVar 0
+
 data TypeConstructor
   = TSyntax
   | TInteger
@@ -46,7 +53,7 @@ data TypeConstructor
   | TIO
   | TType
   | TDatatype Datatype
-  | TSchemaVar Natural
+  | TSchemaVar SchemeVar
   | TMetaVar MetaPtr
   deriving (Data, Eq, Show)
 makePrisms ''TypeConstructor
@@ -115,7 +122,7 @@ class TyLike a arg | a -> arg where
   tIO         :: arg -> a
   tType       :: a
   tDatatype   :: Datatype -> [arg] -> a
-  tSchemaVar  :: Natural -> [arg] -> a
+  tSchemaVar  :: SchemeVar -> [arg] -> a
   tMetaVar    :: MetaPtr -> a
 
 instance TyLike (TyF a) a where

src/Util/Key.hs

@@ -4,10 +4,16 @@ module Util.Key
   (HasKey(..)
   ) where
 
+import Numeric.Natural
+
 class HasKey a where
   getKey :: a -> Int
   fromKey :: Int -> a
 
 instance HasKey Int where
   getKey  = id
   fromKey = id
+
+instance HasKey Natural where
+  getKey  = fromIntegral
+  fromKey = fromIntegral

src/Util/Store.hs

@@ -15,9 +15,11 @@
 
 module Util.Store
   ( lookup
+  , (!)
   , singleton
   , insert
   , toList
+  , toAscList
   , fromList
   , Store
   , unionWith
@@ -64,12 +66,17 @@
   at k f s = alterF f k s
   {-# INLINE at #-}
 
 instance (c ~ d) => Each (Store c a) (Store d b) a b where
   each = traversed
 
 lookup :: HasKey p => p -> Store p v -> Maybe v
 lookup ptr graph = getKey ptr `IM.lookup` unStore graph
 
+(!) :: HasKey p => Store p v -> p -> v
+graph ! ptr = case lookup ptr graph of
+  Just v -> v
+  Nothing -> error "Store.!!: key not found"
+
 singleton :: HasKey p => p -> v -> Store p v
 singleton ptr val = Store $! IM.singleton (getKey ptr) val
 
@@ -79,6 +86,9 @@
 toList :: HasKey p => Store p v -> [(p,v)]
 toList str = map (first fromKey) $ IM.toList (unStore str)
 
+toAscList :: HasKey p => Store p v -> [(p,v)]
+toAscList str = map (first fromKey) $ IM.toAscList (unStore str)
+
 fromList :: HasKey p => [(p,v)] -> Store p v
 fromList ps = Store $! IM.fromList $ map (first getKey) ps