writer.go 5.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289
  1. // Copyright (c) 2012-2020 Ugorji Nwoke. All rights reserved.
  2. // Use of this source code is governed by a MIT license found in the LICENSE file.
  3. package codec
  4. import "io"
  5. // encWriter abstracts writing to a byte array or to an io.Writer.
  6. type encWriter interface {
  7. writeb([]byte)
  8. writestr(string)
  9. writeqstr(string) // write string wrapped in quotes ie "..."
  10. writen1(byte)
  11. // add convenience functions for writing 2,4
  12. writen2(byte, byte)
  13. writen4([4]byte)
  14. writen8([8]byte)
  15. end()
  16. }
  17. // ---------------------------------------------
  18. type bufioEncWriter struct {
  19. w io.Writer
  20. buf []byte
  21. n int
  22. b [16]byte // scratch buffer and padding (cache-aligned)
  23. }
  24. func (z *bufioEncWriter) reset(w io.Writer, bufsize int, blist *bytesFreelist) {
  25. z.w = w
  26. z.n = 0
  27. if bufsize <= 0 {
  28. bufsize = defEncByteBufSize
  29. }
  30. // bufsize must be >= 8, to accomodate writen methods (where n <= 8)
  31. if bufsize <= 8 {
  32. bufsize = 8
  33. }
  34. if cap(z.buf) < bufsize {
  35. if len(z.buf) > 0 && &z.buf[0] != &z.b[0] {
  36. blist.put(z.buf)
  37. }
  38. if len(z.b) > bufsize {
  39. z.buf = z.b[:]
  40. } else {
  41. z.buf = blist.get(bufsize)
  42. }
  43. }
  44. z.buf = z.buf[:cap(z.buf)]
  45. }
  46. func (z *bufioEncWriter) flushErr() (err error) {
  47. n, err := z.w.Write(z.buf[:z.n])
  48. z.n -= n
  49. if z.n > 0 {
  50. if err == nil {
  51. err = io.ErrShortWrite
  52. }
  53. if n > 0 {
  54. copy(z.buf, z.buf[n:z.n+n])
  55. }
  56. }
  57. return err
  58. }
  59. func (z *bufioEncWriter) flush() {
  60. halt.onerror(z.flushErr())
  61. }
  62. func (z *bufioEncWriter) writeb(s []byte) {
  63. LOOP:
  64. a := len(z.buf) - z.n
  65. if len(s) > a {
  66. z.n += copy(z.buf[z.n:], s[:a])
  67. s = s[a:]
  68. z.flush()
  69. goto LOOP
  70. }
  71. z.n += copy(z.buf[z.n:], s)
  72. }
  73. func (z *bufioEncWriter) writestr(s string) {
  74. // z.writeb(bytesView(s)) // inlined below
  75. LOOP:
  76. a := len(z.buf) - z.n
  77. if len(s) > a {
  78. z.n += copy(z.buf[z.n:], s[:a])
  79. s = s[a:]
  80. z.flush()
  81. goto LOOP
  82. }
  83. z.n += copy(z.buf[z.n:], s)
  84. }
  85. func (z *bufioEncWriter) writeqstr(s string) {
  86. // z.writen1('"')
  87. // z.writestr(s)
  88. // z.writen1('"')
  89. if z.n+len(s)+2 > len(z.buf) {
  90. z.flush()
  91. }
  92. z.buf[z.n] = '"'
  93. z.n++
  94. LOOP:
  95. a := len(z.buf) - z.n
  96. if len(s)+1 > a {
  97. z.n += copy(z.buf[z.n:], s[:a])
  98. s = s[a:]
  99. z.flush()
  100. goto LOOP
  101. }
  102. z.n += copy(z.buf[z.n:], s)
  103. z.buf[z.n] = '"'
  104. z.n++
  105. }
  106. func (z *bufioEncWriter) writen1(b1 byte) {
  107. if 1 > len(z.buf)-z.n {
  108. z.flush()
  109. }
  110. z.buf[z.n] = b1
  111. z.n++
  112. }
  113. func (z *bufioEncWriter) writen2(b1, b2 byte) {
  114. if 2 > len(z.buf)-z.n {
  115. z.flush()
  116. }
  117. z.buf[z.n+1] = b2
  118. z.buf[z.n] = b1
  119. z.n += 2
  120. }
  121. func (z *bufioEncWriter) writen4(b [4]byte) {
  122. if 4 > len(z.buf)-z.n {
  123. z.flush()
  124. }
  125. copy(z.buf[z.n:], b[:])
  126. z.n += 4
  127. }
  128. func (z *bufioEncWriter) writen8(b [8]byte) {
  129. if 8 > len(z.buf)-z.n {
  130. z.flush()
  131. }
  132. copy(z.buf[z.n:], b[:])
  133. z.n += 8
  134. }
  135. func (z *bufioEncWriter) endErr() (err error) {
  136. if z.n > 0 {
  137. err = z.flushErr()
  138. }
  139. return
  140. }
  141. // ---------------------------------------------
  142. // bytesEncAppender implements encWriter and can write to an byte slice.
  143. type bytesEncAppender struct {
  144. b []byte
  145. out *[]byte
  146. }
  147. func (z *bytesEncAppender) writeb(s []byte) {
  148. z.b = append(z.b, s...)
  149. }
  150. func (z *bytesEncAppender) writestr(s string) {
  151. z.b = append(z.b, s...)
  152. }
  153. func (z *bytesEncAppender) writeqstr(s string) {
  154. z.b = append(append(append(z.b, '"'), s...), '"')
  155. // z.b = append(z.b, '"')
  156. // z.b = append(z.b, s...)
  157. // z.b = append(z.b, '"')
  158. }
  159. func (z *bytesEncAppender) writen1(b1 byte) {
  160. z.b = append(z.b, b1)
  161. }
  162. func (z *bytesEncAppender) writen2(b1, b2 byte) {
  163. z.b = append(z.b, b1, b2)
  164. }
  165. func (z *bytesEncAppender) writen4(b [4]byte) {
  166. z.b = append(z.b, b[:]...)
  167. // z.b = append(z.b, b[0], b[1], b[2], b[3]) // prevents inlining encWr.writen4
  168. }
  169. func (z *bytesEncAppender) writen8(b [8]byte) {
  170. z.b = append(z.b, b[:]...)
  171. // z.b = append(z.b, b[0], b[1], b[2], b[3], b[4], b[5], b[6], b[7]) // prevents inlining encWr.writen4
  172. }
  173. func (z *bytesEncAppender) endErr() error {
  174. *(z.out) = z.b
  175. return nil
  176. }
  177. func (z *bytesEncAppender) reset(in []byte, out *[]byte) {
  178. z.b = in[:0]
  179. z.out = out
  180. }
  181. // --------------------------------------------------
  182. type encWr struct {
  183. bytes bool // encoding to []byte
  184. js bool // is json encoder?
  185. be bool // is binary encoder?
  186. c containerState
  187. calls uint16
  188. seq uint16 // sequencer (e.g. used by binc for symbols, etc)
  189. wb bytesEncAppender
  190. wf *bufioEncWriter
  191. }
  192. // MARKER: manually inline bytesEncAppender.writenx/writeqstr methods,
  193. // as calling them causes encWr.writenx/writeqstr methods to not be inlined (cost > 80).
  194. //
  195. // i.e. e.g. instead of writing z.wb.writen2(b1, b2), use z.wb.b = append(z.wb.b, b1, b2)
  196. func (z *encWr) writeb(s []byte) {
  197. if z.bytes {
  198. z.wb.writeb(s)
  199. } else {
  200. z.wf.writeb(s)
  201. }
  202. }
  203. func (z *encWr) writeqstr(s string) {
  204. if z.bytes {
  205. // MARKER: z.wb.writeqstr(s)
  206. z.wb.b = append(append(append(z.wb.b, '"'), s...), '"')
  207. } else {
  208. z.wf.writeqstr(s)
  209. }
  210. }
  211. func (z *encWr) writestr(s string) {
  212. if z.bytes {
  213. z.wb.writestr(s)
  214. } else {
  215. z.wf.writestr(s)
  216. }
  217. }
  218. func (z *encWr) writen1(b1 byte) {
  219. if z.bytes {
  220. z.wb.writen1(b1)
  221. } else {
  222. z.wf.writen1(b1)
  223. }
  224. }
  225. func (z *encWr) writen2(b1, b2 byte) {
  226. if z.bytes {
  227. // MARKER: z.wb.writen2(b1, b2)
  228. z.wb.b = append(z.wb.b, b1, b2)
  229. } else {
  230. z.wf.writen2(b1, b2)
  231. }
  232. }
  233. func (z *encWr) writen4(b [4]byte) {
  234. if z.bytes {
  235. z.wb.writen4(b)
  236. } else {
  237. z.wf.writen4(b)
  238. }
  239. }
  240. func (z *encWr) writen8(b [8]byte) {
  241. if z.bytes {
  242. z.wb.writen8(b)
  243. } else {
  244. z.wf.writen8(b)
  245. }
  246. }
  247. func (z *encWr) endErr() error {
  248. if z.bytes {
  249. return z.wb.endErr()
  250. }
  251. return z.wf.endErr()
  252. }
  253. func (z *encWr) end() {
  254. halt.onerror(z.endErr())
  255. }
  256. var _ encWriter = (*encWr)(nil)