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

View as plain text