Merge adjacent formatting entities of same kind

2 files changed, 131 insertions, 2 deletions

diff --git a/telegram/formatter/formatter.go b/telegram/formatter/formatter.go
index 999d659..5ea5cdb 100644
--- a/telegram/formatter/formatter.go
+++ b/telegram/formatter/formatter.go
@@ -22,7 +22,7 @@ var boldRunesXEP0393 = []rune("*")
 var italicRunes = []rune("_")
 var codeRunes = []rune("\n```\n")
 
-// rebalance pumps all the values at given offset to current stack (growing
+// rebalance pumps all the values until the given offset to current stack (growing
 // from start) from given stack (growing from end); should be called
 // before any insertions to the current stack at the given offset
 func (s InsertionStack) rebalance(s2 InsertionStack, offset int32) (InsertionStack, InsertionStack) {
@@ -66,6 +66,54 @@ func SortEntities(entities []*client.TextEntity) []*client.TextEntity {
 	return sortedEntities
 }
 
+// MergeAdjacentEntities merges entities of a same kind
+func MergeAdjacentEntities(entities []*client.TextEntity) []*client.TextEntity {
+	mergedEntities := make([]*client.TextEntity, 0, len(entities))
+	excludedIndices := make(map[int]bool)
+
+	for i, entity := range entities {
+		if excludedIndices[i] {
+			continue
+		}
+
+		typ := entity.Type.TextEntityTypeType()
+		start := entity.Offset
+		end := start + entity.Length
+		ei := make(map[int]bool)
+
+		// collect continuations
+		for j, entity2 := range entities[i+1:] {
+			if entity2.Type.TextEntityTypeType() == typ && entity2.Offset == end {
+				end += entity2.Length
+				ei[j+i+1] = true
+			}
+		}
+
+		// check for intersections with other entities
+		var isIntersecting bool
+		if len(ei) > 0 {
+			for _, entity2 := range entities {
+				entity2End := entity2.Offset + entity2.Length
+				if (entity2.Offset < start && entity2End > start && entity2End < end) ||
+					(entity2.Offset > start && entity2.Offset < end && entity2End > end) {
+					isIntersecting = true
+					break
+				}
+			}
+		}
+
+		if !isIntersecting {
+			entity.Length = end - start
+			for j := range ei {
+				excludedIndices[j] = true
+			}
+		}
+		mergedEntities = append(mergedEntities, entity)
+	}
+
+	return mergedEntities
+}
+
 func markupBraces(entity *client.TextEntity, lbrace, rbrace []rune) (*Insertion, *Insertion) {
 	return &Insertion{
 			Offset: entity.Offset,
@@ -132,12 +180,14 @@ func Format(
 		return sourceText
 	}
 
+	mergedEntities := SortEntities(MergeAdjacentEntities(SortEntities(entities)))
+
 	startStack := make(InsertionStack, 0, len(sourceText))
 	endStack := make(InsertionStack, 0, len(sourceText))
 
 	// convert entities to a stack of brackets
 	var maxEndOffset int32
-	for _, entity := range entities {
+	for _, entity := range mergedEntities {
 		log.Debugf("%#v", entity)
 		if entity.Length <= 0 {
 			continue
diff --git a/telegram/formatter/formatter_test.go b/telegram/formatter/formatter_test.go
index 8453851..5ac6262 100644
--- a/telegram/formatter/formatter_test.go
+++ b/telegram/formatter/formatter_test.go
@@ -269,3 +269,82 @@ func TestFormattingXEP0393AdjacentAndNested(t *testing.T) {
 		t.Errorf("Wrong adjacent&nested formatting: %v", markup)
 	}
 }
+
+func TestFormattingXEP0393AdjacentItalicBoldItalic(t *testing.T) {
+	markup := Format("раса двуногих крысолюдей, которую так редко замечают, что многие отрицают само их существование", []*client.TextEntity{
+		&client.TextEntity{
+			Offset: 0,
+			Length: 26,
+			Type:   &client.TextEntityTypeItalic{},
+		},
+		&client.TextEntity{
+			Offset: 26,
+			Length: 69,
+			Type:   &client.TextEntityTypeBold{},
+		},
+		&client.TextEntity{
+			Offset: 26,
+			Length: 69,
+			Type:   &client.TextEntityTypeItalic{},
+		},
+	}, EntityToXEP0393)
+	if markup != "_раса двуногих крысолюдей, *которую так редко замечают, что многие отрицают само их существование*_" {
+		t.Errorf("Wrong adjacent italic/bold-italic formatting: %v", markup)
+	}
+}
+
+func TestFormattingXEP0393MultipleAdjacent(t *testing.T) {
+	markup := Format("abcde", []*client.TextEntity{
+		&client.TextEntity{
+			Offset: 1,
+			Length: 1,
+			Type:   &client.TextEntityTypeBold{},
+		},
+		&client.TextEntity{
+			Offset: 2,
+			Length: 1,
+			Type:   &client.TextEntityTypeBold{},
+		},
+		&client.TextEntity{
+			Offset: 3,
+			Length: 1,
+			Type:   &client.TextEntityTypeBold{},
+		},
+		&client.TextEntity{
+			Offset: 4,
+			Length: 1,
+			Type:   &client.TextEntityTypeItalic{},
+		},
+	}, EntityToXEP0393)
+	if markup != "a*bcd*_e_" {
+		t.Errorf("Wrong multiple adjacent formatting: %v", markup)
+	}
+}
+
+func TestFormattingXEP0393Intersecting(t *testing.T) {
+	markup := Format("abcde", []*client.TextEntity{
+		&client.TextEntity{
+			Offset: 1,
+			Length: 1,
+			Type:   &client.TextEntityTypeBold{},
+		},
+		&client.TextEntity{
+			Offset: 2,
+			Length: 3,
+			Type:   &client.TextEntityTypeItalic{},
+		},
+		&client.TextEntity{
+			Offset: 2,
+			Length: 1,
+			Type:   &client.TextEntityTypeBold{},
+		},
+		&client.TextEntity{
+			Offset: 3,
+			Length: 1,
+			Type:   &client.TextEntityTypeBold{},
+		},
+	}, EntityToXEP0393)
+	if markup != "a*b*_*cd*e_" {
+		t.Errorf("Wrong intersecting formatting: %v", markup)
+	}
+}