github.com/cloudwego/dynamicgo@v0.2.6-0.20240519101509-707f41b6b834/thrift/generic/iter.go (about)

     1  /**
     2   * Copyright 2023 CloudWeGo Authors.
     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 generic
    18  
    19  import (
    20  	"fmt"
    21  
    22  	"github.com/cloudwego/dynamicgo/internal/native/types"
    23  	"github.com/cloudwego/dynamicgo/meta"
    24  	"github.com/cloudwego/dynamicgo/thrift"
    25  )
    26  
    27  var (
    28  	_SkipMaxDepth = types.TB_SKIP_STACK_SIZE - 1
    29  )
    30  
    31  func (self Node) iterFields() (fi structIterator) {
    32  	// fi = structIteratorPool.Get().(*StructIterator)
    33  	fi.p.Buf = self.raw()
    34  	// if _, err := fi.p.ReadStructBegin(); err != nil {
    35  	// 	fi.Err = meta.NewError(meta.ErrReadInput, "", err)
    36  	// 	return
    37  	// }
    38  	return
    39  }
    40  
    41  func (self Node) iterPairs() (fi mapIterator) {
    42  	buf := self.raw()
    43  	// fi = pairIteratorPool.Get().(*PairIterator)
    44  	fi.p.Buf = buf
    45  	if kt, et, size, err := fi.p.ReadMapBegin(); err != nil {
    46  		fi.Err = meta.NewError(meta.ErrRead, "", err)
    47  		return
    48  	} else {
    49  		fi.size = size
    50  		fi.et = thrift.Type(et)
    51  		fi.kt = thrift.Type(kt)
    52  	}
    53  	return
    54  }
    55  
    56  func (self Node) iterElems() (fi listIterator) {
    57  	buf := self.raw()
    58  	// fi = listIteratorPool.Get().(*ListIterator)
    59  	fi.p.Buf = buf
    60  	if et, size, err := fi.p.ReadListBegin(); err != nil {
    61  		fi.Err = meta.NewError(meta.ErrRead, "", err)
    62  		return
    63  	} else {
    64  		fi.size = size
    65  		fi.et = thrift.Type(et)
    66  	}
    67  	return
    68  }
    69  
    70  type structIterator struct {
    71  	Err error
    72  	p   thrift.BinaryProtocol
    73  }
    74  
    75  func (it structIterator) HasNext() bool {
    76  	return it.Err == nil && it.p.Left() > 0 && (it.p.Buf)[it.p.Read] != byte(thrift.STOP)
    77  }
    78  
    79  func (it *structIterator) Next(useNative bool) (id thrift.FieldID, typ thrift.Type, start int, end int) {
    80  	_, typeId, ID, err := it.p.ReadFieldBegin()
    81  	if err != nil {
    82  		it.Err = meta.NewError(meta.ErrRead, "", err)
    83  		return
    84  	}
    85  	if typeId == thrift.STOP {
    86  		return
    87  	}
    88  	typ = thrift.Type(typeId)
    89  	start = it.p.Read
    90  	err = it.p.Skip(typeId, useNative)
    91  	if err != nil {
    92  		it.Err = meta.NewError(meta.ErrRead, "", err)
    93  		return
    94  	}
    95  	end = it.p.Read
    96  	err = it.p.ReadFieldEnd()
    97  	if err != nil {
    98  		it.Err = meta.NewError(meta.ErrRead, "", err)
    99  		return
   100  	}
   101  	id = thrift.FieldID(ID)
   102  	return
   103  }
   104  
   105  type listIterator struct {
   106  	Err  error
   107  	size int
   108  	k    int
   109  	et   thrift.Type
   110  	p    thrift.BinaryProtocol
   111  }
   112  
   113  func (it listIterator) HasNext() bool {
   114  	return it.Err == nil && it.p.Left() > 0 && it.k < it.size
   115  }
   116  
   117  func (it listIterator) Size() int {
   118  	return it.size
   119  }
   120  
   121  func (it listIterator) Pos() int {
   122  	return it.k
   123  }
   124  
   125  func (it *listIterator) Next(useNative bool) (start int, end int) {
   126  	start = it.p.Read
   127  	err := it.p.Skip(it.et, useNative)
   128  	if err != nil {
   129  		it.Err = meta.NewError(meta.ErrRead, "", err)
   130  		return
   131  	}
   132  	end = it.p.Read
   133  	it.k++
   134  	return
   135  }
   136  
   137  type mapIterator struct {
   138  	Err  error
   139  	size int
   140  	i    int
   141  	kt   thrift.Type
   142  	et   thrift.Type
   143  	p    thrift.BinaryProtocol
   144  }
   145  
   146  func (it mapIterator) HasNext() bool {
   147  	return it.Err == nil && it.p.Left() > 0 && it.i < it.size
   148  }
   149  
   150  func (it mapIterator) Size() int {
   151  	return it.size
   152  }
   153  
   154  func (it mapIterator) Pos() int {
   155  	return it.i
   156  }
   157  
   158  func (it *mapIterator) NextBin(useNative bool) (keyStart int, keyBin []byte, start int, end int) {
   159  	keyStart = it.p.Read
   160  	var err error
   161  	ks := it.p.Read
   162  	if err = it.p.Skip(it.kt, useNative); err != nil {
   163  		it.Err = meta.NewError(meta.ErrRead, "", err)
   164  		return
   165  	}
   166  	keyBin = it.p.Buf[ks:it.p.Read]
   167  	start = it.p.Read
   168  	err = it.p.Skip(it.et, useNative)
   169  	if err != nil {
   170  		it.Err = meta.NewError(meta.ErrRead, "", err)
   171  		return
   172  	}
   173  	end = it.p.Read
   174  	it.i++
   175  	return
   176  }
   177  
   178  func (it *mapIterator) NextStr(useNative bool) (keyStart int, keyString string, start int, end int) {
   179  	keyStart = it.p.Read
   180  	var err error
   181  	if it.kt == thrift.STRING {
   182  		keyString, err = it.p.ReadString(false)
   183  		if err != nil {
   184  			it.Err = meta.NewError(meta.ErrRead, "", err)
   185  			return
   186  		}
   187  	} else {
   188  		it.Err = wrapError(meta.ErrUnsupportedType, "MapIterator.nextStr: key type is not string", nil)
   189  		return
   190  	}
   191  	start = it.p.Read
   192  	err = it.p.Skip(it.et,
   193  		 useNative)
   194  	if err != nil {
   195  		it.Err = meta.NewError(meta.ErrRead, "", err)
   196  		return
   197  	}
   198  	end = it.p.Read
   199  	it.i++
   200  	return
   201  }
   202  
   203  func (it *mapIterator) NextInt(useNative bool) (keyStart int, keyInt int, start int, end int) {
   204  	keyStart = it.p.Read
   205  	switch it.kt {
   206  	case thrift.I08:
   207  		n, err := it.p.ReadByte()
   208  		keyInt = int(n)
   209  		if err != nil {
   210  			it.Err = meta.NewError(meta.ErrRead, "", err)
   211  			return
   212  		}
   213  	case thrift.I16:
   214  		n, err := it.p.ReadI16()
   215  		if err != nil {
   216  			it.Err = meta.NewError(meta.ErrRead, "", err)
   217  			return
   218  		}
   219  		keyInt = int(n)
   220  	case thrift.I32:
   221  		n, err := it.p.ReadI32()
   222  		if err != nil {
   223  			it.Err = meta.NewError(meta.ErrRead, "", err)
   224  			return
   225  		}
   226  		keyInt = int(n)
   227  	case thrift.I64:
   228  		n, err := it.p.ReadI64()
   229  		if err != nil {
   230  			it.Err = meta.NewError(meta.ErrRead, "", err)
   231  			return
   232  		}
   233  		keyInt = int(n)
   234  	default:
   235  		it.Err = wrapError(meta.ErrUnsupportedType, "MapIterator.nextInt: key type is not int", nil)
   236  	}
   237  	start = it.p.Read
   238  	if err := it.p.Skip(it.et,
   239  		 useNative); err != nil {
   240  		it.Err = meta.NewError(meta.ErrRead, "", err)
   241  		return
   242  	}
   243  	end = it.p.Read
   244  	it.i++
   245  	return
   246  }
   247  
   248  // Foreach scan each element of a complex type (LIST/SET/MAP/STRUCT),
   249  // and call handler sequentially with corresponding path and value
   250  func (self Value) Foreach(handler func(path Path, val Value) bool, opts *Options) error {
   251  	switch self.t {
   252  	case thrift.STRUCT:
   253  		it := self.iterFields()
   254  		if it.Err != nil {
   255  			return it.Err
   256  		}
   257  		for it.HasNext() {
   258  			id, typ, start, end := it.Next(opts.UseNativeSkip)
   259  			if it.Err != nil {
   260  				return it.Err
   261  			}
   262  			f := self.Desc.Struct().FieldById(id)
   263  			if f == nil {
   264  				if opts.DisallowUnknow {
   265  					return wrapError(meta.ErrUnknownField, fmt.Sprintf("unknown field %d", id), nil)
   266  				}
   267  				continue
   268  			}
   269  			if f.Type().Type() != typ {
   270  				return wrapError(meta.ErrDismatchType, fmt.Sprintf("field %d type %v mismatch read type %v", id, f.Type().Type(), typ), nil)
   271  			}
   272  			v := self.slice(start, end, f.Type())
   273  			var p Path
   274  			if opts.IterateStructByName && f != nil {
   275  				p = NewPathFieldName(f.Name())
   276  			} else {
   277  				p = NewPathFieldId(id)
   278  			}
   279  			if !handler(p, v) {
   280  				return nil
   281  			}
   282  		}
   283  	case thrift.LIST, thrift.SET:
   284  		it := self.iterElems()
   285  		if it.Err != nil {
   286  			return it.Err
   287  		}
   288  		et := self.Desc.Elem()
   289  		for i := 0; it.HasNext(); i++ {
   290  			start, end := it.Next(opts.UseNativeSkip)
   291  			if it.Err != nil {
   292  				return it.Err
   293  			}
   294  			v := self.slice(start, end, et)
   295  			p := NewPathIndex(i)
   296  			if !handler(p, v) {
   297  				return nil
   298  			}
   299  		}
   300  	case thrift.MAP:
   301  		it := self.iterPairs()
   302  		if it.Err != nil {
   303  			return it.Err
   304  		}
   305  		kt := self.Desc.Key()
   306  		if kt.Type() != it.kt {
   307  			return wrapError(meta.ErrDismatchType, "dismatched descriptor key type and binary type", nil)
   308  		}
   309  		et := self.Desc.Elem()
   310  		if et.Type() != it.et {
   311  			return wrapError(meta.ErrDismatchType, "dismatched descriptor elem type and binary type", nil)
   312  		}
   313  		if kt.Type() == thrift.STRING {
   314  			for i := 0; it.HasNext(); i++ {
   315  				_, keyString, start, end := it.NextStr(opts.UseNativeSkip)
   316  				if it.Err != nil {
   317  					return it.Err
   318  				}
   319  				v := self.slice(start, end, et)
   320  				p := NewPathStrKey(keyString)
   321  				if !handler(p, v) {
   322  					return nil
   323  				}
   324  			}
   325  		} else if kt.Type().IsInt() {
   326  			for i := 0; it.HasNext(); i++ {
   327  				_, keyInt, start, end := it.NextInt(opts.UseNativeSkip)
   328  				if it.Err != nil {
   329  					return it.Err
   330  				}
   331  				v := self.slice(start, end, et)
   332  				p := NewPathIntKey(keyInt)
   333  				if !handler(p, v) {
   334  					return nil
   335  				}
   336  			}
   337  		} else {
   338  			for i := 0; it.HasNext(); i++ {
   339  				_, keyBin, start, end := it.NextBin(opts.UseNativeSkip)
   340  				if it.Err != nil {
   341  					return it.Err
   342  				}
   343  				v := self.slice(start, end, et)
   344  				p := NewPathBinKey(keyBin)
   345  				if !handler(p, v) {
   346  					return nil
   347  				}
   348  			}
   349  		}
   350  	}
   351  	return nil
   352  }
   353  
   354  // ForeachKV scan each element of a MAP type, and call handler sequentially with corresponding key and value
   355  func (self Value) ForeachKV(handler func(key Value, val Value) bool, opts *Options) error {
   356  	if err := self.should("MAP", thrift.MAP); err != "" {
   357  		return wrapError(meta.ErrUnsupportedType, err, nil)
   358  	}
   359  	switch self.t {
   360  	case thrift.MAP:
   361  		p := thrift.BinaryProtocol{Buf: self.raw()}
   362  		kt, et, size, e := p.ReadMapBegin()
   363  		if e != nil {
   364  			return errNode(meta.ErrRead, "", nil)
   365  		}
   366  		kd := self.Desc.Key()
   367  		if kd.Type() != kt {
   368  			return wrapError(meta.ErrDismatchType, "dismatched descriptor key type and binary type", nil)
   369  		}
   370  		ed := self.Desc.Elem()
   371  		if ed.Type() != et {
   372  			return wrapError(meta.ErrDismatchType, "dismatched descriptor elem type and binary type", nil)
   373  		}
   374  		for i := 0; i < size; i++ {
   375  			ks := p.Read
   376  			if err := p.Skip(kt, opts.UseNativeSkip); err != nil {
   377  				return errNode(meta.ErrRead, "", nil)
   378  			}
   379  			key := self.slice(ks, p.Read, kd)
   380  			es := p.Read
   381  			if err := p.Skip(et, opts.UseNativeSkip); err != nil {
   382  				return errNode(meta.ErrRead, "", nil)
   383  			}
   384  			val := self.slice(es, p.Read, ed)
   385  			if !handler(key, val) {
   386  				return nil
   387  			}
   388  		}
   389  		return nil
   390  	default:
   391  		return errNode(meta.ErrUnsupportedType, "", nil)
   392  	}
   393  }
   394  
   395  // Foreach scan each element of a complex type (LIST/SET/MAP/STRUCT),
   396  // and call handler sequentially with corresponding path and node
   397  func (self Node) Foreach(handler func(path Path, node Node) bool, opts *Options) error {
   398  	switch self.t {
   399  	case thrift.STRUCT:
   400  		it := self.iterFields()
   401  		if it.Err != nil {
   402  			return it.Err
   403  		}
   404  		for it.HasNext() {
   405  			id, typ, start, end := it.Next(opts.UseNativeSkip)
   406  			if it.Err != nil {
   407  				return it.Err
   408  			}
   409  			v := self.slice(start, end, typ)
   410  			var p Path
   411  			p = NewPathFieldId(id)
   412  			if !handler(p, v) {
   413  				return nil
   414  			}
   415  		}
   416  	case thrift.LIST, thrift.SET:
   417  		it := self.iterElems()
   418  		if it.Err != nil {
   419  			return it.Err
   420  		}
   421  		for i := 0; it.HasNext(); i++ {
   422  			start, end := it.Next(opts.UseNativeSkip)
   423  			if it.Err != nil {
   424  				return it.Err
   425  			}
   426  			v := self.slice(start, end, self.et)
   427  			p := NewPathIndex(i)
   428  			if !handler(p, v) {
   429  				return nil
   430  			}
   431  		}
   432  	case thrift.MAP:
   433  		it := self.iterPairs()
   434  		if it.Err != nil {
   435  			return it.Err
   436  		}
   437  		kt := self.kt
   438  		if kt == thrift.STRING {
   439  			for i := 0; it.HasNext(); i++ {
   440  				_, keyString, start, end := it.NextStr(opts.UseNativeSkip)
   441  				if it.Err != nil {
   442  					return it.Err
   443  				}
   444  				v := self.slice(start, end, self.et)
   445  				p := NewPathStrKey(keyString)
   446  				if !handler(p, v) {
   447  					return nil
   448  				}
   449  			}
   450  		} else if kt.IsInt() {
   451  			for i := 0; it.HasNext(); i++ {
   452  				_, keyInt, start, end := it.NextInt(opts.UseNativeSkip)
   453  				if it.Err != nil {
   454  					return it.Err
   455  				}
   456  				v := self.slice(start, end, self.et)
   457  				p := NewPathIntKey(keyInt)
   458  				if !handler(p, v) {
   459  					return nil
   460  				}
   461  			}
   462  		} else {
   463  			for i := 0; it.HasNext(); i++ {
   464  				_, keyBin, start, end := it.NextBin(opts.UseNativeSkip)
   465  				if it.Err != nil {
   466  					return it.Err
   467  				}
   468  				v := self.slice(start, end, self.et)
   469  				p := NewPathBinKey(keyBin)
   470  				if !handler(p, v) {
   471  					return nil
   472  				}
   473  			}
   474  		}
   475  	}
   476  	return nil
   477  }
   478  
   479  // ForeachKV scan each element of a MAP type, and call handler sequentially with corresponding key and value
   480  func (self Node) ForeachKV(handler func(key Node, val Node) bool, opts *Options) error {
   481  	switch self.t {
   482  	case thrift.MAP:
   483  		p := thrift.BinaryProtocol{Buf: self.raw()}
   484  		kt, et, size, e := p.ReadMapBegin()
   485  		if e != nil {
   486  			return errNode(meta.ErrRead, "", nil)
   487  		}
   488  		for i := 0; i < size; i++ {
   489  			ks := p.Read
   490  			if err := p.Skip(kt, opts.UseNativeSkip); err != nil {
   491  				return errNode(meta.ErrRead, "", nil)
   492  			}
   493  			key := self.slice(ks, p.Read, kt)
   494  			es := p.Read
   495  			if err := p.Skip(et, opts.UseNativeSkip); err != nil {
   496  				return errNode(meta.ErrRead, "", nil)
   497  			}
   498  			val := self.slice(es, p.Read, et)
   499  			if !handler(key, val) {
   500  				return nil
   501  			}
   502  		}
   503  		return nil
   504  	default:
   505  		return errNode(meta.ErrUnsupportedType, "", nil)
   506  	}
   507  }