Simplify binary lambdas with intermediate decls

We used to make a simplifying assumption that the binary lambda
arguments to effect handlers have no decls in between the first and
second lambda binder. However, it's been easy to violate that
assumption, because inference likes to insert decls while unpacking
patterns. This meant that accumulating over a monoid such as
`AddMonoid Complex` caused a simplification error, because of `Complex`
argument unpacking.

This change relaxes the constraint slightly, to allow for the
problematic decls. However, it's still not a complete solution as it
assumes that the second binder and the eventual reconstruction don't
depend on those decls. This seems sufficient for now and we can always
revise that in the future.

src/lib/Builder.hs

@@ -35,7 +35,7 @@ module Builder (
   TopEnvFrag (..), emitPartialTopEnvFrag, emitLocalModuleEnv,
   inlineLastDecl, fabricateEmitsEvidence, fabricateEmitsEvidenceM,
   singletonBinderNest, varsAsBinderNest, typesAsBinderNest,
-  liftBuilder, liftEmitBuilder, makeBlock,
+  liftBuilder, liftEmitBuilder, makeBlock, makeBlockOfType,
   indexToInt, indexSetSize, intToIndex,
   getIxImpl, IxImpl (..),
   litValToPointerlessAtom, emitPtrLit,
@@ -438,8 +438,12 @@ buildBlock cont = do
 makeBlock :: EnvReader m => Nest Decl n l -> Expr l -> m l (Block n)
 makeBlock decls expr = do
   ty <- {-# SCC blockTypeNormalization #-} cheapNormalize =<< getType expr
-  let ty' = ignoreHoistFailure $ hoist decls ty
-  return $ Block (BlockAnn ty') decls expr
+  return $ makeBlockOfType decls expr ty
+{-# INLINE makeBlock #-}
+
+makeBlockOfType :: Nest Decl n l -> Expr l -> Type l -> Block n
+makeBlockOfType decls expr ty = Block (BlockAnn ty') decls expr
+  where ty' = ignoreHoistFailure $ hoist decls ty
 
 inlineLastDecl :: Nest Decl n l -> Expr l -> Abs (Nest Decl) Expr n
 inlineLastDecl Empty result = Abs Empty result

src/lib/Simplify.hs

@@ -381,17 +381,27 @@ type BinaryLamBinder = (PairB LamBinder LamBinder)
 simplifyBinaryLam :: Emits o => Atom i
   -> SimplifyM i o (Atom o, Abs BinaryLamBinder ReconstructAtom o)
 simplifyBinaryLam atom = case atom of
-  Lam (LamExpr b1 (AtomicBlock (Lam (LamExpr b2 body)))) -> doSimpBinaryLam b1 b2 body
+  Lam (LamExpr b1 (Block _ body1 (Atom (Lam (LamExpr b2 body2))))) -> doSimpBinaryLam b1 body1 b2 body2
   _ -> simplifyAtom atom >>= \case
-    Lam (LamExpr b1 (AtomicBlock (Lam (LamExpr b2 body)))) -> dropSubst $ doSimpBinaryLam b1 b2 body
+    Lam (LamExpr b1 (Block _ body1 (Atom (Lam (LamExpr b2 body2))))) -> dropSubst $ doSimpBinaryLam b1 body1 b2 body2
     _ -> error "Not a binary lambda expression"
   where
-    doSimpBinaryLam :: LamBinder i i' -> LamBinder i' i'' -> Block i''
+    doSimpBinaryLam :: LamBinder i i' -> Nest Decl i' i'' -> LamBinder i'' i''' -> Block i'''
       -> SimplifyM i o (Atom o, Abs BinaryLamBinder ReconstructAtom o)
-    doSimpBinaryLam b1 b2 body = do
-      (Abs (b1' `PairB` b2') body', recon) <- simplifyAbs $ Abs (b1 `PairB` b2) body
-      let binaryLam' = Lam $ LamExpr b1' $ AtomicBlock $ Lam $ LamExpr b2' body'
-      return (binaryLam', recon)
+    doSimpBinaryLam b1 body1 b2 body2 =
+      substBinders b1 \b1' -> do
+        Abs decls (lam2 `PairE` lam2Ty `PairE` (Abs b2' recon')) <- buildScoped $
+          simplifyDecls body1 do
+            (Abs b2' body2', recon) <- simplifyAbs $ Abs b2 body2
+            let lam2' = Lam (LamExpr b2' body2')
+            lam2Ty' <- getType lam2'
+            return (lam2' `PairE` lam2Ty' `PairE` recon)
+        return $ case hoist decls $ Abs b2' recon' of
+          HoistSuccess (Abs b2'' recon'') -> do
+            let binBody = makeBlockOfType decls (Atom lam2) lam2Ty
+            let binRecon = Abs (b1' `PairB` b2'') recon''
+            (Lam (LamExpr b1' binBody), binRecon)
+          HoistFailure _ -> error "Binary lambda simplification failed: binder/recon depends on intermediate decls"
 
 data SplitDataNonData n = SplitDataNonData
   { dataTy    :: Type n