go-lextwt-encoding/twt-fbs/helper.go

package twt_fbs

import (
	"io"
	"time"

	flatbuffers "github.com/google/flatbuffers/go"
	"go.yarn.social/lextwt"
	"go.yarn.social/types"

	fbs "go.sour.is/lextwt-encoding/twt-fbs/lextwt"
)

type TwtRegistry interface {
	Twters() []*types.Twter
	Preamble() lextwt.Comments
	Twts() types.Twts
	WriteTo(w io.Writer) (int64, error)
}

func EncodeTwt(twt types.Twt) []byte {
	lx, ok := twt.(*lextwt.Twt)
	if !ok {
		return nil
	}

	return EncodeLextwt(lx)
}

func EncodeLextwt(twt *lextwt.Twt) []byte {
	b := flatbuffers.NewBuilder(0)
	b.Finish(createTwt(b, twt))
	return b.FinishedBytes()
}

func EncodeRegistry(lis TwtRegistry) []byte {
	b := flatbuffers.NewBuilder(len(lis.Twts()))
	b.Finish(createRegistry(b, lis))
	return b.FinishedBytes()
}

func createRegistry(b *flatbuffers.Builder, r TwtRegistry) flatbuffers.UOffsetT {
	preamble := make([]flatbuffers.UOffsetT, len(r.Preamble()))
	twters := make([]flatbuffers.UOffsetT, len(r.Twters()))
	twts := make([]flatbuffers.UOffsetT, len(r.Twts()))

	for i, twt := range r.Twts() {
		if twt, ok := twt.(*lextwt.Twt); ok {
			twts[i] = createTwt(b, twt)
		}
	}

	fbs.TwtRegistryStartTwtsVector(b, len(twts))
	for _, offset := range twts {
		b.PrependUOffsetT(offset)
	}
	oTwts := b.EndVector(len(twts))

	for _, twter := range r.Twters() {
		twters = append(twters, createTwter(b, twter))
	}

	fbs.TwtRegistryStartTwtersVector(b, len(twters))
	for _, offset := range twters {
		b.PrependUOffsetT(offset)
	}
	oTwters := b.EndVector(len(r.Twters()))

	for _, comment := range r.Preamble() {
		preamble = append(preamble, createComment(b, comment))
	}

	fbs.TwtRegistryStartPreambleVector(b, len(preamble))
	for _, offset := range preamble {
		b.PrependUOffsetT(offset)
	}
	oPreamble := b.EndVector(len(preamble))

	fbs.TwtRegistryStart(b)
	fbs.TwtRegistryAddTwts(b, oTwts)
	fbs.TwtRegistryAddTwters(b, oTwters)
	fbs.TwtRegistryAddPreamble(b, oPreamble)
	return fbs.TwtRegistryEnd(b)
}

func createTwt(b *flatbuffers.Builder, twt *lextwt.Twt) flatbuffers.UOffsetT {
	twter := twt.Twter()

	oTwter := createTwter(b, &twter)
	oCreated := createDateTime(b, twt.Created())
	oElems := createMsg(b, twt.Elems())

	fbs.TwtStart(b)
	fbs.TwtAddTwter(b, oTwter)
	fbs.TwtAddCreated(b, oCreated)
	fbs.TwtAddMsg(b, oElems)

	return fbs.TwtEnd(b)
}

func createTwter(b *flatbuffers.Builder, twter *types.Twter) flatbuffers.UOffsetT {
	nick := b.CreateString(twter.Nick)
	uri := b.CreateString(twter.URI)

	fbs.TwterStart(b)
	fbs.TwterAddNick(b, nick)
	fbs.TwterAddUri(b, uri)

	return fbs.TwterEnd(b)
}

func createDateTime(b *flatbuffers.Builder, ts time.Time) flatbuffers.UOffsetT {
	_, offset := ts.Zone()

	fbs.DateTimeStart(b)
	fbs.DateTimeAddTimestamp(b, ts.UnixMilli())
	fbs.DateTimeAddOffset(b, int16(offset))

	return fbs.DateTimeEnd(b)
}

func createMsg(b *flatbuffers.Builder, msgs []lextwt.Elem) flatbuffers.UOffsetT {
	offsets := make([]flatbuffers.UOffsetT, len(msgs))
	for i, msg := range msgs {
		offsets[i] = createElem(b, msg)
	}

	fbs.TwtStartMsgVector(b, len(msgs))
	for _, offset := range offsets {
		b.PrependUOffsetT(offset)
	}

	return b.EndVector(len(msgs))
}

func createElem(b *flatbuffers.Builder, elem lextwt.Elem) flatbuffers.UOffsetT {

	if elem == lextwt.LineSeparator {
		offset := createLineSeparator(b)

		fbs.MsgStart(b)
		fbs.MsgAddElemType(b, fbs.ElemLineSeparator)
		fbs.MsgAddElem(b, offset)

		return fbs.MsgEnd(b)
	}

	var oElem flatbuffers.UOffsetT
	var oType fbs.Elem

	switch e := elem.(type) {
	case *lextwt.BangMention:
		oType = fbs.ElemBangMention
		oElem = createBangMention(b, e)
	case *lextwt.Code:
		oType = fbs.ElemCode
		oElem = createCode(b, e)
	case *lextwt.Comment:
		oType = fbs.ElemComment
		oElem = createComment(b, e)
	case *lextwt.Link:
		oType = fbs.ElemLink
		oElem = createLink(b, e)
	case *lextwt.Mention:
		oType = fbs.ElemMention
		oElem = createMention(b, e)
	case *lextwt.Subject:
		oType = fbs.ElemSubject
		oElem = createSubject(b, e)
	case *lextwt.Tag:
		oType = fbs.ElemTag
		oElem = createTag(b, e)
	case *lextwt.Text:
		oType = fbs.ElemText
		oElem = createText(b, e)
	default:
		fbs.MsgStart(b)
		fbs.MsgAddElemType(b, fbs.ElemNONE)

		return fbs.MsgEnd(b)
	}

	fbs.MsgStart(b)
	fbs.MsgAddElemType(b, oType)
	fbs.MsgAddElem(b, oElem)

	return fbs.MsgEnd(b)
}

func createLineSeparator(b *flatbuffers.Builder) flatbuffers.UOffsetT {
	fbs.LineSeparatorStart(b)

	return fbs.LineSeparatorEnd(b)
}

func createBangMention(b *flatbuffers.Builder, mention *lextwt.BangMention) flatbuffers.UOffsetT {
	name := b.CreateString(mention.Name())
	target := b.CreateString(mention.Target())

	fbs.BangMentionStart(b)
	fbs.BangMentionAddName(b, name)
	fbs.BangMentionAddTarget(b, target)

	return fbs.BangMentionEnd(b)
}

func createCode(b *flatbuffers.Builder, code *lextwt.Code) flatbuffers.UOffsetT {
	typ := fbs.CodeType(code.CodeType())
	value := b.CreateString(code.Text())

	fbs.CodeStart(b)
	fbs.CodeAddType(b, typ)
	fbs.CodeAddCode(b, value)

	return fbs.CodeEnd(b)
}

func createComment(b *flatbuffers.Builder, comment *lextwt.Comment) flatbuffers.UOffsetT {
	text := b.CreateString(comment.Text())
	var key, value flatbuffers.UOffsetT

	if comment.Key() != "" {
		key = b.CreateString(comment.Key())
		value = b.CreateString(comment.Value())
	}

	fbs.CommentStart(b)
	fbs.CommentAddComment(b, text)
	if comment.Key() != "" {
		fbs.CommentAddKey(b, key)
		fbs.CommentAddValue(b, value)
	}

	return fbs.CommentEnd(b)
}

func createLink(b *flatbuffers.Builder, link *lextwt.Link) flatbuffers.UOffsetT {
	linkType := fbs.LinkType(link.LinkType())
	alt := b.CreateString(link.Text())
	target := b.CreateString(link.Target())
	title := b.CreateString(link.Title())

	fbs.LinkStart(b)
	fbs.LinkAddType(b, linkType)
	fbs.LinkAddAlt(b, alt)
	fbs.LinkAddTarget(b, target)
	fbs.LinkAddTitle(b, title)

	return fbs.LinkEnd(b)
}

func createMention(b *flatbuffers.Builder, mention *lextwt.Mention) flatbuffers.UOffsetT {
	name := b.CreateString(mention.Name())
	target := b.CreateString(mention.Target())

	fbs.MentionStart(b)
	fbs.MentionAddName(b, name)
	fbs.MentionAddTarget(b, target)

	return fbs.MentionEnd(b)
}

func createSubject(b *flatbuffers.Builder, subject *lextwt.Subject) flatbuffers.UOffsetT {
	var oTag flatbuffers.UOffsetT
	if tag, ok := subject.Tag().(*lextwt.Tag); ok && tag != nil {
		oTag = createTag(b, tag)
	}
	oSubject := b.CreateString(subject.Subject())

	fbs.SubjectStart(b)
	fbs.SubjectAddSubject(b, oSubject)
	if oTag != 0 {
		fbs.SubjectAddTag(b, oTag)
	}

	return fbs.SubjectEnd(b)
}

func createTag(b *flatbuffers.Builder, tag *lextwt.Tag) flatbuffers.UOffsetT {
	text := b.CreateString(tag.Text())
	target := b.CreateString(tag.Target())

	fbs.TagStart(b)
	fbs.TagAddTag(b, text)
	fbs.TagAddTarget(b, target)

	return fbs.TagEnd(b)
}

func createText(b *flatbuffers.Builder, text *lextwt.Text) flatbuffers.UOffsetT {
	literal := b.CreateString(text.Literal())

	fbs.TextStart(b)
	fbs.TextAddText(b, literal)

	return fbs.TextEnd(b)
}

func DecodeTwt(in []byte) (types.Twt, error) {
	return DecodeLextwt(in)
}

func DecodeLextwt(in []byte) (*lextwt.Twt, error) {
	o := fbs.GetRootAsTwt(in, 0)
	return fromTwt(o), nil
}

func fromTwt(o *fbs.Twt) *lextwt.Twt {
	elems := make([]lextwt.Elem, o.MsgLength())

	for i := 0; i < o.MsgLength(); i++ {
		var m fbs.Msg

		if o.Msg(&m, i) {
			elems[len(elems)-i-1] = fromMsg(&m)
		}
	}

	return lextwt.NewTwt(
		*fromTwter(o.Twter(nil)),
		fromDateTime(o.Created(nil)),
		elems...,
	)
}

func DecodeRegistry(in []byte) (TwtRegistry, error) {
	root := fbs.GetRootAsTwtRegistry(in, 0)
	return fromRegistry(root), nil
}
func fromRegistry(r *fbs.TwtRegistry) TwtRegistry {
	twts := make([]types.Twt, r.TwtsLength())
	for i := 0; i < r.TwtsLength(); i++ {
		twt := &fbs.Twt{}
		r.Twts(twt, i)
		twts[len(twts)-i-1] = fromTwt(twt)
	}

	twters := make([]*types.Twter, r.TwtersLength())
	for i := 0; i < r.TwtersLength(); i++ {
		twter := &fbs.Twter{}
		if r.Twters(twter, i) {
			twters[len(twters)-i-1] = fromTwter(twter)
		}
	}

	preamble := make(lextwt.Comments, r.PreambleLength())
	for i := 0; i < r.PreambleLength(); i++ {
		comment := &fbs.Comment{}
		if r.Preamble(comment, i) {
			preamble[len(preamble)-i-1] = fromComment(comment)
		}
	}

	return lextwt.NewTwtRegistry(twters, preamble, twts)
}

func fromDateTime(o *fbs.DateTime) *lextwt.DateTime {
	if o == nil {
		return nil
	}
	t := time.UnixMilli(o.Timestamp())
	t.Add(time.Duration(o.Offset()) * time.Second)

	return lextwt.NewDateTime(t, t.Format(time.RFC3339))
}

func fromTwter(o *fbs.Twter) *types.Twter {
	t := &types.Twter{}

	if o != nil {
		t.Nick = string(o.Nick())
		t.URI = string(o.Uri())
	}
	return t
}

func fromMsg(o *fbs.Msg) lextwt.Elem {
	elem := &flatbuffers.Table{}

	if o.Elem(elem) {
		switch o.ElemType() {
		case fbs.ElemText:
			return from(elem, fromText)
		case fbs.ElemSubject:
			return from(elem, fromSubject)
		case fbs.ElemCode:
			return from(elem, fromCode)
		case fbs.ElemMention:
			return from(elem, fromMention)
		case fbs.ElemLink:
			return from(elem, fromLink)
		case fbs.ElemComment:
			return from(elem, fromComment)
		case fbs.ElemBangMention:
			return from(elem, fromBangMention)
		case fbs.ElemTag:
			return from(elem, fromTag)
		case fbs.ElemLineSeparator:
			return lextwt.LineSeparator
		}
	}
	return nil
}

type PA[T any] interface {
	Init([]byte, flatbuffers.UOffsetT)
	*T
}

func from[T, A any, U PA[A]](tbl *flatbuffers.Table, fn func(U) T) T {
	var elem U = new(A)
	elem.Init(tbl.Bytes, tbl.Pos)
	return fn(elem)
}

func fromText(o *fbs.Text) *lextwt.Text {
	if o == nil {
		return nil
	}
	return lextwt.NewText(string(o.Text()))
}

func fromLink(o *fbs.Link) *lextwt.Link {
	if o == nil {
		return nil
	}

	if o.Type() == fbs.LinkTypeMedia {
		return lextwt.NewMedia(string(o.Alt()), string(o.Target()), string(o.Title()))
	}

	return lextwt.NewLink(string(o.Alt()), string(o.Target()), lextwt.LinkType(o.Type()))
}

func fromComment(o *fbs.Comment) *lextwt.Comment {
	if o == nil {
		return nil
	}
	if len(o.Key()) > 0 {
		return lextwt.NewCommentValue(string(o.Comment()), string(o.Key()), string(o.Value()))
	}

	return lextwt.NewComment(string(o.Comment()))
}

func fromTag(o *fbs.Tag) *lextwt.Tag {
	if o == nil {
		return nil
	}
	return lextwt.NewTag(string(o.Tag()), string(o.Target()))
}

func fromSubject(o *fbs.Subject) *lextwt.Subject {
	if o == nil {
		return nil
	}
	if tag := o.Tag(nil); tag != nil {
		return lextwt.NewSubjectTag(string(tag.Tag()), string(tag.Target()))
	}

	return lextwt.NewSubject(string(o.Subject()))
}

func fromCode(o *fbs.Code) *lextwt.Code {
	if o == nil {
		return nil
	}
	return lextwt.NewCode(string(o.Code()), lextwt.CodeType(o.Type()))
}

func fromMention(o *fbs.Mention) *lextwt.Mention {
	if o == nil {
		return nil
	}
	return lextwt.NewMention(string(o.Name()), string(o.Target()))
}

func fromBangMention(o *fbs.BangMention) *lextwt.BangMention {
	if o == nil {
		return nil
	}
	return lextwt.NewBangMention(string(o.Name()), string(o.Target()))
}