sow-tools/internal/gff/binary.go

package gff

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"io"
	"math"
)

const (
	headerSize = 56
)

type header struct {
	FileType           [4]byte
	FileVersion        [4]byte
	StructOffset       uint32
	StructCount        uint32
	FieldOffset        uint32
	FieldCount         uint32
	LabelOffset        uint32
	LabelCount         uint32
	FieldDataOffset    uint32
	FieldDataCount     uint32
	FieldIndicesOffset uint32
	FieldIndicesCount  uint32
	ListIndicesOffset  uint32
	ListIndicesCount   uint32
}

type rawStruct struct {
	Type         uint32
	DataOrOffset uint32
	FieldCount   uint32
}

type rawField struct {
	Type         uint32
	LabelIndex   uint32
	DataOrOffset uint32
}

func Read(r io.Reader) (Document, error) {
	data, err := io.ReadAll(r)
	if err != nil {
		return Document{}, fmt.Errorf("read gff: %w", err)
	}
	if len(data) < headerSize {
		return Document{}, fmt.Errorf("gff file too small: %d bytes", len(data))
	}

	var hdr header
	if err := binary.Read(bytes.NewReader(data[:headerSize]), binary.LittleEndian, &hdr); err != nil {
		return Document{}, fmt.Errorf("decode header: %w", err)
	}

	reader := binaryReader{data: data, hdr: hdr}
	rawStructs, err := reader.readStructs()
	if err != nil {
		return Document{}, err
	}
	rawFields, err := reader.readFields()
	if err != nil {
		return Document{}, err
	}
	labels, err := reader.readLabels()
	if err != nil {
		return Document{}, err
	}

	root, err := reader.decodeStruct(0, rawStructs, rawFields, labels)
	if err != nil {
		return Document{}, err
	}

	return Document{
		FileType:    string(hdr.FileType[:]),
		FileVersion: string(hdr.FileVersion[:]),
		Root:        root,
	}, nil
}

func Write(w io.Writer, doc Document) error {
	encoder := newBinaryEncoder(doc)
	data, err := encoder.encode()
	if err != nil {
		return err
	}
	_, err = w.Write(data)
	return err
}

type binaryReader struct {
	data []byte
	hdr  header
}

func (r binaryReader) readStructs() ([]rawStruct, error) {
	return readTable[rawStruct](r.data, r.hdr.StructOffset, r.hdr.StructCount)
}

func (r binaryReader) readFields() ([]rawField, error) {
	return readTable[rawField](r.data, r.hdr.FieldOffset, r.hdr.FieldCount)
}

func (r binaryReader) readLabels() ([]string, error) {
	start := int(r.hdr.LabelOffset)
	end := start + int(r.hdr.LabelCount)*16
	if end > len(r.data) {
		return nil, fmt.Errorf("label table exceeds file bounds")
	}

	labels := make([]string, 0, r.hdr.LabelCount)
	for offset := start; offset < end; offset += 16 {
		chunk := r.data[offset : offset+16]
		n := bytes.IndexByte(chunk, 0)
		if n == -1 {
			n = len(chunk)
		}
		labels = append(labels, string(chunk[:n]))
	}
	return labels, nil
}

func (r binaryReader) decodeStruct(index uint32, structs []rawStruct, fields []rawField, labels []string) (Struct, error) {
	if index >= uint32(len(structs)) {
		return Struct{}, fmt.Errorf("struct index %d out of range", index)
	}

	raw := structs[index]
	fieldIndices, err := r.fieldIndices(raw)
	if err != nil {
		return Struct{}, err
	}

	out := Struct{
		Type:   raw.Type,
		Fields: make([]Field, 0, len(fieldIndices)),
	}
	for _, fieldIndex := range fieldIndices {
		if fieldIndex >= uint32(len(fields)) {
			return Struct{}, fmt.Errorf("field index %d out of range", fieldIndex)
		}
		field, err := r.decodeField(fields[fieldIndex], structs, fields, labels)
		if err != nil {
			return Struct{}, err
		}
		out.Fields = append(out.Fields, field)
	}

	return out, nil
}

func (r binaryReader) fieldIndices(s rawStruct) ([]uint32, error) {
	switch s.FieldCount {
	case 0:
		return nil, nil
	case 1:
		return []uint32{s.DataOrOffset}, nil
	default:
		start := int(r.hdr.FieldIndicesOffset + s.DataOrOffset)
		end := start + int(s.FieldCount)*4
		if end > len(r.data) {
			return nil, fmt.Errorf("field indices exceed file bounds")
		}

		out := make([]uint32, s.FieldCount)
		reader := bytes.NewReader(r.data[start:end])
		if err := binary.Read(reader, binary.LittleEndian, &out); err != nil {
			return nil, fmt.Errorf("decode field indices: %w", err)
		}
		return out, nil
	}
}

func (r binaryReader) decodeField(raw rawField, structs []rawStruct, fields []rawField, labels []string) (Field, error) {
	if raw.LabelIndex >= uint32(len(labels)) {
		return Field{}, fmt.Errorf("label index %d out of range", raw.LabelIndex)
	}
	fieldType := FieldType(raw.Type)
	value, err := r.decodeValue(fieldType, raw.DataOrOffset, structs, fields, labels)
	if err != nil {
		return Field{}, fmt.Errorf("decode field %q: %w", labels[raw.LabelIndex], err)
	}

	return Field{
		Label: labels[raw.LabelIndex],
		Type:  fieldType,
		Value: value,
	}, nil
}

func (r binaryReader) decodeValue(fieldType FieldType, data uint32, structs []rawStruct, fields []rawField, labels []string) (Value, error) {
	switch fieldType {
	case TypeByte:
		return ByteValue(uint8(data)), nil
	case TypeChar:
		return CharValue(int8(data)), nil
	case TypeWord:
		return WordValue(uint16(data)), nil
	case TypeShort:
		return ShortValue(int16(data)), nil
	case TypeDWord:
		return DWordValue(data), nil
	case TypeInt:
		return IntValue(int32(data)), nil
	case TypeFloat:
		return FloatValue(math.Float32frombits(data)), nil
	case TypeDWord64:
		raw, err := r.sliceFieldData(data, 8)
		if err != nil {
			return nil, err
		}
		return DWord64Value(binary.LittleEndian.Uint64(raw)), nil
	case TypeInt64:
		raw, err := r.sliceFieldData(data, 8)
		if err != nil {
			return nil, err
		}
		return Int64Value(int64(binary.LittleEndian.Uint64(raw))), nil
	case TypeDouble:
		raw, err := r.sliceFieldData(data, 8)
		if err != nil {
			return nil, err
		}
		return DoubleValue(math.Float64frombits(binary.LittleEndian.Uint64(raw))), nil
	case TypeCExoString:
		raw, err := r.prefixedFieldData(data)
		if err != nil {
			return nil, err
		}
		return StringValue(string(raw)), nil
	case TypeResRef:
		raw, err := r.prefixedByteFieldData(data, 1)
		if err != nil {
			return nil, err
		}
		return ResRefValue(string(raw)), nil
	case TypeCExoLocString:
		raw, err := r.locStringFieldData(data)
		if err != nil {
			return nil, err
		}
		return decodeLocString(raw)
	case TypeVoid:
		raw, err := r.prefixedFieldData(data)
		if err != nil {
			return nil, err
		}
		return VoidValue(raw), nil
	case TypeStruct:
		return r.decodeStruct(data, structs, fields, labels)
	case TypeList:
		return r.decodeList(data, structs, fields, labels)
	case TypeOrientation:
		raw, err := r.sliceFieldData(data, 16)
		if err != nil {
			return nil, err
		}
		return Orientation{
			X: math.Float32frombits(binary.LittleEndian.Uint32(raw[0:4])),
			Y: math.Float32frombits(binary.LittleEndian.Uint32(raw[4:8])),
			Z: math.Float32frombits(binary.LittleEndian.Uint32(raw[8:12])),
			W: math.Float32frombits(binary.LittleEndian.Uint32(raw[12:16])),
		}, nil
	case TypeVector:
		raw, err := r.sliceFieldData(data, 12)
		if err != nil {
			return nil, err
		}
		return Vector{
			X: math.Float32frombits(binary.LittleEndian.Uint32(raw[0:4])),
			Y: math.Float32frombits(binary.LittleEndian.Uint32(raw[4:8])),
			Z: math.Float32frombits(binary.LittleEndian.Uint32(raw[8:12])),
		}, nil
	default:
		return nil, fmt.Errorf("unsupported field type %d", fieldType)
	}
}

func (r binaryReader) decodeList(offset uint32, structs []rawStruct, fields []rawField, labels []string) (ListValue, error) {
	start := int(r.hdr.ListIndicesOffset + offset)
	if start+4 > len(r.data) {
		return nil, fmt.Errorf("list data exceeds file bounds")
	}
	count := binary.LittleEndian.Uint32(r.data[start : start+4])
	start += 4
	end := start + int(count)*4
	if end > len(r.data) {
		return nil, fmt.Errorf("list indices exceed file bounds")
	}

	list := make(ListValue, 0, count)
	for pos := start; pos < end; pos += 4 {
		index := binary.LittleEndian.Uint32(r.data[pos : pos+4])
		child, err := r.decodeStruct(index, structs, fields, labels)
		if err != nil {
			return nil, err
		}
		list = append(list, child)
	}
	return list, nil
}

func (r binaryReader) sliceFieldData(offset uint32, size int) ([]byte, error) {
	start := int(r.hdr.FieldDataOffset + offset)
	end := start + size
	if end > len(r.data) {
		return nil, fmt.Errorf("field data exceeds file bounds")
	}
	return r.data[start:end], nil
}

func (r binaryReader) prefixedFieldData(offset uint32) ([]byte, error) {
	sizeRaw, err := r.sliceFieldData(offset, 4)
	if err != nil {
		return nil, err
	}
	size := int(binary.LittleEndian.Uint32(sizeRaw))
	return r.sliceFieldData(offset+4, size)
}

func (r binaryReader) prefixedByteFieldData(offset uint32, prefixBytes uint32) ([]byte, error) {
	header, err := r.sliceFieldData(offset, int(prefixBytes))
	if err != nil {
		return nil, err
	}
	size := int(header[0])
	return r.sliceFieldData(offset+prefixBytes, size)
}

func (r binaryReader) locStringFieldData(offset uint32) ([]byte, error) {
	sizeRaw, err := r.sliceFieldData(offset, 4)
	if err != nil {
		return nil, err
	}
	size := int(binary.LittleEndian.Uint32(sizeRaw))
	return r.sliceFieldData(offset, size+4)
}

func readTable[T any](data []byte, offset uint32, count uint32) ([]T, error) {
	if count == 0 {
		return nil, nil
	}
	var row T
	size := binary.Size(row)
	start := int(offset)
	end := start + int(count)*size
	if end > len(data) {
		return nil, fmt.Errorf("table exceeds file bounds")
	}

	out := make([]T, count)
	reader := bytes.NewReader(data[start:end])
	if err := binary.Read(reader, binary.LittleEndian, &out); err != nil {
		return nil, fmt.Errorf("decode table: %w", err)
	}
	return out, nil
}

func decodeLocString(data []byte) (LocString, error) {
	if len(data) < 12 {
		return LocString{}, fmt.Errorf("locstring too small")
	}
	totalSize := binary.LittleEndian.Uint32(data[0:4])
	if totalSize+4 != uint32(len(data)) {
		return LocString{}, fmt.Errorf("locstring size mismatch")
	}
	stringRef := binary.LittleEndian.Uint32(data[4:8])
	count := binary.LittleEndian.Uint32(data[8:12])
	pos := 12
	entries := make([]LocStringEntry, 0, count)
	for i := uint32(0); i < count; i++ {
		if pos+8 > len(data) {
			return LocString{}, fmt.Errorf("locstring entry header truncated")
		}
		id := binary.LittleEndian.Uint32(data[pos : pos+4])
		length := binary.LittleEndian.Uint32(data[pos+4 : pos+8])
		pos += 8
		if pos+int(length) > len(data) {
			return LocString{}, fmt.Errorf("locstring entry %d truncated", i)
		}
		entries = append(entries, LocStringEntry{
			ID:    id,
			Value: string(data[pos : pos+int(length)]),
		})
		pos += int(length)
	}

	return LocString{
		StringRef: stringRef,
		Entries:   entries,
	}, nil
}

type binaryEncoder struct {
	document     Document
	structs      []rawStruct
	fields       []rawField
	labels       []string
	labelIndex   map[string]uint32
	fieldData    bytes.Buffer
	fieldIndices []uint32
	listIndices  []uint32
}

func newBinaryEncoder(doc Document) *binaryEncoder {
	return &binaryEncoder{
		document:   doc,
		labelIndex: map[string]uint32{},
	}
}

func (e *binaryEncoder) encode() ([]byte, error) {
	if err := e.addStruct(e.document.Root); err != nil {
		return nil, err
	}

	var buf bytes.Buffer
	hdr := header{}
	copy(hdr.FileType[:], []byte(padFour(e.document.FileType)))
	copy(hdr.FileVersion[:], []byte(padFour(defaultString(e.document.FileVersion, "V3.2"))))

	hdr.StructOffset = headerSize
	hdr.StructCount = uint32(len(e.structs))
	hdr.FieldOffset = hdr.StructOffset + uint32(len(e.structs))*12
	hdr.FieldCount = uint32(len(e.fields))
	hdr.LabelOffset = hdr.FieldOffset + uint32(len(e.fields))*12
	hdr.LabelCount = uint32(len(e.labels))
	hdr.FieldDataOffset = hdr.LabelOffset + uint32(len(e.labels))*16
	hdr.FieldDataCount = uint32(e.fieldData.Len())
	hdr.FieldIndicesOffset = hdr.FieldDataOffset + uint32(e.fieldData.Len())
	hdr.FieldIndicesCount = uint32(len(e.fieldIndices) * 4)
	hdr.ListIndicesOffset = hdr.FieldIndicesOffset + uint32(len(e.fieldIndices))*4
	hdr.ListIndicesCount = uint32(len(e.listIndices) * 4)

	if err := binary.Write(&buf, binary.LittleEndian, hdr); err != nil {
		return nil, err
	}
	if err := binary.Write(&buf, binary.LittleEndian, e.structs); err != nil {
		return nil, err
	}
	if err := binary.Write(&buf, binary.LittleEndian, e.fields); err != nil {
		return nil, err
	}
	for _, label := range e.labels {
		var raw [16]byte
		copy(raw[:], []byte(label))
		if _, err := buf.Write(raw[:]); err != nil {
			return nil, err
		}
	}
	if _, err := buf.Write(e.fieldData.Bytes()); err != nil {
		return nil, err
	}
	if err := binary.Write(&buf, binary.LittleEndian, e.fieldIndices); err != nil {
		return nil, err
	}
	if err := binary.Write(&buf, binary.LittleEndian, e.listIndices); err != nil {
		return nil, err
	}

	return buf.Bytes(), nil
}

func (e *binaryEncoder) addStruct(s Struct) error {
	index := uint32(len(e.structs))
	e.structs = append(e.structs, rawStruct{Type: s.Type})

	fieldIndices := make([]uint32, 0, len(s.Fields))
	for _, field := range s.Fields {
		fieldIndex, err := e.addField(field)
		if err != nil {
			return err
		}
		fieldIndices = append(fieldIndices, fieldIndex)
	}

	entry := &e.structs[index]
	entry.FieldCount = uint32(len(fieldIndices))
	switch len(fieldIndices) {
	case 0:
		entry.DataOrOffset = 0
	case 1:
		entry.DataOrOffset = fieldIndices[0]
	default:
		entry.DataOrOffset = uint32(len(e.fieldIndices) * 4)
		e.fieldIndices = append(e.fieldIndices, fieldIndices...)
	}
	return nil
}

func (e *binaryEncoder) addField(field Field) (uint32, error) {
	labelIndex, err := e.indexLabel(field.Label)
	if err != nil {
		return 0, err
	}
	data, err := e.encodeValue(field.Value)
	if err != nil {
		return 0, fmt.Errorf("encode field %q: %w", field.Label, err)
	}

	entry := rawField{
		Type:         uint32(field.Type),
		LabelIndex:   labelIndex,
		DataOrOffset: data,
	}
	index := uint32(len(e.fields))
	e.fields = append(e.fields, entry)
	return index, nil
}

func (e *binaryEncoder) encodeValue(value Value) (uint32, error) {
	switch typed := value.(type) {
	case ByteValue:
		return uint32(uint8(typed)), nil
	case CharValue:
		return uint32(uint8(typed)), nil
	case WordValue:
		return uint32(uint16(typed)), nil
	case ShortValue:
		return uint32(uint16(typed)), nil
	case DWordValue:
		return uint32(typed), nil
	case IntValue:
		return uint32(int32(typed)), nil
	case FloatValue:
		return math.Float32bits(float32(typed)), nil
	case DWord64Value:
		return e.writeFieldData(func(buf *bytes.Buffer) error {
			return binary.Write(buf, binary.LittleEndian, uint64(typed))
		}), nil
	case Int64Value:
		return e.writeFieldData(func(buf *bytes.Buffer) error {
			return binary.Write(buf, binary.LittleEndian, int64(typed))
		}), nil
	case DoubleValue:
		return e.writeFieldData(func(buf *bytes.Buffer) error {
			return binary.Write(buf, binary.LittleEndian, math.Float64bits(float64(typed)))
		}), nil
	case StringValue:
		return e.writeLengthPrefixed([]byte(typed)), nil
	case ResRefValue:
		if len(typed) > 255 {
			return 0, fmt.Errorf("resref exceeds 255 bytes")
		}
		return e.writeFieldData(func(buf *bytes.Buffer) error {
			if err := buf.WriteByte(byte(len(typed))); err != nil {
				return err
			}
			_, err := buf.Write([]byte(typed))
			return err
		}), nil
	case LocString:
		return e.writeFieldData(func(buf *bytes.Buffer) error {
			payload := bytes.Buffer{}
			if err := binary.Write(&payload, binary.LittleEndian, typed.StringRef); err != nil {
				return err
			}
			if err := binary.Write(&payload, binary.LittleEndian, uint32(len(typed.Entries))); err != nil {
				return err
			}
			for _, entry := range typed.Entries {
				if err := binary.Write(&payload, binary.LittleEndian, entry.ID); err != nil {
					return err
				}
				if err := binary.Write(&payload, binary.LittleEndian, uint32(len(entry.Value))); err != nil {
					return err
				}
				if _, err := payload.Write([]byte(entry.Value)); err != nil {
					return err
				}
			}
			if err := binary.Write(buf, binary.LittleEndian, uint32(payload.Len())); err != nil {
				return err
			}
			if _, err := buf.Write(payload.Bytes()); err != nil {
				return err
			}
			return nil
		}), nil
	case VoidValue:
		return e.writeLengthPrefixed([]byte(typed)), nil
	case Struct:
		index := uint32(len(e.structs))
		if err := e.addStruct(typed); err != nil {
			return 0, err
		}
		return index, nil
	case ListValue:
		return e.writeList(typed)
	case Orientation:
		return e.writeFieldData(func(buf *bytes.Buffer) error {
			for _, component := range []float32{typed.X, typed.Y, typed.Z, typed.W} {
				if err := binary.Write(buf, binary.LittleEndian, math.Float32bits(component)); err != nil {
					return err
				}
			}
			return nil
		}), nil
	case Vector:
		return e.writeFieldData(func(buf *bytes.Buffer) error {
			for _, component := range []float32{typed.X, typed.Y, typed.Z} {
				if err := binary.Write(buf, binary.LittleEndian, math.Float32bits(component)); err != nil {
					return err
				}
			}
			return nil
		}), nil
	default:
		return 0, fmt.Errorf("unsupported value type %T", value)
	}
}

func (e *binaryEncoder) writeList(list ListValue) (uint32, error) {
	indices := make([]uint32, 0, len(list))
	for _, item := range list {
		index := uint32(len(e.structs))
		if err := e.addStruct(item); err != nil {
			return 0, err
		}
		indices = append(indices, index)
	}

	offset := uint32(len(e.listIndices) * 4)
	e.listIndices = append(e.listIndices, uint32(len(list)))
	e.listIndices = append(e.listIndices, indices...)
	return offset, nil
}

func (e *binaryEncoder) writeFieldData(write func(*bytes.Buffer) error) uint32 {
	offset := uint32(e.fieldData.Len())
	if err := write(&e.fieldData); err != nil {
		panic(err)
	}
	return offset
}

func (e *binaryEncoder) writeLengthPrefixed(data []byte) uint32 {
	return e.writeFieldData(func(buf *bytes.Buffer) error {
		if err := binary.Write(buf, binary.LittleEndian, uint32(len(data))); err != nil {
			return err
		}
		_, err := buf.Write(data)
		return err
	})
}

func (e *binaryEncoder) indexLabel(label string) (uint32, error) {
	if len(label) > 16 {
		return 0, fmt.Errorf("label %q exceeds 16 bytes", label)
	}
	if index, ok := e.labelIndex[label]; ok {
		return index, nil
	}
	index := uint32(len(e.labels))
	e.labels = append(e.labels, label)
	e.labelIndex[label] = index
	return index, nil
}

func padFour(value string) string {
	if len(value) >= 4 {
		return value[:4]
	}
	return value + string(bytes.Repeat([]byte(" "), 4-len(value)))
}

func defaultString(value, fallback string) string {
	if value == "" {
		return fallback
	}
	return value
}