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"
	"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()
}

// 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) {
	// We can use the context to detect cyclic reference resolution attempts
	// We are allowed to do this outside of the mutex lock because context is immutable
	if ctx.Value(r) != nil {
		// 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
	}
	newCtx := context.WithValue(ctx, r, true)
	desc, e := r.getter(newCtx, 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]
}

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

func (d *referenceDescriptions) ForTarget(target AstNode) iter.Seq[*ReferenceDescription] {
	if d == nil {
		return extiter.Empty[*ReferenceDescription]()
	}
	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,
	}
}