set.go 2.0 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586
  1. // Copyright 2016 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. // Package stringset provides a way to represent a collection of strings
  5. // compactly.
  6. package stringset
  7. import "sort"
  8. // A Set holds a collection of strings that can be looked up by an index number.
  9. type Set struct {
  10. // These fields are exported to allow for code generation.
  11. Data string
  12. Index []uint16
  13. }
  14. // Elem returns the string with index i. It panics if i is out of range.
  15. func (s *Set) Elem(i int) string {
  16. return s.Data[s.Index[i]:s.Index[i+1]]
  17. }
  18. // Len returns the number of strings in the set.
  19. func (s *Set) Len() int {
  20. return len(s.Index) - 1
  21. }
  22. // Search returns the index of the given string or -1 if it is not in the set.
  23. // The Set must have been created with strings in sorted order.
  24. func Search(s *Set, str string) int {
  25. // TODO: optimize this if it gets used a lot.
  26. n := len(s.Index) - 1
  27. p := sort.Search(n, func(i int) bool {
  28. return s.Elem(i) >= str
  29. })
  30. if p == n || str != s.Elem(p) {
  31. return -1
  32. }
  33. return p
  34. }
  35. // A Builder constructs Sets.
  36. type Builder struct {
  37. set Set
  38. index map[string]int
  39. }
  40. // NewBuilder returns a new and initialized Builder.
  41. func NewBuilder() *Builder {
  42. return &Builder{
  43. set: Set{
  44. Index: []uint16{0},
  45. },
  46. index: map[string]int{},
  47. }
  48. }
  49. // Set creates the set created so far.
  50. func (b *Builder) Set() Set {
  51. return b.set
  52. }
  53. // Index returns the index for the given string, which must have been added
  54. // before.
  55. func (b *Builder) Index(s string) int {
  56. return b.index[s]
  57. }
  58. // Add adds a string to the index. Strings that are added by a single Add will
  59. // be stored together, unless they match an existing string.
  60. func (b *Builder) Add(ss ...string) {
  61. // First check if the string already exists.
  62. for _, s := range ss {
  63. if _, ok := b.index[s]; ok {
  64. continue
  65. }
  66. b.index[s] = len(b.set.Index) - 1
  67. b.set.Data += s
  68. x := len(b.set.Data)
  69. if x > 0xFFFF {
  70. panic("Index too > 0xFFFF")
  71. }
  72. b.set.Index = append(b.set.Index, uint16(x))
  73. }
  74. }