Fix MustNotOverwrite when setting slices/indices of slices (#1325)

Closes #1321.

Claude did most of this, I tweaked the tests a fair bit from what it
wrote though.

If anyone is interested, this is the prompt I fed Claude

<details>

# Issue

In VarNamedTuple you can set different indices of the same *slice* of
the same array (i.e. x[1:2][1] and x[1:2][2]):

```julia
julia> using DynamicPPL

julia> vnt = VarNamedTuple()
VarNamedTuple()

julia> x = zeros(2)
2-element Vector{Float64}:
 0.0
 0.0

julia> vnt = DynamicPPL.templated_setindex!!(vnt, 1.0, @varname(x[1:2][1]), x)
VarNamedTuple
└─ x => PartialArray size=(2,) data::Vector{Float64}
        └─ (1,) => 1.0

julia> vnt = DynamicPPL.templated_setindex!!(vnt, 2.0, @varname(x[1:2][2]), x)
VarNamedTuple
└─ x => PartialArray size=(2,) data::Vector{Float64}
        ├─ (1,) => 1.0
        └─ (2,) => 2.0
```

This required some careful thought to make sure that the second set
didn't overwrite the first one.
See e.g. lines 217 to 234 of src/varnamedtuple/getset.jl.

However, I forgot to handle this for the case where we check for
MustNotOverwrite:

```julia
julia> using AbstractPPL: @opticof

julia> using DynamicPPL.VarNamedTuples: _setindex_optic!!, SkipTemplate, MustNotOverwrite

julia> vnt = VarNamedTuple()
VarNamedTuple()

julia> vnt = _setindex_optic!!(vnt, 1.0, @opticof(_.x[1:2][1]), SkipTemplate{1}(x), MustNotOverwrite(@varname(x[1:2][1])))
VarNamedTuple
└─ x => PartialArray size=(2,) data::Vector{Float64}
        └─ (1,) => 1.0

julia> vnt = _setindex_optic!!(vnt, 2.0, @opticof(_.x[1:2][2]), SkipTemplate{1}(x), MustNotOverwrite(@varname(x[1:2][2])))
ERROR: MustNotOverwriteError: Attempted to set a value for x[1:2][2], but a value already existed. This indicates that a value is being set twice (e.g. if the same variable occurs in a model twice).
Stacktrace:
 [1] _setindex_optic!!(pa::DynamicPPL.VarNamedTuples.PartialArray{…}, value::Float64, optic::AbstractPPL.Index{…}, template::Vector{…}, permissions::MustNotOverwrite{…})
   @ DynamicPPL.VarNamedTuples ~/ppl/dppl/src/varnamedtuple/getset.jl:173
 [2] _setindex_optic!!(vnt::VarNamedTuple{…}, value::Float64, optic::AbstractPPL.Property{…}, template::SkipTemplate{…}, permissions::MustNotOverwrite{…})
   @ DynamicPPL.VarNamedTuples ~/ppl/dppl/src/varnamedtuple/getset.jl:276
 [3] top-level scope
   @ REPL[20]:1
Some type information was truncated. Use `show(err)` to see complete types.
```

Two requests:

(1) could you help me fix that?

And after we've done that (but don't launch into this first; just do the
fix first)

(2) could you reason about whether there's a better way to handle the
ordinary case (the first code block)? Right now we do a merge, but I
think it might be possible for us to create a separate method for
setting a (slice of a) PartialArray into a PartialArray, where we make
sure to only copy the indices where the mask is set to true.


</details>

---------

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>

Author: penelopeysm

HISTORY.md

@@ -1,3 +1,7 @@
+# 0.40.14
+
+Fixed `check_model()` erroneously failing for models such as `x[1:2] .~ univariate_dist`.`
+
 # 0.40.13
 
 Fixed `densify!!` not recursing into `VarNamedTuple`s or `ArrayLikeBlock`s inside `PartialArray`s.

Project.toml

@@ -1,6 +1,6 @@
 name = "DynamicPPL"
 uuid = "366bfd00-2699-11ea-058f-f148b4cae6d8"
-version = "0.40.13"
+version = "0.40.14"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

src/varnamedtuple/getset.jl

@@ -168,7 +168,7 @@ function _setindex_optic!!(
 
     is_multiindex = _is_multiindex(template, coptic.ix...; coptic.kw...)
 
-    if permissions isa MustNotOverwrite
+    if permissions isa MustNotOverwrite && optic.child isa AbstractPPL.Iden
         if any(view(pa.mask, coptic.ix...; coptic.kw...))
             throw(MustNotOverwriteError(permissions))
         end
@@ -248,6 +248,16 @@ function _setindex_optic!!(
         sub_value
     end
 
+    # In the merge path, some indices in the slice already have data but not all of them
+    # (haskey returned false but any(mask) was true). If MustNotOverwrite is set, check
+    # that the new sub-value doesn't overlap with existing data at the specific sub-indices.
+    if need_merge && permissions isa MustNotOverwrite && grown_sub_value isa PartialArray
+        existing_mask = view(pa.mask, coptic.ix...; coptic.kw...)
+        if any(existing_mask .& grown_sub_value.mask)
+            throw(MustNotOverwriteError(permissions))
+        end
+    end
+
     return if need_merge
         new_pa = BangBang.setindex!!(copy(pa), grown_sub_value, coptic.ix...; coptic.kw...)
         _merge(pa, new_pa, Val(false))

test/debug_utils.jl

@@ -36,6 +36,34 @@ end
             test_model_can_run_but_fails_check(buggy_demo_model())
         end
 
+        @testset "different sub-indices of the same slice" begin
+            # https://github.com/TuringLang/DynamicPPL.jl/issues/1321
+            @model function demo_slice_subindices()
+                x = Vector{Float64}(undef, 2)
+                x[1:2] .~ Normal()
+                return x
+            end
+            @test check_model(demo_slice_subindices(); error_on_failure=true)
+
+            # Same sub-index twice should still fail
+            @model function buggy_slice_subindices()
+                x = Vector{Float64}(undef, 2)
+                x[1:2][1] ~ Normal()
+                x[1:2][1] ~ Normal()
+                return x
+            end
+            test_model_can_run_but_fails_check(buggy_slice_subindices())
+
+            # Slices of slices
+            @model function buggy_slice_subindices2()
+                x = Vector{Float64}(undef, 3)
+                x[1:3][1:2] ~ MvNormal(zeros(2), I)
+                x[1:3][2:3] ~ MvNormal(zeros(2), I)
+                return x
+            end
+            test_model_can_run_but_fails_check(buggy_slice_subindices2())
+        end
+
         @testset "submodel" begin
             @model ModelInner() = x ~ Normal()
             @model function ModelOuterBroken()

test/varnamedtuple.jl

@@ -1101,6 +1101,8 @@ Base.size(st::SizedThing) = st.size
 
     @testset "_setindex_optic!! with MustNotOverwrite" begin
         function test_must_not_overwrite(vnt, value, vn, template)
+            # Avoid mutating
+            vnt = deepcopy(vnt)
             # Check that calling `_setindex_optic!!` with `MustNotOverwrite` errors if the
             # variable already exists, and that it works if it doesn't.
             @test_throws MustNotOverwriteError templated_setindex_no_overwrite!!(
@@ -1174,6 +1176,45 @@ Base.size(st::SizedThing) = st.size
             )
             test_must_not_overwrite(vnt, 2.0, @varname(x.a), NoTemplate())
         end
+
+        @testset "different sub-indices of the same slice" begin
+            # Setting different sub-indices of the same slice should NOT error.
+            # This is the bug from issue #1321.
+            x = zeros(2)
+            vnt = templated_setindex_no_overwrite!!(
+                VarNamedTuple(), 1.0, @varname(x[1:2][1]), x
+            )
+            vnt = templated_setindex_no_overwrite!!(vnt, 2.0, @varname(x[1:2][2]), x)
+            @test vnt[@varname(x)] == [1.0, 2.0]
+
+            # But setting the SAME sub-index twice should still error.
+            vnt2 = templated_setindex_no_overwrite!!(
+                VarNamedTuple(), 1.0, @varname(x[1:2][1]), x
+            )
+            test_must_not_overwrite(vnt2, 2.0, @varname(x[1:2][1]), x)
+
+            # Also test with a larger array and different slices.
+            y = zeros(4)
+            vnt3 = templated_setindex_no_overwrite!!(
+                VarNamedTuple(), 1.0, @varname(y[1:3][1]), y
+            )
+            vnt3 = templated_setindex_no_overwrite!!(vnt3, 2.0, @varname(y[1:3][2]), y)
+            test_must_not_overwrite(vnt3, [3.0, 4.0], @varname(y[1:3][2:3]), y)
+            vnt3 = templated_setindex_no_overwrite!!(vnt3, 3.0, @varname(y[1:3][3]), y)
+            @test vnt3[@varname(y[1:3])] == [1.0, 2.0, 3.0]
+            # Setting an already-set sub-index, or slice, should error.
+            test_must_not_overwrite(vnt3, 4.0, @varname(y[1:3][2]), y)
+            test_must_not_overwrite(vnt3, [4.0, 5.0], @varname(y[2:3][1:2]), y)
+
+            # Also try with indices of different slices
+            z = zeros(3)
+            vnt4 = templated_setindex_no_overwrite!!(
+                VarNamedTuple(), [1.0, 2.0], @varname(z[1:2]), z
+            )
+            test_must_not_overwrite(vnt4, 3.0, @varname(z[2:3][1]), z)
+            # but you can set it if it's not overlapping
+            templated_setindex_no_overwrite!!(vnt4, 3.0, @varname(z[2:3][2]), z)
+        end
     end
 
     @testset "subset" begin