Remove type/effect annotations from Block.

Blocks are now just a bunch of decls followed by an atom. The type and effects
information is carried by the context -- e.g. a pi type or the other arguments
to a built-in HOF. The reason is that we're about to add decls to binders and
this way we can treat binder-decls block-decls uniformly.

src/lib/Algebra.hs

@@ -137,7 +137,7 @@ type BlockTraverserM i o a = SubstReaderT PolySubstVal (MaybeT1 (BuilderM SimpIR
 blockAsPoly
   :: (EnvExtender m, EnvReader m)
   => Block SimpIR n -> m n (Maybe (Polynomial n))
-blockAsPoly (Block _ decls result) =
+blockAsPoly (Abs decls result) =
   liftBuilder $ runMaybeT1 $ runSubstReaderT idSubst $ blockAsPolyRec decls result
 
 blockAsPolyRec :: Nest (Decl SimpIR) i i' -> Atom SimpIR i' -> BlockTraverserM i o (Polynomial o)

src/lib/Builder.hs

@@ -15,8 +15,9 @@ import Control.Monad.Reader
 import Control.Monad.Writer.Strict hiding (Alt)
 import Control.Monad.State.Strict (MonadState (..), StateT (..), runStateT)
 import qualified Data.Map.Strict as M
+import Data.Foldable (fold)
 import Data.Graph (graphFromEdges, topSort)
-import Data.Text.Prettyprint.Doc (Pretty (..), group, line, nest)
+import Data.Text.Prettyprint.Doc (Pretty (..))
 import Foreign.Ptr
 
 import qualified Unsafe.Coerce as TrulyUnsafe
@@ -28,7 +29,6 @@ import IRVariants
 import MTL1
 import Subst
 import Name
-import PPrint (prettyBlock)
 import QueryType
 import Types.Core
 import Types.Imp
@@ -88,11 +88,11 @@ emitUnOp op x = emitOp $ UnOp op x
 {-# INLINE emitUnOp #-}
 
 emitBlock :: (Builder r m, Emits n) => Block r n -> m n (Atom r n)
-emitBlock (Block _ decls result) = emitDecls decls result
+emitBlock = emitDecls
 
 emitDecls :: (Builder r m, Emits n, RenameE e, SinkableE e)
-          => Nest (Decl r) n l -> e l -> m n (e n)
-emitDecls decls result = runSubstReaderT idSubst $ emitDecls' decls result
+          => WithDecls r e n -> m n (e n)
+emitDecls (Abs decls result) = runSubstReaderT idSubst $ emitDecls' decls result
 
 emitDecls' :: (Builder r m, Emits o, RenameE e, SinkableE e)
            => Nest (Decl r) i i' -> e i' -> SubstReaderT Name m i o (e o)
@@ -278,10 +278,9 @@ emitTopLet hint letAnn expr = do
   v <- emitBinding hint $ AtomNameBinding $ LetBound (DeclBinding letAnn expr)
   return $ AtomVar v ty
 
-emitTopFunBinding :: (Mut n, TopBuilder m) => NameHint -> TopFunDef n -> LamExpr SimpIR n -> m n (TopFunName n)
+emitTopFunBinding :: (Mut n, TopBuilder m) => NameHint -> TopFunDef n -> STopLam n -> m n (TopFunName n)
 emitTopFunBinding hint def f = do
-  ty <- return $ getLamExprType f
-  emitBinding hint $ TopFunBinding $ DexTopFun def ty f Waiting
+  emitBinding hint $ TopFunBinding $ DexTopFun def f Waiting
 
 emitSourceMap :: TopBuilder m => SourceMap n -> m n ()
 emitSourceMap sm = emitLocalModuleEnv $ mempty {envSourceMap = sm}
@@ -334,7 +333,7 @@ extendLinearizationCache s fs =
 
 queryObjCache :: EnvReader m => TopFunName n -> m n (Maybe (FunObjCodeName n))
 queryObjCache v = lookupEnv v >>= \case
-  TopFunBinding (DexTopFun _ _ _ (Finished impl)) -> return $ Just $ topFunObjCode impl
+  TopFunBinding (DexTopFun _ _ (Finished impl)) -> return $ Just $ topFunObjCode impl
   _ -> return Nothing
 
 emitObjFile :: (Mut n, TopBuilder m) => CFunction n -> m n (FunObjCodeName n)
@@ -468,7 +467,7 @@ liftEmitBuilder cont = do
   Distinct <- getDistinct
   let (result, decls, _) = runHardFail $ unsafeRunInplaceT (runBuilderT' cont) env emptyOutFrag
   Emits <- fabricateEmitsEvidenceM
-  emitDecls (unsafeCoerceB $ unRNest decls) result
+  emitDecls $ Abs (unsafeCoerceB $ unRNest decls) result
 
 instance (IRRep r, Fallible m) => ScopableBuilder r (BuilderT r m) where
   buildScoped cont = BuilderT do
@@ -601,56 +600,14 @@ buildBlock
   :: ScopableBuilder r m
   => (forall l. (Emits l, DExt n l) => m l (Atom r l))
   -> m n (Block r n)
-buildBlock cont = buildScoped (cont >>= withType) >>= computeAbsEffects >>= absToBlock
-
-withType :: ((EnvReader m, IRRep r), HasType r e) => e l -> m l ((e `PairE` Type r) l)
-withType e = do
-  ty <- {-# SCC blockTypeNormalization #-} cheapNormalize $ getType e
-  return $ e `PairE` ty
-{-# INLINE withType #-}
-
-makeBlock :: IRRep r => Nest (Decl r) n l -> EffectRow r l -> Atom r l -> Type r l -> Block r n
-makeBlock decls effs atom ty = Block (BlockAnn (EffTy effs' ty')) decls atom where
-  ty' = ignoreHoistFailure $ hoist decls ty
-  effs' = ignoreHoistFailure $ hoist decls effs
-{-# INLINE makeBlock #-}
-
-absToBlockInferringTypes :: (EnvReader m, IRRep r) => Abs (Nest (Decl r)) (Atom r) n -> m n (Block r n)
-absToBlockInferringTypes ab = liftEnvReaderM do
-  abWithEffs <-  computeAbsEffects ab
-  refreshAbs abWithEffs \decls (effs `PairE` result) -> do
-    ty <- cheapNormalize $ getType result
-    return $ ignoreExcept $
-      absToBlock $ Abs decls (effs `PairE` (result `PairE` ty))
-{-# INLINE absToBlockInferringTypes #-}
-
-absToBlock
-  :: (Fallible m, IRRep r)
-  => Abs (Nest (Decl r)) (EffectRow r `PairE` (Atom r `PairE` Type r)) n -> m (Block r n)
-absToBlock (Abs decls (effs `PairE` (result `PairE` ty))) = do
-  let msg = "Block:" <> nest 1 (prettyBlock decls result) <> line
-            <> group ("Of type:" <> nest 2 (line <> pretty ty)) <> line
-            <> group ("With effects:" <> nest 2 (line <> pretty effs))
-  ty' <- liftHoistExcept' (docAsStr msg) $ hoist decls ty
-  effs' <- liftHoistExcept' (docAsStr msg) $ hoist decls effs
-  return $ Block (BlockAnn (EffTy effs' ty')) decls result
-{-# INLINE absToBlock #-}
-
-makeBlockFromDecls :: (EnvReader m, IRRep r) => Abs (Nest (Decl r)) (Atom r) n -> m n (Block r n)
-makeBlockFromDecls (Abs Empty result) = return $ AtomicBlock result
-makeBlockFromDecls ab = liftEnvReaderM $ refreshAbs ab \decls result -> do
-  ty <- return $ getType result
-  effs <- declNestEffects decls
-  PairE ty' effs' <- return $ ignoreHoistFailure $ hoist decls $ PairE ty effs
-  return $ Block (BlockAnn (EffTy effs' ty')) decls result
-{-# INLINE makeBlockFromDecls #-}
+buildBlock = buildScoped
 
 coreLamExpr :: EnvReader m => AppExplicitness
             -> Abs (Nest (WithExpl CBinder)) (PairE (EffectRow CoreIR) CBlock) n
             -> m n (CoreLamExpr n)
 coreLamExpr appExpl ab = liftEnvReaderM do
   refreshAbs ab \bs' (PairE effs' body') -> do
-    resultTy <- return $ getType body'
+    EffTy _ resultTy <- blockEffTy body'
     let bs'' = fmapNest withoutExpl bs'
     return $ CoreLamExpr (CorePiType appExpl bs' (EffTy effs' resultTy)) (LamExpr bs'' body')
 
@@ -736,12 +693,13 @@ buildLamExpr (Abs bs UnitE) cont = case bs of
       buildLamExpr rest' \vs -> cont $ sink v : vs
     return $ LamExpr (Nest b' bs') body'
 
-buildLamExprFromPi
+buildTopLamFromPi
   :: ScopableBuilder r m
   => PiType r n
   -> (forall l. (Emits l, Distinct l, DExt n l) => [AtomVar r l] -> m l (Atom r l))
-  -> m n (LamExpr r n)
-buildLamExprFromPi (PiType bs _) cont = buildLamExpr (EmptyAbs bs) cont
+  -> m n (TopLam r n)
+buildTopLamFromPi piTy@(PiType bs _) cont =
+  TopLam False piTy <$> buildLamExpr (EmptyAbs bs) cont
 
 buildAlt
   :: ScopableBuilder r m
@@ -789,7 +747,7 @@ buildCase' scrut resultTy indexedAltBody = do
       (alts, effs) <- unzip <$> forM (enumerate altBinderTys) \(i, bTy) -> do
         (Abs b' (body `PairE` eff')) <- buildAbs noHint bTy \x -> do
           blk <- buildBlock $ indexedAltBody i $ Var $ sink x
-          eff <- return $ getEffects blk
+          EffTy eff _ <- blockEffTy blk
           return $ blk `PairE` eff
         return (Abs b' body, ignoreHoistFailure $ hoist b' eff')
       return $ Case scrut alts $ EffTy (mconcat effs) resultTy
@@ -1157,9 +1115,17 @@ mkDictAtom d = do
   ty <- typeOfDictExpr d
   return $ DictCon ty d
 
+mkCase :: (EnvReader m, IRRep r) => Atom r n -> Type r n -> [Alt r n] -> m n (Expr r n)
+mkCase scrut resultTy alts = liftEnvReaderM do
+  eff' <- fold <$> forM alts \alt -> refreshAbs alt \b body -> do
+    EffTy eff _ <- blockEffTy body
+    return $ ignoreHoistFailure $ hoist b eff
+  return $ Case scrut alts (EffTy eff' resultTy)
+
 mkCatchException :: EnvReader m => CBlock n -> m n (Hof CoreIR n)
 mkCatchException body = do
-  resultTy <- makePreludeMaybeTy $ getType body
+  EffTy _ bodyTy <- blockEffTy body
+  resultTy <- makePreludeMaybeTy bodyTy
   return $ CatchException resultTy body
 
 app :: (CBuilder m, Emits n) => CAtom n -> CAtom n -> m n (CAtom n)
@@ -1177,7 +1143,7 @@ naryTopAppInlined :: (Builder SimpIR m, Emits n) => TopFunName n -> [SAtom n] ->
 naryTopAppInlined f xs = do
   TopFunBinding f' <- lookupEnv f
   case f' of
-    DexTopFun _ _ (LamExpr bs body) _ ->
+    DexTopFun _ (TopLam _ _ (LamExpr bs body)) _ ->
       applySubst (bs@@>(SubstVal<$>xs)) body >>= emitBlock
     _ -> naryTopApp f xs
 {-# INLINE naryTopAppInlined #-}
@@ -1237,7 +1203,7 @@ applyIxMethod dict method args = case dict of
   IxDictSpecialized _ d params -> do
     SpecializedDict _ maybeFs <- lookupSpecDict d
     Just fs <- return maybeFs
-    LamExpr bs body <- return $ fs !! fromEnum method
+    TopLam _ _ (LamExpr bs body) <- return $ fs !! fromEnum method
     emitBlock =<< applySubst (bs @@> fmap SubstVal (params ++ args)) body
 
 unsafeFromOrdinal :: (SBuilder m, Emits n) => IxType SimpIR n -> Atom SimpIR n -> m n (Atom SimpIR n)
@@ -1551,15 +1517,10 @@ type ExprVisitorNoEmits2 m r = forall i o. ExprVisitorNoEmits (m i o) r i o
 visitLamNoEmits
   :: (ExprVisitorNoEmits2 m r, IRRep r, AtomSubstReader v m, EnvExtender2 m)
   => LamExpr r i -> m i o (LamExpr r o)
-visitLamNoEmits (LamExpr bs body) =
-  visitBinders bs \bs' -> LamExpr bs' <$> visitBlockNoEmits body
-
-visitBlockNoEmits
-  :: (ExprVisitorNoEmits2 m r, IRRep r, AtomSubstReader v m, EnvExtender2 m)
-  => Block r i -> m i o (Block r o)
-visitBlockNoEmits (Block _ decls result) =
-  absToBlockInferringTypes =<< visitDeclsNoEmits decls \decls' -> do
-    Abs decls' <$> visitAtom result
+visitLamNoEmits (LamExpr bs (Abs decls result)) =
+  visitBinders bs \bs' -> LamExpr bs' <$>
+    visitDeclsNoEmits decls \decls' -> Abs decls' <$> do
+      visitAtom result
 
 visitDeclsNoEmits
   :: (ExprVisitorNoEmits2 m r, IRRep r, AtomSubstReader v m, EnvExtender2 m)
@@ -1602,7 +1563,7 @@ visitLamEmits (LamExpr bs body) = visitBinders bs \bs' -> LamExpr bs' <$>
 visitBlockEmits
   :: (ExprVisitorEmits2 m r, SubstReader AtomSubstVal m, EnvExtender2 m, IRRep r, Emits o)
   => Block r i -> m i o (Atom r o)
-visitBlockEmits (Block _ decls result) = visitDeclsEmits decls $ visitAtom result
+visitBlockEmits (Abs decls result) = visitDeclsEmits decls $ visitAtom result
 
 visitDeclsEmits
   :: (ExprVisitorEmits2 m r, SubstReader AtomSubstVal m, EnvExtender2 m, IRRep r, Emits o)

src/lib/CheapReduction.hs

@@ -15,7 +15,8 @@ module CheapReduction
   , unwrapLeadingNewtypesType, wrapNewtypesData, liftSimpAtom, liftSimpType
   , liftSimpFun, makeStructRepVal, NonAtomRenamer (..), Visitor (..), VisitGeneric (..)
   , visitAtomPartial, visitTypePartial, visitAtomDefault, visitTypeDefault, Visitor2
-  , visitBinders, visitPiDefault, visitAlt, toAtomVar)
+  , visitBinders, visitPiDefault, visitAlt, toAtomVar, instantiatePiTy
+  , bindersToVars, bindersToAtoms)
   where
 
 import Control.Applicative
@@ -261,9 +262,6 @@ instance CheaplyReducibleE CoreIR TyConParams TyConParams where
 instance (CheaplyReducibleE r e e', NiceE r e') => CheaplyReducibleE r (Abs (Nest (Decl r)) e) e' where
   cheapReduceE (Abs decls result) = cheapReduceWithDeclsB decls $ cheapReduceE result
 
-instance (CheaplyReducibleE r (Atom r) e', NiceE r e') => CheaplyReducibleE r (Block r) e' where
-  cheapReduceE (Block _ decls result) = cheapReduceE $ Abs decls result
-
 instance IRRep r => CheaplyReducibleE r (Expr r) (Atom r) where
   cheapReduceE expr = confuseGHC >>= \_ -> case expr of
     Atom atom -> cheapReduceE atom
@@ -472,6 +470,10 @@ instantiateTyConDef (TyConDef _ bs conDefs) (TyConParams _ xs) = do
   applySubst (bs @@> (SubstVal <$> xs)) conDefs
 {-# INLINE instantiateTyConDef #-}
 
+instantiatePiTy :: (EnvReader m, IRRep r) => PiType r n -> [Atom r n] -> m n (EffTy r n)
+instantiatePiTy (PiType bs effTy) xs = do
+  applySubst (bs @@> (SubstVal <$> xs)) effTy
+
 -- Returns a representation type (type of an TypeCon-typed Newtype payload)
 -- given a list of instantiated DataConDefs.
 dataDefRep :: DataConDefs n -> CType n
@@ -517,10 +519,10 @@ instance VisitGeneric (Type    r) r where visitGeneric = visitType
 instance VisitGeneric (LamExpr r) r where visitGeneric = visitLam
 instance VisitGeneric (PiType  r) r where visitGeneric = visitPi
 
-instance VisitGeneric (Block r) r where
-  visitGeneric b = visitGeneric (LamExpr Empty b) >>= \case
-    LamExpr Empty b' -> return b'
-    _ -> error "not a block"
+visitBlock :: Visitor m r i o => Block r i -> m (Block r o)
+visitBlock b = visitGeneric (LamExpr Empty b) >>= \case
+  LamExpr Empty b' -> return b'
+  _ -> error "not a block"
 
 visitAlt :: Visitor m r i o => Alt r i -> m (Alt r o)
 visitAlt (Abs b body) = do
@@ -615,17 +617,11 @@ instance IRRep r => VisitGeneric (Expr r) r where
     -- TODO: should we reuse the original effects? Whether it's valid depends on
     -- the type-preservation requirements for a visitor. We should clarify what
     -- those are.
-    Case x alts (EffTy _ t) -> do
+    Case x alts effTy -> do
       x' <- visitGeneric x
-      t' <- visitGeneric t
       alts' <- mapM visitAlt alts
-      let effs' = foldMap altEffects alts'
-      return $ Case x' alts' $ EffTy effs' t'
-      where
-        altEffects :: Alt r n -> EffectRow r n
-        altEffects (Abs bs (Block ann _ _)) = case ann of
-          NoBlockAnn -> Pure
-          BlockAnn (EffTy effs _) -> ignoreHoistFailure $ hoist bs effs
+      effTy' <- visitGeneric effTy
+      return $ Case x' alts' effTy'
     Atom x -> Atom <$> visitGeneric x
     TabCon Nothing t xs -> TabCon Nothing <$> visitGeneric t <*> mapM visitGeneric xs
     TabCon (Just (WhenIRE d)) t xs -> TabCon <$> (Just . WhenIRE <$> visitGeneric d) <*> visitGeneric t <*> mapM visitGeneric xs
@@ -654,10 +650,10 @@ instance IRRep r => VisitGeneric (Hof r) r where
     RunReader x body -> RunReader <$> visitGeneric x <*> visitGeneric body
     RunWriter dest bm body -> RunWriter <$> mapM visitGeneric dest <*> visitGeneric bm <*> visitGeneric body
     RunState  dest s body ->  RunState  <$> mapM visitGeneric dest <*> visitGeneric s <*> visitGeneric body
-    While          b -> While          <$> visitGeneric b
-    RunIO          b -> RunIO          <$> visitGeneric b
-    RunInit        b -> RunInit        <$> visitGeneric b
-    CatchException t b -> CatchException <$> visitType t <*> visitGeneric b
+    While          b -> While          <$> visitBlock b
+    RunIO          b -> RunIO          <$> visitBlock b
+    RunInit        b -> RunInit        <$> visitBlock b
+    CatchException t b -> CatchException <$> visitType t <*> visitBlock b
     Linearize      lam x -> Linearize <$> visitGeneric lam <*> visitGeneric x
     Transpose      lam x -> Transpose <$> visitGeneric lam <*> visitGeneric x
 
@@ -674,7 +670,7 @@ instance IRRep r => VisitGeneric (DAMOp r) r where
 
 instance VisitGeneric UserEffectOp CoreIR where
   visitGeneric = \case
-    Handle name xs body -> Handle <$> renameN name <*> mapM visitGeneric xs <*> visitGeneric body
+    Handle name xs body -> Handle <$> renameN name <*> mapM visitGeneric xs <*> visitBlock body
     Resume t x -> Resume <$> visitGeneric t <*> visitGeneric x
     Perform x i -> Perform <$> visitGeneric x <*> pure i
 
@@ -798,6 +794,15 @@ toAtomVar v = do
   ty <- getType <$> lookupAtomName v
   return $ AtomVar v ty
 
+bindersToVars :: (EnvReader m,  IRRep r) => Nest (Binder r) n' n -> m n [AtomVar r n]
+bindersToVars bs = do
+  withExtEvidence bs do
+    Distinct <- getDistinct
+    mapM toAtomVar $ nestToNames bs
+
+bindersToAtoms :: (EnvReader m,  IRRep r) => Nest (Binder r) n' n -> m n [Atom r n]
+bindersToAtoms bs = liftM (Var <$>) $ bindersToVars bs
+
 newtype SubstVisitor i o a = SubstVisitor { runSubstVisitor :: Reader (Env o, Subst AtomSubstVal i o) a }
         deriving (Functor, Applicative, Monad, MonadReader (Env o, Subst AtomSubstVal i o))
 
@@ -889,7 +894,6 @@ instance IRRep r => SubstE AtomSubstVal (PrimOp r)
 instance IRRep r => SubstE AtomSubstVal (RefOp r)
 instance IRRep r => SubstE AtomSubstVal (EffTy r)
 instance IRRep r => SubstE AtomSubstVal (Expr r)
-instance IRRep r => SubstE AtomSubstVal (Block r)
 instance IRRep r => SubstE AtomSubstVal (GenericOpRep const r)
 instance SubstE AtomSubstVal InstanceBody
 instance SubstE AtomSubstVal DictType