intern.go 4.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238
  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 len(s) >= minInternedStringLen {
  48. if idx, ok := e.dict[s]; ok {
  49. return e.encodeInternedStringIndex(idx)
  50. }
  51. if intern && len(e.dict) < maxDictLen {
  52. if e.dict == nil {
  53. e.dict = make(map[string]int)
  54. }
  55. idx := len(e.dict)
  56. e.dict[s] = idx
  57. }
  58. }
  59. return e.encodeNormalString(s)
  60. }
  61. func (e *Encoder) encodeInternedStringIndex(idx int) error {
  62. if idx <= math.MaxUint8 {
  63. if err := e.writeCode(msgpcode.FixExt1); err != nil {
  64. return err
  65. }
  66. return e.write1(byte(internedStringExtID), uint8(idx))
  67. }
  68. if idx <= math.MaxUint16 {
  69. if err := e.writeCode(msgpcode.FixExt2); err != nil {
  70. return err
  71. }
  72. return e.write2(byte(internedStringExtID), uint16(idx))
  73. }
  74. if uint64(idx) <= math.MaxUint32 {
  75. if err := e.writeCode(msgpcode.FixExt4); err != nil {
  76. return err
  77. }
  78. return e.write4(byte(internedStringExtID), uint32(idx))
  79. }
  80. return fmt.Errorf("msgpack: interned string index=%d is too large", idx)
  81. }
  82. //------------------------------------------------------------------------------
  83. func decodeInternedInterfaceValue(d *Decoder, v reflect.Value) error {
  84. s, err := d.decodeInternedString(true)
  85. if err == nil {
  86. v.Set(reflect.ValueOf(s))
  87. return nil
  88. }
  89. if err != nil {
  90. if _, ok := err.(unexpectedCodeError); !ok {
  91. return err
  92. }
  93. }
  94. if err := d.s.UnreadByte(); err != nil {
  95. return err
  96. }
  97. return decodeInterfaceValue(d, v)
  98. }
  99. func decodeInternedStringValue(d *Decoder, v reflect.Value) error {
  100. s, err := d.decodeInternedString(true)
  101. if err != nil {
  102. return err
  103. }
  104. v.SetString(s)
  105. return nil
  106. }
  107. func (d *Decoder) decodeInternedString(intern bool) (string, error) {
  108. c, err := d.readCode()
  109. if err != nil {
  110. return "", err
  111. }
  112. if msgpcode.IsFixedString(c) {
  113. n := int(c & msgpcode.FixedStrMask)
  114. return d.decodeInternedStringWithLen(n, intern)
  115. }
  116. switch c {
  117. case msgpcode.Nil:
  118. return "", nil
  119. case msgpcode.FixExt1, msgpcode.FixExt2, msgpcode.FixExt4:
  120. typeID, extLen, err := d.extHeader(c)
  121. if err != nil {
  122. return "", err
  123. }
  124. if typeID != internedStringExtID {
  125. err := fmt.Errorf("msgpack: got ext type=%d, wanted %d",
  126. typeID, internedStringExtID)
  127. return "", err
  128. }
  129. idx, err := d.decodeInternedStringIndex(extLen)
  130. if err != nil {
  131. return "", err
  132. }
  133. return d.internedStringAtIndex(idx)
  134. case msgpcode.Str8, msgpcode.Bin8:
  135. n, err := d.uint8()
  136. if err != nil {
  137. return "", err
  138. }
  139. return d.decodeInternedStringWithLen(int(n), intern)
  140. case msgpcode.Str16, msgpcode.Bin16:
  141. n, err := d.uint16()
  142. if err != nil {
  143. return "", err
  144. }
  145. return d.decodeInternedStringWithLen(int(n), intern)
  146. case msgpcode.Str32, msgpcode.Bin32:
  147. n, err := d.uint32()
  148. if err != nil {
  149. return "", err
  150. }
  151. return d.decodeInternedStringWithLen(int(n), intern)
  152. }
  153. return "", unexpectedCodeError{
  154. code: c,
  155. hint: "interned string",
  156. }
  157. }
  158. func (d *Decoder) decodeInternedStringIndex(extLen int) (int, error) {
  159. switch extLen {
  160. case 1:
  161. n, err := d.uint8()
  162. if err != nil {
  163. return 0, err
  164. }
  165. return int(n), nil
  166. case 2:
  167. n, err := d.uint16()
  168. if err != nil {
  169. return 0, err
  170. }
  171. return int(n), nil
  172. case 4:
  173. n, err := d.uint32()
  174. if err != nil {
  175. return 0, err
  176. }
  177. return int(n), nil
  178. }
  179. err := fmt.Errorf("msgpack: unsupported ext len=%d decoding interned string", extLen)
  180. return 0, err
  181. }
  182. func (d *Decoder) internedStringAtIndex(idx int) (string, error) {
  183. if idx >= len(d.dict) {
  184. err := fmt.Errorf("msgpack: interned string at index=%d does not exist", idx)
  185. return "", err
  186. }
  187. return d.dict[idx], nil
  188. }
  189. func (d *Decoder) decodeInternedStringWithLen(n int, intern bool) (string, error) {
  190. if n <= 0 {
  191. return "", nil
  192. }
  193. s, err := d.stringWithLen(n)
  194. if err != nil {
  195. return "", err
  196. }
  197. if intern && len(s) >= minInternedStringLen && len(d.dict) < maxDictLen {
  198. d.dict = append(d.dict, s)
  199. }
  200. return s, nil
  201. }