Fix: careful with linenumber removal in order to improve code coverage

Project.toml

@@ -1,6 +1,6 @@
 name = "TensorOperations"
 uuid = "6aa20fa7-93e2-5fca-9bc0-fbd0db3c71a2"
-version = "5.6.2"
+version = "5.6.3"
 authors = ["Lukas Devos <lukas.devos@ugent.be>", "Maarten Van Damme <maartenvd1994@gmail.com>", "Jutho Haegeman <jutho.haegeman@ugent.be>"]
 
 [deps]

docs/src/man/implementation.md

@@ -77,13 +77,17 @@ objects.
 
 ```@docs
 TensorOperations._flatten
-TensorOperations.removelinenumbernode
 TensorOperations.addtensoroperations
 TensorOperations.insertargument
 TensorOperations.insertbackend
 TensorOperations.insertallocator
 ```
 
+Finally, after all postprocessors have run, the parser strips the `LineNumberNode`s that it
+synthesized itself, while preserving the ones that came from the user's code. As a result the
+generated code stays attributable to the lines of the original `@tensor` expression, which is
+what makes those lines show up in code coverage reports and in stacktraces.
+
 ## Analysis of contraction graphs and optimizing contraction order
 
 The macro [`@tensoropt`](@ref) or the combination of [`@tensor`](@ref) with the keyword

ext/TensorOperationsBumperExt.jl

@@ -27,16 +27,14 @@ function TensorOperations._butensor(src, ex...)
     buf_sym = gensym("buffer")
 
     # TODO: there is no check for doubled tensor kwargs
-    newex = quote
-        $buf_sym = $(Expr(:call, GlobalRef(Bumper, :default_buffer)))
-        $(
-            Expr(
-                :macrocall, GlobalRef(TensorOperations, Symbol("@tensor")),
-                src, :(allocator = $buf_sym), ex...
-            )
+    return Expr(
+        :block,
+        Expr(:(=), buf_sym, Expr(:call, GlobalRef(Bumper, :default_buffer))),
+        Expr(
+            :macrocall, GlobalRef(TensorOperations, Symbol("@tensor")),
+            src, :(allocator = $buf_sym), ex...
         )
-    end
-    return Base.remove_linenums!(newex)
+    )
 end
 
 end

src/indexnotation/contractiontrees.jl

@@ -161,7 +161,7 @@ function insertcontractiontrees!(
             end
         )
     end
-    push!(postexprs, removelinenumbernode(costcompareex))
+    push!(postexprs, costcompareex)
     return treeex
 end
 

src/indexnotation/parser.jl

@@ -13,7 +13,7 @@ mutable struct TensorParser
         contractiontreebuilder = defaulttreebuilder
         contractiontreesorter = defaulttreesorter
         contractioncostcheck = nothing
-        postprocessors = [_flatten, removelinenumbernode, addtensoroperations]
+        postprocessors = [_flatten, addtensoroperations]
         return new(
             preprocessors,
             contractiontreebuilder, contractiontreesorter, contractioncostcheck,
@@ -24,6 +24,9 @@ end
 
 function (parser::TensorParser)(ex::Expr)
     verifytensorexpr(ex)
+    # any `LineNumberNode` present here belongs to the user's code: record its file so that the
+    # ones synthesized further down can be told apart and removed again at the very end
+    userfiles = linenumberfiles(ex)
     for p in parser.preprocessors
         ex = p(ex)::Expr
     end
@@ -35,6 +38,8 @@ function (parser::TensorParser)(ex::Expr)
     for p in parser.postprocessors
         ex = p(ex)::Expr
     end
+    # this has to happen after all (possibly user-supplied) postprocessors have run
+    ex = removeinternallinenumbernodes(ex, userfiles)::Expr
     return ex
 end
 

src/indexnotation/postprocessors.jl

@@ -31,19 +31,60 @@ function _flatten(ex)
 end
 
 """
-    removelinenumbernode(ex)
+    linenumberfiles(ex, files = Set{Symbol}())
 
-Remove all `LineNumberNode`s from an expression.
+Collect the set of files referenced by the `LineNumberNode`s in `ex`.
+
+This is used on the expression that is handed to a [`TensorParser`](@ref), before any
+processing takes place, to determine which `LineNumberNode`s belong to the user's code:
+see [`removeinternallinenumbernodes`](@ref).
+"""
+function linenumberfiles(ex, files = Set{Symbol}())
+    if ex isa LineNumberNode
+        push!(files, ex.file)
+    elseif ex isa Expr
+        foreach(e -> linenumberfiles(e, files), ex.args)
+    end
+    return files
+end
+
+"""
+    removeinternallinenumbernodes(ex, userfiles)
+
+Remove the `LineNumberNode`s that were synthesized by the parser, i.e. the ones whose file is
+not in `userfiles`, as obtained from [`linenumberfiles`](@ref) on the original expression.
+
+`LineNumberNode`s originating from user code are kept, so that the generated code remains
+attributable to the user's source lines. This matters for code coverage: Julia only emits a
+coverage counter for a line that a `LineNumberNode` points at, so stripping the user's
+`LineNumberNode`s leaves every statement of a `@tensor begin ... end` block after the first one
+without any coverage information at all. Conversely, a synthesized `LineNumberNode` would
+re-attribute all statements that follow it to a line in the parser's own source, so both halves
+are needed.
 """
-function removelinenumbernode(ex)
+function removeinternallinenumbernodes(ex, userfiles)
     if isexpr(ex, :block)
-        args = [removelinenumbernode(e) for e in ex.args if !(e isa LineNumberNode)]
+        # within a block, `LineNumberNode`s are statement markers: drop the internal ones
+        args = Any[
+            removeinternallinenumbernodes(e, userfiles) for e in ex.args
+                if !_isinternallinenumber(e, userfiles)
+        ]
         return Expr(:block, args...)
+    elseif isa(ex, Expr)
+        # elsewhere a `LineNumberNode` may be structurally required -- most notably as the
+        # mandatory 2nd argument of a `:macrocall` -- so keep all positions here and only
+        # recurse into nested blocks
+        return Expr(
+            ex.head, Any[removeinternallinenumbernodes(e, userfiles) for e in ex.args]...
+        )
     else
         return ex
     end
 end
 
+_isinternallinenumber(@nospecialize(x), userfiles) =
+    x isa LineNumberNode && x.file ∉ userfiles
+
 # list of functions that are used in expressions produced by `@tensor`
 const tensoroperationsfunctions = (
     :tensoralloc, :tensorfree!,

test/butensor.jl

@@ -6,6 +6,20 @@
 end
 
 using Bumper
+@testset "@butensor preserves user line numbers (issue #280)" begin
+    # `@butensor` wraps the block in an inner `@tensor`, so the expansion has to be recursive
+    # here to reach the code that the parser generated. This also covers the extension itself:
+    # `_butensor` must not introduce `LineNumberNode`s pointing into `ext/`.
+    firstline = @__LINE__() + 2
+    block = @macroexpand @butensor begin
+        T[a, b] := X[a, c] * Y[c, b]
+        Z[a, b] := T[a, c] * W[c, b]
+    end
+    lnns = statementlinenumbernodes(block)
+    @test all(l -> l.file === Symbol(@__FILE__), lnns)
+    @test sort!(unique(l.line for l in lnns)) == collect(firstline .+ (0:1))
+end
+
 @testset "Bumper tests with eltype $T" for T in (Float32, ComplexF64)
     D1, D2, D3 = 30, 40, 20
     d1, d2 = 2, 3

test/macro_kwargs.jl

@@ -92,6 +92,91 @@ end
     end
 end
 
+# https://github.com/QuantumKitHub/TensorOperations.jl/issues/280: the generated code must keep
+# the user's `LineNumberNode`s (so `@tensor` lines show up in code coverage) while dropping the
+# parser's own, which would steal that attribution. See `statementlinenumbernodes` in
+# `runtests.jl` for why only block-statement positions are inspected. `@macroexpand1` is used
+# throughout so that we look at what `@tensor` itself produced, and not at the expansion of
+# macros that it merely passes through (e.g. the `@warn` of the `costcheck` path).
+@testset "line numbers (issue #280)" begin
+    thisfile = Symbol(@__FILE__)
+
+    @testset "single-statement expressions carry no LineNumberNodes" begin
+        # nothing to preserve here: the input has no `LineNumberNode`s of its own, and the
+        # statement is attributed to the line of the `@tensor` call by the surrounding scope
+        exprs = [
+            @macroexpand1(@tensor T[a, b] := A[a, c] * B[c, b]),
+            @macroexpand1(@tensor R[a, b] := A[a, c] * B[c, d] * C[d, e] * E[e, f] * F[f, b]),
+            @macroexpand1(@tensor s = X[a, b] * Y[a, b]),
+            @macroexpand1(@tensoropt R[a, b] := A[a, c] * B[c, d] * C[d, e] * E[e, b]),
+            @macroexpand1(@tensor allocator = alloc R[a, b] := A[a, c] * B[c, d] * C[d, b]),
+            @macroexpand1(@tensor costcheck = warn R[a, b] := A[a, c] * B[c, d] * C[d, b]),
+            @macroexpand1(@tensor contractcheck = true R[a, b] := A[a, c] * B[c, b]),
+        ]
+        for ex in exprs
+            @test isempty(statementlinenumbernodes(ex))
+        end
+    end
+
+    @testset "one user LineNumberNode per statement of a block" begin
+        # multi-statement block, including a scalar assignment; `s = ...` additionally
+        # exercises the `_flatten` path that hoists a block into the right hand side
+        firstline = @__LINE__() + 2
+        block = @macroexpand1 @tensor begin
+            T[a, e] := A[a, c] * B[c, d] * C[d, e]
+            D[a, b] := T[a, e] * E[e, b]
+            s = D[a, b] * F[a, b]
+        end
+        lnns = statementlinenumbernodes(block)
+        @test all(l -> l.file === thisfile, lnns)
+        @test sort!(unique(l.line for l in lnns)) == collect(firstline .+ (0:2))
+
+        # dst-reuse: `tensorify` wraps this in a `quote` of its own
+        reuseline = @__LINE__() + 2
+        reuseblock = @macroexpand1 @tensor begin
+            T[a, b] := A[a, c] * B[c, b]
+            T[a, b] := T[a, c] * C[c, b]
+        end
+        reuselnns = statementlinenumbernodes(reuseblock)
+        @test all(l -> l.file === thisfile, reuselnns)
+        @test sort!(unique(l.line for l in reuselnns)) == collect(reuseline .+ (0:1))
+
+        optline = @__LINE__() + 2
+        optblock = @macroexpand1 @tensoropt begin
+            T[a, e] := A[a, c] * B[c, d] * C[d, e]
+            D[a, b] := T[a, e] * E[e, b]
+        end
+        optlnns = statementlinenumbernodes(optblock)
+        @test all(l -> l.file === thisfile, optlnns)
+        @test sort!(unique(l.line for l in optlnns)) == collect(optline .+ (0:1))
+    end
+
+    @testset "kwargs that generate extra code preserve user line numbers" begin
+        # `allocator` inserts a checkpoint `quote`, `costcheck` inserts a `@notensor` block
+        # containing a `@warn` whose own `LineNumberNode` is structurally required
+        allocline = @__LINE__() + 2
+        allocblock = @macroexpand1 @tensor allocator = alloc begin
+            T[a, b] := A[a, c] * B[c, b]
+            Z[a, b] := T[a, c] * C[c, b]
+        end
+        alloclnns = statementlinenumbernodes(allocblock)
+        @test all(l -> l.file === thisfile, alloclnns)
+        @test sort!(unique(l.line for l in alloclnns)) == collect(allocline .+ (0:1))
+
+        costline = @__LINE__() + 2
+        costblock = @macroexpand1 @tensor costcheck = warn begin
+            T[a, b] := A[a, c] * B[c, b]
+            Z[a, b] := T[a, c] * C[c, b]
+        end
+        costlnns = statementlinenumbernodes(costblock)
+        @test all(l -> l.file === thisfile, costlnns)
+        @test sort!(unique(l.line for l in costlnns)) == collect(costline .+ (0:1))
+        # the `@warn` macrocall must have kept its (TensorOperations-internal) LineNumberNode,
+        # otherwise the expression cannot be expanded at all
+        @test macroexpand(@__MODULE__, costblock; recursive = true) isa Expr
+    end
+end
+
 @testset "opt" begin
     A = randn(5, 5, 5, 5)
     B = randn(5, 5, 5)

test/runtests.jl

@@ -11,6 +11,24 @@ using TensorOperations: DefaultAllocator, ManualAllocator, BufferAllocator
 precision(::Type{<:Union{Float32, Complex{Float32}}}) = 1.0e-2
 precision(::Type{<:Union{Float64, Complex{Float64}}}) = 1.0e-8
 
+# https://github.com/QuantumKitHub/TensorOperations.jl/issues/280: the generated code has to
+# keep the user's `LineNumberNode`s -- Julia only emits a coverage counter for lines that a
+# `LineNumberNode` points at -- and drop the ones synthesized by the parser, which would
+# otherwise re-attribute the statements that follow them to a line in TensorOperations itself.
+# Only `LineNumberNode`s in block-statement position matter for that: elsewhere (most notably
+# the mandatory 2nd argument of a `:macrocall`) they are structurally required and left alone.
+function statementlinenumbernodes(ex, acc = LineNumberNode[])
+    if ex isa Expr
+        if ex.head === :block
+            for e in ex.args
+                e isa LineNumberNode && push!(acc, e)
+            end
+        end
+        foreach(e -> statementlinenumbernodes(e, acc), ex.args)
+    end
+    return acc
+end
+
 # don't run all tests on GPU, only the GPU
 # specific ones
 is_buildkite = get(ENV, "BUILDKITE", "false") == "true"