isFunctionType in PlutusIR.Transform.ThunkRecursions misses a case for polymorphic functions

[edsko]

Summary

isFunctionType in PlutusIR.Transform.ThunkRecursions misses a case for polymorphic functions. This means that the thunkRecursions pass will create a non-strict let-binding for such functions, which in the subsequent compileNonStrictBindings True will then give an unnecessary (Scott-encoded) unit argument. This matters, because this unit argument then makes it all the way through to the untyped Plutus core, resulting in larger than necessary core (indeed, Section 7 of "Formal Specification of the Plutus Core Language" explicitly states that the reason force/delay were introduced in the first place is the size of these unit arguments).

Steps to reproduce the behavior

Consider

newtype Reverse = Reverse (forall a. [a] -> [a])

reversePoly :: [a] -> [a]
reversePoly = reversePoly' []

reversePoly' :: [a] -> [a] -> [a]
reversePoly' acc []     = acc
reversePoly' acc (x:xs) = reversePoly' (x:acc) xs

reverseMono :: [Integer] -> [Integer]
reverseMono = reverseMono' []

reverseMono' :: [Integer] -> [Integer] -> [Integer]
reverseMono' acc []     = acc
reverseMono' acc (x:xs) = reverseMono' (x:acc) xs

compiledReversePoly :: CompiledCode (Reverse)
compiledReversePoly = $$(compile [|| Reverse reversePoly ||])

compiledReverseMono :: CompiledCode ([Integer] -> [Integer])
compiledReverseMono = $$(compile [|| reverseMono ||])

(the newtype is merely there to give CompiledCode a monomorphic type). If we trace the compilation of the monomorphic reverse through the pipeline, just before thunkRecursions we have

let rec data List :: * -> * a = Nil  | Cons a (List a)
in let rec !reverseMono = λ(acc :: List Integer) (ds :: List Integer) -> Nil_match {Integer} ds {∀ _. List Integer} (λ{_} -> acc) (λ(x :: Integer) (xs :: List Integer) {_} -> reverseMono (Cons {Integer} x acc) xs) {_}
   in reverseMono (Nil {Integer})

(this is using my custom pretty-printer, excuse the non-standard syntax); the subsequent thunkRecursions pass then does not affect reverseMono binding at all, and consequently compileNonStrictBindings True doesn't do anything either. In the polymorphic case we see instead

let rec data List :: * -> * a = Nil  | Cons a (List a)
in let rec !reversePoly = λ{a} (acc :: List a) (ds :: List a) -> Nil_match {a} ds {∀ _. List a} (λ{_} -> acc) (λ(x :: a) (xs :: List a) {_} -> reversePoly {a} (Cons {a} x acc) xs) {_}
   in λ{a} -> reversePoly {a} (Nil {a})

which then by thunkRecursions is turned into

let rec data List :: * -> * a = Nil  | Cons a (List a)
in let rec reversePoly = λ{a} (acc :: List a) (ds :: List a) -> Nil_match {a} ds {∀ _. List a} (λ{_} -> acc) (λ(x :: a) (xs :: List a) {_} -> reversePoly {a} (Cons {a} x acc) xs) {_}
   in λ{a} -> reversePoly {a} (Nil {a})

and then by compileNonStrictBindings True into

let rec data List :: * -> * a = Nil  | Cons a (List a)
in let rec !reversePoly = λ(arg :: ∀ a. a -> a) {a} (acc :: List a) (ds :: List a) -> Nil_match {a} ds {∀ _. List a} (λ{_} -> acc) (λ(x :: a) (xs :: List a) {_} -> reversePoly (λ{a} (x :: a) -> x) {a} (Cons {a} x acc) xs) {_}
   in λ{a} -> reversePoly (λ{a} (x :: a) -> x) {a} (Nil {a})

where we see that Scott-encoded unit argument being introduced. This then makes it all the way to the core.

Actual Result

The untyped Plutus core for reversePoly is

let* s = \ s_0 x -> let* reversePoly = s_0 s_0 in \ ~ acc ds -> ds ! (\ ~ -> acc) (\ x_0 xs ~ -> reversePoly (\ ~ x_1 -> x_1) ! (\ ~ case_Nil case_Cons -> case_Cons x_0 acc) xs) !
     reversePoly_0 = s s
in \ ~ -> reversePoly_0 (\ ~ x_2 -> x_2) ! (\ ~ case_Nil_0 case_Cons_0 -> case_Nil_0)

Expected Result

The untyped Plutus core for reverseMono is

reverseMono =
let* s = \ s_0 x -> let* reverseMono = s_0 s_0 in \ ds -> ds ! (\ ~ -> x) (\ x_0 xs ~ -> reverseMono (\ ~ case_Nil case_Cons -> case_Cons x_0 x) xs) !
     reverseMono_0 = s s
in reverseMono_0 (\ ~ case_Nil_0 case_Cons_0 -> case_Nil_0)

which is significantly smaller. Some size difference is perhaps expected (polymorphism leading to force/delay constructors elsewhere), but that is orthogonal to this ticket. The extra (\ ~ x_2 -> x_2) argument in the polymorphic version should not be needed.

Describe the approach you would take to fix this

No response

System info

The above is tested with 4127e9c , though I see that isFunctionType has not changed in latest master.

[michaelpj]

Unfortunately it really is the case that our fixpoint combinators can only handle things of function type (ultimately they're variants on the z combinator), and polymorphic types are not function types. See https://github.com/input-output-hk/plutus/blob/master/plutus-core/plutus-ir/src/PlutusIR/Transform/ThunkRecursions.hs#L16 .

We could implement the other approach described in that note, namely trying to pull the polymorphism out of the recursion. That could work as a standalone optimization. But it's unclear how generically we'd be able to do that (what about mutually recursive functions? etc. etc.), so I don't know if it's worth it.

[effectfully]

Unfortunately it really is the case that our fixpoint combinators can only handle things of function type (ultimately they're variants on the z combinator)

It's fixable, though! We didn't realize that when we were writing the Unraveling paper, but I did update the doc on mutual term-level recursion some time afterwards.

[effectfully]

Since the behavior of the function is intentional, I've removed the bug label. However it would be great to handle polymorphism properly for performance reasons, so I've added the Performance label.

[michaelpj]

There's an internal JIRA ticket for this also, PLT-717.