fix order in serialisation of keyword indices

Author: penelopeysm

Project.toml

@@ -13,6 +13,7 @@ DensityInterface = "b429d917-457f-4dbc-8f4c-0cc954292b1d"
 JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
+OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 
@@ -31,6 +32,7 @@ Distributions = "0.25"
 JSON = "0.19 - 0.21, 1"
 LinearAlgebra = "<0.0.1, 1"
 MacroTools = "0.5"
+OrderedCollections = "1.8.1"
 Random = "1.6"
 StatsBase = "0.32, 0.33, 0.34"
 julia = "1.10.8"

src/varname/serialize.jl

@@ -1,6 +1,7 @@
 ### Serialisation to JSON / string
 
 using JSON: JSON
+using OrderedCollections: OrderedDict
 
 # String constants for each index type that we support serialisation /
 # deserialisation of
@@ -70,6 +71,10 @@ then implement the following two methods:
    type (perhaps prefixed with module qualifiers) as this value to avoid
    clashes. The remainder of the dictionary can have any structure you like.
 
+   Note that `Base.Dict` does not guarantee key order, so if you need to preserve
+   key order for fidelity in serialisation, consider returning an
+   `OrderedCollections.OrderedDict` instead.
+
 2. Suppose the value of `index_to_dict(i)["type"]` is `"MyModule.MyIndexType"`.
    You should then implement the corresponding method
    `AbstractPPL.dict_to_index(::Val{Symbol("MyModule.MyIndexType")}, dict)`,
@@ -119,8 +124,7 @@ function optic_to_dict(i::Index)
         # For some reason if you don't do the isempty check, it gets serialised as `{}`
         # rather than `[]`
         "ix" => isempty(i.ix) ? [] : collect(map(index_to_dict, i.ix)),
-        # TODO(penelopeysm): This is potentially lossy since order is not guaranteed
-        "kw" => Dict(String(x) => index_to_dict(y) for (x, y) in pairs(i.kw)),
+        "kw" => OrderedDict(String(x) => index_to_dict(y) for (x, y) in pairs(i.kw)),
         "child" => optic_to_dict(i.child),
     )
 end
@@ -144,7 +148,7 @@ function varname_to_dict(vn::VarName)
     return Dict("sym" => getsym(vn), "optic" => optic_to_dict(getoptic(vn)))
 end
 
-function dict_to_varname(dict::Dict{<:AbstractString,Any})
+function dict_to_varname(dict::AbstractDict{<:AbstractString,Any})
     return VarName{Symbol(dict["sym"])}(dict_to_optic(dict["optic"]))
 end
 
@@ -178,4 +182,4 @@ Convert a string representation of a `VarName` back to a `VarName`. The string
 should have been generated by `varname_to_string`.
 """
 string_to_varname(str::AbstractString) =
-    dict_to_varname(JSON.parse(str; dicttype=Dict{String,Any}))
+    dict_to_varname(JSON.parse(str; dicttype=OrderedDict{String,Any}))

test/varname/serialize.jl

@@ -38,6 +38,8 @@ using Test
             @varname(z[2:5, :], false),
             @varname(z[2:5, :], true),
             @varname(x[i=1]),
+            @varname(x[j=2, i=1]),
+            @varname(x[i=1, j=2]),
             @varname(x[].a[j=2].b[3, 4, 5, [6]]),
         ]
         for vn in vns
@@ -76,6 +78,15 @@ using Test
 
         # Serialisation should now work
         @test string_to_varname(varname_to_string(vn)) == vn
+
+        # Delete the methods to avoid side effects when running tests again.
+        Base.delete_method(which(AbstractPPL.index_to_dict, (InvertedIndex{Int},)))
+        Base.delete_method(
+            which(
+                AbstractPPL.dict_to_index,
+                (Val{Symbol("InvertedIndices.InvertedIndex")}, Dict),
+            ),
+        )
     end
 end