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

Author: axch

src/lib/CheapReduction.hs

@@ -15,7 +15,7 @@ module CheapReduction
   , unwrapLeadingNewtypesType, wrapNewtypesData, liftSimpAtom, liftSimpType
   , liftSimpFun, makeStructRepVal, NonAtomRenamer (..), Visitor (..), VisitGeneric (..)
   , visitAtomPartial, visitTypePartial, visitAtomDefault, visitTypeDefault, Visitor2
-  , visitBinders, visitPiDefault)
+  , visitBinders, visitPiDefault, visitAlt)
   where
 
 import Control.Applicative

src/lib/Inline.hs

@@ -16,6 +16,7 @@ import IRVariants
 import Name
 import Subst
 import Occurrence hiding (Var)
+import Optimize
 import Types.Core
 import Types.Primitives
 
@@ -90,7 +91,7 @@ inlineDeclsSubst = \case
       s <- getSubst
       extendSubst (b @> SubstVal (SuspEx expr s)) $ inlineDeclsSubst rest
     else do
-      expr' <- inlineExpr Stop expr
+      expr' <- inlineExpr Stop expr >>= (liftEnvReaderM . peepholeExpr)
       -- If the inliner starts moving effectful expressions, it may become
       -- necessary to query the effects of the new expression here.
       let presInfo = resolveWorkConservation ann expr'
@@ -248,6 +249,9 @@ data Context (from::E) (to::E) (o::S) where
   Stop :: Context e e o
   TabAppCtx :: [SAtom i] -> Subst InlineSubstVal i o
             -> Context SExpr e o -> Context SExpr e o
+  CaseCtx :: [SAlt i] -> SType i -> EffectRow SimpIR i
+          -> Subst InlineSubstVal i o
+          -> Context SExpr e o -> Context SExpr e o
   EmitToAtomCtx :: Context SAtom e o -> Context SExpr e o
   EmitToNameCtx :: Context SAtomName e o -> Context SAtom e o
 
@@ -271,6 +275,9 @@ inlineExpr ctx = \case
   TabApp tbl ixs -> do
     s <- getSubst
     inlineAtom (TabAppCtx ixs s ctx) tbl
+  Case scrut alts resultTy effs -> do
+    s <- getSubst
+    inlineAtom (CaseCtx alts resultTy effs s ctx) scrut
   expr -> visitGeneric expr >>= reconstruct ctx
 
 inlineAtom :: Emits o => Context SExpr e o -> SAtom i -> InlineM i o (e o)
@@ -340,12 +347,18 @@ instance Inlinable SBlock where
       effs' <- inline Stop effs  -- TODO Really?
       reconstruct ctx $ Block (BlockAnn ty' effs') decls' ans'
 
+inlineBlockEmits :: Emits o => Context SExpr e2 o -> SBlock i -> InlineM i o (e2 o)
+inlineBlockEmits ctx (Block _ decls ans) = do
+  inlineDecls decls $ inlineAtom ctx ans
+
 -- Still using InlineM because we may call back into inlining, and we wish to
 -- retain our output binding environment.
 reconstruct :: Emits o => Context e1 e2 o -> e1 o -> InlineM i o (e2 o)
 reconstruct ctx e = case ctx of
   Stop -> return e
   TabAppCtx ixs s ctx' -> withSubst s $ reconstructTabApp ctx' e ixs
+  CaseCtx alts resultTy effs s ctx' ->
+    withSubst s $ reconstructCase ctx' e alts resultTy effs
   EmitToAtomCtx ctx' -> emitExprToAtom e >>= reconstruct ctx'
   EmitToNameCtx ctx' -> emit (Atom e) >>= reconstruct ctx'
 {-# INLINE reconstruct #-}
@@ -404,5 +417,23 @@ reconstructTabApp ctx expr ixs =
       ixs' <- mapM (inline Stop) ixs
       reconstruct ctx $ TabApp array' 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
+      resultTy' <- inline Stop resultTy
+      effs' <- inline Stop effs
+      reconstruct ctx $ Case scrut alts' resultTy' effs'
+
 instance Inlinable (EffectRow SimpIR)
 instance Inlinable (EffectAndType SimpIR)

src/lib/Optimize.hs

@@ -7,7 +7,7 @@
 {-# LANGUAGE UndecidableInstances #-}
 
 module Optimize
-  ( optimize, peepholeOp
+  ( optimize, peepholeOp, peepholeExpr
   , hoistLoopInvariant, hoistLoopInvariantDest
   , dceTop, dceTopDest
   , foldCast ) where

src/lib/Types/Core.hs

@@ -425,6 +425,7 @@ type SAtom  = Atom SimpIR
 type SType  = Type SimpIR
 type SExpr  = Expr SimpIR
 type SBlock = Block SimpIR
+type SAlt   = Alt   SimpIR
 type SDecl  = Decl  SimpIR
 type SDecls = Decls SimpIR
 type SAtomName  = AtomName SimpIR

tests/inline-tests.dx

@@ -87,3 +87,22 @@ def id'(x:Nat) -> Nat = x
   sum (for i:(Fin 2) j:(..i). ordinal j)[ix]
 -- CHECK: 1
 -- CHECK-NOT: Compiler bug
+
+-- CHECK-LABEL: Inlining simplifies case-of-known-constructor
+"Inlining simplifies case-of-known-constructor"
+
+-- Inlining xs exposes a case-of-known-constructor opportunity here;
+-- the first inlining pass doesn't take it (yet) because it's
+-- conservative about inlining `i` into the body of `xs`, but the
+-- second pass does.
+%passes inline
+:pp
+  xs = for i:(Either (Fin 3) (Fin 4)).
+    case i of
+      Left k -> 1
+      Right k -> 2
+  for j:(Fin 3). xs[Left j]
+-- CHECK: === inline ===
+-- CHECK: case
+-- CHECK: === inline ===
+-- CHECK-NOT: case