[Nonlinear] add amontoison's Forest type and remove DataStructures.jl

Author: odow

Project.toml

@@ -6,7 +6,6 @@ version = "1.46.0"
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 CodecBzip2 = "523fee87-0ab8-5b00-afb7-3ecf72e48cfd"
 CodecZlib = "944b1d66-785c-5afd-91f1-9de20f533193"
-DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 JSON3 = "0f8b85d8-7281-11e9-16c2-39a750bddbf1"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
@@ -23,7 +22,6 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 BenchmarkTools = "1"
 CodecBzip2 = "0.6, 0.7, 0.8"
 CodecZlib = "0.6, 0.7"
-DataStructures = "0.18, 0.19"
 ForwardDiff = "0.10, 1"
 JSON3 = "1"
 JSONSchema = "1"

src/FileFormats/MPS/MPS.jl

@@ -8,7 +8,7 @@ module MPS
 
 import ..FileFormats
 import MathOptInterface as MOI
-import DataStructures: OrderedDict
+import OrderedCollections: OrderedDict
 
 const IndicatorLessThanTrue{T} =
     MOI.Indicator{MOI.ACTIVATE_ON_ONE,MOI.LessThan{T}}

src/Nonlinear/ReverseAD/Coloring/Coloring.jl

@@ -6,8 +6,7 @@
 
 module Coloring
 
-import DataStructures
-
+include("Forest.jl")
 include("topological_sort.jl")
 
 """
@@ -154,7 +153,7 @@ function _prevent_cycle(
     forbiddenColors,
     color,
 )
-    er = DataStructures.find_root!(S, e_idx2)
+    er = _find_root!(S, e_idx2)
     @inbounds first = firstVisitToTree[er]
     p = first.source # but this depends on the order?
     q = first.target
@@ -172,29 +171,18 @@ function _grow_star(v, w, e_idx, firstNeighbor, color, S)
     @inbounds if p != v
         firstNeighbor[color[w]] = _Edge(e_idx, v, w)
     else
-        union!(S, e_idx, e.index)
+        _root_union!(S, e_idx, e.index)
     end
     return
 end
 
-function _merge_trees(eg, eg1, S)
-    e1 = DataStructures.find_root!(S, eg)
-    e2 = DataStructures.find_root!(S, eg1)
-    if e1 != e2
-        union!(S, eg, eg1)
+function _merge_trees(S::_Forest, eg::Int, eg1::Int)
+    if _find_root!(S, eg) != _find_root!(S, eg1)
+        _root_union!(S, eg, eg1)
     end
     return
 end
 
-# Work-around a deprecation in DataStructures@0.19
-function _IntDisjointSet(n)
-    @static if isdefined(DataStructures, :IntDisjointSet)
-        return DataStructures.IntDisjointSet(n)
-    else
-        return DataStructures.IntDisjointSets(n)                 # COV_EXCL_LINE
-    end
-end
-
 """
     acyclic_coloring(g::UndirectedGraph)
 
@@ -214,8 +202,7 @@ function acyclic_coloring(g::UndirectedGraph)
     firstNeighbor = _Edge[]
     firstVisitToTree = fill(_Edge(0, 0, 0), _num_edges(g))
     color = fill(0, _num_vertices(g))
-    # disjoint set forest of edges in the graph
-    S = _IntDisjointSet(_num_edges(g))
+    S = _Forest(_num_edges(g))
     @inbounds for v in 1:_num_vertices(g)
         n_neighbor = _num_neighbors(v, g)
         start_neighbor = _start_neighbors(v, g)
@@ -293,7 +280,7 @@ function acyclic_coloring(g::UndirectedGraph)
                     continue
                 end
                 if color[x] == color[v]
-                    _merge_trees(e_idx, e2_idx, S)
+                    _merge_trees(S, e_idx, e2_idx)
                 end
             end
         end

src/Nonlinear/ReverseAD/Coloring/Forest.jl

@@ -0,0 +1,51 @@
+# Copyright (c) 2017: Miles Lubin and contributors
+# Copyright (c) 2017: Google Inc.
+# Copyright (c) 2024: Guillaume Dalle and Alexis Montoison
+#
+# Use of this source code is governed by an MIT-style license that can be found
+# in the LICENSE.md file or at https://opensource.org/licenses/MIT.
+
+# The code in this file was taken from
+# https://github.com/gdalle/SparseMatrixColorings.jl/blob/main/src/forest.jl
+#
+# It was copied at the suggestion of Alexis in his JuMP-dev 2025 talk.
+#
+# @odow made minor changes to match MOI coding styles.
+#
+# x-ref https://github.com/gdalle/SparseMatrixColorings.jl/pull/190
+
+mutable struct _Forest
+    # current number of distinct trees in the forest
+    number_of_trees::Int
+    # vector storing the index of a parent in the tree for each edge, used in
+    # union-find operations
+    parents::Vector{Int}
+    # vector approximating the depth of each tree to optimize path compression
+    ranks::Vector{Int}
+
+    _Forest(n::Integer) = new(n, collect(Base.OneTo(n)), zeros(Int, n))
+end
+
+function _find_root!(parents::Vector{Int}, index_edge::Integer)
+    p = parents[index_edge]
+    if parents[p] != p
+        parents[index_edge] = p = _find_root!(parents, p)
+    end
+    return p
+end
+
+function _find_root!(forest::_Forest, index_edge::Integer)
+    return _find_root!(forest.parents, index_edge)
+end
+
+function _root_union!(forest::_Forest, index_edge1::Int, index_edge2::Int)
+    rank1, rank2 = forest.ranks[index_edge1], forest.ranks[index_edge2]
+    if rank1 < rank2
+        index_edge1, index_edge2 = index_edge2, index_edge1
+    elseif rank1 == rank2
+        forest.ranks[index_edge1] += 1
+    end
+    forest.parents[index_edge2] = index_edge1
+    forest.number_of_trees -= 1
+    return
+end

test/FileFormats/MPS/MPS.jl

@@ -10,7 +10,7 @@ using Test
 
 import MathOptInterface as MOI
 import MathOptInterface.FileFormats: MPS
-import DataStructures: OrderedDict
+import OrderedCollections: OrderedDict
 
 function runtests()
     for name in names(@__MODULE__; all = true)