Source file src/pkg/index/suffixarray/suffixarray.go
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 }