libgo/go/strings/replace.go

   1 // Copyright 2011 The Go Authors. All rights reserved.
   2 // Use of this source code is governed by a BSD-style
   3 // license that can be found in the LICENSE file.
   4
   5 package strings
   6
   7 import "io"
   8
   9 // A Replacer replaces a list of strings with replacements.
  10 type Replacer struct {
  11         r replacer
  12 }
  13
  14 // replacer is the interface that a replacement algorithm needs to implement.
  15 type replacer interface {
  16         Replace(s string) string
  17         WriteString(w io.Writer, s string) (n int, err error)
  18 }
  19
  20 // byteBitmap represents bytes which are sought for replacement.
  21 // byteBitmap is 256 bits wide, with a bit set for each old byte to be
  22 // replaced.
  23 type byteBitmap [256 / 32]uint32
  24
  25 func (m *byteBitmap) set(b byte) {
  26         m[b>>5] |= uint32(1 << (b & 31))
  27 }
  28
  29 // NewReplacer returns a new Replacer from a list of old, new string pairs.
  30 // Replacements are performed in order, without overlapping matches.
  31 func NewReplacer(oldnew ...string) *Replacer {
  32         if len(oldnew)%2 == 1 {
  33                 panic("strings.NewReplacer: odd argument count")
  34         }
  35
  36         if len(oldnew) == 2 && len(oldnew[0]) > 1 {
  37                 return &Replacer{r: makeSingleStringReplacer(oldnew[0], oldnew[1])}
  38         }
  39
  40         allNewBytes := true
  41         for i := 0; i < len(oldnew); i += 2 {
  42                 if len(oldnew[i]) != 1 {
  43                         return &Replacer{r: makeGenericReplacer(oldnew)}
  44                 }
  45                 if len(oldnew[i+1]) != 1 {
  46                         allNewBytes = false
  47                 }
  48         }
  49
  50         if allNewBytes {
  51                 bb := &byteReplacer{}
  52                 for i := 0; i < len(oldnew); i += 2 {
  53                         o, n := oldnew[i][0], oldnew[i+1][0]
  54                         if bb.old[o>>5]&uint32(1<<(o&31)) != 0 {
  55                                 // Later old->new maps do not override previous ones with the same old string.
  56                                 continue
  57                         }
  58                         bb.old.set(o)
  59                         bb.new[o] = n
  60                 }
  61                 return &Replacer{r: bb}
  62         }
  63
  64         bs := &byteStringReplacer{}
  65         for i := 0; i < len(oldnew); i += 2 {
  66                 o, new := oldnew[i][0], oldnew[i+1]
  67                 if bs.old[o>>5]&uint32(1<<(o&31)) != 0 {
  68                         // Later old->new maps do not override previous ones with the same old string.
  69                         continue
  70                 }
  71                 bs.old.set(o)
  72                 bs.new[o] = []byte(new)
  73         }
  74         return &Replacer{r: bs}
  75 }
  76
  77 // Replace returns a copy of s with all replacements performed.
  78 func (r *Replacer) Replace(s string) string {
  79         return r.r.Replace(s)
  80 }
  81
  82 // WriteString writes s to w with all replacements performed.
  83 func (r *Replacer) WriteString(w io.Writer, s string) (n int, err error) {
  84         return r.r.WriteString(w, s)
  85 }
  86
  87 // trieNode is a node in a lookup trie for prioritized key/value pairs. Keys
  88 // and values may be empty. For example, the trie containing keys "ax", "ay",
  89 // "bcbc", "x" and "xy" could have eight nodes:
  90 //
  91 //  n0  -
  92 //  n1  a-
  93 //  n2  .x+
  94 //  n3  .y+
  95 //  n4  b-
  96 //  n5  .cbc+
  97 //  n6  x+
  98 //  n7  .y+
  99 //
 100 // n0 is the root node, and its children are n1, n4 and n6; n1's children are
 101 // n2 and n3; n4's child is n5; n6's child is n7. Nodes n0, n1 and n4 (marked
 102 // with a trailing "-") are partial keys, and nodes n2, n3, n5, n6 and n7
 103 // (marked with a trailing "+") are complete keys.
 104 type trieNode struct {
 105         // value is the value of the trie node's key/value pair. It is empty if
 106         // this node is not a complete key.
 107         value string
 108         // priority is the priority (higher is more important) of the trie node's
 109         // key/value pair; keys are not necessarily matched shortest- or longest-
 110         // first. Priority is positive if this node is a complete key, and zero
 111         // otherwise. In the example above, positive/zero priorities are marked
 112         // with a trailing "+" or "-".
 113         priority int
 114
 115         // A trie node may have zero, one or more child nodes:
 116         //  * if the remaining fields are zero, there are no children.
 117         //  * if prefix and next are non-zero, there is one child in next.
 118         //  * if table is non-zero, it defines all the children.
 119         //
 120         // Prefixes are preferred over tables when there is one child, but the
 121         // root node always uses a table for lookup efficiency.
 122
 123         // prefix is the difference in keys between this trie node and the next.
 124         // In the example above, node n4 has prefix "cbc" and n4's next node is n5.
 125         // Node n5 has no children and so has zero prefix, next and table fields.
 126         prefix string
 127         next   *trieNode
 128
 129         // table is a lookup table indexed by the next byte in the key, after
 130         // remapping that byte through genericReplacer.mapping to create a dense
 131         // index. In the example above, the keys only use 'a', 'b', 'c', 'x' and
 132         // 'y', which remap to 0, 1, 2, 3 and 4. All other bytes remap to 5, and
 133         // genericReplacer.tableSize will be 5. Node n0's table will be
 134         // []*trieNode{ 0:n1, 1:n4, 3:n6 }, where the 0, 1 and 3 are the remapped
 135         // 'a', 'b' and 'x'.
 136         table []*trieNode
 137 }
 138
 139 func (t *trieNode) add(key, val string, priority int, r *genericReplacer) {
 140         if key == "" {
 141                 if t.priority == 0 {
 142                         t.value = val
 143                         t.priority = priority
 144                 }
 145                 return
 146         }
 147
 148         if t.prefix != "" {
 149                 // Need to split the prefix among multiple nodes.
 150                 var n int // length of the longest common prefix
 151                 for ; n < len(t.prefix) && n < len(key); n++ {
 152                         if t.prefix[n] != key[n] {
 153                                 break
 154                         }
 155                 }
 156                 if n == len(t.prefix) {
 157                         t.next.add(key[n:], val, priority, r)
 158                 } else if n == 0 {
 159                         // First byte differs, start a new lookup table here. Looking up
 160                         // what is currently t.prefix[0] will lead to prefixNode, and
 161                         // looking up key[0] will lead to keyNode.
 162                         var prefixNode *trieNode
 163                         if len(t.prefix) == 1 {
 164                                 prefixNode = t.next
 165                         } else {
 166                                 prefixNode = &trieNode{
 167                                         prefix: t.prefix[1:],
 168                                         next:   t.next,
 169                                 }
 170                         }
 171                         keyNode := new(trieNode)
 172                         t.table = make([]*trieNode, r.tableSize)
 173                         t.table[r.mapping[t.prefix[0]]] = prefixNode
 174                         t.table[r.mapping[key[0]]] = keyNode
 175                         t.prefix = ""
 176                         t.next = nil
 177                         keyNode.add(key[1:], val, priority, r)
 178                 } else {
 179                         // Insert new node after the common section of the prefix.
 180                         next := &trieNode{
 181                                 prefix: t.prefix[n:],
 182                                 next:   t.next,
 183                         }
 184                         t.prefix = t.prefix[:n]
 185                         t.next = next
 186                         next.add(key[n:], val, priority, r)
 187                 }
 188         } else if t.table != nil {
 189                 // Insert into existing table.
 190                 m := r.mapping[key[0]]
 191                 if t.table[m] == nil {
 192                         t.table[m] = new(trieNode)
 193                 }
 194                 t.table[m].add(key[1:], val, priority, r)
 195         } else {
 196                 t.prefix = key
 197                 t.next = new(trieNode)
 198                 t.next.add("", val, priority, r)
 199         }
 200 }
 201
 202 func (r *genericReplacer) lookup(s string, ignoreRoot bool) (val string, keylen int, found bool) {
 203         // Iterate down the trie to the end, and grab the value and keylen with
 204         // the highest priority.
 205         bestPriority := 0
 206         node := &r.root
 207         n := 0
 208         for node != nil {
 209                 if node.priority > bestPriority && !(ignoreRoot && node == &r.root) {
 210                         bestPriority = node.priority
 211                         val = node.value
 212                         keylen = n
 213                         found = true
 214                 }
 215
 216                 if s == "" {
 217                         break
 218                 }
 219                 if node.table != nil {
 220                         index := r.mapping[s[0]]
 221                         if int(index) == r.tableSize {
 222                                 break
 223                         }
 224                         node = node.table[index]
 225                         s = s[1:]
 226                         n++
 227                 } else if node.prefix != "" && HasPrefix(s, node.prefix) {
 228                         n += len(node.prefix)
 229                         s = s[len(node.prefix):]
 230                         node = node.next
 231                 } else {
 232                         break
 233                 }
 234         }
 235         return
 236 }
 237
 238 // genericReplacer is the fully generic algorithm.
 239 // It's used as a fallback when nothing faster can be used.
 240 type genericReplacer struct {
 241         root trieNode
 242         // tableSize is the size of a trie node's lookup table. It is the number
 243         // of unique key bytes.
 244         tableSize int
 245         // mapping maps from key bytes to a dense index for trieNode.table.
 246         mapping [256]byte
 247 }
 248
 249 func makeGenericReplacer(oldnew []string) *genericReplacer {
 250         r := new(genericReplacer)
 251         // Find each byte used, then assign them each an index.
 252         for i := 0; i < len(oldnew); i += 2 {
 253                 key := oldnew[i]
 254                 for j := 0; j < len(key); j++ {
 255                         r.mapping[key[j]] = 1
 256                 }
 257         }
 258
 259         for _, b := range r.mapping {
 260                 r.tableSize += int(b)
 261         }
 262
 263         var index byte
 264         for i, b := range r.mapping {
 265                 if b == 0 {
 266                         r.mapping[i] = byte(r.tableSize)
 267                 } else {
 268                         r.mapping[i] = index
 269                         index++
 270                 }
 271         }
 272         // Ensure root node uses a lookup table (for performance).
 273         r.root.table = make([]*trieNode, r.tableSize)
 274
 275         for i := 0; i < len(oldnew); i += 2 {
 276                 r.root.add(oldnew[i], oldnew[i+1], len(oldnew)-i, r)
 277         }
 278         return r
 279 }
 280
 281 type appendSliceWriter []byte
 282
 283 // Write writes to the buffer to satisfy io.Writer.
 284 func (w *appendSliceWriter) Write(p []byte) (int, error) {
 285         *w = append(*w, p...)
 286         return len(p), nil
 287 }
 288
 289 // WriteString writes to the buffer without string->[]byte->string allocations.
 290 func (w *appendSliceWriter) WriteString(s string) (int, error) {
 291         *w = append(*w, s...)
 292         return len(s), nil
 293 }
 294
 295 type stringWriterIface interface {
 296         WriteString(string) (int, error)
 297 }
 298
 299 type stringWriter struct {
 300         w io.Writer
 301 }
 302
 303 func (w stringWriter) WriteString(s string) (int, error) {
 304         return w.w.Write([]byte(s))
 305 }
 306
 307 func getStringWriter(w io.Writer) stringWriterIface {
 308         sw, ok := w.(stringWriterIface)
 309         if !ok {
 310                 sw = stringWriter{w}
 311         }
 312         return sw
 313 }
 314
 315 func (r *genericReplacer) Replace(s string) string {
 316         buf := make(appendSliceWriter, 0, len(s))
 317         r.WriteString(&buf, s)
 318         return string(buf)
 319 }
 320
 321 func (r *genericReplacer) WriteString(w io.Writer, s string) (n int, err error) {
 322         sw := getStringWriter(w)
 323         var last, wn int
 324         var prevMatchEmpty bool
 325         for i := 0; i <= len(s); {
 326                 // Ignore the empty match iff the previous loop found the empty match.
 327                 val, keylen, match := r.lookup(s[i:], prevMatchEmpty)
 328                 prevMatchEmpty = match && keylen == 0
 329                 if match {
 330                         wn, err = sw.WriteString(s[last:i])
 331                         n += wn
 332                         if err != nil {
 333                                 return
 334                         }
 335                         wn, err = sw.WriteString(val)
 336                         n += wn
 337                         if err != nil {
 338                                 return
 339                         }
 340                         i += keylen
 341                         last = i
 342                         continue
 343                 }
 344                 i++
 345         }
 346         if last != len(s) {
 347                 wn, err = sw.WriteString(s[last:])
 348                 n += wn
 349         }
 350         return
 351 }
 352
 353 // singleStringReplacer is the implementation that's used when there is only
 354 // one string to replace (and that string has more than one byte).
 355 type singleStringReplacer struct {
 356         finder *stringFinder
 357         // value is the new string that replaces that pattern when it's found.
 358         value string
 359 }
 360
 361 func makeSingleStringReplacer(pattern string, value string) *singleStringReplacer {
 362         return &singleStringReplacer{finder: makeStringFinder(pattern), value: value}
 363 }
 364
 365 func (r *singleStringReplacer) Replace(s string) string {
 366         var buf []byte
 367         i := 0
 368         for {
 369                 match := r.finder.next(s[i:])
 370                 if match == -1 {
 371                         break
 372                 }
 373                 buf = append(buf, s[i:i+match]...)
 374                 buf = append(buf, r.value...)
 375                 i += match + len(r.finder.pattern)
 376         }
 377         if buf == nil {
 378                 return s
 379         }
 380         buf = append(buf, s[i:]...)
 381         return string(buf)
 382 }
 383
 384 func (r *singleStringReplacer) WriteString(w io.Writer, s string) (n int, err error) {
 385         sw := getStringWriter(w)
 386         var i, wn int
 387         for {
 388                 match := r.finder.next(s[i:])
 389                 if match == -1 {
 390                         break
 391                 }
 392                 wn, err = sw.WriteString(s[i : i+match])
 393                 n += wn
 394                 if err != nil {
 395                         return
 396                 }
 397                 wn, err = sw.WriteString(r.value)
 398                 n += wn
 399                 if err != nil {
 400                         return
 401                 }
 402                 i += match + len(r.finder.pattern)
 403         }
 404         wn, err = sw.WriteString(s[i:])
 405         n += wn
 406         return
 407 }
 408
 409 // byteReplacer is the implementation that's used when all the "old"
 410 // and "new" values are single ASCII bytes.
 411 type byteReplacer struct {
 412         // old has a bit set for each old byte that should be replaced.
 413         old byteBitmap
 414
 415         // replacement byte, indexed by old byte. only valid if
 416         // corresponding old bit is set.
 417         new [256]byte
 418 }
 419
 420 func (r *byteReplacer) Replace(s string) string {
 421         var buf []byte // lazily allocated
 422         for i := 0; i < len(s); i++ {
 423                 b := s[i]
 424                 if r.old[b>>5]&uint32(1<<(b&31)) != 0 {
 425                         if buf == nil {
 426                                 buf = []byte(s)
 427                         }
 428                         buf[i] = r.new[b]
 429                 }
 430         }
 431         if buf == nil {
 432                 return s
 433         }
 434         return string(buf)
 435 }
 436
 437 func (r *byteReplacer) WriteString(w io.Writer, s string) (n int, err error) {
 438         // TODO(bradfitz): use io.WriteString with slices of s, avoiding allocation.
 439         bufsize := 32 << 10
 440         if len(s) < bufsize {
 441                 bufsize = len(s)
 442         }
 443         buf := make([]byte, bufsize)
 444
 445         for len(s) > 0 {
 446                 ncopy := copy(buf, s[:])
 447                 s = s[ncopy:]
 448                 for i, b := range buf[:ncopy] {
 449                         if r.old[b>>5]&uint32(1<<(b&31)) != 0 {
 450                                 buf[i] = r.new[b]
 451                         }
 452                 }
 453                 wn, err := w.Write(buf[:ncopy])
 454                 n += wn
 455                 if err != nil {
 456                         return n, err
 457                 }
 458         }
 459         return n, nil
 460 }
 461
 462 // byteStringReplacer is the implementation that's used when all the
 463 // "old" values are single ASCII bytes but the "new" values vary in
 464 // size.
 465 type byteStringReplacer struct {
 466         // old has a bit set for each old byte that should be replaced.
 467         old byteBitmap
 468
 469         // replacement string, indexed by old byte. only valid if
 470         // corresponding old bit is set.
 471         new [256][]byte
 472 }
 473
 474 func (r *byteStringReplacer) Replace(s string) string {
 475         newSize := 0
 476         anyChanges := false
 477         for i := 0; i < len(s); i++ {
 478                 b := s[i]
 479                 if r.old[b>>5]&uint32(1<<(b&31)) != 0 {
 480                         anyChanges = true
 481                         newSize += len(r.new[b])
 482                 } else {
 483                         newSize++
 484                 }
 485         }
 486         if !anyChanges {
 487                 return s
 488         }
 489         buf := make([]byte, newSize)
 490         bi := buf
 491         for i := 0; i < len(s); i++ {
 492                 b := s[i]
 493                 if r.old[b>>5]&uint32(1<<(b&31)) != 0 {
 494                         n := copy(bi[:], r.new[b])
 495                         bi = bi[n:]
 496                 } else {
 497                         bi[0] = b
 498                         bi = bi[1:]
 499                 }
 500         }
 501         return string(buf)
 502 }
 503
 504 // WriteString maintains one buffer that's at most 32KB.  The bytes in
 505 // s are enumerated and the buffer is filled.  If it reaches its
 506 // capacity or a byte has a replacement, the buffer is flushed to w.
 507 func (r *byteStringReplacer) WriteString(w io.Writer, s string) (n int, err error) {
 508         // TODO(bradfitz): use io.WriteString with slices of s instead.
 509         bufsize := 32 << 10
 510         if len(s) < bufsize {
 511                 bufsize = len(s)
 512         }
 513         buf := make([]byte, bufsize)
 514         bi := buf[:0]
 515
 516         for i := 0; i < len(s); i++ {
 517                 b := s[i]
 518                 var new []byte
 519                 if r.old[b>>5]&uint32(1<<(b&31)) != 0 {
 520                         new = r.new[b]
 521                 } else {
 522                         bi = append(bi, b)
 523                 }
 524                 if len(bi) == cap(bi) || (len(bi) > 0 && len(new) > 0) {
 525                         nw, err := w.Write(bi)
 526                         n += nw
 527                         if err != nil {
 528                                 return n, err
 529                         }
 530                         bi = buf[:0]
 531                 }
 532                 if len(new) > 0 {
 533                         nw, err := w.Write(new)
 534                         n += nw
 535                         if err != nil {
 536                                 return n, err
 537                         }
 538                 }
 539         }
 540         if len(bi) > 0 {
 541                 nw, err := w.Write(bi)
 542                 n += nw
 543                 if err != nil {
 544                         return n, err
 545                 }
 546         }
 547         return n, nil
 548 }