feat(alias_elimination): detect x ~ -y perfect aliases

.gitignore

@@ -7,3 +7,4 @@ Manifest.toml
 docs/src/assets/Project.toml
 docs/src/assets/Manifest.toml
 .CondaPkg
+.claude/settings.local.json

Project.toml

@@ -118,7 +118,7 @@ StateSelection = "1.9.1"
 StaticArrays = "1.9.14"
 StochasticDiffEq = "6.82.0, 7"
 SymbolicIndexingInterface = "0.3.39"
-SymbolicUtils = "4.28"
+SymbolicUtils = "4.31"
 Symbolics = "7"
 UnPack = "0.1, 1.0"
 julia = "1.9"

lib/ModelingToolkitBase/src/ModelingToolkitBase.jl

@@ -38,12 +38,6 @@ using SymbolicIndexingInterface
 using LinearAlgebra, SparseArrays
 using InteractiveUtils
 using DataStructures
-@static if pkgversion(DataStructures) >= v"0.19"
-    import DataStructures: IntDisjointSet
-else
-    import DataStructures: IntDisjointSets
-    const IntDisjointSet = IntDisjointSets
-end
 using Base.Threads
 using ArrayInterface
 using Setfield, ConstructionBase

src/ModelingToolkit.jl

@@ -33,12 +33,6 @@ using SymbolicIndexingInterface
 using LinearAlgebra, SparseArrays
 using InteractiveUtils
 using DataStructures
-@static if pkgversion(DataStructures) >= v"0.19"
-    import DataStructures: IntDisjointSet
-else
-    import DataStructures: IntDisjointSets
-    const IntDisjointSet = IntDisjointSets
-end
 using Base.Threads
 using Setfield, ConstructionBase
 import Libdl

src/systems/alias_elimination.jl

@@ -40,6 +40,64 @@ function eliminate_perfect_aliases!(state::TearingState)
     return nothing
 end
 
+"""
+    $TYPEDSIGNATURES
+
+Weighted union-find augmented with sign tracking. Merges the components of `v1`
+and `v2` under the relation `v1 ~ edge_sign · v2` (`edge_sign ∈ ±1`)
+maintains the invariant that `parent[v]` is always the current root and `parity[v]`
+is the sign of `v` relative to that root so finds are O(1).
+
+Merges are smaller-into-larger using `members` (root → member list)
+giving O(E log V) across all unions.
+A contradiction (e.g. `x ~ y, x ~ -y`) zeros out every member's `parity` in the
+affected component; downstream code treats `parity[v] == 0` as "v is forced to 0".
+"""
+function union_with_sign!(
+        parent::Dict{Int, Int}, parity::Dict{Int, Int8},
+        members::Dict{Int, Vector{Int}},
+        v1::Int, v2::Int, edge_sign::Int8
+    )
+    for v in (v1, v2)
+        if !haskey(parent, v)
+            parent[v] = v
+            parity[v] = Int8(1)
+            members[v] = [v]
+        end
+    end
+    r1 = parent[v1]; s1 = parity[v1]
+    r2 = parent[v2]; s2 = parity[v2]
+    if r1 == r2
+        if s1 != Int8(edge_sign * s2)
+            for m in members[r1]
+                parity[m] = Int8(0)
+            end
+        end
+        return
+    end
+    if length(members[r1]) < length(members[r2])
+        r1, r2 = r2, r1
+        s1, s2 = s2, s1
+    end
+    # Edge `v1 = edge_sign · v2` ⇒ sign of r2 relative to r1 = s1 · edge_sign · s2.
+    # If either side is already zero (conflict-tainted), `r2_to_r1` is 0 and we
+    # propagate the zero through both components.
+    r2_to_r1 = Int8(s1 * edge_sign * s2)
+    r2_members = members[r2]
+    for m in r2_members
+        parent[m] = r1
+        parity[m] = Int8(parity[m] * r2_to_r1)
+    end
+    if r2_to_r1 == 0
+        for m in members[r1]
+            parity[m] = Int8(0)
+        end
+    end
+    append!(members[r1], r2_members)
+    delete!(members, r2)
+    return
+end
+
 """
     $TYPEDSIGNATURES
 
@@ -90,12 +148,20 @@ function find_perfect_aliases!(
     original_eqs = state.original_eqs
     irreducibles = get_irreducibles(sys)
 
-    # Not `IntDisjointSet` because we don't want singleton sets for every single variable
-    alias_groups = DisjointSet{Int}()
-    # Candidate alias equations `(ieq, v1_idx, v2_idx)`. Removal is decided below once
-    # each group's target is known: equations with a non-target irreducible endpoint
-    # must stay so the remaining irreducibles are still constrained to the target.
-    candidate_eqs = Tuple{Int, Int, Int}[]
+    # Weighted union-find tracking each variable's sign relative to its
+    # component's current root. `parent[v]` is always the root.
+    # `parity[v] ∈ ±1` is the sign of `v` relative to that root, or `0` if the
+    # component contains a sign contradiction (forcing every member to 0).
+    # `members` maps root → member list (used for smaller-to-larger merges).
+    parent = Dict{Int, Int}()
+    parity = Dict{Int, Int8}()
+    members = Dict{Int, Vector{Int}}()
+    # Candidate alias equations `(ieq, v1_idx, v2_idx, edge_sign)`. 
+    # `edge_sign == ±1` encodes `v1 ~ edge_sign*v2`.
+    # Removal is decided below once each group's target is known:
+    # equations with a non-target irreducible endpoint must stay so
+    # the remaining irreducibles are still constrained to the target.
+    candidate_eqs = Tuple{Int, Int, Int, Int8}[]
 
     for ieq in 1:nsrcs(graph)
         snbors = 𝑠neighbors(graph, ieq)
@@ -106,66 +172,99 @@ function find_perfect_aliases!(
         eq = eqs[ieq]
         v1 = fullvars[snbors[1]]
         v2 = fullvars[snbors[2]]
+        local edge_sign::Int8
         Moshi.Match.@match eq.rhs begin
             BSImpl.AddMul(; variant, coeff, dict) && if variant == SU.AddMulVariant.ADD end => begin
                 SU._iszero(coeff) || continue
                 length(dict) == 2 || continue
-                haskey(dict, v1) && haskey(dict, v2) && SU._iszero(dict[v1] + dict[v2]) || continue
+                haskey(dict, v1) && haskey(dict, v2) || continue
+                if SU._iszero(dict[v1] + dict[v2])
+                    edge_sign = Int8(1)   # v1 ~ v2
+                elseif SU._iszero(dict[v1] - dict[v2])
+                    edge_sign = Int8(-1)  # v1 ~ -v2
+                else
+                    continue
+                end
             end
             _ => continue
         end
-        push!(candidate_eqs, (ieq, snbors[1], snbors[2]))
-        push!(alias_groups, snbors[1])
-        push!(alias_groups, snbors[2])
-        union!(alias_groups, snbors[1], snbors[2])
+        push!(candidate_eqs, (ieq, snbors[1], snbors[2], edge_sign))
+        union_with_sign!(parent, parity, members, snbors[1], snbors[2], edge_sign)
     end
-
-    alias_sets = Dict{Int, Vector{Int}}()
-    sizehint!(alias_sets, DataStructures.num_groups(alias_groups))
-    sizehint!(aliases, length(alias_groups) - DataStructures.num_groups(alias_groups))
-    for var in alias_groups
-        set = get!(() -> Int[], alias_sets, DataStructures.find_root!(alias_groups, var))
-        push!(set, var)
-    end
-
+    sizehint!(aliases, length(parent) - length(members))
+
+    # After the scan above, every variable in any alias group has
+    # `parent[v] == root` and `parity[v] ∈ {-1, 0, +1}`
+    # (sign relative to root, or 0 if the component is conflict-tainted).
+    # The main rewriting loop below picks a target per non-conflict group
+    # and uses `s = parity[v] * parity[target]` (±1) as `v`'s sign relative to the target.
+    # For Conflict groups: every non-irreducible gets substituted with the
+    # symbolic `0`, and one alias eq per irreducible is rewritten to `irr ~ 0`.
     group_target = Dict{Int, Int}()
-    for (root, group_vars) in alias_sets
-        group_target[root] = pick_alias_target(fullvars, group_vars, state_priorities, irreducibles, graph)
-    end
+    sizehint!(group_target, length(members))
 
-    # Queue an alias equation for removal only if both endpoints collapse onto the
-    # target after non-irreducibles are substituted -- i.e. the equation becomes
-    # `T ~ T`. Any equation with a non-target irreducible endpoint is kept; when the
-    # other endpoint is a non-irreducible, the existing substitution machinery below
-    # rewrites the kept equation into `I ~ T` form automatically.
-    for (ieq, v1, v2) in candidate_eqs
-        target = group_target[DataStructures.find_root!(alias_groups, v1)]
-        c1 = contains_possibly_indexed_element(irreducibles, fullvars[v1]) || state_priorities[v1] > 0 ? v1 : target
-        c2 = contains_possibly_indexed_element(irreducibles, fullvars[v2]) || state_priorities[v2] > 0 ? v2 : target
-        c1 == c2 && push!(eqs_to_rm, ieq)
-    end
+    is_sticky = (v) -> contains_possibly_indexed_element(irreducibles, fullvars[v]) || state_priorities[v] > 0
+    is_irreducible_v = (v) -> contains_possibly_indexed_element(irreducibles, fullvars[v])
+
+    # Symbolic zero used as substitution target for variables in conflict groups.
+    zero_sym = Symbolics.COMMON_ZERO
+    signed_sym = (s::Int8, x::SymbolicT) -> s == 1 ? x : (-1) * x
 
     eqs_to_substitute = Int[]
-    for (root, group_vars) in alias_sets
-        target = group_target[root]
-        state.always_present[target] = true
+    irrs_by_root = Dict{Int, Vector{Int}}()
+    for (root, group_vars) in members
+        if parity[root] == 0 # is conflict group
+            # collect the irreducibles that need an equation forcing them to `0`.
+            irrs_by_root[root] = filter(is_irreducible_v, group_vars)
+            for v in group_vars
+                # In conflict groups, ONLY irreducible vars survive as unknowns.
+                # The per-equation cleanup below rewrites one alias eq per
+                # irreducible to `v ~ 0`; here we just mark them present.
+                if is_irreducible_v(v)
+                    state.always_present[v] = true
+                    continue
+                end
+                push!(vars_to_rm, v)
+                subs[fullvars[v]] = zero_sym
+                push!(state.additional_observed, fullvars[v] ~ zero_sym)
+
+                append!(eqs_to_substitute, 𝑑neighbors(graph, v))
+                BipartiteGraphs.set_neighbors!(BipartiteGraphs.invview(graph), v, ())
+
+                dv = var_to_diff[v]
+                while dv isa Int
+                    push!(vars_to_rm, dv)
+                    subs[fullvars[dv]] = zero_sym
+
+                    append!(eqs_to_substitute, 𝑑neighbors(graph, dv))
+                    BipartiteGraphs.set_neighbors!(BipartiteGraphs.invview(graph), dv, ())
+
+                    dv = var_to_diff[dv]
+                end
+            end
+            continue
+        end
+        # For consistent groups pick a target (survives as unknown) and rebase
+        # parities relative to it via `target_p`.
+        target = pick_alias_target(fullvars, group_vars, state_priorities, irreducibles, graph)
+        group_target[root] = target
+        target_p = parity[target]
         for v in group_vars
-            v == target && continue
-            # Irreducibles other than the target stay as unknowns; only non-irreducibles
-            # are eliminated in favor of the target.
-            if contains_possibly_indexed_element(irreducibles, fullvars[v]) || state_priorities[v] > 0
+            # In consistent groups sticky vars (irreducible OR priority>0) other
+            # than the target stay as unknowns.
+            if is_sticky(v) || v == target
                 state.always_present[v] = true
                 continue
             end
 
+            s = Int8(parity[v] * target_p)
             push!(vars_to_rm, v)
-            subs[fullvars[v]] = fullvars[target]
-            push!(state.additional_observed, fullvars[v] ~ fullvars[target])
+            rhs_sym = signed_sym(s, fullvars[target])
+            subs[fullvars[v]] = rhs_sym
+            push!(state.additional_observed, fullvars[v] ~ rhs_sym)
             aliases[v] = target
 
-            dnbors = copy(𝑑neighbors(graph, v))
-            for e in dnbors
-                snbors = 𝑠neighbors(graph, e)
+            for e in copy(𝑑neighbors(graph, v))
                 push!(eqs_to_substitute, e)
                 Graphs.rem_edge!(graph, e, v)
                 Graphs.add_edge!(graph, e, target)
@@ -179,12 +278,10 @@ function find_perfect_aliases!(
                 end
                 push!(vars_to_rm, dv)
 
-                subs[fullvars[dv]] = fullvars[dtarget]
+                subs[fullvars[dv]] = signed_sym(s, fullvars[dtarget])
                 aliases[dv] = dtarget
 
-                dnbors = copy(𝑑neighbors(graph, dv))
-                for e in dnbors
-                    snbors = 𝑠neighbors(graph, e)
+                for e in copy(𝑑neighbors(graph, dv))
                     push!(eqs_to_substitute, e)
                     Graphs.rem_edge!(graph, e, dv)
                     Graphs.add_edge!(graph, e, dtarget)
@@ -196,6 +293,37 @@ function find_perfect_aliases!(
         end
     end
 
+    # Per-equation cleanup. Two cases:
+    #
+    # Conflict groups: every variable is forced to `0`.
+    # Take the first `length(irrs)` eq in the group to force the irrs to 0
+    # (overwriting the eq and replacing its graph row with the edge to `irr`). 
+    # Remaining conflict eqs become `0 ~ 0` and are queued for removal.
+    # Pinned eqs end up in `eqs_to_substitute` from the var-elim pass above
+    # but `subber` leaves `irr ~ 0` untouched (irreducibles aren't in `subs`).
+    #
+    # Consistent groups: queue for removal iff both endpoints collapse. An eq
+    # with a non-target irreducible endpoint must stay so the remaining
+    # irreducible is still pinned to the target after substitution.
+    for (ieq, v1, v2, _) in candidate_eqs
+        if parity[v1] == 0 # if conflict group
+            irrs = irrs_by_root[parent[v1]]
+            if isempty(irrs)
+                push!(eqs_to_rm, ieq)
+            else
+                v_pin = pop!(irrs) # irreducible var to pin to 0
+                BipartiteGraphs.set_neighbors!(graph, ieq, [v_pin])
+                eqs[ieq] = fullvars[v_pin] ~ zero_sym
+                original_eqs[ieq] = fullvars[v_pin] ~ zero_sym
+            end
+        else
+            target = group_target[parent[v1]]
+            c1 = is_sticky(v1) ? v1 : target
+            c2 = is_sticky(v2) ? v2 : target
+            c1 == c2 && push!(eqs_to_rm, ieq)
+        end
+    end
+
     # We need to handle unscalarized array variables
     for k in keys(subs)
         k, isarr = split_indexed_var(k)
@@ -205,12 +333,23 @@ function find_perfect_aliases!(
     end
 
     ir = get_irstructure(sys)
-    subber = SU.IRSubstituter{false}(ir, subs)
+    subber = SU.IRSubstituter{true}(ir, subs)
+    new_vars = SU.IRStructureSearchBuffer(ir, Set{SymbolicT}())
     for e in eqs_to_substitute
         # Double substitute to handle unscalarized array variables. First one
         # substitutes `x` to `[x[1], x[2]]`. The second substitutes `x[1]` and `x[2]`.
         eqs[e] = subber(subber(eqs[e]))
         original_eqs[e] = subber(subber(original_eqs[e]))
+        # Substitution can drop vars beyond the one substituted: zero
+        # substitution annihilates multiplicative cofactors (`v*w` with `v→0`
+        # also removes `w`); alias substitution can cancel the target
+        # (`v - target` with `v→target` leaves neither). Substitution never
+        # introduces new vars, so the new incidence is a subset of the old —
+        # prune the row to entries actually present in the simplified RHS.
+        empty!(new_vars)
+        SU.search_variables!(new_vars, eqs[e])
+        new_row = filter(v_idx -> fullvars[v_idx] in new_vars, 𝑠neighbors(graph, e))
+        BipartiteGraphs.set_neighbors!(graph, e, new_row)
     end
 
     # After substitution, alias equations that connected a sticky non-target
@@ -221,7 +360,7 @@ function find_perfect_aliases!(
     let seen = Set{Tuple{Int,Int}}()
         eqs_rm_set = Set(eqs_to_rm)
         removed_additional_eqs = false
-        for (ieq, _, _) in candidate_eqs
+        for (ieq, _, _, _) in candidate_eqs
             ieq in eqs_rm_set && continue
             snbors = 𝑠neighbors(graph, ieq)
             length(snbors) == 2 || continue