From 7801fe1c0afe694cda5c645e5b9846459abd10b8 Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Tue, 14 Nov 2023 11:11:55 +0000
Subject: [PATCH] move functionality related to `push!` for `VarNameVector`
 into `push!`

---
 src/varnamevector.jl | 85 +++++++++++++++++++++++---------------------
 1 file changed, 44 insertions(+), 41 deletions(-)

diff --git a/src/varnamevector.jl b/src/varnamevector.jl
index 93272a063..4f6a3b7f9 100644
--- a/src/varnamevector.jl
+++ b/src/varnamevector.jl
@@ -121,6 +121,7 @@ Base.length(vnv::VarNameVector) =
     end
 Base.size(vnv::VarNameVector) = (length(vnv),)
 
+# TODO: We should probably remove this
 Base.IndexStyle(::Type{<:VarNameVector}) = IndexLinear()
 
 # Dictionary interface.
@@ -229,59 +230,61 @@ end
 function Base.push!(vnv::VarNameVector, vn::VarName, val, transform=FromVec(val))
     # Error if we already have the variable.
     haskey(vnv, vn) && throw(ArgumentError("variable name $vn already exists"))
-    return update!(vnv, vn, val, transform)
+    # NOTE: We need to compute the `nextrange` BEFORE we start mutating
+    # the underlying; otherwise we might get some strange behaviors.
+    val_vec = tovec(val)
+    r_new = nextrange(vnv, val_vec)
+    vnv.varname_to_index[vn] = length(vnv.varname_to_index) + 1
+    push!(vnv.varnames, vn)
+    push!(vnv.ranges, r_new)
+    append!(vnv.vals, val_vec)
+    push!(vnv.transforms, transform)
+    return nothing
 end
 
 # `update!` and `update!!`: update a variable in the varname vector.
 function update!(vnv::VarNameVector, vn::VarName, val, transform=FromVec(val))
-    val_vec = tovec(val)
     if !haskey(vnv, vn)
         # Here we just add a new entry.
-        # NOTE: We need to compute the `nextrange` BEFORE we start mutating
-        # the underlying; otherwise we might get some strange behaviors.
+        return push!(vnv, vn, val, transform)
+    end
+
+    # Here we update an existing entry.
+    val_vec = tovec(val)
+    idx = getidx(vnv, vn)
+    r_old = getrange(vnv, idx)
+    n_old = length(r_old)
+    n_new = length(val_vec)
+    # Existing keys needs to be handled differently depending on
+    # whether the size of the value is increasing or decreasing.
+    if n_new > n_old
+        # Add the new range.
         r_new = nextrange(vnv, val_vec)
-        vnv.varname_to_index[vn] = length(vnv.varname_to_index) + 1
-        push!(vnv.varnames, vn)
-        push!(vnv.ranges, r_new)
-        append!(vnv.vals, val_vec)
-        push!(vnv.transforms, transform)
+        vnv.ranges[idx] = r_new
+        # Grow the underlying vector to accomodate the new value.
+        resize!(vnv.vals, r_new[end])
+        # Keep track of the deleted ranges.
+        push!(vnv.inactive_ranges, r_old)
     else
-        # Here we update the existing entry.
-        idx = getidx(vnv, vn)
-        r_old = getrange(vnv, idx)
-        n_old = length(r_old)
-        n_new = length(val_vec)
-        # Existing keys needs to be handled differently depending on
-        # whether the size of the value is increasing or decreasing.
-        if n_new > n_old
-            # Add the new range.
-            r_new = nextrange(vnv, val_vec)
-            vnv.ranges[idx] = r_new
-            # Grow the underlying vector to accomodate the new value.
-            resize!(vnv.vals, r_new[end])
-            # Keep track of the deleted ranges.
-            push!(vnv.inactive_ranges, r_old)
-        else
-            # `n_new <= n_old`
-            # Just decrease the current range.
-            r_new = r_old[1]:(r_old[1] + n_new - 1)
-            vnv.ranges[idx] = r_new
-            # And mark the rest as inactive if needed.
-            if n_new < n_old
-                push!(vnv.inactive_ranges, r_old[n_new]:r_old[end])
-            end
+        # `n_new <= n_old`
+        # Just decrease the current range.
+        r_new = r_old[1]:(r_old[1]+n_new-1)
+        vnv.ranges[idx] = r_new
+        # And mark the rest as inactive if needed.
+        if n_new < n_old
+            push!(vnv.inactive_ranges, r_old[n_new]:r_old[end])
         end
+    end
 
-        # Update the value.
-        vnv.vals[r_new] = val_vec
-        # Update the transform.
-        vnv.transforms[idx] = transform
+    # Update the value.
+    vnv.vals[r_new] = val_vec
+    # Update the transform.
+    vnv.transforms[idx] = transform
 
-        # TODO: Should we maybe sweep over inactive ranges and re-contiguify
-        # if we the total number of inactive elements is "large" in some sense?
-    end
+    # TODO: Should we maybe sweep over inactive ranges and re-contiguify
+    # if we the total number of inactive elements is "large" in some sense?
 
-    return vnv
+    return nothing
 end
 
 function recontiguify_ranges!(ranges::AbstractVector{<:AbstractRange})