github.com/ugorji/go/codec@v1.2.13-0.20240307214044-07c54c229a5a/writer.go (about)

     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  
     4  package codec
     5  
     6  import "io"
     7  
     8  // encWriter abstracts writing to a byte array or to an io.Writer.
     9  type encWriter interface {
    10  	writeb([]byte)
    11  	writestr(string)
    12  	writeqstr(string) // write string wrapped in quotes ie "..."
    13  	writen1(byte)
    14  
    15  	// add convenience functions for writing 2,4
    16  	writen2(byte, byte)
    17  	writen4([4]byte)
    18  	writen8([8]byte)
    19  
    20  	end()
    21  }
    22  
    23  // ---------------------------------------------
    24  
    25  type bufioEncWriter struct {
    26  	w io.Writer
    27  
    28  	buf []byte
    29  
    30  	n int
    31  
    32  	b [16]byte // scratch buffer and padding (cache-aligned)
    33  }
    34  
    35  func (z *bufioEncWriter) reset(w io.Writer, bufsize int, blist *bytesFreelist) {
    36  	z.w = w
    37  	z.n = 0
    38  	if bufsize <= 0 {
    39  		bufsize = defEncByteBufSize
    40  	}
    41  	// bufsize must be >= 8, to accomodate writen methods (where n <= 8)
    42  	if bufsize <= 8 {
    43  		bufsize = 8
    44  	}
    45  	if cap(z.buf) < bufsize {
    46  		if len(z.buf) > 0 && &z.buf[0] != &z.b[0] {
    47  			blist.put(z.buf)
    48  		}
    49  		if len(z.b) > bufsize {
    50  			z.buf = z.b[:]
    51  		} else {
    52  			z.buf = blist.get(bufsize)
    53  		}
    54  	}
    55  	z.buf = z.buf[:cap(z.buf)]
    56  }
    57  
    58  func (z *bufioEncWriter) flushErr() (err error) {
    59  	n, err := z.w.Write(z.buf[:z.n])
    60  	z.n -= n
    61  	if z.n > 0 {
    62  		if err == nil {
    63  			err = io.ErrShortWrite
    64  		}
    65  		if n > 0 {
    66  			copy(z.buf, z.buf[n:z.n+n])
    67  		}
    68  	}
    69  	return err
    70  }
    71  
    72  func (z *bufioEncWriter) flush() {
    73  	halt.onerror(z.flushErr())
    74  }
    75  
    76  func (z *bufioEncWriter) writeb(s []byte) {
    77  LOOP:
    78  	a := len(z.buf) - z.n
    79  	if len(s) > a {
    80  		z.n += copy(z.buf[z.n:], s[:a])
    81  		s = s[a:]
    82  		z.flush()
    83  		goto LOOP
    84  	}
    85  	z.n += copy(z.buf[z.n:], s)
    86  }
    87  
    88  func (z *bufioEncWriter) writestr(s string) {
    89  	// z.writeb(bytesView(s)) // inlined below
    90  LOOP:
    91  	a := len(z.buf) - z.n
    92  	if len(s) > a {
    93  		z.n += copy(z.buf[z.n:], s[:a])
    94  		s = s[a:]
    95  		z.flush()
    96  		goto LOOP
    97  	}
    98  	z.n += copy(z.buf[z.n:], s)
    99  }
   100  
   101  func (z *bufioEncWriter) writeqstr(s string) {
   102  	// z.writen1('"')
   103  	// z.writestr(s)
   104  	// z.writen1('"')
   105  
   106  	if z.n+len(s)+2 > len(z.buf) {
   107  		z.flush()
   108  	}
   109  	setByteAt(z.buf, uint(z.n), '"')
   110  	// z.buf[z.n] = '"'
   111  	z.n++
   112  LOOP:
   113  	a := len(z.buf) - z.n
   114  	if len(s)+1 > a {
   115  		z.n += copy(z.buf[z.n:], s[:a])
   116  		s = s[a:]
   117  		z.flush()
   118  		goto LOOP
   119  	}
   120  	z.n += copy(z.buf[z.n:], s)
   121  	setByteAt(z.buf, uint(z.n), '"')
   122  	// z.buf[z.n] = '"'
   123  	z.n++
   124  }
   125  
   126  func (z *bufioEncWriter) writen1(b1 byte) {
   127  	if 1 > len(z.buf)-z.n {
   128  		z.flush()
   129  	}
   130  	setByteAt(z.buf, uint(z.n), b1)
   131  	// z.buf[z.n] = b1
   132  	z.n++
   133  }
   134  func (z *bufioEncWriter) writen2(b1, b2 byte) {
   135  	if 2 > len(z.buf)-z.n {
   136  		z.flush()
   137  	}
   138  	setByteAt(z.buf, uint(z.n+1), b2)
   139  	setByteAt(z.buf, uint(z.n), b1)
   140  	// z.buf[z.n+1] = b2
   141  	// z.buf[z.n] = b1
   142  	z.n += 2
   143  }
   144  
   145  func (z *bufioEncWriter) writen4(b [4]byte) {
   146  	if 4 > len(z.buf)-z.n {
   147  		z.flush()
   148  	}
   149  	// setByteAt(z.buf, uint(z.n+3), b4)
   150  	// setByteAt(z.buf, uint(z.n+2), b3)
   151  	// setByteAt(z.buf, uint(z.n+1), b2)
   152  	// setByteAt(z.buf, uint(z.n), b1)
   153  	copy(z.buf[z.n:], b[:])
   154  	z.n += 4
   155  }
   156  
   157  func (z *bufioEncWriter) writen8(b [8]byte) {
   158  	if 8 > len(z.buf)-z.n {
   159  		z.flush()
   160  	}
   161  	copy(z.buf[z.n:], b[:])
   162  	z.n += 8
   163  }
   164  
   165  func (z *bufioEncWriter) endErr() (err error) {
   166  	if z.n > 0 {
   167  		err = z.flushErr()
   168  	}
   169  	return
   170  }
   171  
   172  // ---------------------------------------------
   173  
   174  // bytesEncAppender implements encWriter and can write to an byte slice.
   175  type bytesEncAppender struct {
   176  	b   []byte
   177  	out *[]byte
   178  }
   179  
   180  func (z *bytesEncAppender) writeb(s []byte) {
   181  	z.b = append(z.b, s...)
   182  }
   183  func (z *bytesEncAppender) writestr(s string) {
   184  	z.b = append(z.b, s...)
   185  }
   186  func (z *bytesEncAppender) writeqstr(s string) {
   187  	z.b = append(append(append(z.b, '"'), s...), '"')
   188  	// z.b = append(z.b, '"')
   189  	// z.b = append(z.b, s...)
   190  	// z.b = append(z.b, '"')
   191  }
   192  func (z *bytesEncAppender) writen1(b1 byte) {
   193  	z.b = append(z.b, b1)
   194  }
   195  func (z *bytesEncAppender) writen2(b1, b2 byte) {
   196  	z.b = append(z.b, b1, b2)
   197  }
   198  
   199  func (z *bytesEncAppender) writen4(b [4]byte) {
   200  	z.b = append(z.b, b[:]...)
   201  	// z.b = append(z.b, b1, b2, b3, b4) // prevents inlining encWr.writen4
   202  }
   203  
   204  func (z *bytesEncAppender) writen8(b [8]byte) {
   205  	z.b = append(z.b, b[:]...)
   206  	// 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
   207  }
   208  
   209  func (z *bytesEncAppender) endErr() error {
   210  	*(z.out) = z.b
   211  	return nil
   212  }
   213  func (z *bytesEncAppender) reset(in []byte, out *[]byte) {
   214  	z.b = in[:0]
   215  	z.out = out
   216  }
   217  
   218  // --------------------------------------------------
   219  
   220  type encWr struct {
   221  	wb bytesEncAppender
   222  	wf *bufioEncWriter
   223  
   224  	bytes bool // encoding to []byte
   225  
   226  	// MARKER: these fields below should belong directly in Encoder.
   227  	// we pack them here for space efficiency and cache-line optimization.
   228  
   229  	js bool // is json encoder?
   230  	be bool // is binary encoder?
   231  
   232  	c containerState
   233  
   234  	calls uint16
   235  	seq   uint16 // sequencer (e.g. used by binc for symbols, etc)
   236  }
   237  
   238  // MARKER: manually inline bytesEncAppender.writenx/writeqstr methods,
   239  // as calling them causes encWr.writenx/writeqstr methods to not be inlined (cost > 80).
   240  //
   241  // i.e. e.g. instead of writing z.wb.writen2(b1, b2), use z.wb.b = append(z.wb.b, b1, b2)
   242  
   243  func (z *encWr) writeb(s []byte) {
   244  	if z.bytes {
   245  		z.wb.writeb(s)
   246  	} else {
   247  		z.wf.writeb(s)
   248  	}
   249  }
   250  func (z *encWr) writestr(s string) {
   251  	if z.bytes {
   252  		z.wb.writestr(s)
   253  	} else {
   254  		z.wf.writestr(s)
   255  	}
   256  }
   257  
   258  // MARKER: Add WriteStr to be called directly by generated code without a genHelper forwarding function.
   259  // Go's inlining model adds cost for forwarding functions, preventing inlining (cost goes above 80 budget).
   260  
   261  func (z *encWr) WriteStr(s string) {
   262  	if z.bytes {
   263  		z.wb.writestr(s)
   264  	} else {
   265  		z.wf.writestr(s)
   266  	}
   267  }
   268  
   269  func (z *encWr) writen1(b1 byte) {
   270  	if z.bytes {
   271  		z.wb.writen1(b1)
   272  	} else {
   273  		z.wf.writen1(b1)
   274  	}
   275  }
   276  
   277  func (z *encWr) writen2(b1, b2 byte) {
   278  	if z.bytes {
   279  		// MARKER: z.wb.writen2(b1, b2)
   280  		z.wb.b = append(z.wb.b, b1, b2)
   281  	} else {
   282  		z.wf.writen2(b1, b2)
   283  	}
   284  }
   285  
   286  func (z *encWr) writen4(b [4]byte) {
   287  	if z.bytes {
   288  		// MARKER: z.wb.writen4(b1, b2, b3, b4)
   289  		z.wb.b = append(z.wb.b, b[:]...)
   290  		// z.wb.writen4(b)
   291  	} else {
   292  		z.wf.writen4(b)
   293  	}
   294  }
   295  func (z *encWr) writen8(b [8]byte) {
   296  	if z.bytes {
   297  		// z.wb.b = append(z.wb.b, b[:]...)
   298  		z.wb.writen8(b)
   299  	} else {
   300  		z.wf.writen8(b)
   301  	}
   302  }
   303  
   304  func (z *encWr) writeqstr(s string) {
   305  	if z.bytes {
   306  		// MARKER: z.wb.writeqstr(s)
   307  		z.wb.b = append(append(append(z.wb.b, '"'), s...), '"')
   308  	} else {
   309  		z.wf.writeqstr(s)
   310  	}
   311  }
   312  
   313  func (z *encWr) endErr() error {
   314  	if z.bytes {
   315  		return z.wb.endErr()
   316  	}
   317  	return z.wf.endErr()
   318  }
   319  
   320  func (z *encWr) end() {
   321  	halt.onerror(z.endErr())
   322  }
   323  
   324  var _ encWriter = (*encWr)(nil)