Span bounds check elision doesn't work properly when method calls occur between bounds check and access

[GrabYourPitchforks]

If the caller performs a bounds check on a span, the runtime will not elide future bounds checks satisfied by that condition if a method call occurs between the original check and the subsequent access. The following repro code demonstrates this issue.

using System;

public class MyClass {
    public int SomeField;
    public int SomeProperty => 42;
    public int SomeMethod() => 42;
    
    private void Foo(Span<int> destination)
    {
        if (!destination.IsEmpty)
        {
            destination[0] = SomeProperty;
        }    
    }
    
    private void Bar(Span<int> destination)
    {
        if (!destination.IsEmpty)
        {
            destination[0] = SomeField;
        }    
    }
    
    private void Baz(Span<int> destination)
    {
        if (!destination.IsEmpty)
        {
            destination[0] = SomeMethod();
        }    
    }
}

x64 codegen, courtesy SharpLab:

MyClass.Foo(System.Span`1<Int32>)
    L0000: sub rsp, 0x28
    L0004: mov rax, [rdx]
    L0007: mov edx, [rdx+8]
    L000a: test edx, edx  ; call to destination.IsEmpty
    L000c: jbe short L0019
    L000e: cmp edx, 0  ; bounds check should've been elided
    L0011: jbe short L001e
    L0013: mov dword ptr [rax], 0x2a
    L0019: add rsp, 0x28
    L001d: ret
    L001e: call 0x00007ff9f75fbab0
    L0023: int3

MyClass.Bar(System.Span`1<Int32>)
    L0000: mov rax, [rdx]
    L0003: mov edx, [rdx+8]
    L0006: test edx, edx  ; call to destination.IsEmpty
    L0008: jbe short L000f
    L000a: mov edx, [rcx+8]
    L000d: mov [rax], edx  ; no bounds check before writing field value to destination
    L000f: ret

MyClass.Baz(System.Span`1<Int32>)
    L0000: sub rsp, 0x28
    L0004: mov rax, [rdx]
    L0007: mov edx, [rdx+8]
    L000a: test edx, edx  ; call to destination.IsEmpty
    L000c: jbe short L0019
    L000e: cmp edx, 0  ; bounds check should've been elided
    L0011: jbe short L001e
    L0013: mov dword ptr [rax], 0x2a
    L0019: add rsp, 0x28
    L001d: ret
    L001e: call 0x00007ff9f75fbab0
    L0023: int3

[SingleAccretion]

A quick glance at the trees suggests that this is the result of inlining breaking them up.
This is for the field case:

***** BB03
STMT00003 (IL 0x009...0x016)
N008 ( 16, 16) [000023] -A-XG-------              *  ASG       int    $VN.Void
N006 ( 14, 14) [000020] ---XG--N----              +--*  COMMA     int    $VN.Void
N003 (  8, 10) [000015] ---XG-------              |  +--*  ARR_BOUNDS_CHECK_Rng void   <l:$2c5, c:$2c4>
N001 (  1,  1) [000010] ------------              |  |  +--*  CNS_INT   int    0 $40
N002 (  3,  2) [000014] ------------              |  |  \--*  LCL_VAR   int    V03 tmp2         u:1 (last use) <l:$200, c:$240>
N005 (  6,  4) [000047] *--X---N----              |  \--*  IND       int    $44
N004 (  3,  2) [000018] ------------              |     \--*  LCL_VAR   byref  V02 tmp1         u:1 (last use) <l:$100, c:$140>
N007 (  1,  1) [000021] ------------              \--*  CNS_INT   int    42 $44

This is for the method case:

------------ BB03 [009..017), preds={BB02} succs={BB04}

***** BB03
STMT00004 (IL 0x009...0x016)
N003 (  8, 10) [000015] ---XG-------              *  ARR_BOUNDS_CHECK_Rng void   <l:$2c5, c:$2c4>
N001 (  1,  1) [000010] ------------              +--*  CNS_INT   int    0 $40
N002 (  3,  2) [000014] ------------              \--*  LCL_VAR   int    V04 tmp3         u:1 (last use) <l:$200, c:$240>

***** BB03
STMT00005 (IL   ???...  ???)
N004 (  8,  6) [000027] -A-XG-------              *  ASG       int    $VN.Void
N002 (  6,  4) [000026] *--X---N----              +--*  IND       int    $44
N001 (  3,  2) [000024] ------------              |  \--*  LCL_VAR   byref  V03 tmp2         u:1 (last use) <l:$100, c:$140>
N003 (  1,  1) [000036] ------------              \--*  CNS_INT   int    42 $44

A source-level workaround would be the use of a local:

if (!destination.IsEmpty)
{
    var tmp = SomeMethod();
    destination[0] = tmp;
}

A Jit-level fix could be stitching these trees together (fairly early) or recognizing them where relevant. A few phases remove checks right now, so stitching would be more effective. It would also be a more complex solution (the early trees are fairly complex). An easier solution would be to update assertion propagation to understand this case:

runtime/src/coreclr/jit/assertionprop.cpp

Lines 4066 to 4068 in bf3789d

	// Defer actually removing the tree until processing reaches its parent comma, since
	// optRemoveRangeCheck needs to rewrite the whole comma tree.
	arrBndsChk->gtFlags |= GTF_ARR_BOUND_INBND;

The flag is already being set for the problematic case, but the machinery to act on it is missing, as pretty much everything assumes bounds checks are GT_COMMA-based.

[JulieLeeMSFT]

CC @AndyAyersMS

[AndyAyersMS]

@SingleAccretion's analysis is spot on.

We break trees for inline candidates and don't glue them back, and the bounds check elimination that happens later depends on intact trees.

We can either not break trees in the first place, break them and then re-form them (forward substitution #6973), or update subsequent optimizations to be more tolerant of broken trees.

[AndyAyersMS]

but the machinery to act on it is missing,

Seems like if we passed stmt into optAssertionProp_BndsChk it could check for the case where the bounds check is the top level statement, and just clean it up on the spot, running a variant of optRemoveRangeCheck.

I don't know for sure if this would catch all the cases but seems like it might; the bounds check doesn't return a value so can't be incorporated into trees other than within commas.

@SingleAccretion want to give this a try?

[SingleAccretion]

@SingleAccretion want to give this a try?

Yep, and I already have a version of the proposed approach working locally. It needs some cleanup and the addition of the same handling in Optimize index checks (so that it works when we have if (span.Length) > 1)). Hopefully that won't take long to be sorted out.

[AndyAyersMS]

Yep, and I already have a version of the proposed approach working locally...

Excellent. I have assigned this issue to you.