notarius/internal/modules/generic/chunk/units/chunker.go

package units

import (
	"context"
	"encoding/json"
	"fmt"
	"math"
	"strconv"

	"gitea.maximumdirect.net/eric/notarius/internal/core/source"
	"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
	"gitea.maximumdirect.net/eric/notarius/internal/framework/pipeline"
)

const Key = "generic"

const (
	defaultMaxUnits     = 50
	defaultOverlapUnits = 0
)

var _ contracts.Chunker = (*Chunker)(nil)

type Options struct {
	MaxUnits     int
	OverlapUnits int
}

type Chunker struct {
	options Options
}

func New(options Options) *Chunker {
	return &Chunker{options: options}
}

func (c *Chunker) Key() string {
	return Key
}

func (*Chunker) ExecutionClass() contracts.ExecutionClass {
	return contracts.ExecutionClassDeterministic
}

func (c *Chunker) ReferenceSlots() []contracts.ReferenceSlot {
	return nil
}

func (c *Chunker) Plan(ctx context.Context, req contracts.ChunkRequest) (contracts.ChunkPlanResult, error) {
	if c == nil {
		return contracts.ChunkPlanResult{}, chunkerErrorf("chunker must not be nil")
	}
	if c.options.MaxUnits <= 0 || c.options.OverlapUnits < 0 || c.options.OverlapUnits >= c.options.MaxUnits {
		return contracts.ChunkPlanResult{}, chunkerErrorf("chunker options must be initialized by construction")
	}
	if ctx == nil {
		return contracts.ChunkPlanResult{}, chunkerErrorf("context must not be nil")
	}
	if err := ctx.Err(); err != nil {
		return contracts.ChunkPlanResult{}, chunkerErrorf("context error before chunking: %w", err)
	}
	if req.Source == nil {
		return contracts.ChunkPlanResult{}, chunkerErrorf("source must not be nil")
	}
	if len(req.Source.Units) == 0 {
		return contracts.ChunkPlanResult{}, chunkerErrorf("source units must not be empty")
	}
	if err := source.ValidateDocument(req.Source); err != nil {
		return contracts.ChunkPlanResult{}, chunkerErrorf("validate source document: %w", err)
	}

	step := c.options.MaxUnits - c.options.OverlapUnits
	ranges := make([]source.ChunkRange, 0, (len(req.Source.Units)+step-1)/step)
	for start := 0; start < len(req.Source.Units); start += step {
		end := start + c.options.MaxUnits
		if end > len(req.Source.Units) {
			end = len(req.Source.Units)
		}
		ranges = append(ranges, source.ChunkRange{
			StartUnitID: req.Source.Units[start].ID,
			EndUnitID:   req.Source.Units[end-1].ID,
		})
		if end == len(req.Source.Units) {
			break
		}
	}

	return contracts.ChunkPlanResult{Plan: source.ChunkPlan{SourceDigest: req.Source.Digest, Ranges: ranges}}, nil
}

func ModuleSpec() pipeline.ModuleSpec {
	return pipeline.ModuleSpec{
		Key:            Key,
		Stage:          pipeline.StageChunk,
		ExecutionClass: contracts.ExecutionClassDeterministic,
		Provides:       []string{"chunks"},
	}
}

func Register(registry *pipeline.ChunkerRegistry) error {
	return registry.RegisterBuilderWithSpec(ModuleSpec(), validateOptions, func(request pipeline.BuildRequest) (contracts.Chunker, error) {
		options, err := DecodeOptions(request.Options)
		if err != nil {
			return nil, err
		}
		return New(options), nil
	})
}

func validateOptions(options map[string]any) error {
	_, err := DecodeOptions(options)
	return err
}

func DecodeOptions(options map[string]any) (Options, error) {
	if err := pipeline.RejectUnknownOptions(options, "max_units", "overlap_units"); err != nil {
		return Options{}, chunkerErrorf("%w", err)
	}
	opts := Options{
		MaxUnits:     defaultMaxUnits,
		OverlapUnits: defaultOverlapUnits,
	}
	var err error
	if value, ok := options["max_units"]; ok {
		opts.MaxUnits, err = positiveIntOption("max_units", value)
		if err != nil {
			return Options{}, err
		}
	}
	if value, ok := options["overlap_units"]; ok {
		opts.OverlapUnits, err = nonNegativeIntOption("overlap_units", value)
		if err != nil {
			return Options{}, err
		}
	}
	if opts.OverlapUnits >= opts.MaxUnits {
		return Options{}, chunkerErrorf("overlap_units must be less than max_units")
	}
	return opts, nil
}

func positiveIntOption(name string, value any) (int, error) {
	got, err := intOption(name, value)
	if err != nil {
		return 0, err
	}
	if got <= 0 {
		return 0, chunkerErrorf("%s must be positive", name)
	}
	return got, nil
}

func nonNegativeIntOption(name string, value any) (int, error) {
	got, err := intOption(name, value)
	if err != nil {
		return 0, err
	}
	if got < 0 {
		return 0, chunkerErrorf("%s must be non-negative", name)
	}
	return got, nil
}

func intOption(name string, value any) (int, error) {
	switch typed := value.(type) {
	case int:
		return typed, nil
	case int8:
		return int(typed), nil
	case int16:
		return int(typed), nil
	case int32:
		return int(typed), nil
	case int64:
		if typed > maxInt() || typed < minInt() {
			return 0, chunkerErrorf("%s is outside supported integer range", name)
		}
		return int(typed), nil
	case uint:
		if uint64(typed) > uint64(maxInt()) {
			return 0, chunkerErrorf("%s is outside supported integer range", name)
		}
		return int(typed), nil
	case uint8:
		return int(typed), nil
	case uint16:
		return int(typed), nil
	case uint32:
		if uint64(typed) > uint64(maxInt()) {
			return 0, chunkerErrorf("%s is outside supported integer range", name)
		}
		return int(typed), nil
	case uint64:
		if typed > uint64(maxInt()) {
			return 0, chunkerErrorf("%s is outside supported integer range", name)
		}
		return int(typed), nil
	case float64:
		if typed != math.Trunc(typed) {
			return 0, chunkerErrorf("%s must be an integer", name)
		}
		if typed > float64(maxInt()) || typed < float64(minInt()) {
			return 0, chunkerErrorf("%s is outside supported integer range", name)
		}
		return int(typed), nil
	case json.Number:
		parsed, err := typed.Int64()
		if err != nil {
			return 0, chunkerErrorf("%s must be an integer", name)
		}
		if parsed > maxInt() || parsed < minInt() {
			return 0, chunkerErrorf("%s is outside supported integer range", name)
		}
		return int(parsed), nil
	default:
		return 0, chunkerErrorf("%s must be an integer", name)
	}
}

func maxInt() int64 {
	return int64(1<<(strconv.IntSize-1) - 1)
}

func minInt() int64 {
	return -maxInt() - 1
}

func chunkerErrorf(format string, args ...any) error {
	return fmt.Errorf("generic chunker: "+format, args...)
}