Refine Xrefstyle Rules in Disco

disco/baller.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"log/slog"
+	"strings"
 
 	"github.com/project-chip/alchemy/asciidoc"
 	"github.com/project-chip/alchemy/asciidoc/parse"
@@ -133,30 +134,63 @@ func (b *Baller) disco(cxt context.Context, doc *asciidoc.Document) error {
 }
 
 func (b *Baller) discoBallTopLevelSection(dc *discoContext, top *asciidoc.Section, docType matter.DocType) error {
-	if b.options.AddXrefstyle {
-		var xrefstyleEntry *asciidoc.AttributeEntry
-		var topIndex int = -1
-
-		for i, el := range dc.doc.Elements {
-			if el == top {
-				topIndex = i
-				break
-			}
-			if ae, ok := el.(*asciidoc.AttributeEntry); ok && ae.Name == "xrefstyle" {
-				xrefstyleEntry = ae
+	if b.options.XrefStyleOnlyInRoot {
+		// Logic:
+		// 1. For non-root, no xrefstyle at all anywhere in doc.
+		// 2. For root, if it is present somewhere in file, it's ok. If not present, add it after Copyright Notice or before first section.
+
+		if dc.doc == dc.library.Root {
+			found := false
+			for _, el := range dc.doc.Elements {
+				if ae, ok := el.(*asciidoc.AttributeEntry); ok && ae.Name == "xrefstyle" {
+					found = true
+					break
+				}
 			}
-		}
+			if !found {
+				ae := asciidoc.NewAttributeEntry("xrefstyle")
+				ae.Elements = asciidoc.Elements{asciidoc.NewString("basic")}
+
+				copyrightIndex := -1
+				topIndex := -1
+				for i, el := range dc.doc.Elements {
+					if s, ok := el.(*asciidoc.Section); ok {
+						if topIndex == -1 {
+							topIndex = i
+						}
+						name := ""
+						if len(s.Title) > 0 {
+							if str, ok := s.Title[0].(*asciidoc.String); ok {
+								name = str.Value
+							}
+						}
+						if strings.Contains(strings.ToLower(name), "copyright notice") {
+							copyrightIndex = i
+							break
+						}
+					}
+				}
 
-		if xrefstyleEntry == nil {
-			ae := asciidoc.NewAttributeEntry("xrefstyle")
-			ae.Elements = asciidoc.Elements{asciidoc.NewString("basic")}
-			if topIndex != -1 {
-				dc.doc.Elements = append(dc.doc.Elements, nil, nil)
-				copy(dc.doc.Elements[topIndex+2:], dc.doc.Elements[topIndex:])
-				dc.doc.Elements[topIndex] = ae
-				dc.doc.Elements[topIndex+1] = &asciidoc.NewLine{}
-			} else {
-				dc.doc.Elements = append(asciidoc.Elements{ae, &asciidoc.NewLine{}}, dc.doc.Elements...)
+				if copyrightIndex != -1 {
+					dc.doc.Elements = append(dc.doc.Elements[:copyrightIndex+1], append([]asciidoc.Element{&asciidoc.NewLine{}, ae}, dc.doc.Elements[copyrightIndex+1:]...)...)
+				} else if topIndex != -1 {
+					dc.doc.Elements = append(dc.doc.Elements[:topIndex], append([]asciidoc.Element{ae, &asciidoc.NewLine{}}, dc.doc.Elements[topIndex:]...)...)
+				} else {
+					dc.doc.Elements = append(asciidoc.Elements{ae, &asciidoc.NewLine{}}, dc.doc.Elements...)
+				}
+			}
+		} else {
+			for i := 0; i < len(dc.doc.Elements); i++ {
+				el := dc.doc.Elements[i]
+				if ae, ok := el.(*asciidoc.AttributeEntry); ok && ae.Name == "xrefstyle" {
+					dc.doc.Elements = append(dc.doc.Elements[:i], dc.doc.Elements[i+1:]...)
+					i--
+					if i+1 < len(dc.doc.Elements) {
+						if _, ok := dc.doc.Elements[i+1].(*asciidoc.NewLine); ok {
+							dc.doc.Elements = append(dc.doc.Elements[:i+1], dc.doc.Elements[i+2:]...)
+						}
+					}
+				}
 			}
 		}
 	}
@@ -191,3 +225,19 @@ func (b *Baller) discoBallTopLevelSection(dc *discoContext, top *asciidoc.Sectio
 	b.postCleanUpStrings(dc.doc, top)
 	return nil
 }
+
+func (b *Baller) TestHelperDiscoBall(doc *asciidoc.Document, isRoot bool) error {
+	lib := &spec.Library{}
+	if isRoot {
+		lib.Root = doc
+	}
+	dc := &discoContext{
+		Context:            context.Background(),
+		doc:                doc,
+		library:            lib,
+		potentialDataTypes: make(map[string][]*DataTypeEntry),
+	}
+	docType := matter.DocTypeCluster
+	top := parse.FindFirst[*asciidoc.Section](doc, asciidoc.RawReader, doc)
+	return b.discoBallTopLevelSection(dc, top, docType)
+}

disco/option.go

@@ -22,7 +22,7 @@ type DiscoOptions struct {
 	NormalizeAnchors              bool `default:"false" aliases:"normalizeAnchors" help:"rewrite anchors and references without labels" group:"Discoballing:"`
 	RemoveMandatoryFallbacks      bool `default:"true" aliases:"removeMandatoryFallbacks" help:"remove fallback values for mandatory fields" group:"Discoballing:"`
 	RenameSections                bool `default:"false" help:"rename sections to disco-ball standard names" group:"Discoballing:"`
-	AddXrefstyle                  bool `default:"true" aliases:"addXrefstyle" help:"add xrefstyle attribute before the first section" group:"Discoballing:"`
+	XrefStyleOnlyInRoot           bool `default:"true" aliases:"xrefStyleOnlyInRoot" help:"enforce xrefstyle: basic only in root" group:"Discoballing:"`
 }
 
 var DefaultOptions = DiscoOptions{
@@ -45,5 +45,5 @@ var DefaultOptions = DiscoOptions{
 	NormalizeAnchors:              false,
 	RemoveMandatoryFallbacks:      true,
 	RenameSections:                false,
-	AddXrefstyle:                  true,
+	XrefStyleOnlyInRoot:           true,
 }

tests/baller_test.go

@@ -0,0 +1,141 @@
+package tests
+
+import (
+	"strings"
+	"testing"
+
+	"github.com/project-chip/alchemy/asciidoc"
+	"github.com/project-chip/alchemy/disco"
+)
+
+func TestXrefStyleOnlyInRoot(t *testing.T) {
+	// Create a Baller with the option enabled
+	b := disco.NewBaller(nil, disco.DiscoOptions{XrefStyleOnlyInRoot: true})
+
+	// Test Case 1: Root doc without xrefstyle, with Copyright Notice
+	{
+		doc := &asciidoc.Document{
+			Elements: asciidoc.Elements{
+				&asciidoc.Section{
+					Title: asciidoc.Elements{asciidoc.NewString("Copyright Notice")},
+					Level: 1,
+				},
+				&asciidoc.Section{
+					Title: asciidoc.Elements{asciidoc.NewString("First Section")},
+					Level: 1,
+				},
+			},
+		}
+
+		err := b.TestHelperDiscoBall(doc, true)
+		if err != nil {
+			t.Fatalf("unexpected error: %v", err)
+		}
+
+		// Verify added after Copyright Notice
+		found := false
+		for i, el := range doc.Elements {
+			if s, ok := el.(*asciidoc.Section); ok && strings.Contains(s.Title[0].(*asciidoc.String).Value, "Copyright Notice") {
+				// Check next element
+				if i+1 < len(doc.Elements) {
+					if _, ok := doc.Elements[i+1].(*asciidoc.NewLine); ok {
+						if i+2 < len(doc.Elements) {
+							if ae, ok := doc.Elements[i+2].(*asciidoc.AttributeEntry); ok && ae.Name == "xrefstyle" {
+								found = true
+								break
+							}
+						}
+					}
+				}
+			}
+		}
+		if !found {
+			t.Error("expected xrefstyle to be added after Copyright Notice section")
+		}
+	}
+
+	// Test Case 2: Root doc without xrefstyle, without Copyright Notice
+	{
+		doc := &asciidoc.Document{
+			Elements: asciidoc.Elements{
+				&asciidoc.Section{
+					Title: asciidoc.Elements{asciidoc.NewString("First Section")},
+					Level: 1,
+				},
+			},
+		}
+
+		err := b.TestHelperDiscoBall(doc, true)
+		if err != nil {
+			t.Fatalf("unexpected error: %v", err)
+		}
+
+		// Verify added before first section
+		if len(doc.Elements) < 3 {
+			t.Fatalf("expected at least 3 elements, got %d", len(doc.Elements))
+		}
+		if ae, ok := doc.Elements[0].(*asciidoc.AttributeEntry); !ok || ae.Name != "xrefstyle" {
+			t.Error("expected xrefstyle to be added at the beginning")
+		}
+	}
+
+	// Test Case 3: Root doc with existing xrefstyle somewhere
+	{
+		ae := asciidoc.NewAttributeEntry("xrefstyle")
+		ae.Elements = asciidoc.Elements{asciidoc.NewString("full")}
+
+		doc := &asciidoc.Document{
+			Elements: asciidoc.Elements{
+				&asciidoc.Section{
+					Title: asciidoc.Elements{asciidoc.NewString("First Section")},
+					Level: 1,
+				},
+				ae,
+			},
+		}
+
+		err := b.TestHelperDiscoBall(doc, true)
+		if err != nil {
+			t.Fatalf("unexpected error: %v", err)
+		}
+
+		// Verify not changed
+		if s, ok := ae.Elements[0].(*asciidoc.String); !ok || s.Value != "full" {
+			t.Errorf("expected xrefstyle to remain full, got %v", ae.Elements[0])
+		}
+	}
+
+	// Test Case 4: Non-root doc with xrefstyle anywhere (should be removed)
+	{
+		ae1 := asciidoc.NewAttributeEntry("xrefstyle")
+		ae2 := asciidoc.NewAttributeEntry("xrefstyle")
+
+		doc := &asciidoc.Document{
+			Elements: asciidoc.Elements{
+				ae1,
+				&asciidoc.Section{
+					Title: asciidoc.Elements{asciidoc.NewString("First Section")},
+					Level: 1,
+				},
+				ae2,
+			},
+		}
+
+		err := b.TestHelperDiscoBall(doc, false)
+		if err != nil {
+			t.Fatalf("unexpected error: %v", err)
+		}
+
+		// Verify no xrefstyle anywhere
+		found := false
+		for _, el := range doc.Elements {
+			if ae, ok := el.(*asciidoc.AttributeEntry); ok && ae.Name == "xrefstyle" {
+				found = true
+				break
+			}
+		}
+		if found {
+			t.Error("expected all xrefstyle to be removed from non-root document")
+		}
+	}
+}