Query rather than check types of scalar operations in Imp.

Author: axch

src/lib/Imp.hs

@@ -29,7 +29,7 @@ import MTL1
 import Name
 import Builder
 import Syntax
-import CheckType (CheckableE (..), checkUnOp, checkBinOp)
+import CheckType (CheckableE (..))
 import Simplify
 import LabeledItems
 import QueryType
@@ -1590,9 +1590,10 @@ impInstrTypes instr = case instr of
 -- TODO: reuse type rules in Type.hs
 impOpType :: IPrimOp n -> IType
 impOpType pop = case pop of
-  ScalarBinOp op x y -> runHardFail $ checkBinOp op (getIType x) (getIType y)
-  ScalarUnOp  op x   -> runHardFail $ checkUnOp  op (getIType x)
-  VectorBinOp op x y -> runHardFail $ checkBinOp op (getIType x) (getIType y)
+  ScalarBinOp op x _ -> typeBinOp op (getIType x)
+  -- All unary ops preserve the type of their input
+  ScalarUnOp  _ x    -> getIType x
+  VectorBinOp op x _ -> typeBinOp op (getIType x)
   Select  _ x  _     -> getIType x
   VectorPack xs      -> Vector ty  where Scalar ty = getIType $ head xs
   VectorIndex x _    -> Scalar ty  where Vector ty = getIType x

src/lib/QueryType.hs

@@ -8,7 +8,7 @@ module QueryType (
   instantiateDataDef, instantiateDepPairTy, instantiatePi, instantiateTabPi,
   litType, lamExprTy,
   numNaryPiArgs, naryLamExprType,
-  oneEffect, projectLength, sourceNameType, typeAsBinderNest
+  oneEffect, projectLength, sourceNameType, typeAsBinderNest, typeBinOp
   ) where
 
 import Control.Monad
@@ -203,6 +203,20 @@ typeAsBinderNest ty = do
   return $ Abs (Nest (ignored:>ty) Empty) body
 {-# INLINE typeAsBinderNest #-}
 
+typeBinOp :: BinOp -> BaseType -> BaseType
+typeBinOp binop xTy = case binop of
+  IAdd   -> xTy;  ISub   -> xTy
+  IMul   -> xTy;  IDiv   -> xTy
+  IRem   -> xTy;
+  ICmp _ -> Scalar Word8Type
+  FAdd   -> xTy;  FSub   -> xTy
+  FMul   -> xTy;  FDiv   -> xTy;
+  FPow   -> xTy
+  FCmp _ -> Scalar Word8Type
+  BAnd   -> xTy;  BOr    -> xTy
+  BXor   -> xTy
+  BShL   -> xTy;  BShR   -> xTy
+
 -- === computing effects ===
 
 computeAbsEffects :: (EnvExtender m, SubstE Name e)
@@ -459,19 +473,7 @@ getTypePrimOp op = case op of
   TabCon ty _ -> substM ty
   ScalarBinOp binop x _ -> do
     xTy <- getTypeBaseType x
-    resTy <- return $ case binop of
-      IAdd   -> xTy;  ISub   -> xTy
-      IMul   -> xTy;  IDiv   -> xTy
-      IRem   -> xTy;
-      ICmp _ -> Scalar Word8Type
-      FAdd   -> xTy;  FSub   -> xTy
-      FMul   -> xTy;  FDiv   -> xTy;
-      FPow   -> xTy
-      FCmp _ -> Scalar Word8Type
-      BAnd   -> xTy;  BOr    -> xTy
-      BXor   -> xTy
-      BShL   -> xTy;  BShR   -> xTy
-    return $ TC $ BaseType resTy
+    return $ TC $ BaseType $ typeBinOp binop xTy
   -- All unary ops preserve the type of the input
   ScalarUnOp _ x -> getTypeE x
   Select _ x _ -> getTypeE x