Expand to cover positional args as well

Author: penelopeysm

docs/src/index.md

@@ -28,5 +28,5 @@ An opt-in mechanism marks functions that might contain `Libtask.produce` stateme
 
 ```@docs; canonical=true
 Libtask.might_produce(::Type{<:Tuple})
-Libtask.@might_produce_kwargs
+Libtask.@might_produce
 ```

src/copyable_task.jl

@@ -360,11 +360,12 @@ that, by default, we assume that calls do not contain `Libtask.produce` statemen
 might_produce(::Type{<:Tuple}) = false
 
 """
-    @might_produce_kwargs(f)
+    @might_produce(f)
 
-If `f` is a function that has keyword arguments and may call `Libtask.produce` inside it,
-then `@might_produce_kwargs(f)` will generate the appropriate methods needed to ensure that
-`Libtask.might_produce` returns `true` for the relevant signatures of `f`.
+If `f` is a function that may call `Libtask.produce` inside it, then `@might_produce(f)`
+will generate the appropriate methods needed to ensure that `Libtask.might_produce` returns
+`true` for all relevant signatures of `f`. This works even if `f` has methods with keyword
+arguments.
 
 ```jldoctest kwargs
 julia> # For this demonstration we need to mark `g` as not being inlineable.
@@ -381,15 +382,18 @@ f (generic function with 1 method)
 julia> # This returns nothing because `g` isn't yet marked as being able to `produce`.
        consume(Libtask.TapedTask(nothing, f))
 
-julia> Libtask.@might_produce_kwargs(g)
+julia> Libtask.@might_produce(g)
 
 julia> # Now it works!
        consume(Libtask.TapedTask(nothing, f))
 6
 """
-macro might_produce_kwargs(f)
+macro might_produce(f)
     # See https://github.com/TuringLang/Libtask.jl/issues/197 for discussion of this macro.
     quote
+        function $(Libtask).might_produce(::Type{<:Tuple{typeof($(esc(f))),Vararg}})
+            return true
+        end
         possible_n_kwargs = unique(map(length ∘ Base.kwarg_decl, methods($(esc(f)))))
         if possible_n_kwargs != [0]
             # Oddly we need to interpolate the module and not the function: either

test/copyable_task.jl

@@ -1,3 +1,15 @@
+module Functions
+using Libtask: produce
+@noinline g3(x) = produce(x)
+@noinline g3(x, y; z) = produce(x + y + z)
+@noinline g3(x, y, z; p, q) = produce(x + y + z + p + q)
+function f3(x)
+    g3(x)
+    g3(x, 1; z=2)
+    g3(x, 1, 2; p=3, q=4)
+end
+end
+
 @testset "copyable_task" begin
     for case in Libtask.TestUtils.test_cases()
         case()
@@ -251,4 +263,46 @@
         @test Libtask.consume(tt) === :a
         @test Libtask.consume(tt) === nothing
     end
+
+    @testset "@might_produce macro" begin
+        # Positional arguments only
+        @noinline g1(x) = produce(x)
+        f1(x) = g1(x)
+        # Without marking it as might_produce
+        tt = Libtask.TapedTask(nothing, f1, 0)
+        @test Libtask.consume(tt) === nothing
+        # Now marking it
+        Libtask.@might_produce(g1)
+        tt = Libtask.TapedTask(nothing, f1, 0)
+        @test Libtask.consume(tt) === 0
+        @test Libtask.consume(tt) === nothing
+
+        # Keyword arguments only
+        @noinline g2(x; y=1, z=2) = produce(x + y + z)
+        f2(x) = g2(x)
+        # Without marking it as might_produce
+        tt = Libtask.TapedTask(nothing, f2, 0)
+        @test Libtask.consume(tt) === nothing
+        # Now marking it
+        Libtask.@might_produce(g2)
+        tt = Libtask.TapedTask(nothing, f2, 0)
+        @test Libtask.consume(tt) === 3
+        @test Libtask.consume(tt) === nothing
+
+        # A function with multiple methods.
+        # Note: f3 and g3 are defined in the module at the top of this file. If
+        # they are defined directly in this testset, for reasons that are
+        # unclear, the `produce` calls are picked up even without using the
+        # `@might_produce` macro, meaning that it's impossible to test that the
+        # macro is having the intended behaviour.
+        tt = Libtask.TapedTask(nothing, Functions.f3, 0)
+        @test Libtask.consume(tt) === nothing
+        # Now marking it
+        Libtask.@might_produce(Functions.g3)
+        tt = Libtask.TapedTask(nothing, Functions.f3, 0)
+        @test Libtask.consume(tt) === 0
+        @test Libtask.consume(tt) === 3
+        @test Libtask.consume(tt) === 10
+        @test Libtask.consume(tt) === nothing
+    end
 end