feature_stream.go 6.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308
  1. package jsoniter
  2. import (
  3. "io"
  4. )
  5. // stream is a io.Writer like object, with JSON specific write functions.
  6. // Error is not returned as return value, but stored as Error member on this stream instance.
  7. type Stream struct {
  8. cfg *frozenConfig
  9. out io.Writer
  10. buf []byte
  11. n int
  12. Error error
  13. indention int
  14. Attachment interface{} // open for customized encoder
  15. }
  16. // NewStream create new stream instance.
  17. // cfg can be jsoniter.ConfigDefault.
  18. // out can be nil if write to internal buffer.
  19. // bufSize is the initial size for the internal buffer in bytes.
  20. func NewStream(cfg API, out io.Writer, bufSize int) *Stream {
  21. return &Stream{
  22. cfg: cfg.(*frozenConfig),
  23. out: out,
  24. buf: make([]byte, bufSize),
  25. n: 0,
  26. Error: nil,
  27. indention: 0,
  28. }
  29. }
  30. // Pool returns a pool can provide more stream with same configuration
  31. func (stream *Stream) Pool() StreamPool {
  32. return stream.cfg
  33. }
  34. // Reset reuse this stream instance by assign a new writer
  35. func (stream *Stream) Reset(out io.Writer) {
  36. stream.out = out
  37. stream.n = 0
  38. }
  39. // Available returns how many bytes are unused in the buffer.
  40. func (stream *Stream) Available() int {
  41. return len(stream.buf) - stream.n
  42. }
  43. // Buffered returns the number of bytes that have been written into the current buffer.
  44. func (stream *Stream) Buffered() int {
  45. return stream.n
  46. }
  47. // Buffer if writer is nil, use this method to take the result
  48. func (stream *Stream) Buffer() []byte {
  49. return stream.buf[:stream.n]
  50. }
  51. // Write writes the contents of p into the buffer.
  52. // It returns the number of bytes written.
  53. // If nn < len(p), it also returns an error explaining
  54. // why the write is short.
  55. func (stream *Stream) Write(p []byte) (nn int, err error) {
  56. for len(p) > stream.Available() && stream.Error == nil {
  57. if stream.out == nil {
  58. stream.growAtLeast(len(p))
  59. } else {
  60. var n int
  61. if stream.Buffered() == 0 {
  62. // Large write, empty buffer.
  63. // Write directly from p to avoid copy.
  64. n, stream.Error = stream.out.Write(p)
  65. } else {
  66. n = copy(stream.buf[stream.n:], p)
  67. stream.n += n
  68. stream.Flush()
  69. }
  70. nn += n
  71. p = p[n:]
  72. }
  73. }
  74. if stream.Error != nil {
  75. return nn, stream.Error
  76. }
  77. n := copy(stream.buf[stream.n:], p)
  78. stream.n += n
  79. nn += n
  80. return nn, nil
  81. }
  82. // WriteByte writes a single byte.
  83. func (stream *Stream) writeByte(c byte) {
  84. if stream.Error != nil {
  85. return
  86. }
  87. if stream.Available() < 1 {
  88. stream.growAtLeast(1)
  89. }
  90. stream.buf[stream.n] = c
  91. stream.n++
  92. }
  93. func (stream *Stream) writeTwoBytes(c1 byte, c2 byte) {
  94. if stream.Error != nil {
  95. return
  96. }
  97. if stream.Available() < 2 {
  98. stream.growAtLeast(2)
  99. }
  100. stream.buf[stream.n] = c1
  101. stream.buf[stream.n+1] = c2
  102. stream.n += 2
  103. }
  104. func (stream *Stream) writeThreeBytes(c1 byte, c2 byte, c3 byte) {
  105. if stream.Error != nil {
  106. return
  107. }
  108. if stream.Available() < 3 {
  109. stream.growAtLeast(3)
  110. }
  111. stream.buf[stream.n] = c1
  112. stream.buf[stream.n+1] = c2
  113. stream.buf[stream.n+2] = c3
  114. stream.n += 3
  115. }
  116. func (stream *Stream) writeFourBytes(c1 byte, c2 byte, c3 byte, c4 byte) {
  117. if stream.Error != nil {
  118. return
  119. }
  120. if stream.Available() < 4 {
  121. stream.growAtLeast(4)
  122. }
  123. stream.buf[stream.n] = c1
  124. stream.buf[stream.n+1] = c2
  125. stream.buf[stream.n+2] = c3
  126. stream.buf[stream.n+3] = c4
  127. stream.n += 4
  128. }
  129. func (stream *Stream) writeFiveBytes(c1 byte, c2 byte, c3 byte, c4 byte, c5 byte) {
  130. if stream.Error != nil {
  131. return
  132. }
  133. if stream.Available() < 5 {
  134. stream.growAtLeast(5)
  135. }
  136. stream.buf[stream.n] = c1
  137. stream.buf[stream.n+1] = c2
  138. stream.buf[stream.n+2] = c3
  139. stream.buf[stream.n+3] = c4
  140. stream.buf[stream.n+4] = c5
  141. stream.n += 5
  142. }
  143. // Flush writes any buffered data to the underlying io.Writer.
  144. func (stream *Stream) Flush() error {
  145. if stream.out == nil {
  146. return nil
  147. }
  148. if stream.Error != nil {
  149. return stream.Error
  150. }
  151. if stream.n == 0 {
  152. return nil
  153. }
  154. n, err := stream.out.Write(stream.buf[0:stream.n])
  155. if n < stream.n && err == nil {
  156. err = io.ErrShortWrite
  157. }
  158. if err != nil {
  159. if n > 0 && n < stream.n {
  160. copy(stream.buf[0:stream.n-n], stream.buf[n:stream.n])
  161. }
  162. stream.n -= n
  163. stream.Error = err
  164. return err
  165. }
  166. stream.n = 0
  167. return nil
  168. }
  169. func (stream *Stream) ensure(minimal int) {
  170. available := stream.Available()
  171. if available < minimal {
  172. stream.growAtLeast(minimal)
  173. }
  174. }
  175. func (stream *Stream) growAtLeast(minimal int) {
  176. if stream.out != nil {
  177. stream.Flush()
  178. if stream.Available() >= minimal {
  179. return
  180. }
  181. }
  182. toGrow := len(stream.buf)
  183. if toGrow < minimal {
  184. toGrow = minimal
  185. }
  186. newBuf := make([]byte, len(stream.buf)+toGrow)
  187. copy(newBuf, stream.Buffer())
  188. stream.buf = newBuf
  189. }
  190. // WriteRaw write string out without quotes, just like []byte
  191. func (stream *Stream) WriteRaw(s string) {
  192. stream.ensure(len(s))
  193. if stream.Error != nil {
  194. return
  195. }
  196. n := copy(stream.buf[stream.n:], s)
  197. stream.n += n
  198. }
  199. // WriteNil write null to stream
  200. func (stream *Stream) WriteNil() {
  201. stream.writeFourBytes('n', 'u', 'l', 'l')
  202. }
  203. // WriteTrue write true to stream
  204. func (stream *Stream) WriteTrue() {
  205. stream.writeFourBytes('t', 'r', 'u', 'e')
  206. }
  207. // WriteFalse write false to stream
  208. func (stream *Stream) WriteFalse() {
  209. stream.writeFiveBytes('f', 'a', 'l', 's', 'e')
  210. }
  211. // WriteBool write true or false into stream
  212. func (stream *Stream) WriteBool(val bool) {
  213. if val {
  214. stream.WriteTrue()
  215. } else {
  216. stream.WriteFalse()
  217. }
  218. }
  219. // WriteObjectStart write { with possible indention
  220. func (stream *Stream) WriteObjectStart() {
  221. stream.indention += stream.cfg.indentionStep
  222. stream.writeByte('{')
  223. stream.writeIndention(0)
  224. }
  225. // WriteObjectField write "field": with possible indention
  226. func (stream *Stream) WriteObjectField(field string) {
  227. stream.WriteString(field)
  228. if stream.indention > 0 {
  229. stream.writeTwoBytes(':', ' ')
  230. } else {
  231. stream.writeByte(':')
  232. }
  233. }
  234. // WriteObjectEnd write } with possible indention
  235. func (stream *Stream) WriteObjectEnd() {
  236. stream.writeIndention(stream.cfg.indentionStep)
  237. stream.indention -= stream.cfg.indentionStep
  238. stream.writeByte('}')
  239. }
  240. // WriteEmptyObject write {}
  241. func (stream *Stream) WriteEmptyObject() {
  242. stream.writeByte('{')
  243. stream.writeByte('}')
  244. }
  245. // WriteMore write , with possible indention
  246. func (stream *Stream) WriteMore() {
  247. stream.writeByte(',')
  248. stream.writeIndention(0)
  249. }
  250. // WriteArrayStart write [ with possible indention
  251. func (stream *Stream) WriteArrayStart() {
  252. stream.indention += stream.cfg.indentionStep
  253. stream.writeByte('[')
  254. stream.writeIndention(0)
  255. }
  256. // WriteEmptyArray write []
  257. func (stream *Stream) WriteEmptyArray() {
  258. stream.writeTwoBytes('[', ']')
  259. }
  260. // WriteArrayEnd write ] with possible indention
  261. func (stream *Stream) WriteArrayEnd() {
  262. stream.writeIndention(stream.cfg.indentionStep)
  263. stream.indention -= stream.cfg.indentionStep
  264. stream.writeByte(']')
  265. }
  266. func (stream *Stream) writeIndention(delta int) {
  267. if stream.indention == 0 {
  268. return
  269. }
  270. stream.writeByte('\n')
  271. toWrite := stream.indention - delta
  272. stream.ensure(toWrite)
  273. for i := 0; i < toWrite && stream.n < len(stream.buf); i++ {
  274. stream.buf[stream.n] = ' '
  275. stream.n++
  276. }
  277. }