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 }