Implement destination passing for Case expressions.

This way we won't have to worry about SoA-converting the return value
of a Case if we should later pull SoA-conversion out of Imp into its
own pass.

Author: axch

src/lib/Builder.hs

@@ -864,6 +864,15 @@ emitRunReader hint r body = do
   lam <- buildEffLam hint rTy \h ref -> body h ref
   emitExpr $ PrimOp $ Hof $ RunReader r lam
 
+buildRememberDest :: (Emits n, ScopableBuilder SimpIR m)
+  => NameHint -> SAtom n
+  -> (forall l. (Emits l, Distinct l, DExt n l) => SAtomName l -> m l (SAtom l))
+  -> m n (SAtom n)
+buildRememberDest hint dest cont = do
+  ty <- getType dest
+  doit <- buildUnaryLamExpr hint ty cont
+  emitExpr $ PrimOp $ DAMOp $ RememberDest dest doit
+
 -- === vector space (ish) type class ===
 
 zeroAt :: (Emits n, SBuilder m) => SType n -> m n (SAtom n)

src/lib/Imp.hs

@@ -304,8 +304,11 @@ translateExpr expr = confuseGHC >>= \_ -> case expr of
     repValAtom =<< naryIndexRepVal f (toList xs)
   Atom x -> substM x
   PrimOp op -> toImpOp op
-  Case e alts ty _ -> do
+  Case e alts unitResultTy _ -> do
     e' <- substM e
+    case unitResultTy of
+      UnitTy -> return ()
+      _ -> error $ "Unexpected returning Case in Imp " ++ pprint expr
     case trySelectBranch e' of
       Just (con, arg) -> do
         Abs b body <- return $ alts !! con
@@ -319,12 +322,11 @@ translateExpr expr = confuseGHC >>= \_ -> case expr of
         where
           go tag xss = do
             tag' <- fromScalarAtom tag
-            dest <- allocDest =<< substM ty
             emitSwitch tag' (zip xss alts) $
               \(xs, Abs b body) ->
                  extendSubst (b @> SubstVal (sink xs)) $
-                   translateBlock body >>= storeAtom (sink dest)
-            loadAtom dest
+                   void $ translateBlock body
+            return UnitVal
   TabCon _ _ _ -> error "Unexpected `TabCon` in Imp pass."
 
 toImpRefOp :: Emits o

src/lib/Lower.hs

@@ -35,7 +35,7 @@ import PPrint
 import QueryType
 import Types.Core
 import Types.Primitives
-import Util (enumerate)
+import Util (enumerate, foldMapM)
 
 -- === For loop resolution ===
 
@@ -109,6 +109,7 @@ lowerExpr :: Emits o => SExpr i -> LowerM i o (SAtom o)
 lowerExpr expr = emitExpr =<< case expr of
   TabCon Nothing ty els -> lowerTabCon Nothing ty els
   PrimOp (Hof (For dir ixDict body)) -> lowerFor Nothing dir ixDict body
+  Case e alts ty _ -> lowerCase Nothing e alts ty
   _ -> visitGeneric expr
 
 lowerBlock :: Emits o => SBlock i -> LowerM i o (SAtom o)
@@ -151,7 +152,6 @@ lowerTabCon maybeDest tabTy elems = do
   dest <- case maybeDest of
     Just  d -> return d
     Nothing -> emitExpr $ PrimOp $ DAMOp $ AllocDest tabTy'
-  destTy <- getType dest
   Abs bord ufoBlock <- buildAbs noHint IdxRepTy \ord -> do
     buildBlock $ unsafeFromOrdinal (sink ixTy') $ Var $ sink ord
   -- This is emitting a chain of RememberDest ops to force `dest` to be used
@@ -163,16 +163,36 @@ lowerTabCon maybeDest tabTy elems = do
   let go incoming_dest [] = return incoming_dest
       go incoming_dest ((ord, e):rest) = do
         i <- dropSubst $ extendSubst (bord@>SubstVal (IdxRepVal (fromIntegral ord))) $
-               lowerBlock ufoBlock
-        do_one_elt <- buildUnaryLamExpr "dest" destTy \local_dest -> do
+          lowerBlock ufoBlock
+        carried_dest <- buildRememberDest "dest" incoming_dest \local_dest -> do
           idest <- indexRef (Var local_dest) (sink i)
           place (FullDest idest) =<< visitAtom e
           return UnitVal
-        carried_dest <- emitExpr $ PrimOp $ DAMOp $ RememberDest incoming_dest do_one_elt
         go carried_dest rest
   dest' <- go dest (enumerate elems)
   return $ PrimOp $ DAMOp $ Freeze dest'
 
+lowerCase :: Emits o
+  => Maybe (Dest SimpIR o) -> SAtom i -> [Alt SimpIR i] -> SType i
+  -> LowerM i o (SExpr o)
+lowerCase maybeDest scrut alts resultTy = do
+  resultTy' <- substM resultTy
+  dest <- case maybeDest of
+    Just  d -> return d
+    Nothing -> emitExpr $ PrimOp $ DAMOp $ AllocDest resultTy'
+  scrut' <- visitAtom scrut
+  dest' <- buildRememberDest "case_dest" dest \local_dest -> do
+    alts' <- forM alts \(Abs (b:>ty) body) -> do
+      ty' <- substM ty
+      buildAbs (getNameHint b) ty' \b' ->
+        extendSubst (b @> Rename b') $
+          buildBlock do
+            lowerBlockWithDest (Var $ sink $ local_dest) body $> UnitVal
+    eff' <- foldMapM getEffects alts'
+    void $ emitExpr $ Case (sink scrut') alts' UnitTy eff'
+    return UnitVal
+  return $ PrimOp $ DAMOp $ Freeze dest'
+
 -- Destination-passing traversals
 --
 -- The idea here is to try to reuse the memory already allocated for outputs of surrounding
@@ -199,6 +219,9 @@ data ProjDest o
   = FullDest (Dest SimpIR o)
   | ProjDest (NE.NonEmpty Projection) (Dest SimpIR o)  -- dest corresponds to the projection applied to name
 
+instance SinkableE ProjDest where
+  sinkingProofE = todoSinkableProof
+
 lookupDest :: DestAssignment i' o -> SAtomName i' -> Maybe (ProjDest o)
 lookupDest = flip lookupNameMap
 
@@ -258,6 +281,13 @@ lowerExprWithDest dest expr = case expr of
   PrimOp (Hof (RunState Nothing s body)) -> traverseRWS body \ref' body' -> do
     s' <- visitAtom s
     return $ RunState ref' s' body'
+  Case e alts ty _ -> case dest of
+    Nothing -> lowerCase Nothing e alts ty
+    Just (FullDest d) -> lowerCase (Just d) e alts ty
+    Just d -> do
+      ans <- lowerCase Nothing e alts ty >>= emitExprToAtom
+      place d ans
+      return $ Atom ans
   _ -> generic
   where
     tabDest = dest <&> \case FullDest d -> d; ProjDest _ _ -> error "unexpected projection"