Unit test for constant folding success

nsajko
Fri, 10 Jan 2025 08:49:05 +0000
General Usage

Here’s a proposal for how to test for successful constant folding:

using Test

function test_constant_folding(f, r)
    vec = code_typed(f, Tuple{})
    p = only(vec)
    code_info = first(p)
    code = try
        code_info.code
    catch
        nothing
    end
    @test repr(only(code)) == repr(:(return $r)) skip=(code isa Nothing)
end

# some example functions to test for foldability
function f(value = 9)
    it_fil = Iterators.filter((_ -> true), value)
    it_map = Iterators.map((x -> x + 1), it_fil)
    v = collect(it_map)
    only(v)
end
function g(collection = (2, 6))
    it_fil = Iterators.filter((_ -> true), collection)
    it_map = Iterators.map((x -> x + 1), it_fil)
    v = collect(it_map)
    sum(v)
end

@testset "constant folding" begin
    test_constant_folding((() -> 7), 7)
    test_constant_folding(f, 10)
    test_constant_folding(g, 10)
end

In the REPL:

julia> @testset "constant folding" begin
           test_constant_folding((() -> 7), 7)
           test_constant_folding(f, 10)
           test_constant_folding(g, 10)
       end
constant folding: Error During Test at REPL[2]:11
  Test threw exception
  Expression: repr(only(code)) == repr(:(return $r))
  ArgumentError: Collection has multiple elements, must contain exactly 1 element
  Stacktrace:
   [1] _only
     @ ./iterators.jl:1545 [inlined]
   [2] only(x::Vector{Any})
     @ Base.Iterators ./iterators.jl:1536
   [3] macro expansion
     @ ~/tmp/jl/jl/nightly_normal/julia-3d85309e80/share/julia/stdlib/v1.12/Test/src/Test.jl:676 [inlined]
   [4] test_constant_folding(f::Function, r::Int64)
     @ Main ./REPL[2]:11
constant folding: Error During Test at REPL[2]:11
  Test threw exception
  Expression: repr(only(code)) == repr(:(return $r))
  ArgumentError: Collection has multiple elements, must contain exactly 1 element
  Stacktrace:
   [1] _only
     @ ./iterators.jl:1545 [inlined]
   [2] only(x::Vector{Any})
     @ Base.Iterators ./iterators.jl:1536
   [3] macro expansion
     @ ~/tmp/jl/jl/nightly_normal/julia-3d85309e80/share/julia/stdlib/v1.12/Test/src/Test.jl:676 [inlined]
   [4] test_constant_folding(f::Function, r::Int64)
     @ Main ./REPL[2]:11
Test Summary:    | Pass  Error  Total  Time
constant folding |    1      2      3  2.1s
RNG of the outermost testset: Random.Xoshiro(0x2bc1148ebb7c1f60, 0x8ecccae6baa3e4fe, 0xc9fc8bc1405b0303, 0xd75d65d894272b1b, 0xada3f8f53a410b26)
ERROR: Some tests did not pass: 1 passed, 0 failed, 2 errored, 0 broken.

I think this test is supposed to work as expected as long as CodeInfo has an iterator of instructions in its :code field. When the property access fails, the test should be skipped gracefully, though.

I wonder if anyone has comments, in particular regarding reliability of the test. Are there any unexpected ways this could fail? cc @Sukera

Regarding possible alternatives, a weaker test would be to check, e.g., with Base.infer_effects or InteractiveUtils.@infer_effects (neither is public ), whether the call infers as :foldable. However that doesn’t guarantee that it’ll actually be constant folded, AFAIK. Also, there’s no public API to check whether an Effects is :foldable, even though we have the non-public Core.Compiler.is_foldable.

7 posts - 2 participants