peachy-go/csexp/sexp.go

package csexp

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"slices"
	"strconv"

	"git.codemonkeysoftware.net/b/gigaparsec"
	pbytes "git.codemonkeysoftware.net/b/gigaparsec/bytes"
	"git.codemonkeysoftware.net/b/gigaparsec/cursor"
	"git.codemonkeysoftware.net/b/peachy-go/shortcircuit"
)

type Sexp interface {
	isSexp()
	WriteTo(w io.Writer) (n int64, err error)
	String() string
	Equal(Sexp) bool
	Clone() Sexp
}

type Atom []byte

func (a Atom) isSexp() {}

func (a Atom) WriteTo(w io.Writer) (int64, error) {
	if scw, ok := w.(*shortcircuit.Writer); ok && scw.Failed() {
		return 0, nil
	}

	n, err := fmt.Fprintf(w, "%d:%s", len(a), []byte(a))
	return int64(n), err
}

func (a Atom) String() string {
	var buf bytes.Buffer
	a.WriteTo(&buf)
	return buf.String()
}

func (a Atom) Equal(s Sexp) bool {
	a2, ok := s.(Atom)
	return ok && bytes.Equal([]byte(a), []byte(a2))
}

func (a Atom) Clone() Sexp {
	return Atom(bytes.Clone([]byte(a)))
}

type List []Sexp

func (l List) isSexp() {}

func (l List) WriteTo(w io.Writer) (int64, error) {
	scw := shortcircuit.EnsureWriter(w)
	if scw.Failed() {
		return 0, nil
	}

	io.WriteString(scw, "(")
	for _, child := range l {
		child.WriteTo(scw)
	}
	io.WriteString(scw, ")")
	return scw.Status()
}

func (l List) String() string {
	var buf bytes.Buffer
	l.WriteTo(&buf)
	return buf.String()
}

func (l List) Equal(s Sexp) bool {
	l2, ok := s.(List)
	return ok && slices.EqualFunc(l, l2, Sexp.Equal)
}

func (l List) Clone() Sexp {
	l2 := make(List, len(l))
	for i, child := range l {
		l2[i] = child.Clone()
	}
	return l2
}

var parseLength = gigaparsec.Map(pbytes.Regexp(`0|[1-9]\d*`), func(s string) uint64 {
	n, err := strconv.ParseUint(s, 10, 64)
	if err != nil {
		panic(err)
	}
	return n
})

func acceptN(n uint64) gigaparsec.Parser[byte, []byte] {
	expected := fmt.Sprintf("%d bytes", n)
	return func(s gigaparsec.State[byte]) (gigaparsec.Result[byte, []byte], error) {
		if n == 0 {
			return gigaparsec.Succeed[byte, []byte](true, nil, s, gigaparsec.MessageOK(s.Pos())), nil
		}
		dst := make([]byte, n)
		_, next, err := s.Read(dst)
		if errors.Is(err, io.EOF) {
			return gigaparsec.Fail[byte, []byte](false, gigaparsec.MessageEnd(s.Pos(), expected)), nil
		}
		if err != nil {
			return gigaparsec.Result[byte, []byte]{}, err
		}
		return gigaparsec.Succeed(true, dst, next, gigaparsec.MessageOK(s.Pos())), nil
	}
}

var parseAtom = gigaparsec.Map(gigaparsec.Bind2(
	parseLength,
	gigaparsec.Pipe[byte, byte, uint64](gigaparsec.Match[byte](':')),
	acceptN,
), func(s []byte) Sexp { return Atom(s) })

func parseSexp(input gigaparsec.State[byte]) (gigaparsec.Result[byte, Sexp], error) {
	return gigaparsec.Choose(parseAtom, parseList)(input)
}

func parseRestOfList(input gigaparsec.State[byte]) (gigaparsec.Result[byte, []Sexp], error) {
	return gigaparsec.Choose(
		gigaparsec.Map(gigaparsec.Match[byte](')'), func(byte) []Sexp { return nil }),
		gigaparsec.Bind(
			parseSexp,
			func(s Sexp) gigaparsec.Parser[byte, []Sexp] {
				return gigaparsec.Map(
					parseRestOfList,
					func(rest []Sexp) []Sexp { return append([]Sexp{s}, rest...) },
				)
			},
		),
	)(input)
}

func parseList(input gigaparsec.State[byte]) (gigaparsec.Result[byte, Sexp], error) {
	return gigaparsec.Map(gigaparsec.Seq2(
		gigaparsec.Match[byte]('('),
		parseRestOfList,
		func(_ byte, rest []Sexp) []Sexp { return rest },
	), func(sexps []Sexp) Sexp { return List(sexps) })(input)
}

func Parse(data []byte) (Sexp, error) {
	parser := gigaparsec.Seq2(
		parseSexp,
		gigaparsec.End[byte](),
		func(s Sexp, _ struct{}) Sexp { return s },
	)
	result, err := parser(gigaparsec.MakeState(cursor.NewSlice(data)))
	if err != nil {
		return nil, fmt.Errorf("csexp.Parse: %w", err)
	}
	if failed, _, msg := result.Failed(); failed {
		return nil, fmt.Errorf("csexp.Parse: %v", msg)
	}
	_, _, sexp, _, _ := result.Succeeded()
	return sexp, nil
}