Fix computeArity.

plutus-core/plutus-ir/src/PlutusIR/Transform/Inline/CallSiteInline.hs

@@ -14,6 +14,7 @@ module PlutusIR.Transform.Inline.CallSiteInline where
 
 import Control.Lens (forMOf)
 import Control.Monad.State
+import Data.Bifunctor (first)
 import PlutusIR.Core
 import PlutusIR.Transform.Inline.Utils
 
@@ -149,8 +150,8 @@ computeArity ::
     Term tyname name uni fun ann
     -> (Arity, Term tyname name uni fun ann)
 computeArity = \case
-    LamAbs _ _ _ body -> (MkTerm : fst (computeArity body), body)
-    TyAbs _ _ _ body  -> (MkType : fst (computeArity body), body)
+    LamAbs _ _ _ body -> (first ((:) MkTerm) (computeArity body))
+    TyAbs _ _ _ body  -> (first ((:) MkType) (computeArity body))
     -- Whenever we encounter a body that is not a lambda or type abstraction, we are done counting
     tm                -> ([],tm)
 
@@ -237,9 +238,10 @@ inlineSat appTerm@(Apply _varAnn _fun _arg) = do
           -- Have checked that with the current tests it doesn't happen
           Nothing -> forMOf termSubterms appTerm inlineSat
           Just varInfo -> do
-            if enoughArgs (arity varInfo) (map fst args) then
+            if enoughArgs (arity varInfo) (map fst args) then do
               -- if the `Var` is fully applied (over-application is allowed) then inline it
-              pure $ mkApps (calledVarDef varInfo) args
+              let inlinedTm = mkApps (calledVarDef varInfo) args
+              forMOf termSubterms inlinedTm inlineSat
             -- otherwise just keep going
             else forMOf termSubterms appTerm inlineSat
           -- if the term being applied is not a `Var`, don't inline, but keep checking
@@ -255,33 +257,18 @@ inlineSat tyInstTerm@(TyInst _varAnn _fun _arg) = do
         case maybeVarInfo of
           Nothing -> forMOf termSubterms tyInstTerm inlineSat
           Just varInfo -> do
-            if enoughArgs (arity varInfo) (map fst args) then
+            if enoughArgs (arity varInfo) (map fst args) then do
               -- if the `Var` is fully applied (over-application is allowed) then inline it
-              pure $ mkApps (calledVarDef varInfo) args
+              let inlinedTm = mkApps (calledVarDef varInfo) args
+              forMOf termSubterms inlinedTm inlineSat
             -- otherwise just keep going
             else forMOf termSubterms tyInstTerm inlineSat
           -- if the term being applied is not a `Var`, don't inline but keep checking the subterms
       v -> forMOf termSubterms v inlineSat
-inlineSat letTm@(Let _ _ bds _letBody) = do
+inlineSat letTm@(Let _ _ _bds _letBody) =
     -- recursive or not, the bindings of this let term *may* contain a saturated function,
     -- so we need to check all the bindings and also the body
     -- `PlutusIR.Core.Plated.termSubterms` gives all that
     forMOf termSubterms letTm inlineSat
-inlineSat (TyAbs _ _ _ tyAbsBody) =
-  -- start count in the body of the type lambda abstraction
-  inlineSat tyAbsBody
-inlineSat (LamAbs _ _ _ fnBody) =
-  -- start the count in the body of the term lambda abstraction
-  inlineSat fnBody
-inlineSat con@(Constant _ _) =
-    -- constants cannot call the variable
-    pure con
-inlineSat bi@(Builtin _ _) =
-  -- default builtin functions in `PlutusCore/Default/Builtins.hs`
-  -- cannot call the variable
-  pure bi
-inlineSat v@(Var _ _) =
-  -- variables being applied should have been checked already, these ones aren't fully applied.
-  -- We don't inline them.
-  pure v
-inlineSat others = pure others
+inlineSat tm =
+  forMOf termSubterms tm inlineSat

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

@@ -214,7 +214,7 @@ inline =
 computeArityTest :: TestNested
 computeArityTest = testNested "computeArityTest" $
         map
-            (goldenPir (fst . computeArity . runQuote . PLC.rename) pTerm)
+            (goldenPir (computeArity . runQuote . PLC.rename) pTerm)
             [ "var" -- from inline tests, testing let terms
             , "tyvar"
             , "single"

plutus-core/plutus-ir/test/transform/computeArityTest/absapp.golden

@@ -1 +1 @@
-[]
+([], { (abs a (type) (lam x a x)) (con integer) })

plutus-core/plutus-ir/test/transform/computeArityTest/ifError.golden

@@ -1 +1,16 @@
-[MkTerm]
+( [MkTerm]
+, [
+  [
+    { b (con integer) }
+    (lam
+      z
+      (con integer)
+      (let
+        (nonrec)
+        (termbind (strict) (vardecl err (con integer)) (error (con integer)))
+        err
+      )
+    )
+  ]
+  (lam z (con integer) (con integer 1))
+] )

plutus-core/plutus-ir/test/transform/computeArityTest/immediateVar.golden

@@ -1 +1,4 @@
-[]
+( []
+, (let
+  (nonrec) (termbind (strict) (vardecl x (con integer)) (error (con integer))) x
+) )

plutus-core/plutus-ir/test/transform/computeArityTest/lamAbs3.golden

@@ -1 +1 @@
-[MkTerm, MkTerm, MkTerm]
+([MkTerm, MkTerm, MkTerm], x)

plutus-core/plutus-ir/test/transform/computeArityTest/lamAbsApp.golden

@@ -1 +1 @@
-[MkTerm, MkTerm, MkTerm]
+([MkTerm, MkTerm, MkTerm], [ (lam k (con integer) x) (con integer 3) ])

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

@@ -1 +1,15 @@
-[]
+( []
+, [
+  [
+    [
+      (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/computeArityTest/multilet.golden

@@ -1 +1,14 @@
-[]
+( []
+, [
+  (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/computeArityTest/single.golden

@@ -1 +1,62 @@
-[]
+( []
+, (let
+  (nonrec)
+  (termbind (strict) (vardecl x (con integer)) (con integer 42))
+  (termbind
+    (strict)
+    (vardecl simple (con integer))
+    (let
+      (nonrec)
+      (termbind
+        (strict)
+        (vardecl f (fun (con integer) (con integer)))
+        (lam y (con integer) y)
+      )
+      [ f (con integer 1) ]
+    )
+  )
+  (termbind
+    (strict)
+    (vardecl insideLambda (fun (con integer) (con integer)))
+    (let
+      (nonrec)
+      (termbind
+        (strict)
+        (vardecl f (fun (con integer) (con integer)))
+        (lam y (con integer) y)
+      )
+      (lam y (con integer) [ f y ])
+    )
+  )
+  (termbind
+    (strict)
+    (vardecl trivialLambda (con integer))
+    (let
+      (nonrec)
+      (termbind
+        (strict)
+        (vardecl f (fun (con integer) (con integer)))
+        (lam y (con integer) y)
+      )
+      [ f [ f (con integer 1) ] ]
+    )
+  )
+  (termbind
+    (strict)
+    (vardecl variableCapture (con integer))
+    (let
+      (nonrec)
+      (termbind
+        (strict)
+        (vardecl f (fun (con integer) (con integer)))
+        (lam y (con integer) x)
+      )
+      (let
+        (nonrec)
+        (termbind (strict) (vardecl x (con integer)) (con integer 24))
+        [ f x ]
+      )
+    )
+  )
+  x
+) )

plutus-core/plutus-ir/test/transform/computeArityTest/tyAbs.golden

@@ -1 +1 @@
-[MkType, MkTerm]
+([MkType, MkTerm], x)

plutus-core/plutus-ir/test/transform/computeArityTest/tyAbs2.golden

@@ -1 +1 @@
-[MkType, MkType, MkTerm, MkTerm]
+([MkType, MkType, MkTerm, MkTerm], y)

plutus-core/plutus-ir/test/transform/computeArityTest/tyAbs2Arrow.golden

@@ -1 +1 @@
-[MkType, MkType, MkTerm, MkTerm]
+([MkType, MkType, MkTerm, MkTerm], [ y x ])

plutus-core/plutus-ir/test/transform/computeArityTest/tyAbsInterleaved.golden

@@ -1 +1 @@
-[MkType, MkTerm, MkType, MkTerm, MkTerm]
+([MkType, MkTerm, MkType, MkTerm, MkTerm], y)

plutus-core/plutus-ir/test/transform/computeArityTest/tyvar.golden

@@ -1 +1,11 @@
-[]
+( []
+, (let
+  (nonrec)
+  (typebind (tyvardecl a (type)) (con integer))
+  (typebind (tyvardecl b (type)) (con integer))
+  (let
+    (nonrec)
+    (termbind (strict) (vardecl y b) (con integer 1))
+    (lam z a [ (lam x b x) y ])
+  )
+) )

plutus-core/plutus-ir/test/transform/computeArityTest/var.golden

@@ -1 +1 @@
-[MkTerm]
+([MkTerm], (let (nonrec) (termbind (strict) (vardecl x (con integer)) y) x))