reference.go

// Copyright 2025 TypeFox GmbH
// This program and the accompanying materials are made available under the
// terms of the MIT License, which is available in the project root.

package fastbelt

import (
	"context"
	"encoding/json/jsontext"
	"encoding/json/v2"
	"errors"
	"fmt"
	"iter"
	"reflect"
	"slices"
	"sync"
	"sync/atomic"

	"typefox.dev/fastbelt/util/collections"
	"typefox.dev/fastbelt/util/extiter"
)

// Reference represents a reference to another AST node of type T.
//
// Resolving a reference is thread safe.
// The resolution is triggered when [Reference.Ref], [Reference.RefNode] or [Reference.Resolve]
// are called for the first time.
type Reference[T AstNode] struct {
	unit        StringUnit
	owner       AstNode
	description *SymbolDescription
	err         *ReferenceError
	ref         T
	mu          sync.Mutex
	getter      ReferenceGetter[T]
	resolved    atomic.Bool
}

// Reset clears all cached resolution results so the reference can be resolved again.
func (r *Reference[T]) Reset() {
	if r == nil {
		return
	}
	r.mu.Lock()
	defer r.mu.Unlock()
	r.resolved.Store(false)
	r.description = nil
	r.err = nil
	var zero T
	r.ref = zero
}

// Unit returns the [StringUnit] that contains the textual reference.
func (r *Reference[T]) Unit() StringUnit {
	if r == nil {
		return nil
	}
	return r.unit
}

// Text returns the textual value of the reference, or an empty string if unavailable.
func (r *Reference[T]) Text() string {
	if r == nil || r.unit == nil {
		return ""
	}
	return r.unit.String()
}

// Owner returns the AST node that owns this reference.
func (r *Reference[T]) Owner() AstNode {
	if r == nil {
		return nil
	}
	return r.owner
}

// Description returns the resolved symbol description, or nil when unresolved.
func (r *Reference[T]) Description() *SymbolDescription {
	if r == nil {
		return nil
	}
	return r.description
}

// RefNode returns the resolved target node as [AstNode].
func (r *Reference[T]) RefNode(ctx context.Context) AstNode {
	return r.Ref(ctx)
}

// Ref resolves the reference and returns the typed target node.
// It returns the zero value of T when resolution fails.
func (r *Reference[T]) Ref(ctx context.Context) T {
	var zero T
	if r == nil {
		return zero
	}
	r.Resolve(ctx)
	return r.ref
}

// Error returns the resolution error for this reference, if any.
func (r *Reference[T]) Error() *ReferenceError {
	if r == nil {
		return nil
	}
	return r.err
}

// Range returns the text range of the reference in the source document.
func (r *Reference[T]) TextRange() TextRange {
	return r.unit.TextRange()
}

// resolutionChainKey is the context key under which a [resolutionChain] is stored.
type resolutionChainKey struct{}

// resolutionChain tracks the references currently being resolved, for cycle
// detection. It is mutated without synchronization and must therefore only be
// used by a single goroutine at a time.
type resolutionChain struct {
	refs []UntypedReference
}

// WithResolutionChain returns a context that carries a reusable resolution
// chain for reference cycle detection.
//
// The chain is mutated without synchronization: the returned context must not
// be shared by goroutines that resolve references concurrently. Create one
// chain context per goroutine instead.
func WithResolutionChain(ctx context.Context) context.Context {
	return context.WithValue(ctx, resolutionChainKey{}, &resolutionChain{})
}

func getResolveChain(ctx context.Context) *resolutionChain {
	chain, _ := ctx.Value(resolutionChainKey{}).(*resolutionChain)
	return chain
}

// Resolve resolves the reference exactly once for this instance.
// It is safe for concurrent use.
func (r *Reference[T]) Resolve(ctx context.Context) {
	// Fast path: check if already resolved without locking
	if r == nil || r.resolved.Load() {
		return
	}
	// Slow path (outlined so that the fast path can be inlined)
	r.resolveSlow(ctx)
}

func (r *Reference[T]) resolveSlow(ctx context.Context) {
	// Cyclic resolution attempts are detected via the resolution chain
	chain := getResolveChain(ctx)
	if chain != nil && slices.Contains(chain.refs, UntypedReference(r)) {
		// Note that we write directly to r.err without locking here
		// This is safe, because the reference is already locked by the caller
		// Attempting to lock it again would cause a deadlock anyway
		r.err = NewReferenceError("Cyclic reference resolution detected")
		// Return directly, do not set the resolved flag
		return
	}
	r.mu.Lock()
	defer r.mu.Unlock()
	if r.resolved.Load() {
		// Another goroutine might have resolved it while we were waiting for the lock
		return
	}
	if chain == nil {
		// No chain provided, so create a context with a new chain
		chain = &resolutionChain{}
		ctx = context.WithValue(ctx, resolutionChainKey{}, chain)
	}
	// Append the current reference to the chain for cycle detection
	chain.refs = append(chain.refs, r)
	defer func() {
		// Pop the reference from the chain after resolution
		chain.refs = chain.refs[:len(chain.refs)-1]
	}()
	desc, e := r.getter(ctx, r)
	r.description = desc
	if r.err == nil {
		// Do not overwrite existing errors
		r.err = e
	}
	if desc != nil {
		if node, ok := desc.Node.(T); ok {
			r.ref = node
		} else if r.err == nil {
			expectedType := reflect.TypeFor[T]().String()
			actualType := "nil"
			if desc.Node != nil {
				actualType = reflect.TypeOf(desc.Node).String()
			}
			r.err = NewReferenceError("Reference resolution type mismatch: expected " + expectedType + ", got " + actualType)
		}
	}
	r.resolved.Store(true)
}

// NewReference creates a lazy, typed reference owned by owner and backed by getter.
func NewReference[T AstNode](owner AstNode, unit StringUnit, getter ReferenceGetter[T]) *Reference[T] {
	return &Reference[T]{
		owner:  owner,
		unit:   unit,
		getter: getter,
	}
}

// UntypedReference is an untyped representation of a [Reference] that can be used when the target
// type is not known at compile time.
//
// Used throughout the fastbelt codebase to generically deal with different references.
type UntypedReference interface {
	Owner() AstNode
	Description() *SymbolDescription
	RefNode(ctx context.Context) AstNode
	Resolve(ctx context.Context)
	Reset()
	Error() *ReferenceError
	Unit() StringUnit
	TextRange() TextRange
	Text() string
}

// ReferenceGetter resolves a [Reference] and returns its symbol description.
// The returned [ReferenceError] is stored on the reference when not nil.
type ReferenceGetter[T AstNode] func(context.Context, *Reference[T]) (*SymbolDescription, *ReferenceError)

// Computes the reference that this token represents.
// Returns nil if the token does not represent a reference.
func ReferenceOfToken(token *Token) UntypedReference {
	owner := token.Element
	rng := token.TextRange()
	if composite, ok := owner.(CompositeNode); ok {
		// If the token is part of a composite node,
		// we first have to retrieve its parent node, which is the actual owner of the reference.
		owner = composite.Container()
		rng = composite.TextRange()
	}
	if owner == nil {
		return nil
	}
	var ref UntypedReference = nil
	// We don't have a direct reference from token -> reference, so we need to search for it
	// We have to iterate over all references of the owner node
	// This might seem inefficient, but in practice the number of references per node is usually very small
	// Also, we only do this in select LSP requests, so the performance impact is negligible
	for ur := range References(owner) {
		// Simply compare the text indices to find the matching reference
		if ur.TextRange() == rng {
			ref = ur
			break
		}
	}
	return ref
}

// ReferenceDescription describes one concrete use of a symbol in source text.
type ReferenceDescription struct {
	// SourceNode is the node that contains the reference.
	SourceNode AstNode
	// TargetNode is the node that is being referenced.
	TargetNode AstNode
	// Range is the text range of the reference in the source document (usually the symbol name).
	Range TextRange
}

// NewReferenceDescription creates a [ReferenceDescription] for a source-to-target link.
func NewReferenceDescription(source, target AstNode, rng TextRange) *ReferenceDescription {
	return &ReferenceDescription{
		SourceNode: source,
		TargetNode: target,
		Range:      rng,
	}
}

// SourceURI returns the URI of the document containing the source node, or nil if not available.
func (d *ReferenceDescription) SourceURI() URI {
	if d.SourceNode != nil {
		doc := d.SourceNode.Document()
		if doc != nil {
			return doc.URI
		}
	}
	return nil
}

// TargetURI returns the URI of the document containing the target node, or nil if not available.
func (d *ReferenceDescription) TargetURI() URI {
	if d.TargetNode != nil {
		doc := d.TargetNode.Document()
		if doc != nil {
			return doc.URI
		}
	}
	return nil
}

// ReferenceDescriptions is a collection of reference descriptions for a document, indexed by target node.
// It allows efficient retrieval of all references to a given target node.
type ReferenceDescriptions interface {
	// All returns an iterator over all reference descriptions in the document.
	All() iter.Seq[*ReferenceDescription]
	// ForTarget returns an iterator over all reference descriptions that point to the given target node.
	ForTarget(target AstNode) iter.Seq[*ReferenceDescription]
}

type referenceDescriptions struct {
	descriptions collections.MultiMap[AstNode, *ReferenceDescription]
}

var emptyReferenceDescriptions = extiter.Empty[*ReferenceDescription]()

func (d *referenceDescriptions) All() iter.Seq[*ReferenceDescription] {
	if d == nil || d.descriptions == nil {
		return emptyReferenceDescriptions
	}
	return d.descriptions.Values()
}

func (d *referenceDescriptions) ForTarget(target AstNode) iter.Seq[*ReferenceDescription] {
	if d == nil || d.descriptions == nil {
		return emptyReferenceDescriptions
	}
	return slices.Values(d.descriptions.Get(target))
}

// NewReferenceDescriptionsFromMap wraps precomputed descriptions keyed by target node.
func NewReferenceDescriptionsFromMap(descriptions collections.MultiMap[AstNode, *ReferenceDescription]) ReferenceDescriptions {
	return &referenceDescriptions{
		descriptions: descriptions,
	}
}

// MarshalJSONTo serializes the reference as a JSON object containing the URI of the resolved target ast node,
// the cross ref text, as well as the err msg if the reference is unresolvable.
// Returns an error if the reference has not been attempted to resolve (r.resolved == false).
// With this function Reference[T] implements [json.MarshalerTo] and the method is called by `json.Marshal()`.
func (r *Reference[T]) MarshalJSONTo(encoder *jsontext.Encoder) error {
	// A nil reference marshals as JSON null.
	if r == nil {
		return encoder.WriteToken(jsontext.Null)
	}
	// We expect at this point that all references have been attempted to resolve.
	if !r.resolved.Load() {
		if path, err := PathOf(r.Owner()); err == nil {
			return errors.New("Reference.MarshalJSONTo(): reference not resolved in node '" + path.String() + "'")
		}
		return errors.New("Reference.MarshalJSONTo(): reference not resolved")
	}

	var uri, errMsg string
	if r.err != nil {
		errMsg = r.err.Msg
	} else if referenced := r.ref; any(referenced) != nil {
		refPath, err := PathOf(referenced)
		if err != nil {
			referencerPath, _ := PathOf(r.Owner())
			return fmt.Errorf("Reference.MarshalJSONTo(): failed to compute path fragment of node of type %T referenced by '%s'", referenced, referencerPath)
		}

		refDoc := referenced.Document()
		if refDoc == nil {
			referencerPath, _ := PathOf(r.Owner())
			return fmt.Errorf("Reference.MarshalJSONTo(): node of type %T referenced by '%s' is not contained in any document", referenced, referencerPath)
		}

		if r.Owner() != nil && r.Owner().Document() == refDoc {
			uri = "#" + refPath.String()
		} else {
			uri = refDoc.URI.WithFragment(refPath.String()).String()
		}
	} else {
		return errors.New("Reference.MarshalJSONTo(): unexpected state")
	}

	return json.MarshalEncode(encoder, struct {
		RefText string `json:"$refText"`
		Ref     string `json:"$ref,omitempty"`
		Err     string `json:"$error,omitempty"`
	}{
		RefText: r.unit.String(),
		Ref:     uri,
		Err:     errMsg,
	})
}

// UnmarshalJSON revives the reference as a JSON object. It inspects properties named '$ref', '$refText', and '$error'.
// '$ref' is expected to contain a URI string denoting the target node within the same or another document, may be absent if the reference was unresolvable before serializing.
// '$refText' denotes the cross reference string.
// '$error' contains the error msg if the reference was unresolvable before serializing
// With this method Reference[T] implements [json.Unmarshaler], and instances can be populated by [json.Unmarshal()].
// However, we typically call it directly via [typefox.dev/fastbelt/util.UnmarshalReference],
// which the code generated for [json.UnmarshalerFrom] implementations of AST node types calls in turn.
//
// Note: Since we're calling this method directly within the unmarshaling implementations of the AST nodes based on a raw [jsontext.Value],
// implementing [json.UnmarshalerFrom].UnmarshalJSONFrom(*jsontext.Decoder) doesn't bring any benefit here, since the value chunks ([]byte) are rather small.
// Therefore, we stick to implementing the plain [json.Unmarshaler] interface here.
func (r *Reference[T]) UnmarshalJSON(value []byte) error {
	aux := &struct {
		RefText string `json:"$refText"`
		Ref     string `json:"$ref"`
		Err     string `json:"$error"`
	}{}
	if err := json.Unmarshal(value, aux); err != nil {
		return err
	}

	// requirement: r.owner is initialized
	r.unit = NewSyntheticToken(aux.RefText, r.owner)
	r.getter = newJsonReferenceGetter[T](aux.Ref)
	if aux.Err != "" {
		r.err = NewReferenceError(aux.Err)
	}
	return nil
}

// JsonLinkingHelper instances are required for properly linking cross references in ASTs
// being created by unmarshaling JSON data via the [json.Unmarshal] API.
// A default implementation is provided by [typefox.dev/fastbelt/util/NewJsonLinkingHelper]().
//
// The helper is expected to be present even for references that turn out to target the same document:
// whether a reference is same- or cross-document is only known once its '$ref' URI is inspected during
// resolution, and by design that resolution fails fast if no helper is present in the context,
// rather than silently succeeding for some references and not others depending on their target.
// Consequently, any [context.Context] used to (re)build a workspace that may contain JSON-loaded
// documents must carry a helper via [JsonLinkingHelperKey] - including on rebuilds triggered by
// unrelated document edits, which commonly build with a fresh context.Background()-derived one.
//
// It is to be provided as part of the context argument of typefox.dev/fastbelt/workspace.Builder.Build](Context, ...), like
//
//	documents, err := service.Get[workspace.DocumentManager](sc)
//	err = builder.Build(
//	    context.WithValue(
//	        ctx,
//	        fastbelt.JsonLinkingHelperKey(),
//	        util.NewJsonLinkingHelper(documents),
//	    ),
//	    ...,
//	    ...,
//	);
type JsonLinkingHelper interface {
	GetDocument(uri URI) *Document
}

var jsonLinkingHelperKey = reflect.TypeFor[JsonLinkingHelper]()

// JsonLinkingHelperKey returns a singleton identifier key for registering a [JsonLinkingHelper] instance
// in the [context.Context] object used by [typefox.dev/fastbelt/workspace.Builder.Build](Context, ...),
// which is required for linking references in ASTs obtained by unmarshaling JSON data.
// See also [JsonLinkingHelper].
func JsonLinkingHelperKey() any {
	return jsonLinkingHelperKey
}

func newJsonReferenceGetter[T AstNode](uriString string) ReferenceGetter[T] {
	return func(ctx context.Context, ref *Reference[T]) (*SymbolDescription, *ReferenceError) {
		// ref.err is set from the imported '$error' by UnmarshalJSON and only survives here on the
		// very first resolution; Reset() clears it before any later resolution attempt. That's
		// intentional: a reference that failed to link in the workspace it was exported from may
		// well resolve in the workspace it is imported into (or after a rebuild), so re-resolution
		// deliberately re-derives the error (e.g. "absent" below) instead of replaying the old one.
		if ref.err != nil {
			return nil, ref.err
		}
		if uriString == "" {
			return nil, NewReferenceError("Reviving reference in Json document failed, '$ref' is absent or empty")
		}
		helper, ok := ctx.Value(jsonLinkingHelperKey).(JsonLinkingHelper)
		if !ok {
			return nil, NewReferenceError("JsonLinkingHelper unavailable")
		}
		var document *Document
		uri := ParseURI(uriString)
		if uri.Path() == "" {
			document = ref.Owner().Document()
		} else if doc := helper.GetDocument(uri.WithFragment("") /* zeros the fragment part */); doc != nil {
			document = doc
		} else {
			return nil, NewReferenceError("Document not found: '" + uri.StringUnencoded() + "'")
		}

		if document.Root == nil {
			return nil, NewReferenceError("Document is empty: '" + document.URI.StringUnencoded() + "'")
		}

		node, err := Resolve(document.Root, uri.Fragment())
		if err != nil {
			return nil, NewReferenceError(err.Error())
		}
		return &SymbolDescription{
			URI:  uri,
			Node: node,
		}, nil
	}
}