Merge pull request #1322 from google-research/type-checking-refactor

Update type checker in anticipation of decls-in-types.

Author: dougalm

src/lib/CheckType.hs

@@ -82,6 +82,7 @@ newtype EnvReaderT (m::MonadKind) (n::S) (a:: *) =
            , MonadWriter w, Fallible, Searcher, Alternative)
 
 type EnvReaderM = EnvReaderT Identity
+type FallibleEnvReaderM = EnvReaderT FallibleM
 
 runEnvReaderM :: Distinct n => Env n -> EnvReaderM n a -> a
 runEnvReaderM bindings m = runIdentity $ runEnvReaderT bindings m

src/lib/Core.hs

@@ -8,7 +8,6 @@ module Generalize (generalizeArgs, generalizeIxDict) where
 
 import Control.Monad
 
-import CheckType (isData)
 import Core
 import Err
 import Types.Core

src/lib/Generalize.hs

@@ -1535,7 +1535,7 @@ impInstrTypes instr = case instr of
   where hostPtrTy ty = PtrType (CPU, ty)
 
 instance CheckableE SimpIR ImpFunction where
-  checkE _ = return () -- TODO
+  checkE = renameM -- TODO
 
 -- TODO: Don't use Core Envs for Imp!
 instance BindsEnv ImpDecl where

src/lib/Imp.hs

@@ -11,7 +11,7 @@
 module Inference
   ( inferTopUDecl, checkTopUType, inferTopUExpr
   , trySynthTerm, generalizeDict, asTopBlock
-  , synthTopE, UDeclInferenceResult (..)) where
+  , synthTopE, UDeclInferenceResult (..), asFFIFunType) where
 
 import Prelude hiding ((.), id)
 import Control.Category
@@ -45,6 +45,7 @@ import SourceInfo
 import Subst
 import QueryType
 import Types.Core
+import Types.Imp
 import Types.Primitives
 import Types.Source
 import Util hiding (group)
@@ -3184,6 +3185,43 @@ withFabricatedEmitsInf cont = fromWrapWithEmitsInf
 newtype WrapWithEmitsInf n r =
   WrapWithEmitsInf { fromWrapWithEmitsInf :: EmitsInf n => r }
 
+-- === IFunType ===
+
+asFFIFunType :: EnvReader m => CType n -> m n (Maybe (IFunType, CorePiType n))
+asFFIFunType ty = return do
+  Pi piTy <- return ty
+  impTy <- checkFFIFunTypeM piTy
+  return (impTy, piTy)
+
+checkFFIFunTypeM :: Fallible m => CorePiType n -> m IFunType
+checkFFIFunTypeM (CorePiType appExpl (_:expls) (Nest b bs) effTy) = do
+  argTy <- checkScalar $ binderType b
+  case bs of
+    Empty -> do
+      resultTys <- checkScalarOrPairType (etTy effTy)
+      let cc = case length resultTys of
+                 0 -> error "Not implemented"
+                 1 -> FFICC
+                 _ -> FFIMultiResultCC
+      return $ IFunType cc [argTy] resultTys
+    Nest b' rest -> do
+      let naryPiRest = CorePiType appExpl expls (Nest b' rest) effTy
+      IFunType cc argTys resultTys <- checkFFIFunTypeM naryPiRest
+      return $ IFunType cc (argTy:argTys) resultTys
+checkFFIFunTypeM _ = error "expected at least one argument"
+
+checkScalar :: (IRRep r, Fallible m) => Type r n -> m BaseType
+checkScalar (BaseTy ty) = return ty
+checkScalar ty = throw TypeErr $ pprint ty
+
+checkScalarOrPairType :: (IRRep r, Fallible m) => Type r n -> m [BaseType]
+checkScalarOrPairType (PairTy a b) = do
+  tys1 <- checkScalarOrPairType a
+  tys2 <- checkScalarOrPairType b
+  return $ tys1 ++ tys2
+checkScalarOrPairType (BaseTy ty) = return [ty]
+checkScalarOrPairType ty = throw TypeErr $ pprint ty
+
 -- === instances ===
 
 instance PrettyE e => Pretty (UDeclInferenceResult e l) where
@@ -3197,9 +3235,11 @@ instance SinkableE e => SinkableE (UDeclInferenceResult e) where
 
 instance (RenameE e, CheckableE CoreIR e) => CheckableE CoreIR (UDeclInferenceResult e) where
   checkE = \case
-    UDeclResultDone _ -> return ()
-    UDeclResultBindName _ block _ -> checkE block
-    UDeclResultBindPattern _ block _ -> checkE block
+    UDeclResultDone e -> UDeclResultDone <$> checkE e
+    UDeclResultBindName ann block ab ->
+      UDeclResultBindName ann <$> checkE block <*> renameM ab -- TODO: check result
+    UDeclResultBindPattern hint block recon ->
+      UDeclResultBindPattern hint <$> checkE block <*> renameM recon -- TODO: check recon
 
 instance HasType CoreIR InfEmission where
   getType = \case

src/lib/Inference.hs

@@ -291,7 +291,7 @@ linearize f x = runPrimalMInit $ linearizeLambdaApp f x
 linearizeTopLam :: STopLam n -> [Active] -> DoubleBuilder SimpIR n (STopLam n, STopLam n)
 linearizeTopLam (TopLam False _ (LamExpr bs body)) actives = do
   (primalFun, tangentFun) <- runPrimalMInit $ refreshBinders bs \bs' frag -> extendSubst frag do
-    let allPrimals = nestToAtomVars bs'
+    let allPrimals = bindersVars bs'
     activeVs <- catMaybes <$> forM (zip actives allPrimals) \(active, v) -> case active of
       True  -> return $ Just v
       False -> return $ Nothing

src/lib/Linearize.hs

@@ -213,7 +213,7 @@ getMethodNameType v = liftEnvReaderM $ lookupEnv v >>= \case
   MethodBinding className i -> do
     ClassDef _ _ paramNames _ paramBs scBinders methodTys <- lookupClassDef className
     refreshAbs (Abs paramBs $ Abs scBinders (methodTys !! i)) \paramBs' absPiTy -> do
-      let params = Var <$> nestToAtomVars paramBs'
+      let params = Var <$> bindersVars paramBs'
       dictTy <- DictTy <$> dictType (sink className) params
       withFreshBinder noHint dictTy \dictB -> do
         scDicts <- getSuperclassDicts (Var $ binderVar dictB)
@@ -384,3 +384,47 @@ liftIFunType (IFunType _ argTys resultTys) = liftEnvReaderM $ go argTys where
     t:ts -> withFreshBinder noHint (BaseTy 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 $ checkDataLike ty
+  case runFallibleM result of
+    Success () -> return True
+    Failure _  -> return False
+
+checkDataLike :: Type CoreIR i -> SubstReaderT Name FallibleEnvReaderM i o ()
+checkDataLike ty = case ty of
+  TyVar _ -> notData
+  TabPi (TabPiType _ b eltTy) -> do
+    renameBinders b \_ ->
+      checkDataLike eltTy
+  DepPairTy (DepPairType _ b@(_:>l) r) -> do
+    recur l
+    renameBinders b \_ -> checkDataLike r
+  NewtypeTyCon nt -> do
+    (_, ty') <- unwrapNewtypeType =<< renameM nt
+    dropSubst $ recur ty'
+  TC con -> case con of
+    BaseType _       -> return ()
+    ProdType as      -> mapM_ recur as
+    SumType  cs      -> mapM_ recur cs
+    RefType _ _      -> return ()
+    HeapType         -> return ()
+    _ -> notData
+  _   -> notData
+  where
+    recur = checkDataLike
+    notData = throw TypeErr $ pprint ty
+
+checkExtends :: (Fallible m, IRRep r) => EffectRow r n -> EffectRow r n -> m ()
+checkExtends allowed (EffectRow effs effTail) = do
+  let (EffectRow allowedEffs allowedEffTail) = allowed
+  case effTail of
+    EffectRowTail _ -> assertEq allowedEffTail effTail ""
+    NoTail -> return ()
+  forM_ (eSetToList effs) \eff -> unless (eff `eSetMember` allowedEffs) $
+    throw CompilerErr $ "Unexpected effect: " ++ pprint eff ++
+                      "\nAllowed: " ++ pprint allowed
+

src/lib/QueryType.hs

@@ -19,7 +19,7 @@ import Data.Text.Prettyprint.Doc (Pretty (..), hardline)
 
 import Builder
 import CheapReduction
-import CheckType (CheckableE (..), isData, checkBlock)
+import CheckType
 import Core
 import Err
 import Generalize
@@ -274,9 +274,12 @@ instance SinkableE SimplifiedBlock
 instance RenameE SimplifiedBlock
 instance HoistableE SimplifiedBlock
 instance CheckableE SimpIR SimplifiedBlock where
-  checkE (SimplifiedBlock block _) =
-    -- TODO: CheckableE instance for the recon too
-    void $ checkBlock block
+  checkE (SimplifiedBlock block recon) = do
+    block' <- renameM block
+    effTy <- blockEffTy block' -- TODO: store this in the simplified block instead
+    block'' <- dropSubst $ checkBlock effTy block'
+    recon' <- renameM recon -- TODO: CheckableE instance for the recon too
+    return $ SimplifiedBlock block'' recon'
 
 instance Pretty (SimplifiedBlock n) where
   pretty (SimplifiedBlock block recon) =
@@ -286,9 +289,9 @@ instance SinkableE SimplifiedTopLam where
   sinkingProofE = todoSinkableProof
 
 instance CheckableE SimpIR SimplifiedTopLam where
-  checkE (SimplifiedTopLam lam _) = do
+  checkE (SimplifiedTopLam lam recon) =
     -- TODO: CheckableE instance for the recon too
-    checkE lam
+    SimplifiedTopLam <$> checkE lam <*> renameM recon
 
 instance Pretty (SimplifiedTopLam n) where
   pretty (SimplifiedTopLam lam recon) = pretty lam <> hardline <> pretty recon

src/lib/Simplify.hs

@@ -43,9 +43,9 @@ import qualified LLVM.AST
 
 import AbstractSyntax
 import Builder
-import CheckType ( CheckableE (..), asFFIFunType, checkHasType)
+import CheckType ( CheckableE (..), checkTypeIs)
 #ifdef DEX_DEBUG
-import CheckType (checkTypesM)
+import CheckType (checkTypes)
 #endif
 import Core
 import ConcreteSyntax
@@ -316,7 +316,7 @@ evalSourceBlock' mname block = case sbContents block of
           _ -> evalUExpr expr
         fType <- getType <$> toAtomVar fname'
         (nimplicit, nexplicit, linFunTy) <- liftExceptEnvReaderM $ getLinearizationType zeros fType
-        impl `checkHasType` linFunTy >>= \case
+        liftEnvReaderT (impl `checkTypeIs` linFunTy) >>= \case
           Failure _ -> do
             let implTy = getType impl
             throw TypeErr $ unlines
@@ -744,7 +744,7 @@ checkPass name cont = do
     return result
 #ifdef DEX_DEBUG
   logTop $ MiscLog $ "Running checks"
-  checkTypesM result
+  checkTypes result
   logTop $ MiscLog $ "Checks passed"
 #else
   logTop $ MiscLog $ "Checks skipped (not a debug build)"

src/lib/TopLevel.hs

@@ -943,12 +943,12 @@ binderType (_:>ty) = ty
 binderVar  :: (IRRep r, DExt n l) => Binder r n l -> AtomVar r l
 binderVar (b:>ty) = AtomVar (binderName b) (sink ty)
 
-nestToAtomVars :: (Distinct l, Ext n l, IRRep r)
-               => Nest (Binder r) n l -> [AtomVar r l]
-nestToAtomVars = \case
+bindersVars :: (Distinct l, Ext n l, IRRep r)
+            => Nest (Binder r) n l -> [AtomVar r l]
+bindersVars = \case
   Empty -> []
   Nest b bs -> withExtEvidence b $ withSubscopeDistinct bs $
-    sink (binderVar b) : nestToAtomVars bs
+    sink (binderVar b) : bindersVars bs
 
 -- === ToBinding ===