Go言語のorderedmapパッケージを改善した

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 (JSONOrderedMap の変換) の実装は、正直あまりきれいではありませんでした。 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トークンを一つずつデコードする関数です。 *DecoderJSONのステートを持っているので、単なる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 (OrderedMapJSONへの変換) にも改善を入れました。 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としてリリースされています。 個人的には実装がかなりクリアになり (まあ自分で書いたし)、パフォーマンスも良くなったので、安心してこのパッケージを使えるようになりました。 また一つ世界が良くなりましたね。 おしまい!