From d8fe4c9b1d35198056fde43a29f3e9842d489aee Mon Sep 17 00:00:00 2001
From: Viacheslav Poturaev <vearutop@gmail.com>
Date: Sat, 24 Feb 2024 00:14:42 +0100
Subject: [PATCH] Optionally expose embedded structs with allOf

---
 reflect.go      | 23 ++++++++++++++++++++++-
 reflect_test.go | 35 +++++++++++++++++++++++++++++++++++
 2 files changed, 57 insertions(+), 1 deletion(-)

diff --git a/reflect.go b/reflect.go
index f1276e8..30e8ef4 100644
--- a/reflect.go
+++ b/reflect.go
@@ -24,6 +24,7 @@ var (
 	typeOfTextMarshaler   = reflect.TypeOf((*encoding.TextMarshaler)(nil)).Elem()
 	typeOfEmptyInterface  = reflect.TypeOf((*interface{})(nil)).Elem()
 	typeOfSchemaInliner   = reflect.TypeOf((*SchemaInliner)(nil)).Elem()
+	typeOfEmbedReferencer = reflect.TypeOf((*EmbedReferencer)(nil)).Elem()
 )
 
 const (
@@ -47,6 +48,11 @@ type SchemaInliner interface {
 	InlineJSONSchema()
 }
 
+// EmbedReferencer is a marker interface to enable reference to embedded struct type.
+type EmbedReferencer interface {
+	ReferEmbedded()
+}
+
 // IgnoreTypeName instructs reflector to keep original type name during mapping.
 func (s Schema) IgnoreTypeName() {}
 
@@ -261,6 +267,10 @@ func checkSchemaSetup(params InterceptSchemaParams) (bool, error) {
 //		ProcessWithoutTags
 //		SkipEmbeddedMapsSlices
 //		SkipUnsupportedProperties
+//
+// Fields from embedded structures are processed as if they were defined in the root structure.
+// Alternatively, if embedded structure has a field tag `refer:"true"` or implements EmbedReferencer,
+// its reference will be added to `allOf` of the parent schema.
 func (r *Reflector) Reflect(i interface{}, options ...func(rc *ReflectContext)) (Schema, error) {
 	rc := ReflectContext{}
 	rc.Context = context.Background()
@@ -914,7 +924,18 @@ func (r *Reflector) walkProperties(v reflect.Value, parent *Schema, rc *ReflectC
 
 		if tag == "" && field.Anonymous &&
 			(field.Type.Kind() == reflect.Struct || deepIndirect.Kind() == reflect.Struct) {
-			if err := r.walkProperties(values[i], parent, rc); err != nil {
+
+			forceReference := (field.Type.Implements(typeOfEmbedReferencer) && field.Tag.Get("refer") == "") ||
+				field.Tag.Get("refer") == "true"
+
+			if forceReference {
+				rc.Path = append(rc.Path, "")
+				if s, err := r.reflect(values[i].Interface(), rc, false, parent); err != nil {
+					return err
+				} else {
+					parent.AllOf = append(parent.AllOf, s.ToSchemaOrBool())
+				}
+			} else if err := r.walkProperties(values[i], parent, rc); err != nil {
 				return err
 			}
 
diff --git a/reflect_test.go b/reflect_test.go
index 000013d..52fb3d7 100644
--- a/reflect_test.go
+++ b/reflect_test.go
@@ -1827,3 +1827,38 @@ func TestReflector_Reflect_multipleTags(t *testing.T) {
 	  "type":"object"
 	}`, s)
 }
+
+func TestReflector_Reflect_embedded(t *testing.T) {
+	type A struct {
+		FieldA int `json:"field_a"`
+	}
+
+	type C struct {
+		jsonschema.EmbedReferencer
+		FieldC int `json:"field_c"`
+	}
+
+	type B struct {
+		A      `refer:"true"`
+		FieldB int `json:"field_b"`
+		C
+	}
+
+	r := jsonschema.Reflector{}
+
+	s, err := r.Reflect(B{}, jsonschema.InterceptProp(func(params jsonschema.InterceptPropParams) error {
+		return nil
+	}))
+	require.NoError(t, err)
+	assertjson.EqMarshal(t, `{
+	  "definitions":{
+		"JsonschemaGoTestA":{"properties":{"field_a":{"type":"integer"}},"type":"object"},
+		"JsonschemaGoTestC":{"properties":{"field_c":{"type":"integer"}},"type":"object"}
+	  },
+	  "properties":{"field_b":{"type":"integer"}},"type":"object",
+	  "allOf":[
+		{"$ref":"#/definitions/JsonschemaGoTestA"},
+		{"$ref":"#/definitions/JsonschemaGoTestC"}
+	  ]
+	}`, s)
+}