Touchups in response to code review of PR 894.

Author: axch

src/lib/Imp.hs

@@ -1591,8 +1591,7 @@ impInstrTypes instr = case instr of
 impOpType :: IPrimOp n -> IType
 impOpType pop = case pop of
   ScalarBinOp op x _ -> typeBinOp op (getIType x)
-  -- All unary ops preserve the type of their input
-  ScalarUnOp  _ x    -> getIType x
+  ScalarUnOp  op x   -> typeUnOp  op (getIType x)
   VectorBinOp op x _ -> typeBinOp op (getIType x)
   Select  _ x  _     -> getIType x
   VectorPack xs      -> Vector ty  where Scalar ty = getIType $ head xs

src/lib/Linearize.hs

@@ -143,14 +143,16 @@ liftTangentM args m = liftSubstReaderT $ lift11 $ runReaderT1 args m
 
 isTrivialForAD :: Expr o -> PrimalM i o Bool
 isTrivialForAD expr = do
-  trivialTy  <- any isSingletonType . maybeTangentType <$> getType expr
+  trivialTy  <- presentAnd isSingletonType . maybeTangentType <$> getType expr
   hasActiveEffs <- getEffects expr >>= \case
                      Pure -> return False
                      -- TODO: Be more precise here, such as checking
                      -- whether the effects are themselves active.
                      _ -> return True
   hasActiveVars <- isActive expr
   return $ not hasActiveEffs && (trivialTy || not hasActiveVars)
+    where presentAnd :: (a -> Bool) -> Maybe a -> Bool
+          presentAnd = any
 
 isActive :: HoistableE e => e o -> PrimalM i o Bool
 isActive e = do

src/lib/QueryType.hs

@@ -8,7 +8,7 @@ module QueryType (
   instantiateDataDef, instantiateDepPairTy, instantiatePi, instantiateTabPi,
   litType, lamExprTy,
   numNaryPiArgs, naryLamExprType,
-  oneEffect, projectLength, sourceNameType, typeAsBinderNest, typeBinOp,
+  oneEffect, projectLength, sourceNameType, typeAsBinderNest, typeBinOp, typeUnOp,
   isSingletonType, singletonTypeVal,
   ) where
 
@@ -219,6 +219,9 @@ typeBinOp binop xTy = case binop of
   BXor   -> xTy
   BShL   -> xTy;  BShR   -> xTy
 
+typeUnOp :: UnOp -> BaseType -> BaseType
+typeUnOp = const id  -- All unary ops preserve the type of the input
+
 -- === computing effects ===
 
 computeAbsEffects :: (EnvExtender m, SubstE Name e)
@@ -476,8 +479,7 @@ getTypePrimOp op = case op of
   ScalarBinOp binop x _ -> do
     xTy <- getTypeBaseType x
     return $ TC $ BaseType $ typeBinOp binop xTy
-  -- All unary ops preserve the type of the input
-  ScalarUnOp _ x -> getTypeE x
+  ScalarUnOp unop x -> TC . BaseType . typeUnOp unop <$> getTypeBaseType x
   Select _ x _ -> getTypeE x
   UnsafeFromOrdinal ty _ -> substM ty
   ToOrdinal _ -> return IdxRepTy