文字列を段階的に組み立てる

プログラムで文字列を組み立てていると、 "(" ^ なにか ^ ")" のように、文字列結合で段階的に結果を組み立てたいことがある。

文字列が小さいことがわかっているうちはこれでもよいのだけれど、文字列が大きくなってくると段々と計算量が気になってくる（特に最適化をしていなければ、文字列の長さを n とすると結合のたびに O(n) のコストがかかる）。

左から右に組み立てることがわかっていれば、OCamlなら Buffer モジュールが使えるが、括弧のようなものがからむとそうも言っていられない。

stringを使うのをやめて、ropeのようなデータ構造を使う手もあるけど、依存ライブラリを増やしたくない場合もある（言語処理系によっては文字列の内部表現がropeになっていて以下のような小細工をする必要がない場合もある）。

ということで、自分でデータ構造をつくってみる。

type t =
  | Leaf of int * string
  | Concat of int * t * t

型 t は ^ による文字列結合の構文木を表現したもので、 Concat が ^、 Leaf が文字列値に対応する。

ついでに、文字列化のときの便利のために結合後の文字列長を持たせておく。

let length = function
  | Leaf (i, _) -> i
  | Concat (i, _, _) -> i

長さを手で与えると間違えるので、smart constructor的な関数を用意する。

let leaf s =
  Leaf (String.length s, s)

let empty = leaf ""

let (^) t1 t2 =
  Concat (length t1 + length t2, t1, t2)

文字列を結合するたびに leaf と書くのも面倒なので、文字列を追加するための演算子も定義する。

let (@+) s t =
  leaf s ^ t

let (+@) t s =
  t ^ leaf s

"a" @+ "b" @+ x は "a" @+ ("b" @+ x) のように解釈されてほしいので右結合になるように演算子名を @ で始めておく（OCamlの演算子の結合順位は、演算子の初めの文字で決まる）。逆に x +@ "a" +@ "b" は (x +@ "a") +@ "b" になってほしいので、 @+ の反対の +@ にしておく。こうしておくと、 + で始まる演算子は右結合で @ で始まる演算子よりも結合が強いので、 "a" @+ x +@ "b" +@ "c" が "a" @+ ((x +@ "b") +@ "c") のようになり、うまい具合いに型がつくようになるし、文字列の側に @ をつけるという規則になって見た目にも覚えやすい。

これで、結合は O(1) でできるようになった。

あとはこれを文字列化できるようにする。

let to_string t =
  let buf = Buffer.create @@ length t in
  let rec loop = function
    | Leaf (_, s) -> Buffer.add_string buf s
    | Concat (_, t1, t2) ->
      loop t1;
      loop t2
  in
  loop t;
  Buffer.contents buf

ここは素直に副作用をつかって Buffer で組み立てる。固定長のバッファを端から埋めているので O(n) だろう。

to_string を見ていると、 loop が末尾再帰になっていないのがちょっと気になる。

^ は結合的な演算なので、 (a ^ b) ^ c = a ^ (b ^ c) だ。さらに、 a ^ (b ^ (c ^ empty)) のように変形していけば list と同型にできるので、 List.iter と同じように末尾再帰でループをまわせそうな気がする。

let to_string' t =
  let buf = Buffer.create @@ length t in
  let rec loop rest = function
    | Leaf (_, s) ->
      Buffer.add_string buf s;
      begin match rest with
      | r::rs ->
        loop rs r
      | [] ->
        Buffer.contents buf
      end
    | Concat (_, t1, t2) ->
      loop (t2::rest) t1
  in loop [] t

構造を一度に変換する必要はないので、 Concat に出逢うたびに右部分木をリストに押し込んでおいて、少しずつ構造を変換していく。

使い方はこんな風になる。

let () =
  ("1" @+ empty) +@ "2" |> to_string' |> print_endline;
  "1" @+ (empty +@ "2") |> to_string' |> print_endline;
  "1" @+ (empty +@ "2" +@ "3") |> to_string' |> print_endline;
  ("1" @+ empty +@ "2") +@ "3" |> to_string' |> print_endline;
  empty +@ "1" +@ "2" +@ "3" +@ "4" |> to_string' |> print_endline;
  let x = leaf "xy" in
  let x = "{" @+ "(" @+ x +@ ")" +@ "}" in
  x |> to_string' |> print_endline;
  ()

ropeや他のデータ構造にしておけば、結合以外に文字列の分割などもサポートできるが、結合と文字列化だけに機能を搾れば比較的簡単に効率のよい実装ができる。

実際に使ってみると、 OCaml の関数適用演算子と優先順位がかぶって使いにくいという問題があったがそれは御愛嬌ということで（関数適用は (@@) でなくて (&) の方がいいと言う人の気持ちがよくわかった）。

おまけ: String.concat 版

Standard MLにはOCamlの Buffer のようなものがないらしいので、Standard MLに移植しやすそうな String.concat を使った to_string も書いてみた。

let to_string'' t =
  let rec loop acc rest = function
    | Leaf (_, s) ->
      begin match rest with
      | [] -> s::acc
      | r::rs -> loop (s::acc) rs r
      end
    | Concat (_, t1, t2) ->
      loop acc (t1::rest) t2
  in
  loop [] [] t |> String.concat ""

末尾再帰版と同じ方針で右部分木から処理していけばリストを後ろから組み立てていける。

String.concat は最初に結果文字列の長さ分の領域を確保し、順番に埋めていくので O(n) だ（OCaml, SML#, MLton調べ）。

追記: Standard MLでもBasisライブラリへのBufferモジュールの追加は議論されているようだ。2015 004 Addition of Buffer module · SMLFamily/BasisLibrary Wiki

参考

• Gaucheではこのような処理がtext.treeモジュールとして提供されている。