From 4e0c2cfc79ab356c39d3bc829df57c243ee225c6 Mon Sep 17 00:00:00 2001
From: Shuhei Kadowaki <aviatesk@gmail.com>
Date: Tue, 20 Aug 2024 12:23:50 +0900
Subject: [PATCH] inference: propagate partially initialized mutable structs
 more

Following up JuliaLang/julia#55297.
A mutable struct can have undefined fields in a non-contiguous manner,
but `PartialStruct` cannot model such a state. So in JuliaLang/julia#55297
`PartialStruct` was used to represent only the mutable objects where the
field following the minimum initialized fields is also defined.

As a follow-up for the PR, this commit implements a minor improvement
that, in cases when the mutable object is already represented as a
`PartialStruct`, allows inference to add one more `isdefined`-field
information on top of those implied by its `fields`.

```julia
mutable struct PartiallyInitialized2
    a; b; c
    PartiallyInitialized2(a) = (@nospecialize; new(a))
    PartiallyInitialized2(a, b) = (@nospecialize; new(a, b))
    PartiallyInitialized2(a, b, c) = (@nospecialize; new(a, b, c))
end

@test Base.infer_effects((PartiallyInitialized2,); optimize=false) do x
    if isdefined(x, :b)
        if isdefined(x, :c)
            return x.c
        end
        return x.b
    end
    return nothing
end |> Core.Compiler.is_nothrow
```
---
 base/compiler/abstractinterpretation.jl | 11 ++++++++++-
 test/compiler/inference.jl              |  9 +++++++++
 2 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/base/compiler/abstractinterpretation.jl b/base/compiler/abstractinterpretation.jl
index 351f241878e7d8..04c91aa9a11d8b 100644
--- a/base/compiler/abstractinterpretation.jl
+++ b/base/compiler/abstractinterpretation.jl
@@ -2057,7 +2057,16 @@ function form_partially_defined_struct(@nospecialize(obj), @nospecialize(name))
     fldidx === nothing && return nothing
     nminfld = datatype_min_ninitialized(objt)
     if ismutabletype(objt)
-        fldidx == nminfld+1 || return nothing
+        if fldidx ≠ nminfld+1
+            # A mutable struct can have gutshot undefined fields, but `PartialStruct` cannot
+            # model such a state. So here `PartialStruct` can be used to represent only the
+            # objects where the field following the minimum initialized fields is also defined.
+            if !(obj isa PartialStruct && fldidx == length(obj.fields)+1)
+                # if it is already represented as a `PartialStruct`, we can add one more
+                # `isdefined`-field information on top of those implied by its `fields`
+                return nothing
+            end
+        end
     else
         fldidx > nminfld || return nothing
     end
diff --git a/test/compiler/inference.jl b/test/compiler/inference.jl
index 75f33a280e245c..0bef73f67838f7 100644
--- a/test/compiler/inference.jl
+++ b/test/compiler/inference.jl
@@ -6033,6 +6033,15 @@ end |> Core.Compiler.is_nothrow
         return x.b
     end
 end |> !Core.Compiler.is_nothrow
+@test Base.infer_effects((PartiallyInitialized2,); optimize=false) do x
+    if isdefined(x, :b)
+        if isdefined(x, :c)
+            return x.c
+        end
+        return x.b
+    end
+    return nothing
+end |> Core.Compiler.is_nothrow
 
 # End to end test case for the partially initialized struct with `PartialStruct`
 @noinline broadcast_noescape1(a) = (broadcast(identity, a); nothing)