Do another round of inlining, and io_bind

A big cost in IO heavy programs is io_bind, and we can often inline it
away and turn it into just sequencing operations. Things have to be
lined up right to do that though - ideally, case inlining and the
newtype optimisation will know just a little bit more to be able to do
it automatically, but for now, the inliner treats io_bind as a special
case.

Also do another round of inlining, since lots more things can become
inlinable (io_bind especially, becoming fully applied to the %World)
after the first pass.

libs/prelude/PrimIO.idr

@@ -29,12 +29,17 @@ prim_io_bind : (1 act : PrimIO a) -> (1 k : a -> PrimIO b) -> PrimIO b
 prim_io_bind fn k w
     = let MkIORes x' w' = fn w in k x' w'
 
+-- There's a special case for inlining this is Compiler.Inline, because
+-- the inliner is cautious about case blocks at the moment. Once it's less
+-- cautious, add an explicit %inline directive and take out the special case.
+-- See also dealing with the newtype optimisation via %World in
+-- Compiler.CompileExpr
 export
 io_bind : (1 act : IO a) -> (1 k : a -> IO b) -> IO b
-io_bind (MkIO fn)
-    = \k => MkIO (\w => let MkIORes x' w' = fn w
-                            MkIO res = k x' in
-                            res w')
+io_bind (MkIO fn) k
+    = MkIO (\w => let MkIORes x' w' = fn w
+                      MkIO res = k x' in
+                      res w')
 
 %extern prim__putStr : String -> (1 x : %World) -> IORes ()
 %extern prim__getStr : (1 x : %World) -> IORes String

src/Compiler/Common.idr

@@ -220,6 +220,12 @@ getCompileData tm_in
          logTime "Inline" $ traverse_ inlineDef cns
          logTime "Merge lambda" $ traverse_ mergeLamDef cns
          logTime "Fix arity" $ traverse_ fixArityDef cns
+         -- Do another round, since merging lambdas might expose more
+         -- optimisation opportunities, especially a really important one
+         -- for io_bind.
+         logTime "Inline" $ traverse_ inlineDef cns
+         logTime "Merge lambda" $ traverse_ mergeLamDef cns
+         logTime "Fix arity" $ traverse_ fixArityDef cns
          logTime "Forget names" $ traverse_ mkForgetDef cns
 
          compiledtm <- fixArityExp !(compileExp tycontags tm)

src/Compiler/CompileExpr.idr

@@ -352,21 +352,28 @@ mutual
                 -- outside world, so we need to evaluate to keep the
                 -- side effect.
                 Just (DCon _ arity (Just (noworld, pos))) =>
-                     if noworld -- just substitute the scrutinee into
-                                -- the RHS
-                        then let env : SubstCEnv args vars
+-- FIXME: We don't need the commented out bit *for now* because io_bind
+-- isn't being inlined, but it does need to be a little bit cleverer to
+-- get the best performance.
+-- I'm (edwinb) keeping it visible here because I plan to put it back in
+-- more or less this form once case inlining works better and the whole thing
+-- works in a nice principled way.
+--                      if noworld -- just substitute the scrutinee into
+--                                 -- the RHS
+--                         then 
+                             let env : SubstCEnv args vars
                                      = mkSubst 0 scr pos args in
                                  pure $ Just (substs env !(toCExpTree tags n sc))
-                        else -- let bind the scrutinee, and substitute the
-                             -- name into the RHS
-                             let env : SubstCEnv args (MN "eff" 0 :: vars)
-                                     = mkSubst 0 (CLocal fc First) pos args in
-                             do sc' <- toCExpTree tags n sc
-                                let scope = thin {outer=args}
-                                                 {inner=vars}
-                                                 (MN "eff" 0) sc'
-                                pure $ Just (CLet fc (MN "eff" 0) False scr
-                                                  (substs env scope))
+--                         else -- let bind the scrutinee, and substitute the
+--                              -- name into the RHS
+--                              let env : SubstCEnv args (MN "eff" 0 :: vars)
+--                                      = mkSubst 0 (CLocal fc First) pos args in
+--                              do sc' <- toCExpTree tags n sc
+--                                 let scope = thin {outer=args}
+--                                                  {inner=vars}
+--                                                  (MN "eff" 0) sc'
+--                                 pure $ Just (CLet fc (MN "eff" 0) False scr
+--                                                   (substs env scope))
                 _ => pure Nothing -- there's a normal match to do
     where
       mkSubst : Nat -> CExp vs ->

src/Compiler/Inline.idr

@@ -132,6 +132,18 @@ mutual
          List Name -> EEnv free vars -> Stack free -> CExp (vars ++ free) ->
          Core (CExp free)
   eval rec env stk (CLocal fc p) = evalLocal fc rec stk env p
+  -- This is hopefully a temporary hack, giving a special case for io_bind.
+  -- Currently the elaborator is a bit cautious about inlining case blocks
+  -- in case they duplicate work. We should fix that, to decide more accurately
+  -- whether they're safe to inline, but until then this gives such a huge
+  -- boost by removing unnecessary lambdas that we'll keep the special case.
+  eval rec env (_ :: _ :: act :: cont :: world :: stk)
+               (CRef fc (NS ["PrimIO"] (UN "io_bind")))
+      = do xn <- genName "act"
+           sc <- eval rec [] [] (CApp fc cont [CRef fc xn, world])
+           pure $ unload stk $
+                CLet fc xn False (CApp fc act [world])
+                     (refToLocal xn xn sc)
   eval rec env stk (CRef fc n)
       = do defs <- get Ctxt
            Just gdef <- lookupCtxtExact n (gamma defs)

src/TTImp/ProcessRecord.idr

@@ -144,7 +144,7 @@ elabRecord {vars} eopts fc env nest newns vis tn params conName_in fields
                    processDecl [] nest env
                        (IClaim fc (if isErased rc
                                       then erased
-                                      else top) vis [] (MkImpTy fc projNameNS projTy))
+                                      else top) vis [Inline] (MkImpTy fc projNameNS projTy))
 
                    -- Define the LHS and RHS
                    let lhs_exp

tests/idris2/basic019/expected

@@ -4,5 +4,5 @@ Main> Prelude.elem : Eq a => a -> List a -> Bool
 elem x [] = False
 elem x (y :: ys) = if x == y then True else elem x ys
 Main> PrimIO.io_bind : (1 act : IO a) -> (1 k : (a -> IO b)) -> IO b
-io_bind (MkIO fn) = \1 k => MkIO (\1 w => (case fn w of { MkIORes x' w' => case k x' of { MkIO res => res w' } }))
+io_bind (MkIO fn) k = MkIO (\1 w => (case fn w of { MkIORes x' w' => case k x' of { MkIO res => res w' } }))
 Main> Bye for now!