ranger.go 3.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179
  1. package jet
  2. import (
  3. "errors"
  4. "fmt"
  5. "reflect"
  6. "sync"
  7. )
  8. // Ranger describes an interface for types that iterate over something.
  9. // Implementing this interface means the ranger will be used when it's
  10. // encountered on the right hand side of a range's "let" expression.
  11. type Ranger interface {
  12. // Range calls should return a key, a value and a done bool to indicate
  13. // whether there are more values to be generated.
  14. //
  15. // When the done flag is true, then the loop ends.
  16. Range() (reflect.Value, reflect.Value, bool)
  17. // ProvidesIndex should return true if keys are produced during Range()
  18. // calls. This call should be idempotent across Range() calls (i.e.
  19. // its return value must not change during an iteration).
  20. ProvidesIndex() bool
  21. }
  22. type intsRanger struct {
  23. i, val, to int64
  24. }
  25. var _ Ranger = &intsRanger{}
  26. func (r *intsRanger) Range() (index, value reflect.Value, end bool) {
  27. r.i++
  28. r.val++
  29. end = r.val == r.to
  30. // The indirection in the ValueOf calls avoids an allocation versus
  31. // using the concrete value of 'i' and 'val'. The downside is having
  32. // to interpret 'r.i' as "the current value" after Range() returns,
  33. // and so it needs to be initialized as -1.
  34. index = reflect.ValueOf(&r.i).Elem()
  35. value = reflect.ValueOf(&r.val).Elem()
  36. return
  37. }
  38. func (r *intsRanger) ProvidesIndex() bool { return true }
  39. func newIntsRanger(from, to int64) *intsRanger {
  40. r := &intsRanger{
  41. to: to,
  42. i: -1,
  43. val: from - 1,
  44. }
  45. return r
  46. }
  47. type pooledRanger interface {
  48. Ranger
  49. Setup(reflect.Value)
  50. }
  51. type sliceRanger struct {
  52. v reflect.Value
  53. i int
  54. }
  55. var _ Ranger = &sliceRanger{}
  56. var _ pooledRanger = &sliceRanger{}
  57. func (r *sliceRanger) Setup(v reflect.Value) {
  58. r.i = 0
  59. r.v = v
  60. }
  61. func (r *sliceRanger) Range() (index, value reflect.Value, end bool) {
  62. if r.i == r.v.Len() {
  63. end = true
  64. return
  65. }
  66. index = reflect.ValueOf(r.i)
  67. value = r.v.Index(r.i)
  68. r.i++
  69. return
  70. }
  71. func (r *sliceRanger) ProvidesIndex() bool { return true }
  72. type mapRanger struct {
  73. iter *reflect.MapIter
  74. hasMore bool
  75. }
  76. var _ Ranger = &mapRanger{}
  77. var _ pooledRanger = &mapRanger{}
  78. func (r *mapRanger) Setup(v reflect.Value) {
  79. r.iter = v.MapRange()
  80. r.hasMore = r.iter.Next()
  81. }
  82. func (r *mapRanger) Range() (key, value reflect.Value, end bool) {
  83. if !r.hasMore {
  84. end = true
  85. return
  86. }
  87. key, value = r.iter.Key(), r.iter.Value()
  88. r.hasMore = r.iter.Next()
  89. return
  90. }
  91. func (r *mapRanger) ProvidesIndex() bool { return true }
  92. type chanRanger struct {
  93. v reflect.Value
  94. }
  95. var _ Ranger = &chanRanger{}
  96. var _ pooledRanger = &chanRanger{}
  97. func (r *chanRanger) Setup(v reflect.Value) {
  98. r.v = v
  99. }
  100. func (r *chanRanger) Range() (_, value reflect.Value, end bool) {
  101. v, ok := r.v.Recv()
  102. value, end = v, !ok
  103. return
  104. }
  105. func (r *chanRanger) ProvidesIndex() bool { return false }
  106. // ranger pooling
  107. var (
  108. poolSliceRanger = &sync.Pool{
  109. New: func() interface{} {
  110. return new(sliceRanger)
  111. },
  112. }
  113. poolsByKind = map[reflect.Kind]*sync.Pool{
  114. reflect.Slice: poolSliceRanger,
  115. reflect.Array: poolSliceRanger,
  116. reflect.Map: &sync.Pool{
  117. New: func() interface{} {
  118. return new(mapRanger)
  119. },
  120. },
  121. reflect.Chan: &sync.Pool{
  122. New: func() interface{} {
  123. return new(chanRanger)
  124. },
  125. },
  126. }
  127. )
  128. func getRanger(v reflect.Value) (r Ranger, cleanup func(), err error) {
  129. if !v.IsValid() {
  130. return nil, nil, errors.New("can't range over invalid value")
  131. }
  132. t := v.Type()
  133. if t.Implements(rangerType) {
  134. return v.Interface().(Ranger), func() { /* no cleanup needed */ }, nil
  135. }
  136. v, isNil := indirect(v)
  137. if isNil {
  138. return nil, nil, fmt.Errorf("cannot range over nil pointer/interface (%s)", t)
  139. }
  140. pool, ok := poolsByKind[v.Kind()]
  141. if !ok {
  142. return nil, nil, fmt.Errorf("value %v (type %s) is not rangeable", v, t)
  143. }
  144. pr := pool.Get().(pooledRanger)
  145. pr.Setup(v)
  146. return pr, func() { pool.Put(pr) }, nil
  147. }