<commit_before>package ubjson

import (
	"encoding/binary"
	"io"
	"math"
	"strconv"

	structform "github.com/urso/go-structform"
)

type Visitor struct {
	w       writer
	scratch [16]byte
}

type writer struct {
	out io.Writer
}

var _ structform.ExtVisitor = &Visitor{}

const (
	maxUint  = ^uint(0)
	maxInt   = int(maxUint >> 1)
	isUint64 = maxUint > math.MaxUint32
	isInt64  = maxInt > math.MaxInt32
)

func (w writer) write(b []byte) error {
	_, err := w.out.Write(b)
	return err
}

func NewVisitor(out io.Writer) *Visitor {
	return &Visitor{w: writer{out}}
}

func (vs *Visitor) writeByte(b byte) error {
	vs.scratch[0] = b
	return vs.w.write(vs.scratch[:1])
}

func (vs *Visitor) optionalCount(l int) error {
	if l <= 0 {
		// don't add size if array is empty or size is unknown
		return nil
	}

	if err := vs.writeByte(countMarker); err != nil {
		return err
	}
	return vs.writeLen(l)
}

func (vs *Visitor) OnObjectStart(l int, _ structform.BaseType) error {
	// TODO: add typed object support in case of values being passed one by one

	// if number of elements is known, add size
	if err := vs.writeByte(objStartMarker); err != nil {
		return err
	}

	return vs.optionalCount(l)
}

func (vs *Visitor) OnObjectFinished() error {
	return vs.writeByte(objEndMarker)
}

func (vs *Visitor) OnKey(s string) error {
	return vs.string(s, false)
}

func (vs *Visitor) OnArrayStart(l int, t structform.BaseType) error {
	// TODO: optimize array by computing type tag

	if err := vs.writeByte(arrStartMarker); err != nil {
		return err
	}

	// if array size is known, add at least size
	return vs.optionalCount(l)
}

func (vs *Visitor) OnArrayFinished() error {
	return vs.writeByte(arrEndMarker)
}

func (vs *Visitor) writeLen(l int) error {
	return vs.onInt(l, true)
}

func (vs *Visitor) OnString(s string) error {
	if len(s) == 0 {
		return vs.OnNil()
	}
	return vs.string(s, true)
}

func (vs *Visitor) string(s string, marker bool) error {
	if marker {
		if err := vs.writeByte(stringMarker); err != nil {
			return err
		}
	}
	if err := vs.writeLen(len(s)); err != nil {
		return err
	}
	return vs.w.write(str2Bytes(s))
}

func (vs *Visitor) OnBool(b bool) error {
	if b {
		return vs.writeByte(trueMarker)
	}
	return vs.writeByte(falseMarker)
}

func (vs *Visitor) OnNil() error {
	return vs.writeByte(nullMarker)
}

// int

func (vs *Visitor) OnInt8(i int8) error {
	return vs.int8(i, true)
}

func (vs *Visitor) int8(i int8, marker bool) error {
	if marker {
		if err := vs.writeByte(int8Marker); err != nil {
			return err
		}
	}
	return vs.writeByte(byte(i))
}

func (vs *Visitor) OnInt16(i int16) error {
	if math.MinInt8 < i && i < math.MaxInt8 {
		return vs.OnInt8(int8(i))
	}
	return vs.int16(i, true)
}

func (vs *Visitor) int16(i int16, marker bool) error {
	if marker {
		if err := vs.writeByte(int16Marker); err != nil {
			return err
		}
	}
	binary.BigEndian.PutUint16(vs.scratch[:2], uint16(i))
	return vs.w.write(vs.scratch[:2])
}

func (vs *Visitor) OnInt32(i int32) error {
	if math.MinInt16 < i && i < math.MaxInt16 {
		return vs.OnInt16(int16(i))
	}
	return vs.int32(i, true)
}

func (vs *Visitor) int32(i int32, marker bool) error {
	if marker {
		if err := vs.writeByte(int32Marker); err != nil {
			return err
		}
	}
	binary.BigEndian.PutUint32(vs.scratch[:4], uint32(i))
	return vs.w.write(vs.scratch[:4])
}

func (vs *Visitor) OnInt64(i int64) error {
	if math.MinInt32 < i && i < math.MaxInt32 {
		return vs.OnInt32(int32(i))
	}
	return vs.int64(i, true)<commit_msg>use <= when checking value target size in ubjson<commit_after>package ubjson

import (
	"encoding/binary"
	"io"
	"math"
	"strconv"

	structform "github.com/urso/go-structform"
)

type Visitor struct {
	w       writer
	scratch [16]byte
}

type writer struct {
	out io.Writer
}

var _ structform.ExtVisitor = &Visitor{}

const (
	maxUint  = ^uint(0)
	maxInt   = int(maxUint >> 1)
	isUint64 = maxUint > math.MaxUint32
	isInt64  = maxInt > math.MaxInt32
)

func (w writer) write(b []byte) error {
	_, err := w.out.Write(b)
	return err
}

func NewVisitor(out io.Writer) *Visitor {
	return &Visitor{w: writer{out}}
}

func (vs *Visitor) writeByte(b byte) error {
	vs.scratch[0] = b
	return vs.w.write(vs.scratch[:1])
}

func (vs *Visitor) optionalCount(l int) error {
	if l <= 0 {
		// don't add size if array is empty or size is unknown
		return nil
	}

	if err := vs.writeByte(countMarker); err != nil {
		return err
	}
	return vs.writeLen(l)
}

func (vs *Visitor) OnObjectStart(l int, _ structform.BaseType) error {
	// TODO: add typed object support in case of values being passed one by one

	// if number of elements is known, add size
	if err := vs.writeByte(objStartMarker); err != nil {
		return err
	}

	return vs.optionalCount(l)
}

func (vs *Visitor) OnObjectFinished() error {
	return vs.writeByte(objEndMarker)
}

func (vs *Visitor) OnKey(s string) error {
	return vs.string(s, false)
}

func (vs *Visitor) OnArrayStart(l int, t structform.BaseType) error {
	// TODO: optimize array by computing type tag

	if err := vs.writeByte(arrStartMarker); err != nil {
		return err
	}

	// if array size is known, add at least size
	return vs.optionalCount(l)
}

func (vs *Visitor) OnArrayFinished() error {
	return vs.writeByte(arrEndMarker)
}

func (vs *Visitor) writeLen(l int) error {
	return vs.onInt(l, true)
}

func (vs *Visitor) OnString(s string) error {
	if len(s) == 0 {
		return vs.OnNil()
	}
	return vs.string(s, true)
}

func (vs *Visitor) string(s string, marker bool) error {
	if marker {
		if err := vs.writeByte(stringMarker); err != nil {
			return err
		}
	}
	if err := vs.writeLen(len(s)); err != nil {
		return err
	}
	return vs.w.write(str2Bytes(s))
}

func (vs *Visitor) OnBool(b bool) error {
	if b {
		return vs.writeByte(trueMarker)
	}
	return vs.writeByte(falseMarker)
}

func (vs *Visitor) OnNil() error {
	return vs.writeByte(nullMarker)
}

// int

func (vs *Visitor) OnInt8(i int8) error {
	return vs.int8(i, true)
}

func (vs *Visitor) int8(i int8, marker bool) error {
	if marker {
		if err := vs.writeByte(int8Marker); err != nil {
			return err
		}
	}
	return vs.writeByte(byte(i))
}

func (vs *Visitor) OnInt16(i int16) error {
	if math.MinInt8 <= i && i <= math.MaxInt8 {
		return vs.int8(int8(i), true)
	}
	return vs.int16(i, true)
}

func (vs *Visitor) int16(i int16, marker bool) error {
	if marker {
		if err := vs.writeByte(int16Marker); err != nil {
			return err
		}
	}
	binary.BigEndian.PutUint16(vs.scratch[:2], uint16(i))
	return vs.w.write(vs.scratch[:2])
}

func (vs *Visitor) OnInt32(i int32) error {
	if math.MinInt16 <= i && i <= math.MaxInt16 {
		return vs.OnInt16(int16(i))
	}
	return vs.int32(i, true)
}

func (vs *Visitor) int32(i int32, marker bool) error {
	if marker {
		if err := vs.writeByte(int32Marker); err != nil {
			return err
		}
	}
	binary.BigEndian.PutUint32(vs.scratch[:4], uint32(i))
	return vs.w.write(vs.scratch[:4])
}

func (vs *Visitor) OnInt64(i int64) error {
	if math.MinInt32 <= i && i <= math.MaxInt32 {
		return vs.OnInt32(int32(i))
	}
	return vs.int64(i, true)
{'fraction_non_alphanumeric': 0.14050291234550363, 'fraction_numerical': 0.024151157834919734, 'mean_word_length': 3.6933333333333334, 'pattern_counts': {'":': 0, '<': 18, '<?xml version=': 0, '>': 11, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6841713', 'n_tokens_mistral': 2865, 'n_tokens_neox': 2723, 'n_words': 757}