https://itchyny.hatenablog.com/entry/2021/01/20/100000

Go言語で書かれたorderedmapというサードパーティパッケージがあります。 github.com

Goのmapには順序がなく、JSONをデコードすると順序が失われ、それをエンコードするとオブジェクトのキーの順序にソートされます。これに困る人はそこそこいるようで、順序を保持するmapはいくつか実装されてきました。その中の一つが、orderedmapというパッケージです。シンプルなインターフェイスが気に入っています。

orderedmapパッケージの利用例

package main

import (
    "encoding/json"
    "fmt"
    "log"

    "github.com/iancoleman/orderedmap"
)

func main() {
    src := `{ "z": 1, "x": 2, "y": 3 }`

    fmt.Println("# map[string]interface{}")
    var v map[string]interface{}
    if err := json.Unmarshal([]byte(src), &v); err != nil {
        log.Fatal(err)
    }
    bs, err := json.MarshalIndent(v, "", "  ")
    if err != nil {
        log.Fatal(err)
    }
    fmt.Printf("%s\n\n", bs)

    fmt.Println("# OrderedMap")
    o := orderedmap.New()
    if err := json.Unmarshal([]byte(src), &o); err != nil {
        log.Fatal(err)
    }
    bs, err = json.MarshalIndent(o, "", "  ")
    if err != nil {
        log.Fatal(err)
    }
    fmt.Printf("%s\n\n", bs)
}

# map[string]interface{}
{
  "x": 2,
  "y": 3,
  "z": 1
}

# OrderedMap
{
  "z": 1,
  "x": 2,
  "y": 3
}

しかし、このパッケージのUnmarshal (JSON → OrderedMap の変換) の実装は、正直あまりきれいではありませんでした。 v0.1.0までの実装は以下のようになっていました。

map[string]interface{} に Unmarshal する
オブジェクトの各キーに対して、JSON文字列の中の出現位置を探す
- キーの " をエスケープして、 strings.LastIndex で探す
- JSON文字列を探したインデックスまでを切り取って } を付け加えたときにvalidなJSONかチェックする
  - validなJSONでなければ、ネストされたJSONのキーにヒットしまっているので、更に前にたどってキーのポジションを探す
オブジェクトの各値に対して
- オブジェクトであれば、OrderedMapへの変換を再帰的に行う
- 配列であれば、各要素の変換を再帰的に行う
オブジェクトのキー一覧を、出現位置でソートする

オブジェクトのキーの順序を保持したい → キーの文字列の出現位置を探索してソートすれば良い、という素朴な発想でこういう実装になっているのだと思いますが、この処理には様々な問題があります。

オブジェクトのキーを文字列から探すときに " しかエスケープしておらず、{ "\n": 1 } のようなものはうまくUnmarshalできない
- { "A\u0041A\u0041A": 0 } のようなキーもありうるので、strings.LastIndex でポジションを探す方針には限界がある (指数関数的な個数のエスケープを試す必要がある)
パフォーマンスが良くない (#2)
- validなJSONかかどうかをチェックするという処理でいちいち json.Unmarshal している (orderedmap.go#L178-L192)。
- 探索範囲を絞ったりスペースをスキップしたりするために、文字列のアロケーションが多い
ネストしたJSONに対処するために実装が複雑で、一見してバグっていない自信がない (主観)
- 重複キーに対しては結構怪しい気がする

個人的には好きなパッケージですが、Unmarshalの実装が複雑であまり使いたくないという気持ちがありました。しかし作者も課題感を持っているようでしたし、解決策も思いついたので、実装し直すことにしました。

標準パッケージのencoding/jsonの*DecoderにはToken() (Token, error)という関数があります。 JSONのトークンを一つずつデコードする関数です。 *DecoderはJSONのステートを持っているので、単なるlexerではなくてvalidなJSONのトークン列を返してくれます。これを使います。

map[string]interface{} に Unmarshal する
JSONから *Decoder を作り、dec.Token() を使ってトークンをたどっていく
- キーはスライスに追加する
  - 重複キーは後ろに移動する
- 値の最初のトークンによって、オブジェクトか配列かその他の場合に分かれる
  - 値がオブジェクトの場合
    - map[string]interface{}またはOrderedMap (重複キー) の場合はOrderedMapに変換する処理を行う
    - それ以外の場合 (型が違う) は重複キーなので、OrderedMapへのデコードを行うが得られた値は捨てる
  - 値が配列の場合
    - 配列の各要素に対してオブジェクトか配列の場合は再帰的にデコードを行う
    - それ以外の場合 (型が違う、配列のインデックスを超えているなど) は重複キーなので、デコードを行うが得られた値は捨てる

この実装によって様々な点が改善されました。

オブジェクトのキーがエスケープされていても問題なく動く
- { "A\u0041A\u0041A": 0 } のようなものでもOK
以前の実装では何度も文字列上を走査する必要があったが、この実装では二回スキャンすればよい
- 実装を突き詰めてJSONパースを自前で持てば一回で済むが、そこはencoding/jsonに任せて小さく実装している
オブジェクトのキー一覧をソートする必要がない
- ただし重複キーが大量にある場合のパフォーマンスは良くない
文字列のアロケーションが大幅に減少し、パフォーマンスも改善される
実装が比較的読みやすい

ベンチマークは以下のように改善されています。アロケーションは半分以下に抑えられ、パフォーマンスも三倍近く改善しています。

benchmark                     old ns/op     new ns/op     delta
BenchmarkUnmarshalJSON-16     125485        37536         -70.09%

benchmark                     old allocs    new allocs    delta
BenchmarkUnmarshalJSON-16     964           389           -59.65%

benchmark                     old bytes     new bytes     delta
BenchmarkUnmarshalJSON-16     49885         13174         -73.59%

一方、Marshal (OrderedMap → JSONへの変換) にも改善を入れました。 Unmarshal ほどの問題はありませんでしたが、オブジェクトのキーのエスケープが足りていない ({ "\n": 1 } 相当を出力すると改行が入りvalidなJSONにならない) とか、文字列のアロケーションが多いといった問題を解決しています。文字列の結合をbytes.Bufferに置き換える単純なお仕事です。

ベンチマークは以下のようになっています。アロケーションはかなり減りましたが、パフォーマンスとしては7%程度の改善にとどまっています。

benchmark                   old ns/op     new ns/op     delta
BenchmarkMarshalJSON-16     15687         14535         -7.34%

benchmark                   old allocs    new allocs    delta
BenchmarkMarshalJSON-16     94            33            -64.89%

benchmark                   old bytes     new bytes     delta
BenchmarkMarshalJSON-16     9309          2195          -76.42%

以上の改善は、v0.2.0としてリリースされています。個人的には実装がかなりクリアになり (まあ自分で書いたし)、パフォーマンスも良くなったので、安心してこのパッケージを使えるようになりました。また一つ世界が良くなりましたね。おしまい！