aego/asset/loader/loader.go

package loader

import (
	"bytes"

	"aego/asset"
	"aego/core/errs"
	"aego/core/ids"
	"aego/core/log"
	"aego/gpu"
	"aego/render/mesh"
	"aego/render/texture"
	"aego/vfs"
)

const ServiceID = "asset.loader"

type pageEntry struct {
	handle gpu.TextureHandle
	width  int
	height int
	filter gpu.FilterMode
}

type regionEntry struct {
	region texture.Region
	page   ids.AssetID
}

type decodedKey struct {
	id  ids.AssetID
	key string
}

type decodedValue struct {
	v   any
	err error
}

type Loader struct {
	dev     gpu.Device
	up      *texture.Uploader
	log     *log.ChannelLogger
	index   *asset.Index
	files   vfs.FS
	pages   map[ids.AssetID]pageEntry
	region  map[ids.AssetID]regionEntry
	blobs   map[ids.AssetID][]byte
	meshes  map[ids.AssetID]*mesh.Mesh
	failed  map[ids.AssetID]bool
	decoded map[decodedKey]decodedValue
	subs    []subscriber
	nextID  uint64
	gen     uint32
}

type subscriber struct {
	id uint64
	fn func(asset.Reloaded)
}

func New(dev gpu.Device, up *texture.Uploader, lg *log.Logger) *Loader {
	if lg == nil {
		lg = log.Nop()
	}
	return &Loader{
		dev:     dev,
		up:      up,
		log:     lg.With(log.Asset),
		pages:   map[ids.AssetID]pageEntry{},
		region:  map[ids.AssetID]regionEntry{},
		blobs:   map[ids.AssetID][]byte{},
		meshes:  map[ids.AssetID]*mesh.Mesh{},
		failed:  map[ids.AssetID]bool{},
		decoded: map[decodedKey]decodedValue{},
		gen:     1,
	}
}

func (l *Loader) Mount(idx *asset.Index, files vfs.FS) error {
	if idx == nil || files == nil {
		return errs.New(errs.AssetNotMounted, "loader needs both an index and an artifact file system")
	}
	l.index = idx.Clone()
	l.files = files
	clear(l.decoded)
	l.gen++
	return nil
}

func (l *Loader) Decoded(id ids.AssetID, key string, build func() (any, error)) (any, error) {
	k := decodedKey{id: id, key: key}
	if d, ok := l.decoded[k]; ok {
		return d.v, d.err
	}
	v, err := build()
	l.decoded[k] = decodedValue{v: v, err: err}
	return v, err
}

func (l *Loader) Generation() uint32 {
	return l.gen
}

func (l *Loader) Resident() int {
	return len(l.pages) + len(l.blobs) + len(l.meshes)
}

func (l *Loader) Subscribe(fn func(asset.Reloaded)) func() {
	l.nextID++
	id := l.nextID
	l.subs = append(l.subs, subscriber{id: id, fn: fn})
	return func() {
		for i := range l.subs {
			if l.subs[i].id == id {
				l.subs = append(l.subs[:i], l.subs[i+1:]...)
				return
			}
		}
	}
}

func (l *Loader) notify(r asset.Reloaded) {
	for i := range l.subs {
		l.subs[i].fn(r)
	}
}

func (l *Loader) entry(id ids.AssetID) (*asset.Entry, bool) {
	if l.index == nil {
		return nil, false
	}
	e, ok := l.index.Lookup(id)
	if !ok || e.Missing {
		return nil, false
	}
	return e, true
}

func (l *Loader) read(a asset.Artifact) ([]byte, error) {
	return l.files.ReadFile(a.Path)
}

func (l *Loader) Scripts() (map[string][]byte, error) {
	out := map[string][]byte{}
	if l.index == nil {
		return out, nil
	}
	var first error
	l.index.Each(func(e *asset.Entry) {
		if e.Kind != asset.ArtScript || e.Missing {
			return
		}
		for _, a := range e.Artifacts {
			if a.Kind != asset.ArtScript || a.Name == "" {
				continue
			}
			data, err := l.read(a)
			if err != nil {
				if first == nil {
					first = errs.Wrap(err, errs.AssetLoadFailed, "cannot read a script artifact").With(log.F("package", a.Name))
				}
				continue
			}
			out[a.Name] = data
		}
	})
	return out, first
}

func (l *Loader) builtinWhite() (texture.Region, bool) {
	d := &texture.Data{Width: 1, Height: 1, Format: gpu.TexSRGBA8, Pixels: []byte{255, 255, 255, 255}, Premultiplied: true}
	h, err := l.up.Upload(d, texture.UploadOptions{Label: "builtin.white"})
	if err != nil {
		l.fail(asset.BuiltinWhite, "cannot upload the builtin white texture", err)
		return texture.Region{}, false
	}
	l.pages[asset.BuiltinWhite] = pageEntry{handle: h, width: 1, height: 1, filter: gpu.FilterNearest}
	reg := texture.Whole(h, 1, 1)
	l.region[asset.BuiltinWhite] = regionEntry{region: reg, page: asset.BuiltinWhite}
	return reg, true
}

func (l *Loader) Texture(id ids.AssetID) (texture.Region, bool) {
	if r, ok := l.region[id]; ok {
		return r.region, true
	}
	if l.failed[id] {
		return texture.Region{}, false
	}
	if id == asset.BuiltinWhite {
		return l.builtinWhite()
	}
	e, ok := l.entry(id)
	if !ok {
		return texture.Region{}, false
	}
	switch e.Kind {
	case asset.ArtSpriteRegion:
		return l.loadRegion(id, e)
	case asset.ArtTexture, asset.ArtAtlasPage:
		page, ok := l.loadPage(id, e)
		if !ok {
			return texture.Region{}, false
		}
		reg := texture.Whole(page.handle, page.width, page.height)
		l.region[id] = regionEntry{region: reg, page: id}
		return reg, true
	}
	return texture.Region{}, false
}

func (l *Loader) loadRegion(id ids.AssetID, e *asset.Entry) (texture.Region, bool) {
	art, ok := e.Artifact(asset.ArtSpriteRegion)
	if !ok {
		l.fail(id, "sprite region artifact is missing", nil)
		return texture.Region{}, false
	}
	data, err := l.read(art)
	if err != nil {
		l.fail(id, "cannot read the sprite region artifact", err)
		return texture.Region{}, false
	}
	ra, err := asset.DecodeRegion(data)
	if err != nil {
		l.fail(id, "malformed sprite region artifact", err)
		return texture.Region{}, false
	}
	page, ok := l.pages[ra.Page]
	if !ok {
		pe, pok := l.entry(ra.Page)
		if !pok {
			l.fail(id, "sprite region points at a missing page", nil)
			return texture.Region{}, false
		}
		page, ok = l.loadPage(ra.Page, pe)
		if !ok {
			return texture.Region{}, false
		}
	}
	whole := texture.Whole(page.handle, page.width, page.height)
	reg := whole.Sub(ra.X, ra.Y, ra.W, ra.H)
	l.region[id] = regionEntry{region: reg, page: ra.Page}
	return reg, true
}

func (l *Loader) loadPage(id ids.AssetID, e *asset.Entry) (pageEntry, bool) {
	if p, ok := l.pages[id]; ok {
		return p, true
	}
	kind := asset.ArtTexture
	if e.Kind == asset.ArtAtlasPage {
		kind = asset.ArtAtlasPage
	}
	art, ok := e.Artifact(kind)
	if !ok {
		l.fail(id, "texture artifact is missing", nil)
		return pageEntry{}, false
	}
	data, err := l.read(art)
	if err != nil {
		l.fail(id, "cannot read the texture artifact", err)
		return pageEntry{}, false
	}
	p, err := l.uploadTexture(data)
	if err != nil {
		l.fail(id, "cannot upload the texture", err)
		return pageEntry{}, false
	}
	l.pages[id] = p
	return p, true
}

func (l *Loader) uploadTexture(data []byte) (pageEntry, error) {
	ta, err := asset.DecodeTexture(data)
	if err != nil {
		return pageEntry{}, err
	}
	d := &texture.Data{
		Width:         ta.Width,
		Height:        ta.Height,
		Format:        ta.Format,
		Pixels:        ta.Pixels,
		Premultiplied: true,
	}
	h, err := l.up.Upload(d, texture.UploadOptions{Mips: ta.Mips, Label: "asset"})
	if err != nil {
		return pageEntry{}, err
	}
	return pageEntry{handle: h, width: ta.Width, height: ta.Height, filter: ta.Filter}, nil
}

func (l *Loader) readMesh(e *asset.Entry) (*mesh.Data, error) {
	art, ok := e.Artifact(asset.ArtMesh)
	if !ok {
		return nil, errs.New(errs.AssetBadArtifact, "mesh artifact is missing", log.F("asset", e.ID.String()))
	}
	raw, err := l.read(art)
	if err != nil {
		return nil, err
	}
	return mesh.Decode(raw)
}

func (l *Loader) Mesh(id ids.AssetID) (*mesh.Mesh, bool) {
	if m, ok := l.meshes[id]; ok {
		return m, true
	}
	if id.IsZero() || l.failed[id] {
		return nil, false
	}
	data, ok := asset.BuiltinMesh(id)
	if !ok {
		e, found := l.entry(id)
		if !found || e.Kind != asset.ArtMesh {
			return nil, false
		}
		var err error
		data, err = l.readMesh(e)
		if err != nil {
			l.fail(id, "cannot load the mesh", err)
			return nil, false
		}
	}
	m, err := mesh.Upload(l.dev, data, "asset.mesh")
	if err != nil {
		l.fail(id, "cannot upload the mesh", err)
		return nil, false
	}
	l.meshes[id] = m
	return m, true
}

func (l *Loader) fail(id ids.AssetID, msg string, cause error) {
	l.failed[id] = true
	fields := []log.Field{log.F("asset", id.String())}
	if cause != nil {
		fields = append(fields, log.F("err", cause.Error()))
	}
	l.log.Error(msg, fields...)
}

func (l *Loader) blob(id ids.AssetID, kind asset.ArtifactKind) ([]byte, bool) {
	e, ok := l.entry(id)
	if !ok || e.Kind != kind {
		return nil, false
	}
	if b, ok := l.blobs[id]; ok {
		return b, true
	}
	if l.failed[id] {
		return nil, false
	}
	art, ok := e.Artifact(kind)
	if !ok {
		l.fail(id, "artifact is missing", nil)
		return nil, false
	}
	data, err := l.read(art)
	if err != nil {
		l.fail(id, "cannot read the artifact", err)
		return nil, false
	}
	l.blobs[id] = data
	return data, true
}

func (l *Loader) Scene(id ids.AssetID) ([]byte, bool) {
	return l.blob(id, asset.ArtScene)
}

func (l *Loader) Prefab(id ids.AssetID) ([]byte, bool) {
	return l.blob(id, asset.ArtPrefab)
}

func (l *Loader) Preload(list ...ids.AssetID) error {
	for _, id := range list {
		if b, ok := asset.LookupBuiltin(id); ok {
			loaded := false
			if b.Kind == asset.ArtTexture {
				_, loaded = l.Texture(id)
			} else {
				_, loaded = l.Mesh(id)
			}
			if !loaded {
				return errs.New(errs.AssetLoadFailed, "builtin asset failed to load", log.F("asset", id.String()))
			}
			continue
		}
		e, ok := l.entry(id)
		if !ok {
			return errs.New(errs.AssetNotFound, "asset is not in the index", log.F("asset", id.String()))
		}
		ok = true
		switch e.Kind {
		case asset.ArtTexture, asset.ArtAtlasPage, asset.ArtSpriteRegion:
			_, ok = l.Texture(id)
		case asset.ArtMesh:
			_, ok = l.Mesh(id)
		case asset.ArtScene:
			_, ok = l.Scene(id)
		case asset.ArtPrefab:
			_, ok = l.Prefab(id)
		}
		if !ok {
			return errs.New(errs.AssetLoadFailed, "asset failed to load", log.F("asset", id.String()), log.F("kind", string(e.Kind)))
		}
	}
	return nil
}

type staged struct {
	id        ids.AssetID
	kind      asset.ArtifactKind
	page      pageEntry
	geo       *mesh.Mesh
	blob      []byte
	isTex     bool
	unchanged bool
}

func (l *Loader) prepare(e *asset.Entry) (staged, bool, error) {
	switch e.Kind {
	case asset.ArtTexture, asset.ArtAtlasPage:
		art, ok := e.Artifact(e.Kind)
		if !ok {
			return staged{}, false, errs.New(errs.AssetBadArtifact, "entry has no artifact of its own kind", log.F("asset", e.ID.String()))
		}
		data, err := l.read(art)
		if err != nil {
			return staged{}, false, err
		}
		p, err := l.uploadTexture(data)
		if err != nil {
			return staged{}, false, err
		}
		return staged{id: e.ID, kind: e.Kind, page: p, isTex: true}, true, nil
	case asset.ArtMesh:
		data, err := l.readMesh(e)
		if err != nil {
			return staged{}, false, err
		}
		m, err := mesh.Upload(l.dev, data, "asset.mesh")
		if err != nil {
			return staged{}, false, err
		}
		return staged{id: e.ID, kind: e.Kind, geo: m}, true, nil
	case asset.ArtScript:
		return staged{id: e.ID, kind: e.Kind}, true, nil
	case asset.ArtSpriteRegion, asset.ArtScene, asset.ArtPrefab:
		art, ok := e.Artifact(e.Kind)
		if !ok {
			return staged{}, false, errs.New(errs.AssetBadArtifact, "entry has no artifact of its own kind", log.F("asset", e.ID.String()))
		}
		data, err := l.read(art)
		if err != nil {
			return staged{}, false, err
		}
		return staged{id: e.ID, kind: e.Kind, blob: data}, true, nil
	}
	return staged{}, false, nil
}

func (l *Loader) Swap(entries []asset.Entry, files vfs.FS) error {
	if files != nil {
		l.files = files
	}
	if l.index == nil {
		return errs.New(errs.AssetNotMounted, "loader has no index mounted")
	}
	prepared := make([]staged, 0, len(entries))
	for i := range entries {
		s, ok, err := l.prepare(&entries[i])
		if err != nil {
			l.rollback(prepared)
			return err
		}
		if ok {
			prepared = append(prepared, s)
		}
	}
	for i := range entries {
		l.index.Put(entries[i])
	}
	for i := range prepared {
		s := &prepared[i]
		delete(l.failed, s.id)
		switch {
		case s.isTex:
			if old, ok := l.pages[s.id]; ok {
				l.dev.DestroyTexture(old.handle)
			}
			l.pages[s.id] = s.page
			l.invalidateRegionsOf(s.id)
		case s.geo != nil:
			if old, ok := l.meshes[s.id]; ok {
				old.Destroy(l.dev)
			}
			l.meshes[s.id] = s.geo
		case s.kind == asset.ArtScript:
		case s.kind == asset.ArtSpriteRegion:
			delete(l.region, s.id)
		default:
			if old, ok := l.blobs[s.id]; ok && bytes.Equal(old, s.blob) {
				s.unchanged = true
				continue
			}
			l.blobs[s.id] = s.blob
		}
	}
	clear(l.decoded)
	l.gen++
	for i := range prepared {
		if prepared[i].unchanged {
			continue
		}
		l.notify(asset.Reloaded{ID: prepared[i].id, Kind: prepared[i].kind})
	}
	return nil
}

func (l *Loader) invalidateRegionsOf(page ids.AssetID) {
	for id, r := range l.region {
		if r.page == page {
			delete(l.region, id)
		}
	}
}

func (l *Loader) rollback(prepared []staged) {
	for _, s := range prepared {
		if s.isTex {
			l.dev.DestroyTexture(s.page.handle)
		}
		if s.geo != nil {
			s.geo.Destroy(l.dev)
		}
	}
}

func (l *Loader) Drop(id ids.AssetID) {
	if p, ok := l.pages[id]; ok {
		l.dev.DestroyTexture(p.handle)
		delete(l.pages, id)
		l.invalidateRegionsOf(id)
	}
	if m, ok := l.meshes[id]; ok {
		m.Destroy(l.dev)
		delete(l.meshes, id)
	}
	delete(l.region, id)
	delete(l.blobs, id)
	delete(l.failed, id)
	clear(l.decoded)
	l.gen++
}

func (l *Loader) Close() {
	for id, p := range l.pages {
		l.dev.DestroyTexture(p.handle)
		delete(l.pages, id)
	}
	for id, m := range l.meshes {
		m.Destroy(l.dev)
		delete(l.meshes, id)
	}
	l.region = map[ids.AssetID]regionEntry{}
	l.blobs = map[ids.AssetID][]byte{}
	l.failed = map[ids.AssetID]bool{}
	l.decoded = map[decodedKey]decodedValue{}
	l.subs = nil
}

var (
	_ asset.Resolver = (*Loader)(nil)
	_ asset.Decoder  = (*Loader)(nil)
)