suffixarray.go

Documentation: index/suffixarray

     1  // Copyright 2010 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 suffixarray implements substring search in logarithmic time using
     6  // an in-memory suffix array.
     7  //
     8  // Example use:
     9  //
    10  //	// create index for some data
    11  //	index := suffixarray.New(data)
    12  //
    13  //	// lookup byte slice s
    14  //	offsets1 := index.Lookup(s, -1) // the list of all indices where s occurs in data
    15  //	offsets2 := index.Lookup(s, 3)  // the list of at most 3 indices where s occurs in data
    16  //
    17  package suffixarray
    18  
    19  import (
    20  	"bytes"
    21  	"encoding/binary"
    22  	"io"
    23  	"regexp"
    24  	"sort"
    25  )
    26  
    27  // Index implements a suffix array for fast substring search.
    28  type Index struct {
    29  	data []byte
    30  	sa   []int // suffix array for data; len(sa) == len(data)
    31  }
    32  
    33  // New creates a new Index for data.
    34  // Index creation time is O(N*log(N)) for N = len(data).
    35  func New(data []byte) *Index {
    36  	return &Index{data, qsufsort(data)}
    37  }
    38  
    39  // writeInt writes an int x to w using buf to buffer the write.
    40  func writeInt(w io.Writer, buf []byte, x int) error {
    41  	binary.PutVarint(buf, int64(x))
    42  	_, err := w.Write(buf[0:binary.MaxVarintLen64])
    43  	return err
    44  }
    45  
    46  // readInt reads an int x from r using buf to buffer the read and returns x.
    47  func readInt(r io.Reader, buf []byte) (int, error) {
    48  	_, err := io.ReadFull(r, buf[0:binary.MaxVarintLen64]) // ok to continue with error
    49  	x, _ := binary.Varint(buf)
    50  	return int(x), err
    51  }
    52  
    53  // writeSlice writes data[:n] to w and returns n.
    54  // It uses buf to buffer the write.
    55  func writeSlice(w io.Writer, buf []byte, data []int) (n int, err error) {
    56  	// encode as many elements as fit into buf
    57  	p := binary.MaxVarintLen64
    58  	for ; n < len(data) && p+binary.MaxVarintLen64 <= len(buf); n++ {
    59  		p += binary.PutUvarint(buf[p:], uint64(data[n]))
    60  	}
    61  
    62  	// update buffer size
    63  	binary.PutVarint(buf, int64(p))
    64  
    65  	// write buffer
    66  	_, err = w.Write(buf[0:p])
    67  	return
    68  }
    69  
    70  // readSlice reads data[:n] from r and returns n.
    71  // It uses buf to buffer the read.
    72  func readSlice(r io.Reader, buf []byte, data []int) (n int, err error) {
    73  	// read buffer size
    74  	var size int
    75  	size, err = readInt(r, buf)
    76  	if err != nil {
    77  		return
    78  	}
    79  
    80  	// read buffer w/o the size
    81  	if _, err = io.ReadFull(r, buf[binary.MaxVarintLen64:size]); err != nil {
    82  		return
    83  	}
    84  
    85  	// decode as many elements as present in buf
    86  	for p := binary.MaxVarintLen64; p < size; n++ {
    87  		x, w := binary.Uvarint(buf[p:])
    88  		data[n] = int(x)
    89  		p += w
    90  	}
    91  
    92  	return
    93  }
    94  
    95  const bufSize = 16 << 10 // reasonable for BenchmarkSaveRestore
    96  
    97  // Read reads the index from r into x; x must not be nil.
    98  func (x *Index) Read(r io.Reader) error {
    99  	// buffer for all reads
   100  	buf := make([]byte, bufSize)
   101  
   102  	// read length
   103  	n, err := readInt(r, buf)
   104  	if err != nil {
   105  		return err
   106  	}
   107  
   108  	// allocate space
   109  	if 2*n < cap(x.data) || cap(x.data) < n {
   110  		// new data is significantly smaller or larger then
   111  		// existing buffers - allocate new ones
   112  		x.data = make([]byte, n)
   113  		x.sa = make([]int, n)
   114  	} else {
   115  		// re-use existing buffers
   116  		x.data = x.data[0:n]
   117  		x.sa = x.sa[0:n]
   118  	}
   119  
   120  	// read data
   121  	if _, err := io.ReadFull(r, x.data); err != nil {
   122  		return err
   123  	}
   124  
   125  	// read index
   126  	for sa := x.sa; len(sa) > 0; {
   127  		n, err := readSlice(r, buf, sa)
   128  		if err != nil {
   129  			return err
   130  		}
   131  		sa = sa[n:]
   132  	}
   133  	return nil
   134  }
   135  
   136  // Write writes the index x to w.
   137  func (x *Index) Write(w io.Writer) error {
   138  	// buffer for all writes
   139  	buf := make([]byte, bufSize)
   140  
   141  	// write length
   142  	if err := writeInt(w, buf, len(x.data)); err != nil {
   143  		return err
   144  	}
   145  
   146  	// write data
   147  	if _, err := w.Write(x.data); err != nil {
   148  		return err
   149  	}
   150  
   151  	// write index
   152  	for sa := x.sa; len(sa) > 0; {
   153  		n, err := writeSlice(w, buf, sa)
   154  		if err != nil {
   155  			return err
   156  		}
   157  		sa = sa[n:]
   158  	}
   159  	return nil
   160  }
   161  
   162  // Bytes returns the data over which the index was created.
   163  // It must not be modified.
   164  //
   165  func (x *Index) Bytes() []byte {
   166  	return x.data
   167  }
   168  
   169  func (x *Index) at(i int) []byte {
   170  	return x.data[x.sa[i]:]
   171  }
   172  
   173  // lookupAll returns a slice into the matching region of the index.
   174  // The runtime is O(log(N)*len(s)).
   175  func (x *Index) lookupAll(s []byte) []int {
   176  	// find matching suffix index range [i:j]
   177  	// find the first index where s would be the prefix
   178  	i := sort.Search(len(x.sa), func(i int) bool { return bytes.Compare(x.at(i), s) >= 0 })
   179  	// starting at i, find the first index at which s is not a prefix
   180  	j := i + sort.Search(len(x.sa)-i, func(j int) bool { return !bytes.HasPrefix(x.at(j+i), s) })
   181  	return x.sa[i:j]
   182  }
   183  
   184  // Lookup returns an unsorted list of at most n indices where the byte string s
   185  // occurs in the indexed data. If n < 0, all occurrences are returned.
   186  // The result is nil if s is empty, s is not found, or n == 0.
   187  // Lookup time is O(log(N)*len(s) + len(result)) where N is the
   188  // size of the indexed data.
   189  //
   190  func (x *Index) Lookup(s []byte, n int) (result []int) {
   191  	if len(s) > 0 && n != 0 {
   192  		matches := x.lookupAll(s)
   193  		if n < 0 || len(matches) < n {
   194  			n = len(matches)
   195  		}
   196  		// 0 <= n <= len(matches)
   197  		if n > 0 {
   198  			result = make([]int, n)
   199  			copy(result, matches)
   200  		}
   201  	}
   202  	return
   203  }
   204  
   205  // FindAllIndex returns a sorted list of non-overlapping matches of the
   206  // regular expression r, where a match is a pair of indices specifying
   207  // the matched slice of x.Bytes(). If n < 0, all matches are returned
   208  // in successive order. Otherwise, at most n matches are returned and
   209  // they may not be successive. The result is nil if there are no matches,
   210  // or if n == 0.
   211  //
   212  func (x *Index) FindAllIndex(r *regexp.Regexp, n int) (result [][]int) {
   213  	// a non-empty literal prefix is used to determine possible
   214  	// match start indices with Lookup
   215  	prefix, complete := r.LiteralPrefix()
   216  	lit := []byte(prefix)
   217  
   218  	// worst-case scenario: no literal prefix
   219  	if prefix == "" {
   220  		return r.FindAllIndex(x.data, n)
   221  	}
   222  
   223  	// if regexp is a literal just use Lookup and convert its
   224  	// result into match pairs
   225  	if complete {
   226  		// Lookup returns indices that may belong to overlapping matches.
   227  		// After eliminating them, we may end up with fewer than n matches.
   228  		// If we don't have enough at the end, redo the search with an
   229  		// increased value n1, but only if Lookup returned all the requested
   230  		// indices in the first place (if it returned fewer than that then
   231  		// there cannot be more).
   232  		for n1 := n; ; n1 += 2 * (n - len(result)) /* overflow ok */ {
   233  			indices := x.Lookup(lit, n1)
   234  			if len(indices) == 0 {
   235  				return
   236  			}
   237  			sort.Ints(indices)
   238  			pairs := make([]int, 2*len(indices))
   239  			result = make([][]int, len(indices))
   240  			count := 0
   241  			prev := 0
   242  			for _, i := range indices {
   243  				if count == n {
   244  					break
   245  				}
   246  				// ignore indices leading to overlapping matches
   247  				if prev <= i {
   248  					j := 2 * count
   249  					pairs[j+0] = i
   250  					pairs[j+1] = i + len(lit)
   251  					result[count] = pairs[j : j+2]
   252  					count++
   253  					prev = i + len(lit)
   254  				}
   255  			}
   256  			result = result[0:count]
   257  			if len(result) >= n || len(indices) != n1 {
   258  				// found all matches or there's no chance to find more
   259  				// (n and n1 can be negative)
   260  				break
   261  			}
   262  		}
   263  		if len(result) == 0 {
   264  			result = nil
   265  		}
   266  		return
   267  	}
   268  
   269  	// regexp has a non-empty literal prefix; Lookup(lit) computes
   270  	// the indices of possible complete matches; use these as starting
   271  	// points for anchored searches
   272  	// (regexp "^" matches beginning of input, not beginning of line)
   273  	r = regexp.MustCompile("^" + r.String()) // compiles because r compiled
   274  
   275  	// same comment about Lookup applies here as in the loop above
   276  	for n1 := n; ; n1 += 2 * (n - len(result)) /* overflow ok */ {
   277  		indices := x.Lookup(lit, n1)
   278  		if len(indices) == 0 {
   279  			return
   280  		}
   281  		sort.Ints(indices)
   282  		result = result[0:0]
   283  		prev := 0
   284  		for _, i := range indices {
   285  			if len(result) == n {
   286  				break
   287  			}
   288  			m := r.FindIndex(x.data[i:]) // anchored search - will not run off
   289  			// ignore indices leading to overlapping matches
   290  			if m != nil && prev <= i {
   291  				m[0] = i // correct m
   292  				m[1] += i
   293  				result = append(result, m)
   294  				prev = m[1]
   295  			}
   296  		}
   297  		if len(result) >= n || len(indices) != n1 {
   298  			// found all matches or there's no chance to find more
   299  			// (n and n1 can be negative)
   300  			break
   301  		}
   302  	}
   303  	if len(result) == 0 {
   304  		result = nil
   305  	}
   306  	return
   307  }
   308
View as plain text