Go - Transformerを紐解く (google.org/x/text/transform) & テキスト置換処理の実装

この記事はWano Group Advent Calendar 2023の18日目の記事となります。

自己紹介

TuneCore Japanでバックエンドエンジニアをしている@_tachi_です。
今年9月より入社し、主にGoやPerlを使って日々開発をしております
HIPHOP（特にtrap/drill系）と服が好きです。

概要

業務としてGoに触れて3ヶ月が経過しました。
Simpleでいい言語だと思います（≠ Easy）。
個人としては2019年ごろに触れていたので、数年ぶりのGoとなりました。
ジェネリクスなどが導入されたおかげで数年前に比べて表現力が高まったように感じています。

前置きはさておき、個人的に気になっているgolang.org/x/text/transformというパッケージがあります。

入力を byte stream として自分で様々な独自処理を挟み込めるようです。
ただし公式ドキュメントを見ての通り、どのように扱うのか分かりにくいです。
（Exampleくらい用意してくれてもいい気がするのですが...）
オンライン上に公開されている使ってみた系の記事では、すでに内容を把握した上で実装されてた内容が多く、基本的な部分について多くは触れられていません。
この記事ではなるべく基本的な部分にフォーカスを当て、どのような動作をするパッケージなのか紐解いてみようと思います。

本記事の最後に、他の先駆者たちに倣って私もテキスト置換処理を実装したので公開します。

GoのTransformerを紐解く

様々な検証を行いTransformerの動作を確認していきます。
以下に登場するサンプルは全て MyTransformer を定義し、Transformerインターフェースを満たすため、Transformメソッドを実装します。

1. 引数について

まずは引数に何が渡されるのか出力して理解するところから始めます。

package main

import (
    "fmt"
)

type MyTransformer struct {
    transformer.NoResetState
}

func (t MyTransformer) Transform(dst, src []byte, atEOF bool) (nDst, nSrc int, err error) {
    fmt.Printf("dst: %v\n", dst)
    fmt.Printf("dst len: %v\n", len(dst))
    fmt.Printf("nDst: %v\n", nDst)
    fmt.Printf("src: %v\n", src)
    fmt.Printf("src len: %v\n", len(src))
    fmt.Printf("nSrc: %v\n", nSrc)
    fmt.Printf("atEOF: %v\n", atEOF)

  return
}

func main() {
    s := "Hello World"
    r := strings.NewReader(s)
    b := make([]byte, len(s))
    t := transform.NewReader(r, MyTransformer{})
    _, err := t.Read(b)
    if err != nil {
        fmt.Println(err.Error())
	}
}

>> go run main.go
#4096個の0が出力される
dst: [0 0 0 0 0 ...]
dst len: 4096
nDst: 0
# Hello World のバイトデータ
src: [72 101 108 108 111 32 87 111 114 108 100]
src len: 11
nSrc: 0
# 
atEOF: false
err: transform: inconsistent byte count returned

引数で受け渡した Hello World が src に渡されていることが確認できます。
dst に関しては、4096byte 分の領域が予め確保されているため len(dst) の出力が 4096 となっています。

このサイズについては transformer の実装で const defaultBuffSize = 4096 と指定されています。
transform.NewReader(io.Reader, transformer.Transform) を呼び出してた時点で、make([]byte, defaultBufSize) で確保されています。

最後に transform: inconsistent byte count returned というエラーが出力されました。
Transform 関数は nDst, nSrc を戻り値として受け取ります。
この2つは出力時点でどちらも0でした。
この2つの変数については次節で検証します。
取り急ぎ以下のコードを加え、エラーが発生しないようにします。

func (t MyTransformer) Transform(dst, src []byte, atEOF bool) (nDst, nSrc int, err error) {
    // 省略
    n := len(src)
    nDst, nSrc = n, n
    return
}

エラーの発生がなくなり、無事 Transform 関数が正常に実行されるようになりました。

2. 戻り値について

Transform 関数は3つの戻り値 (nDst, nSrc int, err error) を受け取ります。
名前付きの戻り値なので、return　nDst,　nSrc,　err と書かずに return のみで値を返すことが可能です。

この3つの戻り値はどれも重要です。
戻り値の値自体は厳密に考えず、適当に値を直打ちして Transform 関数の動作を確認してみます。

func (t myTransformer) Transformer(dst, src []byte, atEOF bool) (nDst, nSrc int, err error) {
    // 全部ゼロ
    return 0, 0, nil
}

>> go run main.go
err: transform: inconsistent byte count returned

引数の値について検証している際にも同じエラーに遭遇しました。
こちらのエラーは、Transform のコメントや、エラー定義のコメントから分かる通り、errがnilならlen(src)の値が入っていないとエラー となります。

※ Transformerのコメントより
// A nil error means that all of the transformed bytes (whether freshly
// transformed from src or left over from previous Transform calls)
// were written to dst. A nil error can be returned regardless of
// whether atEOF is true. ここから>> If err is nil then nSrc must equal len(src);
// the converse is not necessarily true.

つまり前節の後半で追加した len(src) の値を返すことが適切な処理であることがここで明らかになりました。
指示通り、nSrc には len(src) を返します。

func (t MyTransformer) Transform(dst, src []byte, atEOF bool) (nDst, nSrc int, err error) {
    return 0, len(src), nil
}

>> go run main.go
err: EOF

おや、またエラーが出てきてしまいましたね
EOF とだけ出力されました。
コメント通り、nSrc に len(src) の値を返したのですが、他にも何か必要なのでしょうか？
EOFとあるように、何やら atEOF と関連がありそうです。
atEOF の値に注目してみます。

func (t MyTransformer) Transform(dst, src []byte, atEOF bool) (nDst, nSrc int, err error) {
    fmt.Printf("atEOF: %v", atEOF)
    return 0, len(src), nil
}

>> go run main.go
atEOF: false
atEOF: true
err: EOF

実は処理が2回走っていたことが分かりました。
2回目の Transform 関数実行時、atEOF == true となっています。
このあたりの実態は以下の通りです。

つまり atEOF == io.EOF であり、Readメソッドを呼び出すと、src を読み切るか、何かしらエラーを返すまで forループで MyTransformer の Transform メソッドを実行します。　
また、nDst の値も適切に返す必要がありそうです。
src は中身を読み切ったので EOF となったが、nDst の値が変わっていないとエラーとなるようです。

nDst については Transformer のコメント通り、src から dst に対して変換されたbyte数がセットされることを期待しています。

// ErrShortDst means that the destination buffer was too short to
// receive all of the transformed bytes.
ErrShortDst = errors.New("transform: short destination buffer")

各種適切なハンドリングが必要になりますので、以下のコードを追加し、正常終了できる状態に修正します。

func (t MyTransformer) Transform(dst, src []byte, atEOF bool) (nDst, nSrc int, err error) {
    if atEOF {
		return
	}

    n := len(src)
    nDst, nSrc := n, n
    return
}

>> go run main.go

エラーメッセージが表示されなくなりましたね

3. 簡単な実装を通して、入出力の関係を把握する

ここまで引数と戻り値について確認してきました。
ただ src や dst に対してまだ何もデータの操作をしていません。
今度は簡単な実装を通して Transform の実際の動作について確認していきます。

func (t MyTransformer) Transform(dst, src []byte, atEOF bool) (nDst, nSrc int, err error) {
    if atEOF {
		return
	}
    // 残りの長さが5より大きい間は処理を繰り返す
	for len(src[nSrc:]) >= 6 {
        // srcからdstへ、現在処理しているsrcから 5indexコピーする
        // nは移動に成功した数
		n := copy(dst[nDst:], src[nSrc:6])
		nDst += n
		nSrc += n
	}
    // 最終的にはlen(src)でないとエラーになるため
	nSrc += len(src[nSrc:])

    return
}

>> go run main.go
Hello<スペース>

Hello が出力されました。
処理の意図通り src = "Hello World" から5個ずつ取り出し(Hello )、6文字以上余っているなら処理を繰り返します。
残りは "World" の5文字であるため、ループは終了し Transformメソッドの実行も完了します。
Transformの実装時に注意するポイントは適切に nDst を計算しておくことです。
今回の実装ではあまりに移動数が明らかなので、dstへ変換したバイト数を計算するようなロジックはありませんが、もしこの Transform メソッド内で複雑な処理を行った時は注意です。
例えば計算ミスが起こったことを想定して、最後の戻り値の nDst から適当な値を引いてみます。

func (t MyTransformer) Transform(dst, src []byte, atEOF bool) (nDst, nSrc int, err error) {
    // 省略
    // nDstから適当に長さを引いてみる
    return nDst-2, nSrc, err
}

>> go run main.go
Hell

dst へは5文字分コピーしているはずなので、Hello まで変換されたはずです。
そのため nDst = 5 が正しい値ですが、 nDst が 3 になってしまったことで、 実際に書き込まれた数も nDst 分だけ となってしまいました。

4. 適切にエラーを扱う

Transform メソッドの戻り値に err error があるため、このメソッドはエラーを扱うことを期待しています。
transformパッケージではエラーの定義が公開されています。

ErrShortDst は変換先のバッファが少なすぎる場合にこのエラーを返します。
つまり dst へデータを変換した結果、 dst のサイズが 4096byte を超えてしまうことが分かる場合にこのエラーを返してあげる必要があります。

func (t MyTransformer) Transform(dst, src []byte, atEOF bool) (nDst, nSrc int, err error) {
    // 何かしらの処理をしている...
    
    // `dst` へ移動済みのデータのサイズ + これから追加しようとしているデータのサイズ
    dstSize := dst[nDst:] + len(someByteDataToMoveDst)
    // 4096byteより大きくなるなら `ErrShortDst` を返す
    if dstSize > len(dst) {
       err = ErrShortDst
       return
    }
}

ErrShortSrc の使いどきは単純に使う分には無さそうですが、Transform の用途次第では十分活用する機会がありそうです。
コメントにある通り、変換を行う際に十分なサイズの src が渡って来なかった際にこのエラーを返します。

// ErrShortSrc means that the source buffer has insufficient data to
// complete the transformation.
ErrShortSrc = errors.New("transform: short source buffer")

例えば以下のようなパターンだと活用する機会もあるでしょう。

• 最低 len(src) > 100 欲しいのに、src にはそれ以下しか含まれていなかった
• srcにはある区切り文字が含まれる一連のデータが受け渡され、その単位で処理をしたいが現在処理中のsrcには含まれていなかった
  • goreでは \n まで含まれることを期待するチェックが errfilter.go で実装されています
    • https://github.com/x-motemen/gore/blob/f2e73f6a02de3602bc7bad79ffae8a4a483d3349/errfilter.go#L16-L42

func (t MyTransformer) Transform(dst, src []byte, atEOF bool) (nDst, nSrc int, err error) {
    // 100byte以上あることを期待する
    if len(src[nSrc:]) < 100 {
       err = ErrShortSrc
       return
    }


    // 区切り文字 `,` が含まれることを期待する
    if bytes.IndexByte(src[nSrc:], ",") < 0 {
       err = ErrShortSrc
       return
    }

    // 連続した2つの `//` が含まれることを期待する
    if bytes.Index(src[nSrc:], "//") < 0 {
       err = ErrShortSrc
       return
    }
}

これらのエラーを適切に扱うことで、ErrShortDst が返された場合は、dstに対して十分なバッファを確保した上で再度Transformを実行し、ErrShortSrcが返された場合は、十分なサイズの src を再度 Transform に渡して実行します。

テキスト置換処理の実装

ここまで Transformer の様々な動きについて確認してきました。
nDstやnSrc を適切に計算する、そしてエラーを適切に扱えれば問題ないので実は扱いがシンプルなパッケージであることがわかったと思います。
極論byteデータに対してどんな処理をしていても、適切な値が返されてれば大丈夫なのです。(何もしていなくても成り立つことが 1. 引数について の後半のコードから確認できます)

以上のことを踏まえて1つ簡単なテキスト置換を行う Transformer を実装してみます。
他の方と同じ実装にならないような形を意識しました。
置換データのindexにフォーカスしたテキスト置換 Transformer です。

PlayGroundよりお試しいただけます
https://goplay.tools/snippet/dEXUqSq4ARP

以下同様の置換処理を行う標準ライブラリとの比較ベンチマークです。
テキスト置換に関しては標準のstrings.Replacerの方が速度の観点では早いです。
strings.Replacer > 今回実装したReplacer > bytes.ReplaceAll

package main

import (
	"bytes"
	"strings"
	"testing"

	"golang.org/x/text/transform"
)

func Benchmark_Replacer(t *testing.B) {

	var buf bytes.Buffer

	ct := transform.Chain(
		NewReplacer([]byte(`お寿司`), []byte(`お肉`)),
		NewReplacer([]byte(`🍣`), []byte(`🍖🍖🥩`)),
		NewReplacer([]byte(`明日`), []byte(`明々後日`)),
		NewReplacer([]byte(`な`), []byte(`ゼ`)),
	)

	w := transform.NewWriter(
		&buf,
		ct,
	)
	s := bytes.Repeat([]byte("明日はお寿司を沢山食べたいな🍣\n"), 10000)
	w.Write(s)

}

func Benchmark_StringReplacer(t *testing.B) {

	s := strings.Repeat("明日はお寿司を沢山食べたいな🍣\n", 10000)
	r := strings.NewReplacer(
		`お寿司`, `お肉`,
		`🍣`, `🍖🍖🥩`,
		`明日`, `明々後日`,
		`な`, `ゼ`,
	)

	r.Replace(s)

}

func Benchmark_ByteReplaceAll(t *testing.B) {

	s := bytes.Repeat([]byte("明日はお寿司を沢山食べたいな🍣\n"), 10000)
	maps := []string{
		`お寿司`, `お肉`,
		`🍣`, `🍖🍖🥩`,
		`明日`, `明々後日`,
		`な`, `ゼ`,
	}

	for i := 0; i < len(maps)/2; i++ {
		s = bytes.ReplaceAll(s, []byte(maps[i*2+0]), []byte(maps[i*2+1]))
	}

}

>> go test -bench . -benchmem
goos: darwin
goarch: arm64
pkg: github.com/dsktchr/go-scripts/transformer
Benchmark_Replacer-10           1000000000               0.004428 ns/op        0 B/op          0 allocs/op
Benchmark_StringReplacer-10     1000000000               0.002178 ns/op        0 B/op          0 allocs/op
Benchmark_ByteReplaceAll-10     1000000000               0.005672 ns/op        0 B/op          0 allocs/op
PASS
ok      github.com/dsktchr/go-scripts/transformer       0.475s

テキスト置換の処理内容が分かりにくい人のために

PlayGround上で確認したい値をPrintしていただけるとより分かりやすいかなとは思いますが、なるべく可視化したものも記載しておきます

日本語はマルチバイト文字です　
0,1,2...のようなindexに"肉"などは格納されていません
以下の説明は理解を助けるための内容にすぎませんのでご了承ください

// 以下の結果になるようにテキストを変換する
Expect = "明日はお肉を沢山食べたいな🍣"
----------------------------------------------------------------
Old    = "お寿司"　（3文字）
New    = "お酒"　　　　(2文字)
Input  = "明日はお寿司を沢山食べたいな🍣"

置換先がInputから無くなるまで処理を繰り返す
----------------------------------------------------------------
[Loop - 1]

// Dstへの移動数
_nDst = 0
// Srcから既に読み込みが完了したサイズ
_nSrc = 0

// Dstへの書き込み上限が4096
Dst(0/4096):
まだ何もない

Src:
| 　0　　|　 1 |　 2　 |　　　3 |　　4　　　| 5　 |　 6 　| 　7 　| 　8 | 　9　 | 10 | 　11　　| 12 | 13 | 14 |　　
| 明 | 日 | は | お | 寿 | 司 | を | 　沢 | 山 | 食 | 　べ | た  |　 い　　|　 な | 　🍣　 |

// 置換対象データの `開始位置`　を取得
Index: 3 (開始位置(Src[ 0(_nSrc): ], "お寿司"))

// 置換対象（お寿司） `まで` を Dst へ
Dst[ 0(_nDst): ]  　　       <---  Copy ---  Src[ 0(_nSrc) : 6(_nSrc + Index + Oldの長さ) ]
| 　0　　|　 　1 |　 2　 |　 3 | 　4 | 　5 |             | 　0 　|　 1 |　 2　 | 3 | 　4 　| 　5　 |
| 明 | 日 | 　は | お | 寿 | 司 |             | 明 | 日 | は | お | 寿 | 司　 |

// Dstへ '6' つ移動した
// 移動数の記録を残す
N = 6

// 置換前のデータの `開始位置` に置換対象のデータをcopy　
Dst[ 3(_nDst + Index): ]  <--- Copy --- New
| 　3　　|　 4 | 　5 |　　　　　　　　　　　　　　　　　　　　　　　　　　　　　　　　　　　　　　　　　　   　 | 　0　　|　 1 |　　　　　　 　　
| お | 寿 | 司 |　　　　　　　　　　　　　　　　　　　　　　　　　　　　　　　　　　　　　　　　　　    | お | 肉 |

// Dstの中身はこのように変化している
Dst:
| 　0　　|　 　1 |　 2　 |　 3 | 　4 |
| 明 | 日 | 　は | お | 肉 |   

// Dstに対して2回のデータの移動（copy）を行なっている
// Dstへの移動数は
// これまでのDstへの移動数 + 最初の置換対象までの移動数(N) - (置換元の長さ(Old) - 置換先の長さ(New))
// で表現することができる
_nDst = 0(_nDst) + N - (3 - 2)
_nDst: 5

// SrcからDstへ `明日はお寿司` を移動したので `6`
_nSrc = N

----------------------------------------------------------------
[Loop - 2]

_nDst = 5
_nSrc = 6

Dst(5/4096):
| 　0　　|　 　1 |　 2　 |　 3 | 　4 |
| 明 | 日 | 　は | お | 肉 |  


Src:
| 　0　　|　 1 |　 2　 |　　　3 |　　4　　　| 5　 |　 6 　| 　7 　| 　8 | 　9　 | 10 | 　11　　| 12 | 13 | 14 |　　
| 明 | 日 | は | お | 寿 | 司 | を | 　沢 | 山 | 食 | 　べ | た  |　 い　　|　 な | 　🍣　 |
|          前回マッチ         |                   未処理                    |

// 2回目のループの際は、前回処理した位置以降から探す
// Srcの未処理部分には置換元となる `お寿司` がないため開始位置がない(-1)
Index: -1 (開始位置(Src[ 6(_nSrc): ], "お寿司"))

// マッチする内容がないため、Src残りのデータを全てDstへ移動し置換処理を終了とする
// Dstは`Index=5`はまだないため、データの置換ではなく、追加のような動きになる
Dst[ 5(_nDst) : ]  <--- Copy --- Src[ 6(_nSrc): ]
|　　　5 | ... |                  |　 6 　| 　7 　| 　8 | 　9　 | 10 | 　11　　| 12 | 13 | 14 |　
|　　を | ... |                  | を | 沢 | 山 | 食 | 　べ | た |　 い　　|　 な | 　🍣　 |

// Dstへ `9` 移動した
N = 9

// 最後に移動数を確定する
_nDst = 5(_nDst) + 9(N)
_nSrc = 6(_nSrc) + 9(N)

// 長さが14なので_nDstと等しい
// _nDstが正しく計算されたことが分かる
Dst:
| 　0　　|　 1 |　 2　 |　　　3 |　　4　　　| 5　 |　 6 　| 　7 　| 　8 | 　9　 | 10 | 　11　　| 12 | 13 |
| 明 | 日 | は | お | 肉 | を | 　沢 | 山 | 食 | 　べ | た  |　 い　　|　 な | 　🍣　 |

最後に

以上、Goの Transformer を紐解くでした。
byteデータに対する操作は何でもいけるため、テキスト置換以外の活躍の幅は多いと思っています。
皆さんも何かの変換処理を実装する際は是非ともTransformer を活用してみてください

参考文献

人材募集

現在、Wanoグループ / TuneCore Japan では人材募集をしています。興味のある方は下記を参照してください。

Wano | Wano Group JOBS

TuneCore Japan | TuneCore Japan JOBS