plutus
plutus copied to clipboard
`isFunctionType` in `PlutusIR.Transform.ThunkRecursions` misses a case for polymorphic functions
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 4127e9cd6e889824d724c30eae55033cb50cbf3e , though I see that isFunctionType has not changed in latest master.
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.
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.
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.
There's an internal JIRA ticket for this also, PLT-717.