Teach the PIR beta pass to handle multiple arguments at once

This has bothered me for a while. We rely on this to turn function
applications into let-bindings so that we can inline them, but because
of the way it was written it could only ever do this for at most one
argument at a time! So we could only fully-inline a function of arity N
with N rounds of the simplifier! Terrible!

Now it handles any number at once.

I wrote all the functions for this while working on
#4365, since there we
*really* need it, since constructors and destructors appear like
multiple function arguments rather than multiple nested
immediately-applied-lambdas. But it applies well here too, so I thought
I'd just do this improvement on the side.

Author: michaelpj

plutus-core/plutus-ir/src/PlutusIR/Core/Type.hs

@@ -18,7 +18,8 @@ module PlutusIR.Core.Type (
     Binding (..),
     Term (..),
     Program (..),
-    applyProgram
+    applyProgram,
+    termAnn
     ) where
 
 import PlutusPrelude
@@ -158,3 +159,17 @@ applyProgram
     -> Program tyname name uni fun a
     -> Program tyname name uni fun a
 applyProgram (Program a1 t1) (Program a2 t2) = Program (a1 <> a2) (Apply mempty t1 t2)
+
+termAnn :: Term tynam name uni fun a -> a
+termAnn t = case t of
+  Let a _ _ _    -> a
+  Var a _        -> a
+  TyAbs a _ _ _  -> a
+  LamAbs a _ _ _ -> a
+  Apply a _ _    -> a
+  Constant a _   -> a
+  Builtin a _    -> a
+  TyInst a _ _   -> a
+  Error a _      -> a
+  IWrap a _ _ _  -> a
+  Unwrap a _     -> a

plutus-core/plutus-ir/src/PlutusIR/Transform/Beta.hs

@@ -1,39 +1,84 @@
-{-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE LambdaCase   #-}
+{-# LANGUAGE ViewPatterns #-}
 {-|
 A simple beta-reduction pass.
 -}
 module PlutusIR.Transform.Beta (
   beta
   ) where
 
-import PlutusPrelude
-
 import PlutusIR
+import PlutusIR.Core.Type
 
-import Control.Lens (transformOf)
+import Control.Lens.Setter ((%~))
+import Data.Function ((&))
+import Data.List.NonEmpty qualified as NE
 
-{-|
-A single non-recursive application of the beta rule.
+{- Note [Multi-beta]
+Consider two examples where applying beta should be helpful.
+
+1: [(\x . [(\y . t) b]) a]
+2: [[(\x . (\y . t)) a] b]
+
+(1) is the typical "let-binding" pattern: each binding corresponds to an immediately-applied lambda.
+(2) is the typical "function application" pattern: a multi-argument function applied to multiple arguments.
+
+In both cases we would like to produce something like
+
+let
+  x = a
+  y = b
+in t
+
+However, if we naively do a bottom-up pattern-matching transformation on the AST
+to look for immediately-applied lambda abstractions then we will get the following:
+
+1:
+  [(\x . [(\y . t) b]) a]
+  -->
+  [(\x . let y = b in t) a]
+  ->
+  let x = a in let y = b in t
+
+2:
+  [[(\x . (\y . t)) a] b]
+  -->
+  [(let x = a in (\y . t)) b]
+
+Now, if we later lift the let out, then we will be able to see that we can transform (2) further.
+But that means that a) we'd have to do the expensive let-floating pass in every iteration of the simplifier, and
+b) we can only inline one function argument per iteration of the  simplifier, so for a function of
+arity N we *must* do at least N passes.
+
+This isn't great, so the solution is to recognize case (2) properly and handle all the arguments in one go.
+That will also match cases like (1) just fine, since it's just made up of unary function applications.
+
+That does mean that we need to do a manual traversal rather than doing standard bottom-up processing.
 -}
-betaStep
-    :: Term tyname name uni fun a
-    -> Term tyname name uni fun a
-betaStep = \case
-    Apply a (LamAbs _ name typ body) arg ->
-        let varDecl  = VarDecl a name typ
-            binding  = TermBind a Strict varDecl arg
-            bindings = binding :| []
-        in
-            Let a NonRec bindings body
-    -- This case is disabled as it introduces a lot of type inlining (determined from profiling)
-    -- and is currently unsound https://input-output.atlassian.net/browse/SCP-2570.
-    -- TyInst a (TyAbs _ tyname kind body) typ ->
-    --     let tyVarDecl = TyVarDecl a tyname kind
-    --         tyBinding = TypeBind a tyVarDecl typ
-    --         bindings  = tyBinding :| []
-    --     in
-    --         Let a NonRec bindings body
-    t -> t
+
+{-| Extract the list of bindings from a term, a bit like a "multi-beta" reduction.
+
+Some examples will help:
+
+[(\x . t) a] -> Just ([x |-> a], t)
+
+[[[(\x . (\y . (\z . t))) a] b] c] -> Just ([x |-> a, y |-> b, z |-> c]) t)
+
+[[(\x . t) a] b] -> Nothing
+
+When we decide that we want to do beta for types, we will need to extend this to handle type instantiations too.
+-}
+extractBindings :: Term tyname name uni fun a -> Maybe (NE.NonEmpty (Binding tyname name uni fun a), Term tyname name uni fun a)
+extractBindings = collectArgs []
+  where
+      collectArgs argStack (Apply _ f arg) = collectArgs (arg:argStack) f
+      collectArgs argStack t               = matchArgs argStack [] t
+      matchArgs (arg:rest) acc (LamAbs a n ty body) = matchArgs rest (TermBind a Strict (VarDecl a n ty) arg:acc) body
+      matchArgs []         acc t                    =
+          case NE.nonEmpty (reverse acc) of
+              Nothing   -> Nothing
+              Just acc' -> Just (acc', t)
+      matchArgs (_:_)      _   _                    = Nothing
 
 {-|
 Recursively apply the beta transformation on the code, both for the terms
@@ -57,4 +102,17 @@ and types
 beta
     :: Term tyname name uni fun a
     -> Term tyname name uni fun a
-beta = transformOf termSubterms betaStep
+beta = \case
+    -- See Note [Multi-beta]
+    -- This maybe isn't the best annotation for this term, but it will do.
+    (extractBindings -> Just (bs, t)) -> Let (termAnn t) NonRec bs (beta t)
+    -- This case is disabled as it introduces a lot of type inlining (determined from profiling)
+    -- and is currently unsound https://input-output.atlassian.net/browse/SCP-2570.
+    -- TyInst a (TyAbs _ tyname kind body) typ ->
+    --     let tyVarDecl = TyVarDecl a tyname kind
+    --         tyBinding = TypeBind a tyVarDecl typ
+    --         bindings  = tyBinding :| []
+    --     in
+    --         Let a NonRec bindings body
+
+    t                                 -> t & termSubterms %~ beta

plutus-core/plutus-ir/test/TransformSpec.hs

@@ -146,6 +146,8 @@ beta =
     $ map (goldenPir (Beta.beta . runQuote . PLC.rename) $ term @PLC.DefaultUni @PLC.DefaultFun)
     [ "lamapp"
     , "absapp"
+    , "multiapp"
+    , "multilet"
     ]
 
 unwrapCancel :: TestNested

plutus-core/plutus-ir/test/recursion/stupidZero.golden

@@ -84,29 +84,17 @@
                               x_i0
                               Nat_i6
                               [
-                                [
-                                  (lam
-                                    stupidZero_i0
-                                    (fun Nat_i7 Nat_i7)
+                                (lam
+                                  stupidZero_i0
+                                  (fun Nat_i7 Nat_i7)
+                                  [
+                                    [ { [ match_Nat_i4 x_i2 ] Nat_i7 } Zero_i6 ]
                                     (lam
-                                      n_i0
-                                      Nat_i8
-                                      [
-                                        [
-                                          { [ match_Nat_i5 n_i1 ] Nat_i8 }
-                                          Zero_i7
-                                        ]
-                                        (lam
-                                          pred_i0
-                                          Nat_i9
-                                          [ stupidZero_i3 pred_i1 ]
-                                        )
-                                      ]
+                                      pred_i0 Nat_i8 [ stupidZero_i2 pred_i1 ]
                                     )
-                                  )
-                                  [ (unwrap s_i2) s_i2 ]
-                                ]
-                                x_i1
+                                  ]
+                                )
+                                [ (unwrap s_i2) s_i2 ]
                               ]
                             )
                           )

plutus-core/plutus-ir/test/transform/beta/multiapp

@@ -0,0 +1 @@
+[ (lam x (con integer) (lam y (con integer) (lam z (con integer) [ y x z ]))) (con integer 1) (con integer 2) (con integer 3) ]

plutus-core/plutus-ir/test/transform/beta/multiapp.golden

@@ -0,0 +1,7 @@
+(let
+  (nonrec)
+  (termbind (strict) (vardecl x (con integer)) (con integer 1))
+  (termbind (strict) (vardecl y (con integer)) (con integer 2))
+  (termbind (strict) (vardecl z (con integer)) (con integer 3))
+  [ [ y x ] z ]
+)

plutus-core/plutus-ir/test/transform/beta/multilet

@@ -0,0 +1 @@
+[ (lam x (con integer) [ (lam y (con integer) [ (lam z (con integer) [ y x z ]) (con integer 3) ] ) (con integer 2) ]) (con integer 1)]

plutus-core/plutus-ir/test/transform/beta/multilet.golden

@@ -0,0 +1,13 @@
+(let
+  (nonrec)
+  (termbind (strict) (vardecl x (con integer)) (con integer 1))
+  (let
+    (nonrec)
+    (termbind (strict) (vardecl y (con integer)) (con integer 2))
+    (let
+      (nonrec)
+      (termbind (strict) (vardecl z (con integer)) (con integer 3))
+      [ [ y x ] z ]
+    )
+  )
+)