github.com/aacfactory/avro@v1.2.12/internal/base/writer.go (about) 1 package base 2 3 import ( 4 "encoding/binary" 5 "io" 6 "math" 7 ) 8 9 // WriterFunc is a function used to customize the Writer. 10 type WriterFunc func(w *Writer) 11 12 // WithWriterConfig specifies the configuration to use with a writer. 13 func WithWriterConfig(cfg API) WriterFunc { 14 return func(w *Writer) { 15 w.cfg = cfg.(*frozenConfig) 16 } 17 } 18 19 // Writer is an Avro specific io.Writer. 20 type Writer struct { 21 cfg *frozenConfig 22 out io.Writer 23 buf []byte 24 Error error 25 } 26 27 // NewWriter creates a new Writer. 28 func NewWriter(out io.Writer, bufSize int, opts ...WriterFunc) *Writer { 29 writer := &Writer{ 30 cfg: DefaultConfig.(*frozenConfig), 31 out: out, 32 buf: make([]byte, 0, bufSize), 33 Error: nil, 34 } 35 36 for _, opt := range opts { 37 opt(writer) 38 } 39 40 return writer 41 } 42 43 // Reset resets the Writer with a new io.Writer attached. 44 func (w *Writer) Reset(out io.Writer) { 45 w.out = out 46 w.buf = w.buf[:0] 47 } 48 49 // Buffered returns the number of buffered bytes. 50 func (w *Writer) Buffered() int { 51 return len(w.buf) 52 } 53 54 // Buffer gets the Writer buffer. 55 func (w *Writer) Buffer() []byte { 56 return w.buf 57 } 58 59 // Flush writes any buffered data to the underlying io.Writer. 60 func (w *Writer) Flush() error { 61 if w.out == nil { 62 return nil 63 } 64 if w.Error != nil { 65 return w.Error 66 } 67 68 n, err := w.out.Write(w.buf) 69 if n < len(w.buf) && err == nil { 70 err = io.ErrShortWrite 71 } 72 if err != nil { 73 if w.Error == nil { 74 w.Error = err 75 } 76 return err 77 } 78 79 w.buf = w.buf[:0] 80 81 return nil 82 } 83 84 func (w *Writer) writeByte(b byte) { 85 w.buf = append(w.buf, b) 86 } 87 88 // Write writes raw bytes to the Writer. 89 func (w *Writer) Write(b []byte) (int, error) { 90 w.buf = append(w.buf, b...) 91 return len(b), nil 92 } 93 94 // WriteBool writes a Bool to the Writer. 95 func (w *Writer) WriteBool(b bool) { 96 if b { 97 w.writeByte(0x01) 98 return 99 } 100 w.writeByte(0x00) 101 } 102 103 // WriteInt writes an Int to the Writer. 104 func (w *Writer) WriteInt(i int32) { 105 e := uint64((uint32(i) << 1) ^ uint32(i>>31)) 106 w.encodeInt(e) 107 } 108 109 // WriteLong writes a Long to the Writer. 110 func (w *Writer) WriteLong(i int64) { 111 e := (uint64(i) << 1) ^ uint64(i>>63) 112 w.encodeInt(e) 113 } 114 115 func (w *Writer) encodeInt(i uint64) { 116 if i == 0 { 117 w.writeByte(0) 118 return 119 } 120 121 for i > 0 { 122 b := byte(i) & 0x7F 123 i >>= 7 124 125 if i != 0 { 126 b |= 0x80 127 } 128 w.writeByte(b) 129 } 130 } 131 132 // WriteFloat writes a Float to the Writer. 133 func (w *Writer) WriteFloat(f float32) { 134 b := make([]byte, 4) 135 binary.LittleEndian.PutUint32(b, math.Float32bits(f)) 136 137 w.buf = append(w.buf, b...) 138 } 139 140 // WriteDouble writes a Double to the Writer. 141 func (w *Writer) WriteDouble(f float64) { 142 b := make([]byte, 8) 143 binary.LittleEndian.PutUint64(b, math.Float64bits(f)) 144 145 w.buf = append(w.buf, b...) 146 } 147 148 // WriteBytes writes Bytes to the Writer. 149 func (w *Writer) WriteBytes(b []byte) { 150 w.WriteLong(int64(len(b))) 151 w.buf = append(w.buf, b...) 152 } 153 154 // WriteString reads a String to the Writer. 155 func (w *Writer) WriteString(s string) { 156 w.WriteLong(int64(len(s))) 157 w.buf = append(w.buf, s...) 158 } 159 160 // WriteBlockHeader writes a Block Header to the Writer. 161 func (w *Writer) WriteBlockHeader(l, s int64) { 162 if s > 0 && !w.cfg.config.DisableBlockSizeHeader { 163 w.WriteLong(-l) 164 w.WriteLong(s) 165 return 166 } 167 w.WriteLong(l) 168 } 169 170 // WriteBlockCB writes a block using the callback. 171 func (w *Writer) WriteBlockCB(callback func(w *Writer) int64) int64 { 172 var dummyHeader [18]byte 173 headerStart := len(w.buf) 174 175 // Write dummy header 176 _, _ = w.Write(dummyHeader[:]) 177 178 // Write block data 179 capturedAt := len(w.buf) 180 length := callback(w) 181 size := int64(len(w.buf) - capturedAt) 182 183 // Take a reference to the block data 184 captured := w.buf[capturedAt:len(w.buf)] 185 186 // Rewrite the header 187 w.buf = w.buf[:headerStart] 188 w.WriteBlockHeader(length, size) 189 190 // Copy the block data back to its position 191 w.buf = append(w.buf, captured...) 192 193 return length 194 }