feat(id-handling) support for $id syntax

Author: Ryan Lee

document.go

@@ -8,6 +8,7 @@ import (
 	"os"
 	"path/filepath"
 	"regexp"
+	"strings"
 
 	"github.com/ryanolee/go-chaff/internal/util"
 	"github.com/thoas/go-funk"
@@ -25,6 +26,19 @@ type (
 
 		// A map of protocols to their associated document fetchers
 		documentFetchers map[string]*documentFetcherInterface
+
+		// Maps resolved $id URIs to the real document + JSON pointer path
+		// where the sub-schema lives, avoiding document duplication.
+		idAliases map[string]idAlias
+	}
+
+	// idAlias maps a resolved $id URI back to the parent document and the
+	// JSON pointer path within that document. For example, a $defs entry
+	// with "$id": "color" in https://example.com/base.json produces:
+	//   idAlias{documentId: "https://example.com/base.json", path: "#/$defs/color"}
+	idAlias struct {
+		documentId string // The real document containing this sub-schema
+		path       string // JSON pointer within documentId, e.g. "#/$defs/color"
 	}
 
 	documentFetcherInterface interface {
@@ -76,7 +90,7 @@ func newDocumentResolver(opts ParserOptions, rootDocument *schemaNode) (*documen
 		resolvedRootDocumentId = opts.RelativeTo
 	}
 
-	return &documentResolver{
+	resolver := &documentResolver{
 		// Setup the root document
 		documentCurrentlyBeingParsedId: opts.RelativeTo,
 		parsedExternalDocuments: map[string]bool{
@@ -86,7 +100,22 @@ func newDocumentResolver(opts ParserOptions, rootDocument *schemaNode) (*documen
 			resolvedRootDocumentId: rootDocument,
 		},
 		documentFetchers: documentFetchers,
-	}, nil
+		idAliases:        make(map[string]idAlias),
+	}
+
+	// Collect $id aliases from the root document tree so that relative
+	// $ref values (e.g. $ref: "color") can be resolved without I/O.
+	baseURI := resolvedRootDocumentId
+	if rootDocument.Id != nil && isAbsoluteURI(*rootDocument.Id) {
+		baseURI = *rootDocument.Id
+	}
+
+	// Merge into existing idAliases
+	for resolvedId, alias := range collectSubSchemaIds(baseURI, resolvedRootDocumentId, rootDocument) {
+		resolver.idAliases[resolvedId] = alias
+	}
+
+	return resolver, nil
 }
 
 func (r *documentResolver) GetDocumentIdCurrentlyBeingParsed() string {
@@ -115,21 +144,11 @@ func (r *documentResolver) ResolveDocumentIdAndPath(ref string) (string, string,
 	}
 
 	// If we have no document, it is a local or relative reference and can be handled as such
-	if !hasDocument || documentID == "" || r.HasNoFetchers() {
+	if !hasDocument || documentID == "" {
 		return r.GetCurrentScope(), path, nil
 	}
 
-	fetcher, err := r.getFetcherForRef(ref)
-	if err != nil {
-		return "", "", fmt.Errorf("failed to get document fetcher for document '%s': %w", ref, err)
-	}
-
-	resolvedDocumentId, err := fetcher.resolveDocumentId(r.GetCurrentScope(), documentID)
-	if err != nil {
-		return "", "", fmt.Errorf("failed to resolve document ID for reference '%s': %w", ref, err)
-	}
-
-	return resolvedDocumentId, path, nil
+	return r.resolveDocumentRef(documentID, path)
 }
 
 // Traverses a schema node and rewrites any $ref fields to be relative to the given document ID
@@ -223,26 +242,21 @@ func (r *documentResolver) HandleDeferredReferenceResolution(ref string, metadat
 	}
 
 	// If we have no document, it is a local or relative reference and can be handled as such
-	if !hasDocument || documentID == "" || r.HasNoFetchers() {
+	if !hasDocument || documentID == "" {
 		return r.GetCurrentScope(), path, nil
 	}
 
-	fetcher, err := r.getFetcherForRef(documentID)
-	if err != nil {
-		return "", "", fmt.Errorf("failed to get document fetcher for document '%s': %w", ref, err)
-	}
-
-	resolvedDocumentId, err := fetcher.resolveDocumentId(r.GetCurrentScope(), documentID)
+	resolvedDocId, resolvedPath, err := r.resolveDocumentRef(documentID, path)
 	if err != nil {
-		return "", "", fmt.Errorf("failed to resolve document ID for reference '%s': %w", ref, err)
+		return "", "", err
 	}
 
 	// Queue document for parsing if it hasn't already been parsed
-	if !funk.ContainsString(r.externalDocumentsThatNeedParsing, resolvedDocumentId) && r.pathWouldRequireNewDocumentParsed(resolvedDocumentId) {
-		r.externalDocumentsThatNeedParsing = append(r.externalDocumentsThatNeedParsing, resolvedDocumentId)
+	if !funk.ContainsString(r.externalDocumentsThatNeedParsing, resolvedDocId) && r.pathWouldRequireNewDocumentParsed(resolvedDocId) {
+		r.externalDocumentsThatNeedParsing = append(r.externalDocumentsThatNeedParsing, resolvedDocId)
 	}
 
-	return resolvedDocumentId, path, nil
+	return resolvedDocId, resolvedPath, nil
 }
 
 /**
@@ -269,42 +283,37 @@ func (r *documentResolver) ResolvePath(metadata *parserMetadata, ref string) (*s
 	}
 
 	// If we have no document, it is a local or relative reference and can be handled as such
-	if !hasDocument || documentID == "" || r.HasNoFetchers() {
+	if !hasDocument || documentID == "" {
 		subNode, err := resolveSubReferencePath(r.documents[r.GetCurrentScope()], path, "")
 		if err != nil {
 			return nil, "", fmt.Errorf("failed to resolve sub-reference path '%s' in document '%s': %w", path, r.GetCurrentScope(), err)
 		}
 		return subNode, fmt.Sprintf("%s%s", r.GetCurrentScope(), path), nil
 	}
 
-	fetcher, err := r.getFetcherForRef(documentID)
+	resolvedDocId, resolvedPath, err := r.resolveDocumentRef(documentID, path)
 	if err != nil {
-		return nil, "", fmt.Errorf("failed to get document fetcher for document '%s': %w", ref, err)
+		return nil, "", err
 	}
 
-	resolvedDocumentId, err := fetcher.resolveDocumentId(r.GetCurrentScope(), documentID)
-	if resolvedDocumentId == "" {
-		return nil, "", fmt.Errorf("failed to resolve document ID for reference '%s': %w", ref, err)
-	}
-
-	document, documentLoaded := r.documents[resolvedDocumentId]
+	document, documentLoaded := r.documents[resolvedDocId]
 
 	if !documentLoaded {
-		document, err = r.resolveDocument(resolvedDocumentId)
+		document, err = r.resolveDocument(resolvedDocId)
 		if err != nil {
 			return nil, "", fmt.Errorf("failed to resolve document '%s': %w", documentID, err)
 		}
 	}
 
-	subNode, err := resolveSubReferencePath(document, path, "")
+	subNode, err := resolveSubReferencePath(document, resolvedPath, "")
 	if err != nil {
-		return nil, "", fmt.Errorf("failed to resolve sub-reference path '%s' in document '%s': %w", path, documentID, err)
+		return nil, "", fmt.Errorf("failed to resolve sub-reference path '%s' in document '%s': %w", resolvedPath, resolvedDocId, err)
 	}
 
 	// Shift the document scope so base url resolution works correctly
-	r.SetDocumentBeingResolved(resolvedDocumentId)
+	r.SetDocumentBeingResolved(resolvedDocId)
 
-	return subNode, fmt.Sprintf("%s%s", resolvedDocumentId, path), nil
+	return subNode, fmt.Sprintf("%s%s", resolvedDocId, resolvedPath), nil
 }
 
 func (r *documentResolver) addDocument(id string, node *schemaNode) {
@@ -318,6 +327,95 @@ func (r *documentResolver) addDocument(id string, node *schemaNode) {
 	// points into said document
 	r.documents[id] = node
 	r.parsedExternalDocuments[id] = false
+
+	// Collect $id aliases from the new document tree.
+	baseURI := id
+	if node.Id != nil && isAbsoluteURI(*node.Id) {
+		baseURI = *node.Id
+	}
+	for resolvedId, alias := range collectSubSchemaIds(baseURI, id, node) {
+		if _, exists := r.idAliases[resolvedId]; !exists {
+			r.idAliases[resolvedId] = alias
+		}
+	}
+}
+
+// collectSubSchemaIds marshals a schema node to a generic map and recursively
+// walks the entire structure looking for "$id" entries. Each found $id is
+// resolved against its nearest ancestor base URI and recorded as an idAlias.
+func collectSubSchemaIds(baseURI string, documentId string, node *schemaNode) map[string]idAlias {
+	data, err := json.Marshal(node)
+	if err != nil {
+		return nil
+	}
+	var raw map[string]interface{}
+	if err := json.Unmarshal(data, &raw); err != nil {
+		return nil
+	}
+
+	result := make(map[string]idAlias)
+	baseAtPath := map[string]string{"": baseURI}
+
+	WalkSchema(raw, "", func(n map[string]interface{}, path string) {
+		idVal, ok := n["$id"]
+		if !ok {
+			return
+		}
+		idStr, ok := idVal.(string)
+		if !ok || idStr == "" {
+			return
+		}
+		resolved := resolveRelativeURI(nearestBaseURI(baseAtPath, path), idStr)
+		if resolved == "" {
+			return
+		}
+		baseAtPath[path] = resolved
+		if path == "" {
+			return // root $id is the document identity, not an alias
+		}
+		if _, exists := result[resolved]; exists {
+			return
+		}
+		result[resolved] = idAlias{documentId: documentId, path: "#" + path}
+	})
+
+	return result
+}
+
+// nearestBaseURI walks up the JSON pointer path to find the closest ancestor
+// that established a base URI via $id.
+func nearestBaseURI(baseAtPath map[string]string, path string) string {
+	for p := path; ; {
+		if base, ok := baseAtPath[p]; ok {
+			return base
+		}
+		i := strings.LastIndex(p, "/")
+		if i < 0 {
+			break
+		}
+		p = p[:i]
+	}
+	return baseAtPath[""]
+}
+
+// resolveRelativeURI resolves rel against base using standard URL resolution.
+// Returns "" if parsing fails.
+func resolveRelativeURI(base, rel string) string {
+	baseURL, err := url.Parse(base)
+	if err != nil {
+		return ""
+	}
+	relURL, err := url.Parse(rel)
+	if err != nil {
+		return ""
+	}
+	return baseURL.ResolveReference(relURL).String()
+}
+
+// isAbsoluteURI returns true when s looks like an absolute URI (has a scheme).
+func isAbsoluteURI(s string) bool {
+	u, err := url.Parse(s)
+	return err == nil && u.Scheme != ""
 }
 
 // Checks if there are more documents that need to be parsed
@@ -339,10 +437,12 @@ func (r *documentResolver) ParseNextDocument(metadata *parserMetadata) (string,
 		return nextDocumentId, fmt.Errorf("failed to resolve document '%s': %w", nextDocumentId, err)
 	}
 
-	_, err = parseRoot(*documentNode, metadata)
-	// Mark document as parsed
+	// Mark as parsed before parseRoot so that intra-document $id-based
+	// refs resolved during parsing don't re-queue this document.
 	r.parsedExternalDocuments[nextDocumentId] = true
 
+	_, err = parseRoot(*documentNode, metadata)
+
 	if err != nil {
 		return nextDocumentId, fmt.Errorf("failed to parse document '%s': %w", nextDocumentId, err)
 	}
@@ -375,6 +475,49 @@ func (r *documentResolver) resolveDocument(ref string) (*schemaNode, error) {
 	return document, nil
 }
 
+// resolveDocumentRef resolves a document reference to its canonical document
+// ID and path. It checks the $id alias table first (a lightweight map of
+// resolved $id URI → document + JSON pointer), falling back to I/O-based
+// fetchers only if no alias matches.
+func (r *documentResolver) resolveDocumentRef(documentID string, refPath string) (string, string, error) {
+	// Check $id alias table first — resolves bare-name refs like "color"
+	// to the real document + JSON pointer path with no I/O.
+	resolved := resolveRelativeURI(r.GetCurrentScope(), documentID)
+	if resolved != "" {
+		if alias, exists := r.idAliases[resolved]; exists {
+			return alias.documentId, composeJsonPointerPaths(alias.path, refPath), nil
+		}
+	}
+
+	// No alias found — fall back to I/O-based fetchers.
+	if r.HasNoFetchers() {
+		return r.GetCurrentScope(), refPath, nil
+	}
+
+	fetcher, err := r.getFetcherForRef(documentID)
+	if err != nil {
+		return "", "", fmt.Errorf("failed to get document fetcher for document '%s': %w", documentID, err)
+	}
+
+	resolvedDocumentId, err := fetcher.resolveDocumentId(r.GetCurrentScope(), documentID)
+	if err != nil {
+		return "", "", fmt.Errorf("failed to resolve document ID for reference '%s': %w", documentID, err)
+	}
+
+	return resolvedDocumentId, refPath, nil
+}
+
+// composeJsonPointerPaths appends the sub-path from a $ref onto an alias's
+// base path. For example, alias "#/$defs/color" + ref "#/type" = "#/$defs/color/type".
+// If refPath is "#" (root of the aliased resource), the alias path is returned as-is.
+func composeJsonPointerPaths(aliasPath string, refPath string) string {
+	if refPath == "#" || refPath == "" {
+		return aliasPath
+	}
+	// refPath is "#/something" — strip the "#" to get "/something"
+	return aliasPath + refPath[1:]
+}
+
 func (r *documentResolver) getFetcherForRef(documentId string) (documentFetcherInterface, error) {
 	url, err := url.Parse(documentId)
 	if err != nil {

document_test.go

@@ -50,6 +50,20 @@ func TestDeeplyNestedDocumentCluster(t *testing.T) {
 	test.TestJsonSchemaDirWithConfig(t, "test_data/document/file/cluster4_deep_nested", 100, getDocumentChaffConfig(), nil)
 }
 
+func TestDefinitionIdDocumentCluster(t *testing.T) {
+	t.Parallel()
+	test.TestJsonSchemaDirWithConfig(t, "test_data/document/file/cluster5_definition_ids", 100, getDocumentChaffConfig(), nil)
+}
+
+func TestEveryIdPositionCluster(t *testing.T) {
+	t.Parallel()
+	test.TestJsonSchemaDirWithConfig(t, "test_data/document/file/cluster6_every_id_position", 100, getDocumentChaffConfig(), func() *chaff.GeneratorOptions {
+		return &chaff.GeneratorOptions{
+			MaximumOneOfAttempts: 10000,
+		}
+	})
+}
+
 func TestHttpDocuments(t *testing.T) {
 	server, err := http_schema_server.StartTestServer()
 	if err != nil {