Implement and test case-of-case optimization in the inliner.

Author: axch

src/lib/Builder.hs

@@ -768,15 +768,15 @@ injectAltResult sumTys con (Abs b body) = liftBuilder do
 
 -- TODO: consider a version with nonempty list of alternatives where we figure
 -- out the result type from one of the alts rather than providing it explicitly
-buildCase :: (Emits n, ScopableBuilder r m)
-          => Atom r n -> Type r n
-          -> (forall l. (Emits l, DExt n l) => Int -> Atom r l -> m l (Atom r l))
-          -> m n (Atom r n)
-buildCase scrut resultTy indexedAltBody = do
+buildCase' :: (Emits n, ScopableBuilder r m)
+  => Atom r n -> Type r n
+  -> (forall l. (Emits l, DExt n l) => Int -> Atom r l -> m l (Atom r l))
+  -> m n (Expr r n)
+buildCase' scrut resultTy indexedAltBody = do
   case trySelectBranch scrut of
     Just (i, arg) -> do
       Distinct <- getDistinct
-      indexedAltBody i $ sink arg
+      Atom <$> indexedAltBody i (sink arg)
     Nothing -> do
       scrutTy <- getType scrut
       altBinderTys <- caseAltsBinderTys scrutTy
@@ -786,7 +786,13 @@ buildCase scrut resultTy indexedAltBody = do
           eff <- getEffects blk
           return $ blk `PairE` eff
         return (Abs b' body, ignoreHoistFailure $ hoist b' eff')
-      emitExpr $ Case scrut alts resultTy $ mconcat effs
+      return $ Case scrut alts resultTy $ mconcat effs
+
+buildCase :: (Emits n, ScopableBuilder r m)
+  => Atom r n -> Type r n
+  -> (forall l. (Emits l, DExt n l) => Int -> Atom r l -> m l (Atom r l))
+  -> m n (Atom r n)
+buildCase s r b = emitExprToAtom =<< buildCase' s r b
 
 buildEffLam
   :: ScopableBuilder r m

src/lib/Inline.hs

@@ -420,20 +420,35 @@ reconstructTabApp ctx expr ixs =
 reconstructCase :: Emits o
   => Context SExpr e o -> SExpr o -> [SAlt i] -> SType i -> EffectRow SimpIR i
   -> InlineM i o (e o)
-reconstructCase ctx scrutExpr alts resultTy effs = do
-  -- TODO Opportunity here to inspect `scrutExpr` and perform case-of-case
-  -- optimization
-  scrut <- emitExprToAtom scrutExpr
-  case trySelectBranch scrut of
-    Just (i, val) -> do
-      Abs b body <- return $ alts !! i
-      extendSubst (b @> (SubstVal $ DoneEx $ Atom val)) do
-        inlineBlockEmits ctx body
-    Nothing -> do
-      alts' <- mapM visitAlt alts
+reconstructCase ctx scrutExpr alts resultTy effs =
+  case scrutExpr of
+    Case sscrut salts _ _ -> do
+      -- Perform case-of-case optimization
+      -- TODO Add join points to reduce code duplication (and repeated inlining)
+      -- of the arms of the outer case
       resultTy' <- inline Stop resultTy
-      effs' <- inline Stop effs
-      reconstruct ctx $ Case scrut alts' resultTy' effs'
+      reconstruct ctx =<< (buildCase' sscrut resultTy' \i val -> do
+        ans <- applyAbs (sink $ salts !! i) (SubstVal val) >>= emitBlock
+        buildCase ans (sink resultTy') \j jval -> do
+          Abs b body <- return $ alts !! j
+          extendSubst (b @> (SubstVal $ DoneEx $ Atom jval)) do
+            inlineBlockEmits Stop body >>= emitExprToAtom)
+    _ -> do
+      -- Attempt case-of-known-constructor optimization
+      -- I can't use `buildCase` here because I want to propagate the incoming
+      -- context `ctx` into the selected alternative if the optimization fires,
+      -- but leave it around the whole reconstructed `Case` if it doesn't.
+      scrut <- emitExprToAtom scrutExpr
+      case trySelectBranch scrut of
+        Just (i, val) -> do
+          Abs b body <- return $ alts !! i
+          extendSubst (b @> (SubstVal $ DoneEx $ Atom val)) do
+            inlineBlockEmits ctx body
+        Nothing -> do
+          alts' <- mapM visitAlt alts
+          resultTy' <- inline Stop resultTy
+          effs' <- inline Stop effs
+          reconstruct ctx $ Case scrut alts' resultTy' effs'
 
 instance Inlinable (EffectRow SimpIR)
 instance Inlinable (EffectAndType SimpIR)

tests/inline-tests.dx

@@ -106,3 +106,23 @@ def id'(x:Nat) -> Nat = x
 -- CHECK: case
 -- CHECK: === inline ===
 -- CHECK-NOT: case
+
+-- CHECK-LABEL: Inlining carries out the case-of-case optimization
+"Inlining carries out the case-of-case optimization"
+
+-- Before inlining there are two cases, but attempting to inline `x`
+-- reveals a case-of-case opprtunity, which in turn exposes
+-- case-of-known-constructor in each branch, leading to just one case
+-- in the end.
+%passes inline
+:pp
+  x = if id'(3) > 2
+        then Just 4
+        else Nothing
+  case x of
+    Just a -> a * a
+    Nothing -> 0
+-- CHECK: === inline ===
+-- CHECK: case
+-- CHECK-NOT: case
+-- CHECK: === inline ===