intern.go 4.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236
  1. package msgpack
  2. import (
  3. "fmt"
  4. "math"
  5. "reflect"
  6. "github.com/vmihailenco/msgpack/v5/msgpcode"
  7. )
  8. const (
  9. minInternedStringLen = 3
  10. maxDictLen = math.MaxUint16
  11. )
  12. var internedStringExtID = int8(math.MinInt8)
  13. func init() {
  14. extTypes[internedStringExtID] = &extInfo{
  15. Type: stringType,
  16. Decoder: decodeInternedStringExt,
  17. }
  18. }
  19. func decodeInternedStringExt(d *Decoder, v reflect.Value, extLen int) error {
  20. idx, err := d.decodeInternedStringIndex(extLen)
  21. if err != nil {
  22. return err
  23. }
  24. s, err := d.internedStringAtIndex(idx)
  25. if err != nil {
  26. return err
  27. }
  28. v.SetString(s)
  29. return nil
  30. }
  31. //------------------------------------------------------------------------------
  32. func encodeInternedInterfaceValue(e *Encoder, v reflect.Value) error {
  33. if v.IsNil() {
  34. return e.EncodeNil()
  35. }
  36. v = v.Elem()
  37. if v.Kind() == reflect.String {
  38. return e.encodeInternedString(v.String(), true)
  39. }
  40. return e.EncodeValue(v)
  41. }
  42. func encodeInternedStringValue(e *Encoder, v reflect.Value) error {
  43. return e.encodeInternedString(v.String(), true)
  44. }
  45. func (e *Encoder) encodeInternedString(s string, intern bool) error {
  46. // Interned string takes at least 3 bytes. Plain string 1 byte + string len.
  47. if idx, ok := e.dict[s]; ok {
  48. return e.encodeInternedStringIndex(idx)
  49. }
  50. if intern && len(s) >= minInternedStringLen && len(e.dict) < maxDictLen {
  51. if e.dict == nil {
  52. e.dict = make(map[string]int)
  53. }
  54. idx := len(e.dict)
  55. e.dict[s] = idx
  56. }
  57. return e.encodeNormalString(s)
  58. }
  59. func (e *Encoder) encodeInternedStringIndex(idx int) error {
  60. if idx <= math.MaxUint8 {
  61. if err := e.writeCode(msgpcode.FixExt1); err != nil {
  62. return err
  63. }
  64. return e.write1(byte(internedStringExtID), uint8(idx))
  65. }
  66. if idx <= math.MaxUint16 {
  67. if err := e.writeCode(msgpcode.FixExt2); err != nil {
  68. return err
  69. }
  70. return e.write2(byte(internedStringExtID), uint16(idx))
  71. }
  72. if uint64(idx) <= math.MaxUint32 {
  73. if err := e.writeCode(msgpcode.FixExt4); err != nil {
  74. return err
  75. }
  76. return e.write4(byte(internedStringExtID), uint32(idx))
  77. }
  78. return fmt.Errorf("msgpack: interned string index=%d is too large", idx)
  79. }
  80. //------------------------------------------------------------------------------
  81. func decodeInternedInterfaceValue(d *Decoder, v reflect.Value) error {
  82. s, err := d.decodeInternedString(true)
  83. if err == nil {
  84. v.Set(reflect.ValueOf(s))
  85. return nil
  86. }
  87. if err != nil {
  88. if _, ok := err.(unexpectedCodeError); !ok {
  89. return err
  90. }
  91. }
  92. if err := d.s.UnreadByte(); err != nil {
  93. return err
  94. }
  95. return decodeInterfaceValue(d, v)
  96. }
  97. func decodeInternedStringValue(d *Decoder, v reflect.Value) error {
  98. s, err := d.decodeInternedString(true)
  99. if err != nil {
  100. return err
  101. }
  102. v.SetString(s)
  103. return nil
  104. }
  105. func (d *Decoder) decodeInternedString(intern bool) (string, error) {
  106. c, err := d.readCode()
  107. if err != nil {
  108. return "", err
  109. }
  110. if msgpcode.IsFixedString(c) {
  111. n := int(c & msgpcode.FixedStrMask)
  112. return d.decodeInternedStringWithLen(n, intern)
  113. }
  114. switch c {
  115. case msgpcode.Nil:
  116. return "", nil
  117. case msgpcode.FixExt1, msgpcode.FixExt2, msgpcode.FixExt4:
  118. typeID, extLen, err := d.extHeader(c)
  119. if err != nil {
  120. return "", err
  121. }
  122. if typeID != internedStringExtID {
  123. err := fmt.Errorf("msgpack: got ext type=%d, wanted %d",
  124. typeID, internedStringExtID)
  125. return "", err
  126. }
  127. idx, err := d.decodeInternedStringIndex(extLen)
  128. if err != nil {
  129. return "", err
  130. }
  131. return d.internedStringAtIndex(idx)
  132. case msgpcode.Str8, msgpcode.Bin8:
  133. n, err := d.uint8()
  134. if err != nil {
  135. return "", err
  136. }
  137. return d.decodeInternedStringWithLen(int(n), intern)
  138. case msgpcode.Str16, msgpcode.Bin16:
  139. n, err := d.uint16()
  140. if err != nil {
  141. return "", err
  142. }
  143. return d.decodeInternedStringWithLen(int(n), intern)
  144. case msgpcode.Str32, msgpcode.Bin32:
  145. n, err := d.uint32()
  146. if err != nil {
  147. return "", err
  148. }
  149. return d.decodeInternedStringWithLen(int(n), intern)
  150. }
  151. return "", unexpectedCodeError{
  152. code: c,
  153. hint: "interned string",
  154. }
  155. }
  156. func (d *Decoder) decodeInternedStringIndex(extLen int) (int, error) {
  157. switch extLen {
  158. case 1:
  159. n, err := d.uint8()
  160. if err != nil {
  161. return 0, err
  162. }
  163. return int(n), nil
  164. case 2:
  165. n, err := d.uint16()
  166. if err != nil {
  167. return 0, err
  168. }
  169. return int(n), nil
  170. case 4:
  171. n, err := d.uint32()
  172. if err != nil {
  173. return 0, err
  174. }
  175. return int(n), nil
  176. }
  177. err := fmt.Errorf("msgpack: unsupported ext len=%d decoding interned string", extLen)
  178. return 0, err
  179. }
  180. func (d *Decoder) internedStringAtIndex(idx int) (string, error) {
  181. if idx >= len(d.dict) {
  182. err := fmt.Errorf("msgpack: interned string at index=%d does not exist", idx)
  183. return "", err
  184. }
  185. return d.dict[idx], nil
  186. }
  187. func (d *Decoder) decodeInternedStringWithLen(n int, intern bool) (string, error) {
  188. if n <= 0 {
  189. return "", nil
  190. }
  191. s, err := d.stringWithLen(n)
  192. if err != nil {
  193. return "", err
  194. }
  195. if intern && len(s) >= minInternedStringLen && len(d.dict) < maxDictLen {
  196. d.dict = append(d.dict, s)
  197. }
  198. return s, nil
  199. }