exec.go

package starlet

import (
	"bytes"
	"errors"
	"fmt"
	"sort"
	"strings"
	"sync"

	itn "github.com/1set/starlet/internal"
	"go.starlark.net/starlark"
)

// execStarlarkFile executes a Starlark file with the given filename and source, and returns the global environment and any error encountered.
// If the cache is enabled, it will try to load the compiled program from the cache first, and save the compiled program to the cache after compilation.
func (m *Machine) execStarlarkFile(filename string, src interface{}, allowCache bool) (starlark.StringDict, error) {
	// restore the arguments for starlark.ExecFileOptions
	opts := m.getFileOptions()
	thread := m.thread
	predeclared := m.predeclared
	hasCache := m.progCache != nil

	// if cache is not enabled or not allowed, just execute the original source
	if !hasCache || !allowCache {
		return starlark.ExecFileOptions(opts, thread, filename, src, predeclared)
	}

	// for compiled program and cache key; sources that cannot be
	// content-keyed (e.g. an io.Reader) are executed without the cache
	// rather than risking a stale hit on a filename-only key
	key, ok := m.getCacheKey(src)
	if !ok {
		return starlark.ExecFileOptions(opts, thread, filename, src, predeclared)
	}
	var (
		prog *starlark.Program
		err  error
	)

	// try to load compiled program from cache first
	if hasCache {
		// if cache is enabled, try to load compiled bytes from cache first
		if cb, ok := m.progCache.Get(key); ok {
			// load program from compiled bytes
			if prog, err = starlark.CompiledProgram(bytes.NewReader(cb)); err != nil {
				// if failed, remove the result and continue
				prog = nil
			}
		}
	}

	// if program is not loaded from cache, compile and cache it
	if prog == nil {
		// parse, resolve, and compile a Starlark source file.
		if _, prog, err = starlark.SourceProgramOptions(opts, filename, src, predeclared.Has); err != nil {
			return nil, err
		}
		// dump the compiled program to bytes
		buf := new(bytes.Buffer)
		if err = prog.Write(buf); err != nil {
			return nil, err
		}
		// save the compiled bytes to cache
		_ = m.progCache.Set(key, buf.Bytes())
	}

	// execute the compiled program
	g, err := prog.Init(thread, predeclared)
	g.Freeze()
	return g, err
}

// getCacheKey derives the cache key for a compiled program. The key must
// cover everything the compiler consumes, not just the source text:
//
//   - the predeclared NAME SET: resolution binds each identifier as
//     predeclared vs global/undefined via predeclared.Has, so the same
//     source compiled under a different name set is a different program.
//     A stale hit either failed at run time with "internal error:
//     predeclared variable X is uninitialized" or silently validated a
//     source that must not compile;
//   - the dialect bits (recursion / global-reassign FileOptions), which
//     gate parsing and resolution;
//   - the source content. An io.Reader source used to fall back to a
//     filename-only key, so identical names with different content
//     collided — such sources now report ok=false and skip the cache.
//
// The key layout is internal and may change between releases; persisted
// caches then miss once and recompile.
func (m *Machine) getCacheKey(src interface{}) (key string, ok bool) {
	var k string
	switch s := src.(type) {
	case string:
		k = itn.GetStringMD5(s)
	case []byte:
		k = itn.GetBytesMD5(s)
	default:
		return "", false
	}

	// fold in the sorted predeclared names
	names := make([]string, 0, len(m.predeclared))
	for n := range m.predeclared {
		names = append(names, n)
	}
	sort.Strings(names)
	nd := itn.GetStringMD5(strings.Join(names, "\x00"))

	// fold in the dialect bits
	var opt int
	if m.allowRecursion {
		opt |= 1
	}
	if m.allowGlobalReassign {
		opt |= 2
	}

	return fmt.Sprintf("%d:%d:%s:%s", starlark.CompilerVersion, opt, nd, k), true
}

// ByteCache is an interface for caching byte data, used for caching compiled Starlark programs.
type ByteCache interface {
	Get(key string) ([]byte, bool)
	Set(key string, value []byte) error
}

// MemoryCache is a simple in-memory map-based ByteCache, serves as a default cache for Starlark programs.
type MemoryCache struct {
	_ itn.DoNotCompare
	sync.RWMutex
	data map[string][]byte
}

// NewMemoryCache creates a new MemoryCache instance.
func NewMemoryCache() *MemoryCache {
	return &MemoryCache{
		data: make(map[string][]byte),
	}
}

// Get returns the value for the given key, and whether the key exists.
func (c *MemoryCache) Get(key string) ([]byte, bool) {
	c.RLock()
	defer c.RUnlock()

	if c == nil || c.data == nil {
		return nil, false
	}
	v, ok := c.data[key]
	return v, ok
}

// Set sets the value for the given key.
func (c *MemoryCache) Set(key string, value []byte) error {
	c.Lock()
	defer c.Unlock()

	if c == nil || c.data == nil {
		return errors.New("no data map found in the cache")
	}
	c.data[key] = value
	return nil
}