Move singleton-type handling to QueryType.

Also extend isSingletonType to table types and singletonTypeValRec to
lambda types ('cause they were different?), and observe that
isSingletonType is a pure function.

Author: axch

src/lib/CheckType.hs

@@ -11,7 +11,7 @@ module CheckType (
   tryGetType,
   checkUnOp, checkBinOp,
   isData, asFirstOrderFunction, asFFIFunType,
-  isSingletonType, singletonTypeVal, asNaryPiType,
+  asNaryPiType,
   ) where
 
 import Prelude hiding (id)
@@ -1062,52 +1062,6 @@ checkUnOp op x = do
       where
         u = SomeUIntArg; f = SomeFloatArg; sr = SameReturn
 
--- === singleton types ===
-
--- TODO: the following implementation should be valid:
---   isSingletonType :: EnvReader m => Type n -> m n Bool
---   isSingletonType ty =
---     singletonTypeVal ty >>= \case
---       Nothing -> return False
---       Just _  -> return True
--- But we want to be able to query the singleton-ness of types that we haven't
--- implemented tangent types for. So instead we do a separate case analysis.
-isSingletonType :: EnvReader m => Type n -> m n Bool
-isSingletonType topTy =
-  case checkIsSingleton topTy of
-    Just () -> return True
-    Nothing -> return False
-  where
-    checkIsSingleton :: Type n -> Maybe ()
-    checkIsSingleton ty = case ty of
-      Pi (PiType _ _ body) -> checkIsSingleton body
-      StaticRecordTy items -> mapM_ checkIsSingleton items
-      TC con -> case con of
-        ProdType tys -> mapM_ checkIsSingleton tys
-        _ -> Nothing
-      _ -> Nothing
-
-
-singletonTypeVal :: EnvReader m => Type n -> m n (Maybe (Atom n))
-singletonTypeVal ty = liftTyperT do
-  singletonTypeVal' ty
-
--- TODO: TypeCon with a single case?
-singletonTypeVal'
-  :: (MonadFail2 m, SubstReader Name m, EnvReader2 m, EnvExtender2 m)
-  => Type i -> m i o (Atom o)
-singletonTypeVal' ty = case ty of
-  TabPi (TabPiType b body) ->
-    substBinders b \b' -> do
-      body' <- singletonTypeVal' body
-      return $ TabLam $ TabLamExpr b' $ AtomicBlock body'
-  StaticRecordTy items -> Record <$> traverse singletonTypeVal' items
-  TC con -> case con of
-    ProdType tys -> ProdVal <$> traverse singletonTypeVal' tys
-    _            -> notASingleton
-  _ -> notASingleton
-  where notASingleton = fail "not a singleton type"
-
 -- === various helpers for querying types ===
 
 checkedApplyDataDefParams :: (EnvReader m, Fallible1 m) => DataDef n -> [Type n] -> m n [DataConDef n]

src/lib/Linearize.hs

@@ -16,7 +16,6 @@ import GHC.Stack
 import Name
 import Builder
 import Syntax
-import CheckType (isSingletonType)
 import MTL1
 import QueryType
 import Util (bindM2)
@@ -144,9 +143,7 @@ liftTangentM args m = liftSubstReaderT $ lift11 $ runReaderT1 args m
 
 isTrivialForAD :: Expr o -> PrimalM i o Bool
 isTrivialForAD expr = do
-  trivialTy  <- (maybeTangentType <$> getType expr) >>= \case
-    Nothing -> return False
-    Just tTy -> isSingletonType tTy
+  trivialTy  <- any isSingletonType . maybeTangentType <$> getType expr
   hasActiveEffs <- getEffects expr >>= \case
                      Pure -> return False
                      -- TODO: Be more precise here, such as checking

src/lib/QueryType.hs

@@ -8,14 +8,16 @@ module QueryType (
   instantiateDataDef, instantiateDepPairTy, instantiatePi, instantiateTabPi,
   litType, lamExprTy,
   numNaryPiArgs, naryLamExprType,
-  oneEffect, projectLength, sourceNameType, typeAsBinderNest, typeBinOp
+  oneEffect, projectLength, sourceNameType, typeAsBinderNest, typeBinOp,
+  isSingletonType, singletonTypeVal,
   ) where
 
 import Control.Monad
 import Data.Foldable (toList)
 import Data.List (elemIndex)
 import qualified Data.List.NonEmpty as NE
 import qualified Data.Map.Strict as M
+import Data.Maybe (isJust)
 import qualified Data.Set        as S
 
 import CheapReduction (cheapNormalize)
@@ -795,3 +797,53 @@ rwsFunEffects rws f = getTypeSubst f >>= \case
 
 deleteEff :: Effect n -> EffectRow n -> EffectRow n
 deleteEff eff (EffectRow effs t) = EffectRow (S.delete eff effs) t
+
+-- === singleton types ===
+
+-- The following implementation should be valid:
+--   isSingletonType :: EnvReader m => Type n -> m n Bool
+--   isSingletonType ty =
+--     singletonTypeVal ty >>= \case
+--       Nothing -> return False
+--       Just _  -> return True
+-- But a separate implementation doesn't have to go under binders,
+-- because it doesn't have to reconstruct the singleton value (which
+-- may be type annotated and whose type may refer to names).
+
+isSingletonType :: Type n -> Bool
+isSingletonType topTy = isJust $ checkIsSingleton topTy
+  where
+    checkIsSingleton :: Type n -> Maybe ()
+    checkIsSingleton ty = case ty of
+      Pi (PiType _ Pure body) -> checkIsSingleton body
+      TabPi (TabPiType _ body) -> checkIsSingleton body
+      StaticRecordTy items -> mapM_ checkIsSingleton items
+      TC con -> case con of
+        ProdType tys -> mapM_ checkIsSingleton tys
+        _ -> Nothing
+      _ -> Nothing
+
+singletonTypeVal :: EnvReader m => Type n -> m n (Maybe (Atom n))
+singletonTypeVal ty = liftEnvReaderT $
+  runSubstReaderT idSubst $ singletonTypeValRec ty
+{-# INLINE singletonTypeVal #-}
+
+-- TODO: TypeCon with a single case?
+singletonTypeValRec :: Type i
+  -> SubstReaderT Name (EnvReaderT Maybe) i o (Atom o)
+singletonTypeValRec ty = case ty of
+  Pi (PiType b Pure body) ->
+    substBinders b \(PiBinder b' ty' arr) -> do
+      body' <- singletonTypeValRec body
+      return $ Lam $ LamExpr (LamBinder b' ty' arr Pure) $ AtomicBlock body'
+  TabPi (TabPiType b body) ->
+    substBinders b \b' -> do
+      body' <- singletonTypeValRec body
+      return $ TabLam $ TabLamExpr b' $ AtomicBlock body'
+  StaticRecordTy items -> Record <$> traverse singletonTypeValRec items
+  TC con -> case con of
+    ProdType tys -> ProdVal <$> traverse singletonTypeValRec tys
+    _            -> notASingleton
+  _ -> notASingleton
+  where notASingleton = fail "not a singleton type"
+

src/lib/Transpose.hs

@@ -12,7 +12,6 @@ import Control.Monad.Reader
 import qualified Data.Set as S
 
 import MTL1
-import CheckType (singletonTypeVal)
 import Err
 import Name
 import Syntax