feat: add support for $id and $anchor reference resolution (#123)

Author: TristanSpeakEasy

arazzo/criterion/criterion.go

@@ -147,7 +147,7 @@ func (c *CriterionTypeUnion) Populate(source any) error {
 		c.Type = &typ
 	} else if coreCriterionTypeUnion.ExpressionType != nil {
 		c.ExpressionType = &CriterionExpressionType{}
-		if err := marshaller.Populate(*coreCriterionTypeUnion.ExpressionType, c.ExpressionType); err != nil {
+		if err := marshaller.PopulateWithContext(*coreCriterionTypeUnion.ExpressionType, c.ExpressionType, nil); err != nil {
 			return err
 		}
 	}

extensions/extensions.go

@@ -112,7 +112,7 @@ func UnmarshalExtensionModel[H any, L any](ctx context.Context, e *Extensions, e
 
 	var mV H
 
-	if err := marshaller.Populate(*c, &mV); err != nil {
+	if err := marshaller.PopulateWithContext(*c, &mV, nil); err != nil {
 		return nil, err
 	}
 	*m = mV

extensions/extensions_test.go

@@ -122,7 +122,7 @@ func getTestModelWithExtensions(ctx context.Context, t *testing.T, data string)
 	require.Empty(t, validationErrs)
 
 	m := &ModelWithExtensions{}
-	err = marshaller.Populate(c, m)
+	err = marshaller.PopulateWithContext(c, m, nil)
 	require.NoError(t, err)
 
 	return m

jsonschema/oas3/core/jsonschema.go

@@ -35,6 +35,7 @@ type Schema struct {
 	UnevaluatedProperties marshaller.Node[JSONSchema]                                             `key:"unevaluatedProperties"`
 	Items                 marshaller.Node[JSONSchema]                                             `key:"items"`
 	Anchor                marshaller.Node[*string]                                                `key:"$anchor"`
+	ID                    marshaller.Node[*string]                                                `key:"$id"`
 	Not                   marshaller.Node[JSONSchema]                                             `key:"not"`
 	Properties            marshaller.Node[*sequencedmap.Map[string, JSONSchema]]                  `key:"properties"`
 	Defs                  marshaller.Node[*sequencedmap.Map[string, JSONSchema]]                  `key:"$defs"`

jsonschema/oas3/jsonschema.go

@@ -2,10 +2,12 @@ package oas3
 
 import (
 	"context"
+	"strings"
 	"unsafe"
 
 	"github.com/speakeasy-api/openapi/extensions"
 	"github.com/speakeasy-api/openapi/jsonschema/oas3/core"
+	"github.com/speakeasy-api/openapi/marshaller"
 	"github.com/speakeasy-api/openapi/pointer"
 	"github.com/speakeasy-api/openapi/references"
 	"github.com/speakeasy-api/openapi/validation"
@@ -38,6 +40,20 @@ type JSONSchema[T Referenceable | Concrete] struct {
 	//   parent = intermediate reference, topLevelParent = original reference
 	parent         *JSONSchema[Referenceable] // Immediate parent reference in the chain
 	topLevelParent *JSONSchema[Referenceable] // Top-level parent (root of the reference chain)
+
+	// enclosingSchema is the Schema that contains this JSONSchema as a field.
+	// This is used during population to determine the parent's effective base URI
+	// for relative $id resolution. Set during PopulateWithContext.
+	enclosingSchema *Schema
+
+	// schemaRegistry stores $id and $anchor mappings for this document.
+	// Used for standalone JSON Schema documents that are not embedded in an OpenAPI document.
+	// For embedded schemas, the owning OpenAPI document's registry is used instead.
+	schemaRegistry SchemaRegistry
+
+	// documentBaseURI is the base URI for this standalone JSON Schema document.
+	// This is typically derived from the $id keyword or empty if not specified.
+	documentBaseURI string
 }
 
 var _ references.Resolvable[JSONSchema[Concrete]] = (*JSONSchema[Referenceable])(nil)
@@ -281,6 +297,9 @@ func (j *JSONSchema[T]) ShallowCopy() *JSONSchema[T] {
 		resolvedSchemaCache:      j.resolvedSchemaCache,
 		parent:                   j.parent,
 		topLevelParent:           j.topLevelParent,
+		enclosingSchema:          j.enclosingSchema,
+		schemaRegistry:           j.schemaRegistry,
+		documentBaseURI:          j.documentBaseURI,
 	}
 
 	// Shallow copy the EitherValue contents
@@ -295,23 +314,137 @@ func (j *JSONSchema[T]) ShallowCopy() *JSONSchema[T] {
 	return result
 }
 
-// PopulateWithParent implements the ParentAwarePopulator interface to establish parent relationships during population
-func (j *JSONSchema[T]) PopulateWithParent(source any, parent any) error {
-	// If we have a parent that is also a JSONSchema, establish the parent relationship
-	if parent != nil {
-		if parentSchema, ok := parent.(*Schema); ok {
-			j.SetParent(parentSchema.GetParent())
-			// If the parent has a top-level parent, inherit it; otherwise, the parent is the top-level
-			if parentSchema.GetParent().GetTopLevelParent() != nil {
-				j.SetTopLevelParent(parentSchema.GetParent().GetTopLevelParent())
-			} else {
-				j.SetTopLevelParent(parentSchema.GetParent())
-			}
+// GetSchemaRegistry returns the schema registry for this standalone JSON Schema document.
+// The registry stores $id and $anchor mappings for efficient schema resolution.
+// If the registry has not been initialized, it creates one with the document's base URI.
+// This implements the SchemaRegistryProvider interface.
+func (j *JSONSchema[T]) GetSchemaRegistry() SchemaRegistry {
+	if j == nil {
+		return nil
+	}
+
+	// Lazily initialize the registry if needed
+	if j.schemaRegistry == nil {
+		j.schemaRegistry = NewSchemaRegistry(j.GetDocumentBaseURI())
+	}
+
+	return j.schemaRegistry
+}
+
+// GetDocumentBaseURI returns the base URI for this standalone JSON Schema document.
+// This is derived from the $id keyword of the root schema, if present.
+// The returned URI is normalized and stripped of any fragment to align with registry behavior.
+// This implements the SchemaRegistryProvider interface.
+func (j *JSONSchema[T]) GetDocumentBaseURI() string {
+	if j == nil {
+		return ""
+	}
+
+	var uri string
+
+	// If we have an explicit document base URI set, use it
+	if j.documentBaseURI != "" {
+		uri = j.documentBaseURI
+	} else if j.IsSchema() && j.GetSchema() != nil {
+		// Try to get from the root schema's $id
+		uri = j.GetSchema().GetID()
+	}
+
+	if uri == "" {
+		return ""
+	}
+
+	// Strip fragment and normalize to align with registry behavior
+	// Per JSON Schema spec, $id should not contain fragments, but we strip for robustness
+	return normalizeDocumentBaseURI(uri)
+}
+
+// normalizeDocumentBaseURI strips fragments and normalizes a URI for use as a document base.
+func normalizeDocumentBaseURI(uri string) string {
+	if uri == "" {
+		return ""
+	}
+
+	// Strip fragment if present
+	if idx := strings.Index(uri, "#"); idx != -1 {
+		uri = uri[:idx]
+	}
+
+	// Use the same normalization as the registry
+	return normalizeURI(uri)
+}
+
+// SetSchemaRegistry sets the schema registry for this document.
+// This is primarily used during unmarshalling to set a pre-created registry.
+func (j *JSONSchema[T]) SetSchemaRegistry(registry SchemaRegistry) {
+	if j == nil {
+		return
+	}
+	j.schemaRegistry = registry
+}
+
+// SetDocumentBaseURI sets the document base URI for this standalone JSON Schema.
+func (j *JSONSchema[T]) SetDocumentBaseURI(uri string) {
+	if j == nil {
+		return
+	}
+	j.documentBaseURI = uri
+}
+
+// GetEnclosingSchema returns the Schema that contains this JSONSchema as a field.
+// This is used during population to access the parent schema's effective base URI
+// for relative $id resolution. Returns nil if not set.
+func (j *JSONSchema[T]) GetEnclosingSchema() *Schema {
+	if j == nil {
+		return nil
+	}
+	return j.enclosingSchema
+}
+
+// SetEnclosingSchema sets the Schema that contains this JSONSchema as a field.
+// This is used during population to establish parent-child relationships
+// for effective base URI computation.
+func (j *JSONSchema[T]) SetEnclosingSchema(schema *Schema) {
+	if j == nil {
+		return
+	}
+	j.enclosingSchema = schema
+}
+
+// PopulateWithContext implements the ContextAwarePopulator interface for full context-aware population.
+// This method receives the owning document and propagates it to the contained Schema.
+func (j *JSONSchema[T]) PopulateWithContext(source any, ctx *marshaller.PopulationContext) error {
+	// If we have a parent that is a Schema, store it as enclosingSchema.
+	// This is critical for relative $id resolution - children need to access parent's effective base URI.
+	//
+	// Note: We only set enclosingSchema here (document tree relationship).
+	// The parent/topLevelParent fields are for reference chains and are set during
+	// reference resolution, not population. See documentation on those fields (lines 28-41).
+	if ctx != nil && ctx.Parent != nil {
+		if parentSchema, ok := ctx.Parent.(*Schema); ok {
+			j.enclosingSchema = parentSchema
 		}
 	}
 
-	// First, perform the standard population
-	if err := j.EitherValue.PopulateWithParent(source, j); err != nil {
+	// Determine the owning document for context propagation
+	// If we're not nested in any other document, this JSONSchema becomes its own owning document
+	var owningDoc any
+	if ctx != nil && ctx.OwningDocument != nil {
+		owningDoc = ctx.OwningDocument
+	} else {
+		// This is a standalone JSON Schema document - use itself as the owning document
+		owningDoc = j
+	}
+
+	// Create a new context with this JSONSchema as the parent for nested schemas
+	// Note: The Schema's PopulateWithContext will use 'j' (this JSONSchema) as its parent reference
+	childCtx := &marshaller.PopulationContext{
+		Parent:         j,
+		OwningDocument: owningDoc,
+	}
+
+	// Perform the standard population with context through the EitherValue
+	if err := j.EitherValue.PopulateWithContext(source, childCtx); err != nil {
 		return err
 	}