mapiter.go raw

   1  /*
   2   * Copyright 2021 ByteDance Inc.
   3   *
   4   * Licensed under the Apache License, Version 2.0 (the "License");
   5   * you may not use this file except in compliance with the License.
   6   * You may obtain a copy of the License at
   7   *
   8   *     http://www.apache.org/licenses/LICENSE-2.0
   9   *
  10   * Unless required by applicable law or agreed to in writing, software
  11   * distributed under the License is distributed on an "AS IS" BASIS,
  12   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13   * See the License for the specific language governing permissions and
  14   * limitations under the License.
  15   */
  16  
  17  package alg
  18  
  19  import (
  20  	"encoding"
  21  	"reflect"
  22  	"strconv"
  23  	"sync"
  24  	"unsafe"
  25  
  26  	"github.com/bytedance/sonic/internal/encoder/vars"
  27  	"github.com/bytedance/sonic/internal/rt"
  28  )
  29  
  30  type _MapPair struct {
  31      k string  // when the map key is integer, k is pointed to m
  32      v unsafe.Pointer
  33      m [32]byte
  34  }
  35  
  36  type MapIterator struct {
  37      It rt.GoMapIterator     // must be the first field
  38      kv rt.GoSlice           // slice of _MapPair
  39      ki int
  40  }
  41  
  42  var (
  43      iteratorPool = sync.Pool{}
  44      iteratorPair = rt.UnpackType(reflect.TypeOf(_MapPair{}))
  45  )
  46  
  47  func init() {
  48      if unsafe.Offsetof(MapIterator{}.It) != 0 {
  49          panic("_MapIterator.it is not the first field")
  50      }
  51  }
  52  
  53  
  54  func newIterator() *MapIterator {
  55      if v := iteratorPool.Get(); v == nil {
  56          return new(MapIterator)
  57      } else {
  58          return resetIterator(v.(*MapIterator))
  59      }
  60  }
  61  
  62  func resetIterator(p *MapIterator) *MapIterator {
  63      p.ki = 0
  64      p.It = rt.GoMapIterator{}
  65      p.kv.Len = 0
  66      return p
  67  }
  68  
  69  func (self *MapIterator) at(i int) *_MapPair {
  70      return (*_MapPair)(unsafe.Pointer(uintptr(self.kv.Ptr) + uintptr(i) * unsafe.Sizeof(_MapPair{})))
  71  }
  72  
  73  func (self *MapIterator) add() (p *_MapPair) {
  74      p = self.at(self.kv.Len)
  75      self.kv.Len++
  76      return
  77  }
  78  
  79  func (self *MapIterator) data() (p []_MapPair) {
  80      *(*rt.GoSlice)(unsafe.Pointer(&p)) = self.kv
  81      return
  82  }
  83  
  84  func (self *MapIterator) append(t *rt.GoType, k unsafe.Pointer, v unsafe.Pointer) (err error) {
  85      p := self.add()
  86      p.v = v
  87  
  88      /* check for strings */
  89      if tk := t.Kind(); tk != reflect.String {
  90          return self.appendGeneric(p, t, tk, k)
  91      }
  92  
  93      /* fast path for strings */
  94      p.k = *(*string)(k)
  95      return nil
  96  }
  97  
  98  func (self *MapIterator) appendGeneric(p *_MapPair, t *rt.GoType, v reflect.Kind, k unsafe.Pointer) error {
  99      switch v {
 100          case reflect.Int       : p.k = rt.Mem2Str(strconv.AppendInt(p.m[:0], int64(*(*int)(k)), 10))      ; return nil
 101          case reflect.Int8      : p.k = rt.Mem2Str(strconv.AppendInt(p.m[:0], int64(*(*int8)(k)), 10))     ; return nil
 102          case reflect.Int16     : p.k = rt.Mem2Str(strconv.AppendInt(p.m[:0], int64(*(*int16)(k)), 10))    ; return nil
 103          case reflect.Int32     : p.k = rt.Mem2Str(strconv.AppendInt(p.m[:0], int64(*(*int32)(k)), 10))    ; return nil
 104          case reflect.Int64     : p.k = rt.Mem2Str(strconv.AppendInt(p.m[:0], int64(*(*int64)(k)), 10))           ; return nil
 105          case reflect.Uint      : p.k = rt.Mem2Str(strconv.AppendUint(p.m[:0], uint64(*(*uint)(k)), 10))    ; return nil
 106          case reflect.Uint8     : p.k = rt.Mem2Str(strconv.AppendUint(p.m[:0], uint64(*(*uint8)(k)), 10))   ; return nil
 107          case reflect.Uint16    : p.k = rt.Mem2Str(strconv.AppendUint(p.m[:0], uint64(*(*uint16)(k)), 10))  ; return nil
 108          case reflect.Uint32    : p.k = rt.Mem2Str(strconv.AppendUint(p.m[:0], uint64(*(*uint32)(k)), 10))  ; return nil
 109          case reflect.Uint64    : p.k = rt.Mem2Str(strconv.AppendUint(p.m[:0], uint64(*(*uint64)(k)), 10))          ; return nil
 110          case reflect.Uintptr   : p.k = rt.Mem2Str(strconv.AppendUint(p.m[:0], uint64(*(*uintptr)(k)), 10)) ; return nil
 111          case reflect.Interface : return self.appendInterface(p, t, k)
 112          case reflect.Struct, reflect.Ptr : return self.appendConcrete(p, t, k)
 113          default                : panic("unexpected map key type")
 114      }
 115  }
 116  
 117  func (self *MapIterator) appendConcrete(p *_MapPair, t *rt.GoType, k unsafe.Pointer) (err error) {
 118      // compiler has already checked that the type implements the encoding.MarshalText interface
 119      if !t.Indirect() {
 120          k = *(*unsafe.Pointer)(k)
 121      }
 122      eface := rt.GoEface{Value: k, Type: t}.Pack()
 123      out, err := eface.(encoding.TextMarshaler).MarshalText()
 124      if err != nil {
 125          return err
 126      }
 127      p.k = rt.Mem2Str(out)
 128      return
 129  }
 130  
 131  func (self *MapIterator) appendInterface(p *_MapPair, t *rt.GoType, k unsafe.Pointer) (err error) {
 132      if len(rt.IfaceType(t).Methods) == 0 {
 133          panic("unexpected map key type")
 134      } else if p.k, err = asText(k); err == nil {
 135          return nil
 136      } else {
 137          return
 138      }
 139  }
 140  
 141  func IteratorStop(p *MapIterator) {
 142      iteratorPool.Put(p)
 143  }
 144  
 145  func IteratorNext(p *MapIterator) {
 146      i := p.ki
 147      t := &p.It
 148  
 149      /* check for unordered iteration */
 150      if i < 0 {
 151          rt.Mapiternext(t)
 152          return
 153      }
 154  
 155      /* check for end of iteration */
 156      if p.ki >= p.kv.Len {
 157          t.K = nil
 158          t.V = nil
 159          return
 160      }
 161  
 162      /* update the key-value pair, and increase the pointer */
 163      t.K = unsafe.Pointer(&p.at(p.ki).k)
 164      t.V = p.at(p.ki).v
 165      p.ki++
 166  }
 167  
 168  func IteratorStart(t *rt.GoMapType, m unsafe.Pointer, fv uint64) (*MapIterator, error) {
 169      it := newIterator()
 170      rt.Mapiterinit(t, m, &it.It)
 171      count := rt.Maplen(m)
 172  
 173      /* check for key-sorting, empty map don't need sorting */
 174      if count == 0 || (fv & (1<<BitSortMapKeys)) == 0 {
 175          it.ki = -1
 176          return it, nil
 177      }
 178  
 179      /* pre-allocate space if needed */
 180      if count > it.kv.Cap {
 181          it.kv = rt.GrowSlice(iteratorPair, it.kv, count)
 182      }
 183  
 184      /* dump all the key-value pairs */
 185      for ; it.It.K != nil; rt.Mapiternext(&it.It) {
 186          if err := it.append(t.Key, it.It.K, it.It.V); err != nil {
 187              IteratorStop(it)
 188              return nil, err
 189          }
 190      }
 191  
 192      /* sort the keys, map with only 1 item don't need sorting */
 193      if it.ki = 1; count > 1 {
 194          radixQsort(it.data(), 0, maxDepth(it.kv.Len))
 195      }
 196  
 197      /* load the first pair into iterator */
 198      it.It.V = it.at(0).v
 199      it.It.K = unsafe.Pointer(&it.at(0).k)
 200      return it, nil
 201  }
 202  
 203  func asText(v unsafe.Pointer) (string, error) {
 204  	text := rt.AssertI2I(rt.UnpackType(vars.EncodingTextMarshalerType), *(*rt.GoIface)(v))
 205  	r, e := (*(*encoding.TextMarshaler)(unsafe.Pointer(&text))).MarshalText()
 206  	return rt.Mem2Str(r), e
 207  }
 208