<commit_before>package lich

import (
	"bufio"
	"bytes"
	"errors"
	"fmt"
	"io"
	"strconv"
)

type LexTokenType int

func (t LexTokenType) String() string {
	return []string{
		Data:       "Data",
		ArrayOpen:  "ArrayOpen",
		ArrayClose: "ArrayClose",
		DictOpen:   "DictOpen",
		DictClose:  "DictClose",
	}[int(t)]
}

// Lexer token enum
const (
	Data LexTokenType = iota
	ArrayOpen
	ArrayClose
	DictOpen
	DictClose
)

type LexToken struct {
	LexTokenType
	Data []byte
}

func (t LexToken) String() string {
	if t.LexTokenType == Data {
		return fmt.Sprintf("%v<%s>", t.LexTokenType, t.Data)
	}
	return t.LexTokenType.String()
}

// Errors introduced by Lex
var (
	ErrUnexpectedChar     = errors.New("Unexpected character")
	ErrBufferTooShort     = errors.New("Input too short")
	ErrMissingClosingChar = errors.New("Missing closing character")
)

func Lex(r io.Reader) ([]LexToken, error) {
	b := bufio.NewReader(r)
	return lexElement(b)
}

type readScanner interface {
	io.Reader
	io.ByteScanner
}

func lexElement(b readScanner) ([]LexToken, error) {
	// Spec says size will always fit into 20 bytes or less
	const maxSizeLength = 20

	sizeBuf := make([]byte, maxSizeLength)
	i := 0

	for i = range sizeBuf {
		c, err := b.ReadByte()
		if err != nil {
			return nil, err
		}

		if !isDigit(c) {
			b.UnreadByte()
			break
		}

		sizeBuf[i] = c
	}

	// Turn the size into an int
	size, err := strconv.Atoi(string(sizeBuf[:i]))
	if err != nil {
		return nil, ErrUnexpectedChar
	}

	c, err := b.ReadByte()
	if err != nil {
		return nil, err
	}

	// Bail out on bad data before making a possibly large buffer
	switch c {
	case '<', '[', '{':
	default:
		return nil, ErrUnexpectedChar
	}

	buf := make([]byte, size+1) // +1 for terminal character
	if _, err := io.ReadFull(b, buf); err != nil {
		return nil, err
	}

	switch c {
	case '<':
		return lexData(buf)
	case '[':
		return lexArrayOrDict(buf, true)
	case '{':
		return lexArrayOrDict(buf, false)
	}
	panic("unreachable")
}

func isDigit(b byte) bool {
	return b >= '0' && b <= '9'
}

func lexData(buf []byte) ([]LexToken, error) {
	if len(buf) < 1 {
		return nil, ErrBufferTooShort
	}
	if buf[len(buf)-1] != '>' {
		return nil, ErrMissingClosingChar
	}
	tokens := []LexToken{{LexTokenType: Data, Data: buf[:len(buf)-1]}}
	return tokens, nil
}

func lexArrayOrDict(buf []byte, isArray bool) ([]LexToken, error) {
	var (
		closingChar = byte(']')
		openingType = ArrayOpen
		closingType = ArrayClose
	)
	if !isArray {
		closingChar = byte('}')
		openingType = DictOpen
		closingType = DictClose
	}

	// First check for consistency before parsing guts
	if buf[len(buf)-1] != closingChar {
		return nil, ErrMissingClosingChar
	}

	tokens := []LexToken{{LexTokenType: openingType}}

	b := bytes.NewReader(buf)

	for b.Len() > 1 {
		newTokens, err := lexElement(b)
		if err != nil {
			return nil, err
		}
		tokens = append(tokens, newTokens...)
	}

	tokens = append(tokens, LexToken{LexTokenType: closingType})
	return tokens, nil
}<commit_msg>Lexer: Improve efficiency<commit_after>package lich

import (
	"bufio"
	"errors"
	"fmt"
	"io"
	"strconv"
)

type LexTokenType int

func (t LexTokenType) String() string {
	return []string{
		Data:       "Data",
		ArrayOpen:  "ArrayOpen",
		ArrayClose: "ArrayClose",
		DictOpen:   "DictOpen",
		DictClose:  "DictClose",
	}[int(t)]
}

// Lexer token enum
const (
	Data LexTokenType = iota
	ArrayOpen
	ArrayClose
	DictOpen
	DictClose
)

type LexToken struct {
	LexTokenType
	Data []byte
}

func (t LexToken) String() string {
	if t.LexTokenType == Data {
		return fmt.Sprintf("%v<%s>", t.LexTokenType, t.Data)
	}
	return t.LexTokenType.String()
}

// Errors introduced by Lex
var (
	ErrUnexpectedChar     = errors.New("Unexpected character")
	ErrUnreadableSize     = errors.New("Couldn't read size of element")
	ErrMissingClosingChar = errors.New("Missing closing character")
)

func Lex(r io.Reader) ([]LexToken, error) {
	b := bufio.NewReader(r)
	return lexElement(b)
}

func lexElement(b *bufio.Reader) ([]LexToken, error) {
	// Spec says size will always fit into 20 bytes or less
	const maxSizeLength = 20

	sizeBuf := make([]byte, maxSizeLength)
	i := 0

	for i = range sizeBuf {
		c, err := b.ReadByte()
		if err != nil {
			return nil, err
		}

		if !isDigit(c) {
			b.UnreadByte()
			break
		}

		sizeBuf[i] = c
	}

	// Turn the size into an int
	size, err := strconv.ParseInt(string(sizeBuf[:i]), 10, 64)
	if err != nil {
		return nil, ErrUnreadableSize
	}

	c, err := b.ReadByte()
	if err != nil {
		return nil, err
	}

	switch c {
	case '<':
		return lexData(b, size)
	case '[':
		return lexArrayOrDict(b, size, true)
	case '{':
		return lexArrayOrDict(b, size, false)
	}

	return nil, ErrUnexpectedChar
}

func isDigit(b byte) bool {
	return b >= '0' && b <= '9'
}

func lexData(b *bufio.Reader, size int64) ([]LexToken, error) {
	data := make([]byte, int(size))
	_, err := io.ReadFull(b, data)
	if err != nil {
		return nil, err
	}

	c, err := b.ReadByte()
	if err != nil {
		return nil, err
	}
	if c != '>' {
		return nil, ErrMissingClosingChar
	}

	tokens := []LexToken{{LexTokenType: Data, Data: data}}
	return tokens, nil
}

func lexArrayOrDict(b *bufio.Reader, size int64, isArray bool) ([]LexToken, error) {
	var (
		closingChar = byte(']')
		openingType = ArrayOpen
		closingType = ArrayClose
	)
	if !isArray {
		closingChar = byte('}')
		openingType = DictOpen
		closingType = DictClose
	}
	tokens := []LexToken{{LexTokenType: openingType}}

	lb := bufio.NewReader(io.LimitReader(b, size))

	for {
		_, err := lb.ReadByte()
		if err == io.EOF {
			break
		}
		if err != nil {
			return nil, err
		}
		lb.UnreadByte()
		newTokens, err := lexElement(lb)
		if err != nil {
			return nil, err
		}
		tokens = append(tokens, newTokens...)
	}

	c, err := b.ReadByte()
	if err != nil {
		return nil, err
	}
	if c != closingChar {
		return nil, ErrMissingClosingChar
	}
	tokens = append(tokens, LexToken{LexTokenType: closingType})
	return tokens, nil
}
{'fraction_non_alphanumeric': 0.1426653239342061, 'fraction_numerical': 0.004867405169519974, 'mean_word_length': 3.5697852760736195, 'pattern_counts': {'":': 0, '<': 11, '<?xml version=': 0, '>': 10, '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': '3963693', 'n_tokens_mistral': 2388, 'n_tokens_neox': 2194, 'n_words': 606}