ucl/ucl/builtins/lists.go

package builtins

import (
	"context"
	"errors"

	"ucl.lmika.dev/ucl"
)

func Lists() ucl.Module {
	return ucl.Module{
		Name: "lists",
		Builtins: map[string]ucl.BuiltinHandler{
			"append":  listAppend,
			"first":   listFirst,
			"batch":   listBatch,
			"uniq":    listUniq,
			"sublist": listSublist,
		},
	}
}

func listAppend(ctx context.Context, args ucl.CallArgs) (any, error) {
	var (
		what ucl.Object
		item ucl.Object
	)

	if err := args.Bind(&what); err != nil {
		return nil, err
	}

	if what == nil {
		what = ucl.NewListObject()
	}

	t, ok := what.(ucl.ModListable)
	if !ok {
		return nil, errors.New("expected mutable list")
	}

	for args.NArgs() > 0 {
		if err := args.Bind(&item); err != nil {
			return nil, err
		}
		if err := t.Insert(-1, item); err != nil {
			return nil, err
		}
	}

	return t, nil
}

func listFirst(ctx context.Context, args ucl.CallArgs) (any, error) {
	var (
		what  ucl.Object
		count int
	)

	if err := args.Bind(&what, &count); err != nil {
		return nil, err
	}

	if count == 0 {
		return ucl.NewListObject(), nil
	}

	newList := ucl.NewListObject()

	switch t := what.(type) {
	case ucl.Listable:
		if count < 0 {
			count = t.Len() + count
		}

		for i := 0; i < min(count, t.Len()); i++ {
			newList.Append(t.Index(i))
		}
	case ucl.Iterable:
		if count < 0 {
			return nil, errors.New("negative counts not supported on iters")
		}

		for i := 0; t.HasNext() && i < count; i++ {
			v, err := t.Next(ctx)
			if err != nil {
				return nil, err
			}
			newList.Append(v)
		}
	default:
		return nil, errors.New("expected listable")
	}

	return newList, nil
}

func eachListOrIterItem(ctx context.Context, o ucl.Object, f func(int, ucl.Object) error) error {
	switch t := o.(type) {
	case ucl.Listable:
		for i := 0; i < t.Len(); i++ {
			if err := f(i, t.Index(i)); err != nil {
				return err
			}
		}
		return nil
	case ucl.Iterable:
		idx := 0
		for t.HasNext() {
			v, err := t.Next(ctx)
			if err != nil {
				return err
			}
			if err := f(idx, v); err != nil {
				return err
			}
			idx++
		}
		return nil
	}
	return errors.New("expected listable")
}

type uniqKey struct {
	sVal string
	iVal int
}

func listUniq(ctx context.Context, args ucl.CallArgs) (any, error) {
	var (
		what ucl.Object
	)

	if err := args.Bind(&what); err != nil {
		return nil, err
	}

	seen := make(map[uniqKey]bool)
	found := ucl.NewListObject()

	if err := eachListOrIterItem(ctx, what, func(idx int, v ucl.Object) error {
		var key uniqKey
		switch v := v.(type) {
		case ucl.StringObject:
			key = uniqKey{sVal: string(v)}
		case ucl.IntObject:
			key = uniqKey{iVal: int(v)}
		default:
			return errors.New("expected string or int")
		}

		if !seen[key] {
			seen[key] = true
			found.Append(v)
		}

		return nil
	}); err != nil {
		return nil, err
	}

	return found, nil
}

func listBatch(ctx context.Context, args ucl.CallArgs) (any, error) {
	var (
		what      ucl.Object
		groupSize int
	)

	if err := args.Bind(&what, &groupSize); err != nil {
		return nil, err
	}

	if groupSize <= 0 {
		return nil, errors.New("group size must be > 0")
	}

	groups := ucl.NewListObject()
	var thisGroup *ucl.ListObject

	if err := eachListOrIterItem(ctx, what, func(idx int, v ucl.Object) error {
		if thisGroup == nil || thisGroup.Len() == groupSize {
			thisGroup = ucl.NewListObject()
			groups.Append(thisGroup)
		}
		thisGroup.Append(v)
		return nil
	}); err != nil {
		return nil, err
	}

	return groups, nil
}

func listSublist(ctx context.Context, args ucl.CallArgs) (any, error) {
	var (
		l             ucl.Listable
		from, fromIdx int
		to, toIdx     int
	)

	if err := args.Bind(&l, &from); err != nil {
		return nil, err
	}

	if l == nil {
		return nil, nil
	}

	if args.NArgs() >= 1 {
		if err := args.Bind(&to); err != nil {
			return nil, err
		}

		fromIdx, toIdx = listPos(l, from), listPos(l, to)
	} else {
		if from < 0 {
			fromIdx, toIdx = listPos(l, from), l.Len()
		} else {
			fromIdx, toIdx = 0, listPos(l, from)
		}
	}

	if fromIdx > toIdx {
		return ucl.NewListObject(), nil
	}

	newList := ucl.NewListObjectOfLength(toIdx - fromIdx)
	for i := fromIdx; i < toIdx; i++ {
		if err := newList.SetIndex(i-fromIdx, l.Index(i)); err != nil {
			return nil, err
		}
	}
	return newList, nil
}

func listPos(l ucl.Listable, pos int) int {
	if pos < 0 {
		return max(l.Len()+pos, 0)
	}
	return min(pos, l.Len())
}