<commit_before>package blockpool

import (
	"github.com/Datadog/zstd"
	"github.com/go-errors/errors"
	"github.com/itchio/wharf/tlc"
)

/////////////////////////////
// Sink
/////////////////////////////

// A CompressingSink compresses blocks with ztsd-q9 before storing them
// to the underlying sink
type CompressingSink struct {
	Sink Sink

	compressedBuf []byte
}

var _ Sink = (*CompressingSink)(nil)

// Store first compresses the data, then stores it into the underling sink
func (cs *CompressingSink) Store(loc BlockLocation, data []byte) error {
	if cs.compressedBuf == nil {
		// planning for the worst case scenario - that compressing the data
		// actually increased the block size
		cs.compressedBuf = make([]byte, BigBlockSize*2)
	}

	_, err := zstd.CompressLevel(cs.compressedBuf, data, 9)
	if err != nil {
		return errors.Wrap(err, 1)
	}

	return cs.Sink.Store(loc, cs.compressedBuf)
}

// GetContainer returns the underlying source's container
func (cs *CompressingSink) GetContainer() *tlc.Container {
	return cs.Sink.GetContainer()
}
<commit_msg>Don't store full buffer, only size returned by ztsd<commit_after>package blockpool

import (
	"github.com/Datadog/zstd"
	"github.com/go-errors/errors"
	"github.com/itchio/wharf/tlc"
)

/////////////////////////////
// Sink
/////////////////////////////

// A CompressingSink compresses blocks with ztsd-q9 before storing them
// to the underlying sink
type CompressingSink struct {
	Sink Sink

	compressedBuf []byte
}

var _ Sink = (*CompressingSink)(nil)

// Store first compresses the data, then stores it into the underling sink
func (cs *CompressingSink) Store(loc BlockLocation, data []byte) error {
	if cs.compressedBuf == nil {
		// planning for the worst case scenario - that compressing the data
		// actually increased the block size
		cs.compressedBuf = make([]byte, BigBlockSize*2)
	}

	compressedData, err := zstd.CompressLevel(cs.compressedBuf, data, 9)
	if err != nil {
		return errors.Wrap(err, 1)
	}

	return cs.Sink.Store(loc, compressedData)
}

// GetContainer returns the underlying source's container
func (cs *CompressingSink) GetContainer() *tlc.Container {
	return cs.Sink.GetContainer()
}

{'fraction_non_alphanumeric': 0.15374024782010096, 'fraction_numerical': 0.0036714089031665903, 'mean_word_length': 3.2745098039215685, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, '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': '3285438', 'n_tokens_mistral': 705, 'n_tokens_neox': 681, 'n_words': 203}