From df415ddb689023d8234a69823ad213d01c54541a Mon Sep 17 00:00:00 2001
From: Santhosh Kumar Tekuri <santhosh.tekuri@gmail.com>
Date: Wed, 24 Apr 2024 21:50:53 +0530
Subject: [PATCH] separate $schema load from root

---
 draft.go  | 34 ++++++++++++++++++++++++++++++++++
 loader.go | 55 +++++++++++++++++++++++++++++++++++++++++++++++++++++++
 root.go   |  7 -------
 roots.go  | 51 +++++++--------------------------------------------
 4 files changed, 96 insertions(+), 51 deletions(-)

diff --git a/draft.go b/draft.go
index 957cbc9..cc9e109 100644
--- a/draft.go
+++ b/draft.go
@@ -2,6 +2,7 @@ package jsonschema
 
 import (
 	"fmt"
+	"slices"
 	"strings"
 )
 
@@ -222,6 +223,39 @@ func (d *Draft) getID(obj map[string]any) string {
 	return id
 }
 
+func (d *Draft) getVocabs(url url, doc any) ([]string, error) {
+	if d.version < 2019 {
+		return nil, nil
+	}
+	obj, ok := doc.(map[string]any)
+	if !ok {
+		return nil, nil
+	}
+	v, ok := obj["$vocabulary"]
+	if !ok {
+		return nil, nil
+	}
+	obj, ok = v.(map[string]any)
+	if !ok {
+		return nil, nil
+	}
+
+	var vocabs []string
+	for vocab, reqd := range obj {
+		if reqd, ok := reqd.(bool); !ok || !reqd {
+			continue
+		}
+		name, ok := strings.CutPrefix(vocab, d.vocabPrefix)
+		if !ok {
+			return nil, &UnsupportedVocabularyError{url.String(), vocab}
+		}
+		if !slices.Contains(vocabs, name) {
+			vocabs = append(vocabs, name)
+		}
+	}
+	return vocabs, nil
+}
+
 func (d *Draft) validate(up urlPtr, v any, regexpEngine RegexpEngine) error {
 	err := d.sch.validate(v, regexpEngine)
 	if err != nil {
diff --git a/loader.go b/loader.go
index 946f426..3e39089 100644
--- a/loader.go
+++ b/loader.go
@@ -134,6 +134,61 @@ func (l *defaultLoader) load(url url) (any, error) {
 	return doc, nil
 }
 
+func (l *defaultLoader) getDraft(up urlPtr, doc any, defaultDraft *Draft, cycle map[url]struct{}) (*Draft, error) {
+	obj, ok := doc.(map[string]any)
+	if !ok {
+		return defaultDraft, nil
+	}
+	sch, ok := strVal(obj, "$schema")
+	if !ok {
+		return defaultDraft, nil
+	}
+	if draft := draftFromURL(sch); draft != nil {
+		return draft, nil
+	}
+	sch, _ = split(sch)
+	if _, err := gourl.Parse(sch); err != nil {
+		return nil, &InvalidMetaSchemaURLError{up.String(), err}
+	}
+	schUrl := url(sch)
+	if up.ptr.isEmpty() && schUrl == up.url {
+		return nil, &UnsupportedDraftError{schUrl.String()}
+	}
+	if _, ok := cycle[schUrl]; ok {
+		return nil, &MetaSchemaCycleError{schUrl.String()}
+	}
+	cycle[schUrl] = struct{}{}
+	doc, err := l.load(schUrl)
+	if err != nil {
+		return nil, err
+	}
+	return l.getDraft(urlPtr{schUrl, ""}, doc, defaultDraft, cycle)
+}
+
+func (l *defaultLoader) getMetaVocabs(doc any, draft *Draft) ([]string, error) {
+	obj, ok := doc.(map[string]any)
+	if !ok {
+		return nil, nil
+	}
+	sch, ok := strVal(obj, "$schema")
+	if !ok {
+		return nil, nil
+	}
+	if draft := draftFromURL(sch); draft != nil {
+		return nil, nil
+	}
+	sch, _ = split(sch)
+	if _, err := gourl.Parse(sch); err != nil {
+		return nil, &ParseURLError{sch, err}
+	}
+	schUrl := url(sch)
+	doc, err := l.load(schUrl)
+	if err != nil {
+		return nil, err
+	}
+	return draft.getVocabs(schUrl, doc)
+}
+
 // --
 
 type LoadURLError struct {
diff --git a/root.go b/root.go
index 9147a1a..d4b2d18 100644
--- a/root.go
+++ b/root.go
@@ -327,13 +327,6 @@ func newResource(ptr jsonPointer, id url) *resource {
 
 //--
 
-type meta struct {
-	draft  *Draft
-	vocabs []string
-}
-
-// --
-
 type UnsupportedVocabularyError struct {
 	URL        string
 	Vocabulary string
diff --git a/roots.go b/roots.go
index 64bcaa6..f34627e 100644
--- a/roots.go
+++ b/roots.go
@@ -2,7 +2,6 @@ package jsonschema
 
 import (
 	"fmt"
-	gourl "net/url"
 	"strings"
 )
 
@@ -35,60 +34,24 @@ func (rr *roots) orLoad(u url) (*root, error) {
 	if err != nil {
 		return nil, err
 	}
-	return rr.addRoot(u, doc, make(map[url]struct{}))
+	return rr.addRoot(u, doc)
 }
 
-func (rr *roots) getMeta(up urlPtr, doc any, cycle map[url]struct{}) (meta, error) {
-	obj, ok := doc.(map[string]any)
-	if !ok {
-		return meta{rr.defaultDraft, nil}, nil
-	}
-	sch, ok := strVal(obj, "$schema")
-	if !ok {
-		return meta{rr.defaultDraft, nil}, nil
-	}
-	if draft := draftFromURL(sch); draft != nil {
-		return meta{draft, nil}, nil
-	}
-	sch, _ = split(sch)
-	if _, err := gourl.Parse(sch); err != nil {
-		return meta{}, &InvalidMetaSchemaURLError{up.String(), err}
-	}
-	schUrl := url(sch)
-	if r, ok := rr.roots[schUrl]; ok {
-		vocabs, err := r.getReqdVocabs()
-		return meta{r.draft, vocabs}, err
-	}
-	if schUrl == up.url {
-		return meta{}, &UnsupportedDraftError{schUrl.String()}
-	}
-	if _, ok := cycle[schUrl]; ok {
-		return meta{}, &MetaSchemaCycleError{schUrl.String()}
-	}
-	cycle[schUrl] = struct{}{}
-	doc, err := rr.loader.load(schUrl)
-	if err != nil {
-		return meta{}, err
-	}
-	r, err := rr.addRoot(schUrl, doc, cycle)
+func (rr *roots) addRoot(u url, doc any) (*root, error) {
+	draft, err := rr.loader.getDraft(urlPtr{u, ""}, doc, rr.defaultDraft, map[url]struct{}{})
 	if err != nil {
-		return meta{}, err
+		return nil, err
 	}
-	vocabs, err := r.getReqdVocabs()
-	return meta{r.draft, vocabs}, err
-}
-
-func (rr *roots) addRoot(u url, doc any, cycle map[url]struct{}) (*root, error) {
-	meta, err := rr.getMeta(urlPtr{u, ""}, doc, cycle)
+	vocabs, err := rr.loader.getMetaVocabs(doc, draft)
 	if err != nil {
 		return nil, err
 	}
 	r := &root{
 		url:                 u,
 		doc:                 doc,
-		draft:               meta.draft,
+		draft:               draft,
 		resources:           map[jsonPointer]*resource{},
-		metaVocabs:          meta.vocabs,
+		metaVocabs:          vocabs,
 		subschemasProcessed: map[jsonPointer]struct{}{},
 	}
 	if err := r.collectResources(doc, u, ""); err != nil {