Add kinds to distinguish data, refs, functions

src/lib/CheapReduction.hs

@@ -11,7 +11,7 @@ module CheapReduction
   ( reduceWithDecls, reduceExpr
   , instantiateTyConDef, dataDefRep, unwrapNewtypeType, projType
   , NonAtomRenamer (..), Visitor (..), VisitGeneric (..)
-  , visitAtomDefault, visitTypeDefault, Visitor2, mkStuck, mkStuckTy
+  , visitAtomDefault, visitTypeDefault, Visitor2, mkStuck
   , visitBinders, visitPiDefault, visitAlt, toAtomVar, instantiate, withInstantiated
   , bindersToVars, bindersToAtoms, instantiateNames, withInstantiatedNames, assumeConst
   , repValAtom, reduceUnwrap, reduceProj, reduceSuperclassProj, typeOfApp
@@ -163,11 +163,6 @@ mkStuck x = do
   ty <- queryStuckType x
   return $ Stuck ty x
 
-mkStuckTy :: EnvReader m => CStuck n -> m n (CType n)
-mkStuckTy x = do
-  ty <- queryStuckType x
-  return $ StuckTy ty x
-
 queryStuckType :: (IRRep r, EnvReader m) => Stuck r n -> m n (Type r n)
 queryStuckType = \case
   Var v -> return $ getType v
@@ -531,7 +526,7 @@ instance IRRep r => VisitGeneric (TyCon r) r where
     RefType t      -> RefType  <$> visitGeneric t
     TabPi t        -> TabPi     <$> visitGeneric t
     DepPairTy t    -> DepPairTy <$> visitGeneric t
-    TypeKind       -> return TypeKind
+    Kind k         -> return $ Kind k
     DictTy t       -> DictTy <$> visitGeneric t
     Pi t           -> Pi <$> visitGeneric t
     NewtypeTyCon t -> NewtypeTyCon <$> visitGeneric t

src/lib/CheckType.hs

@@ -172,7 +172,9 @@ instance IRRep r => CheckableE r (AtomVar r) where
 
 instance IRRep r => CheckableE r (Type r) where
   checkE = \case
-    StuckTy ty e -> uncurry StuckTy <$> checkStuck ty e
+    StuckTy k e -> do
+      (TyCon (Kind k'), e') <- checkStuck (TyCon (Kind k)) e
+      return $ StuckTy k' e'
     TyCon e -> TyCon <$> checkE e
 
 instance (ToBinding ann c, Color c, CheckableE r ann) => CheckableB r (BinderP c ann) where
@@ -369,7 +371,7 @@ instance IRRep r => CheckableE r (TyCon r) where
     ProdType tys     -> ProdType <$> mapM checkE tys
     SumType  cs      -> SumType  <$> mapM checkE cs
     RefType a        -> RefType  <$> checkE a
-    TypeKind         -> return TypeKind
+    Kind k         -> return $ Kind k
     Pi t           -> Pi           <$> checkE t
     TabPi t        -> TabPi        <$> checkE t
     NewtypeTyCon t -> NewtypeTyCon <$> checkE t

src/lib/Err.hs

@@ -152,6 +152,7 @@ data TypeErr =
  | AmbiguousInferenceVar VarStr TypeStr InfVarDesc
  | FFIResultTyErr    TypeStr
  | FFIArgTyNotScalar TypeStr
+ | ExpectedAType TypeStr
    deriving (Show, Eq)
 
 data MiscErr =
@@ -303,6 +304,7 @@ instance PrintableErr TypeErr where
          "couldn't infer instantiation of type " <> t
        MiscInfVar ->
          "ambiguous type variable: " ++ infVar ++ ": " ++ ty
+    ExpectedAType t -> "Expected a type, got a: " ++ t
 
 instance PrintableErr MiscErr where
   printErr = \case

src/lib/Generalize.hs

@@ -39,16 +39,16 @@ generalizeArgs fTy argsTop = liftGeneralizerM $ runSubstReaderT idSubst do
     go (Nest (WithAttrB expl b) bs) (arg:args) = do
       ty' <- substM $ binderType b
       arg' <- liftSubstReaderT case (ty', expl) of
-        (TyKind, _) -> toAtom <$> generalizeType (fromJust $ toMaybeType arg)
+        (TyCon (Kind TypeKind), _) -> toAtom <$> generalizeType (fromJust $ toMaybeType arg)
         (TyCon (DictTy _), Inferred Nothing (Synth _)) ->
           toAtom <$> generalizeDict ty' (fromJust $ toMaybeDict arg)
-        _ -> isData ty' >>= \case
-          True -> toAtom <$> emitGeneralizationParameter ty' arg
-          False -> do
-            -- Unlike in `inferRoles` in `Inference`, it's ok to have non-data,
-            -- non-type, non-dict arguments (e.g. a function). We just don't
-            -- generalize in that case.
-            return arg
+        -- _ -> isData ty' >>= \case
+        --   True -> toAtom <$> emitGeneralizationParameter ty' arg
+        --   False -> do
+        --     -- Unlike in `inferRoles` in `Inference`, it's ok to have non-data,
+        --     -- non-type, non-dict arguments (e.g. a function). We just don't
+        --     -- generalize in that case.
+        --     return arg
       args'' <- extendSubst (b@>SubstVal arg') $ go bs args
       return $ arg' : args''
     go Empty [] = return []
@@ -127,19 +127,19 @@ traverseTyParams (TyCon ty) f = liftM TyCon $ getDistinct >>= \Distinct -> case
       ClassDef _ _ _ _ _ bs _ _ <- lookupClassDef name
       params' <- traverseRoleBinders f bs params
       return $ DictType sn name params'
-    IxDictType   t -> IxDictType   <$> f' TypeParam TyKind t
+    IxDictType   t -> IxDictType   <$> f' TypeParam (toType $ Kind TypeKind) t
   TabPi (TabPiType d (b:>iTy) resultTy) -> do
-    iTy' <- f' TypeParam TyKind iTy
+    iTy' <- f' TypeParam (toType $ Kind TypeKind) iTy
     let dictTy = toType $ IxDictType iTy'
     d' <- fromJust . toMaybeDict <$> f DictParam dictTy (toAtom d)
     withFreshBinder (getNameHint b) iTy' \(b':>_) -> do
-      resultTy' <- applyRename (b@>binderName b') resultTy >>= (f' TypeParam TyKind)
+      resultTy' <- applyRename (b@>binderName b') resultTy >>= (f' TypeParam (toType $ Kind TypeKind))
       return $ TabPi $ TabPiType d' (b':>iTy') resultTy'
   BaseType b -> return $ BaseType b
-  ProdType tys -> ProdType <$> forM tys \t -> f' TypeParam TyKind t
+  ProdType tys -> ProdType <$> forM tys \t -> f' TypeParam (toType $ Kind TypeKind) t
   RefType _ -> error "not implemented"
-  SumType  tys -> SumType  <$> forM tys \t -> f' TypeParam TyKind t
-  TypeKind     -> return TypeKind
+  SumType  tys -> SumType  <$> forM tys \t -> f' TypeParam (toType $ Kind TypeKind) t
+  Kind k -> return $ Kind k
   NewtypeTyCon con -> NewtypeTyCon <$> case con of
     Nat -> return Nat
     Fin n -> Fin <$> f DataParam NatTy n

src/lib/Inference.hs

@@ -270,7 +270,7 @@ withInferenceVar hint binding cont = diffStateT1 \s -> do
 {-# INLINE withInferenceVar #-}
 
 withFreshUnificationVar
-  :: (Zonkable e, Emits o) => SrcId -> InfVarDesc -> Kind CoreIR o
+  :: (Zonkable e, Emits o) => SrcId -> InfVarDesc -> CType o
   -> (forall o'. (Emits o', DExt o o') => CAtomVar o' -> SolverM i o' (e o'))
   -> SolverM i o (e o)
 withFreshUnificationVar sid desc k cont = do
@@ -283,7 +283,7 @@ withFreshUnificationVar sid desc k cont = do
 {-# INLINE withFreshUnificationVar #-}
 
 withFreshUnificationVarNoEmits
-  :: (Zonkable e) => SrcId -> InfVarDesc -> Kind CoreIR o
+  :: (Zonkable e) => SrcId -> InfVarDesc -> CType o
   -> (forall o'. (DExt o o') => CAtomVar o' -> SolverM i o' (e o'))
   -> SolverM i o (e o)
 withFreshUnificationVarNoEmits sid desc k cont = diffStateT1 \s -> do
@@ -417,6 +417,19 @@ etaExpandPartialPi (PartialPiType appExpl expls bs reqTy) cont = do
     let piTy = CorePiType appExpl expls bs' (getType body)
     return $ CoreLamExpr piTy $ LamExpr bs' body
 
+-- Expects the uexpr to be a type-like thing, including ordinary data types, pi
+-- types, dict types etc.
+topDownType :: forall i o. Emits o => UExpr i -> InfererM i o (CType o)
+topDownType exprWithSrc@(WithSrcE sid expr) = case expr of
+  UPrim UTuple xs -> toType . ProdType <$> mapM checkUType xs
+  _ -> do
+    ty <- bottomUpExplicit exprWithSrc
+    ty' <- maybeInterpretPunsAsTyCons (Check (toType $ Kind TypeKind)) ty
+    ty'' <- instantiateSigma sid Infer ty'
+    case toMaybeType ty'' of
+      Nothing -> throw sid $ ExpectedAType (pprint ty'')
+      Just ty''' -> return ty'''
+
 -- Doesn't introduce implicit pi binders or dependent pairs
 topDownExplicit :: forall i o. Emits o => CType o -> UExpr i -> InfererM i o (CAtom o)
 topDownExplicit reqTy exprWithSrc@(WithSrcE sid expr) = emitExprType sid reqTy >> case expr of
@@ -452,7 +465,7 @@ topDownExplicit reqTy exprWithSrc@(WithSrcE sid expr) = emitExprType sid reqTy >
   UNatLit x -> fromNatLit sid x reqTy
   UIntLit x -> fromIntLit sid x reqTy
   UPrim UTuple xs -> case reqTy of
-    TyKind -> toAtom . ProdType <$> mapM checkUType xs
+    TyCon (Kind TypeKind) -> toAtom . ProdType <$> mapM checkUType xs
     TyCon (ProdType reqTys) -> do
       when (length reqTys /= length xs) $ throw sid $ TupleLengthMismatch (length reqTys) (length xs)
       toAtom <$> ProdCon <$> forM (zip reqTys xs) \(reqTy', x) -> topDown reqTy' x
@@ -644,8 +657,8 @@ withUDecl (WithSrcB _ d) cont = case d of
     bindLetPat p var cont
 
 considerInlineAnn :: LetAnn -> CType n -> LetAnn
-considerInlineAnn PlainLet TyKind = InlineLet
-considerInlineAnn PlainLet (TyCon (Pi (CorePiType _ _ _ TyKind))) = InlineLet
+considerInlineAnn PlainLet (TyCon (Kind _)) = InlineLet
+considerInlineAnn PlainLet (TyCon (Pi (CorePiType _ _ _ (TyCon (Kind _))))) = InlineLet
 considerInlineAnn ann _ = ann
 
 applyFromLiteralMethod
@@ -785,7 +798,7 @@ buildAppConstraints appSrcId reqTy (CorePiType _ _ bs ty) = do
       Check reqTy' -> [TypeConstraint appSrcId (sink reqTy') resultTy]
 
 maybeInterpretPunsAsTyCons :: RequiredTy n -> SigmaAtom n -> InfererM i n (SigmaAtom n)
-maybeInterpretPunsAsTyCons (Check TyKind) (SigmaUVar sn _ (UPunVar v)) = do
+maybeInterpretPunsAsTyCons (Check (TyCon (Kind TypeKind))) (SigmaUVar sn _ (UPunVar v)) = do
   let v' = UTyConVar v
   ty <- getUVarType v'
   return $ SigmaUVar sn ty v'
@@ -1005,7 +1018,7 @@ inferPrimArg :: Emits o => UExpr i -> InfererM i o (CAtom o)
 inferPrimArg x = do
   xBlock <- buildBlock $ bottomUp x
   case getType xBlock of
-    TyKind -> reduceExpr xBlock >>= \case
+    TyCon (Kind TypeKind) -> reduceExpr xBlock >>= \case
       Just reduced -> return reduced
       _ -> throwInternal "Type args to primops must be reducible"
     _ -> emit xBlock
@@ -1020,7 +1033,7 @@ matchPrimApp = \case
    P.ProdType -> \ts -> return $ toAtom $ ProdType $ map (fromJust . toMaybeType) ts
    P.SumType  -> \ts -> return $ toAtom $ SumType  $ map (fromJust . toMaybeType) ts
    P.RefType  -> \case ~[h, a] -> undefined -- return $ toAtom $ RefType h (fromJust $ toMaybeType a)
-   P.TypeKind -> \case ~[] -> return $ Con $ TyConAtom $ TypeKind
+   P.TypeKind -> \case ~[] -> return $ toAtom $ Kind $ TypeKind
  UCon con -> case con of
    P.ProdCon -> \xs -> return $ toAtom $ ProdCon xs
    P.SumCon _ -> error "not supported"
@@ -1050,7 +1063,7 @@ matchPrimApp = \case
    matchGenericOp op xs = do
      (tyArgs, dataArgs) <- partitionEithers <$> forM xs \x -> do
        case getType x of
-         TyKind -> do
+         TyCon (Kind TypeKind) -> do
            Just x' <- return $ toMaybeType x
            return $ Left x'
          _ -> return $ Right x
@@ -1372,7 +1385,8 @@ checkInstanceParams bsTop paramsTop = go bsTop paramsTop
   go :: Nest CBinder o any -> [UExpr i] -> InfererM i o [CAtom o]
   go Empty [] = return []
   go (Nest (b:>ty) bs) (x:xs) = do
-    x' <- checkUParam ty x
+    let msg = CantReduceType $ pprint x
+    x' <- withReducibleEmissions (getSrcId x) msg $ topDown (sink ty) x
     Abs bs' UnitE <- applySubst (b@>SubstVal x') $ Abs bs UnitE
     (x':) <$> go bs' xs
   go _ _ = error "zip error"
@@ -1517,14 +1531,9 @@ bindLetPat (WithSrcB sid pat) v cont = emitExprType sid (getType v) >> case pat
     emitInline atom = emitDecl noHint InlineLet $ Atom atom
 
 checkUType :: UType i -> InfererM i o (CType o)
-checkUType t = do
-  Just t' <- toMaybeType <$> checkUParam TyKind t
-  return t'
-
-checkUParam :: Kind CoreIR o -> UType i -> InfererM i o (CAtom o)
-checkUParam k uty =
-  withReducibleEmissions (getSrcId uty) msg $ topDownExplicit (sink k) uty
-  where msg = CantReduceType $ pprint uty
+checkUType t = withReducibleEmissions (getSrcId t) msg do
+  topDownType t
+  where msg = CantReduceType $ pprint t
 
 inferTabCon :: forall i o. Emits o => SrcId -> [UExpr i] -> InfererM i o (CAtom o)
 inferTabCon sid xs = do
@@ -1662,8 +1671,8 @@ instance Unifiable (TyCon CoreIR) where
   unify t1 t2 = case t1 of
     ( BaseType b ) -> do
     { BaseType b' <- matchit; guard $ b == b'}
-    ( TypeKind ) -> do
-    { TypeKind <- matchit; return () }
+    ( Kind k ) -> do
+    { Kind k'  <- matchit; guard $ k == k' }
     ( Pi piTy ) -> do
     { Pi piTy' <- matchit; unify piTy piTy'}
     ( TabPi piTy) -> do
@@ -1760,7 +1769,7 @@ instance Unifiable (Abs CBinder CType) where
       return UnitE
 
 withFreshSkolemName
-  :: Zonkable e => Kind CoreIR o
+  :: Zonkable e => CType o
   -> (forall o'. DExt o o' => CAtomVar o' -> SolverM i o' (e o'))
   -> SolverM i o (e o)
 withFreshSkolemName ty cont = diffStateT1 \s -> do
@@ -1895,7 +1904,7 @@ generalizeInstanceArg role ty arg cont = case role of
   -- that it's valid to implement `generalizeDict` by synthesizing an entirely
   -- fresh dictionary, and if we were to do that, we would infer this type
   -- parameter exactly as we do here, using inference.
-  TypeParam -> withFreshUnificationVarNoEmits rootSrcId MiscInfVar TyKind \v -> cont $ toAtom v
+  TypeParam -> withFreshUnificationVarNoEmits rootSrcId MiscInfVar (toType $ Kind TypeKind) \v -> cont $ toAtom v
   DictParam -> withFreshDictVarNoEmits ty (
     \ty' -> case toMaybeDict (sink arg) of
               Just d -> liftM toAtom $ lift11 $ generalizeDictRec ty' d

src/lib/QueryType.hs

@@ -137,7 +137,7 @@ getUVarType = \case
   UPunVar     v -> getStructDataConType v
   UClassVar v -> do
     ClassDef _ _ _ _ expls bs _ _ <- lookupClassDef v
-    return $ toType $ CorePiType ExplicitApp expls bs TyKind
+    return $ toType $ CorePiType ExplicitApp expls bs (toType $ Kind TypeKind)
   UMethodVar  v -> getMethodNameType v
 
 getMethodNameType :: EnvReader m => MethodName n -> m n (CType n)
@@ -179,8 +179,8 @@ getTyConNameType :: EnvReader m => TyConName n -> m n (Type CoreIR n)
 getTyConNameType v = do
   TyConDef _ expls bs _ <- lookupTyCon v
   case bs of
-    Empty -> return TyKind
-    _ -> return $ toType $ CorePiType ExplicitApp expls bs TyKind
+    Empty -> return $ toType $ Kind TypeKind
+    _ -> return $ toType $ CorePiType ExplicitApp expls bs $ toType $ Kind TypeKind
 
 getDataConNameType :: EnvReader m => DataConName n -> m n (Type CoreIR n)
 getDataConNameType dataCon = liftEnvReaderM $ withSubstReaderT do
@@ -284,30 +284,3 @@ liftIFunType (IFunType _ argTys resultTys) = liftEnvReaderM $ go argTys where
     t:ts -> withFreshBinder noHint (toType $ BaseType t) \b -> do
       PiType bs effTy <- go ts
       return $ PiType (Nest b bs) effTy
-
--- === Data constraints ===
-
-isData :: EnvReader m => Type CoreIR n -> m n Bool
-isData ty = do
-  result <- liftEnvReaderT $ withSubstReaderT $ go ty
-  case result of
-    Just () -> return True
-    Nothing -> return False
-  where
-    go :: Type CoreIR i -> SubstReaderT Name (EnvReaderT Maybe) i o ()
-    go = \case
-      StuckTy _ _ -> notData
-      TyCon con -> case con of
-        TabPi (TabPiType _ b eltTy) -> renameBinders b \_ -> go eltTy
-        DepPairTy (DepPairType _ b@(_:>l) r) -> go l >> renameBinders b \_ -> go r
-        NewtypeTyCon nt -> do
-          (_, ty') <- unwrapNewtypeType =<< renameM nt
-          dropSubst $ go ty'
-        BaseType _  -> return ()
-        ProdType as -> mapM_ go as
-        SumType  cs -> mapM_ go cs
-        RefType _ -> return ()
-        TypeKind -> notData
-        DictTy _ -> notData
-        Pi _     -> notData
-      where notData = empty

src/lib/QueryTypePure.hs

@@ -66,6 +66,20 @@ typeBinOp binop xTy = case binop of
 typeUnOp :: UnOp -> BaseType -> BaseType
 typeUnOp = const id  -- All unary ops preserve the type of the input
 
+getKind :: Type r n -> Kind
+getKind = \case
+  StuckTy k _ -> k
+  TyCon con -> case con of
+    BaseType _     -> TypeKind
+    ProdType _     -> TypeKind
+    SumType  _     -> TypeKind
+    TabPi _        -> TypeKind
+    DepPairTy _    -> TypeKind
+    NewtypeTyCon _ -> TypeKind
+    RefType _  -> RefKind
+    Pi _       -> FunKind
+    DictTy _   -> DictKind
+    Kind _     -> OtherKind
 
 instance IRRep r => HasType r (AtomVar r) where
   getType (AtomVar _ ty) = ty
@@ -87,13 +101,8 @@ instance HasType CoreIR (DictCon CoreIR) where
     IxFin n -> toType $ IxDictType (FinTy n)
     IxRawFin _ -> toType $ IxDictType IdxRepTy
 
-instance HasType CoreIR CType where
-  getType = \case
-    TyCon _ -> TyKind
-    StuckTy t _ -> t
-
 instance HasType CoreIR NewtypeTyCon where
-  getType _ = TyKind
+  getType _ = TyCon $ Kind TypeKind
 
 getNewtypeType :: NewtypeCon n -> CType n
 getNewtypeType con = case con of
@@ -110,7 +119,7 @@ instance IRRep r => HasType r (Con r) where
     DepPair _ _ ty -> toType $ DepPairTy ty
     DictConAtom d -> getType d
     NewtypeCon con _ -> getNewtypeType con
-    TyConAtom _ -> TyKind
+    TyConAtom k -> TyCon $ Kind $ getKind $ TyCon k
 
 getSuperclassType :: RNest CBinder n l -> Nest CBinder l l' -> Int -> CType n
 getSuperclassType _ Empty = error "bad index"

src/lib/RuntimePrint.hs

@@ -80,7 +80,7 @@ showAnyTyCon tyCon atom = case tyCon of
     xs <- getUnpacked atom
     parens $ sepBy ", " $ map rec xs
   -- TODO: traverse the type and print out data components
-  TypeKind -> printAsConstant
+  Kind _ -> printAsConstant
   Pi _ -> printTypeOnly "function"
   TabPi _ -> brackets $ forEachTabElt atom \iOrd x -> do
     isFirst <- emit $ BinOp (ICmp Equal) iOrd (NatVal 0)

src/lib/Simplify.hs

@@ -181,7 +181,7 @@ getRepType (TyCon con) = case con of
     -- dropSubst $ getRepType ty'
   Pi _     -> error notDataType
   DictTy _ -> error notDataType
-  TypeKind -> error notDataType
+  Kind _ -> error notDataType
   where notDataType = "Not a type of runtime-representable data"
 
 toDataAtomAssumeNoDecls :: CAtom i -> SimplifyM i o (SAtom o)

src/lib/TopLevel.hs

@@ -707,7 +707,7 @@ loadModuleSource config moduleName = do
       LibDirectory dir -> return dir
 {-# SCC loadModuleSource #-}
 
-getDexString :: (MonadIO1 m, EnvReader m, Fallible1 m) => Val CoreIR n -> m n String
+getDexString :: (MonadIO1 m, EnvReader m, Fallible1 m) => CAtom n -> m n String
 getDexString val = do
   -- TODO: use a `ByteString` instead of `String`
   Stuck _ (RepValAtom (RepVal _ tree)) <- return val