Skip to content

Commit

Permalink
Merge pull request #205 from mailgun/maxim/develop
Browse files Browse the repository at this point in the history 
Fix WrapWords excessive memory utilization
  • Loading branch information
@olekukonko
olekukonko authored Sep 25, 2023
2 parents f6b4e4a + f1da30c commit df64c4b
Show file tree
Hide file tree
Showing 5 changed files with 255 additions and 23 deletions.
2 changes: 1 addition & 1 deletion table_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -74,7 +74,7 @@ func ExampleTable() {
// Output: *================================*================================*===============================*==========*
// | NAME | SIGN | RATING | |
// *================================*================================*===============================*==========*
// | Learn East has computers | Some Data | Another Data |
// | Learn East has computers | Some Data | Another Data |
// | with adapted keyboards with | | |
// | enlarged print etc | | |
// | Instead of lining up the | the way across, he splits the | Like most ergonomic keyboards | See Data |
Expand Down
27 changes: 27 additions & 0 deletions testdata/long-text-wrapped.txt
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,27 @@
Я к вам пишу — чего же боле? Что я могу еще сказать? Теперь, я знаю, в вашей
воле Меня презреньем наказать. Но вы, к моей несчастной доле Хоть каплю жалости
храня, Вы не оставите меня. Сначала я молчать хотела; Поверьте: моего стыда Вы
не узнали б никогда, Когда б надежду я имела Хоть редко, хоть в неделю раз В
деревне нашей видеть вас, Чтоб только слышать ваши речи, Вам слово молвить, и
потом Все думать, думать об одном И день и ночь до новой встречи. Но, говорят,
вы нелюдим; В глуши, в деревне всё вам скучно, А мы… ничем мы не блестим,
Хоть вам и рады простодушно. Зачем вы посетили нас? В глуши забытого селенья
Я никогда не знала б вас, Не знала б горького мученья. Души неопытной волненья
Смирив со временем (как знать?), По сердцу я нашла бы друга, Была бы верная
супруга И добродетельная мать. Другой!.. Нет, никому на свете Не отдала бы
сердца я! То в вышнем суждено совете… То воля неба: я твоя; Вся жизнь моя была
залогом Свиданья верного с тобой; Я знаю, ты мне послан богом, До гроба ты
хранитель мой… Ты в сновиденьях мне являлся, Незримый, ты мне был уж мил, Твой
чудный взгляд меня томил, В душе твой голос раздавался Давно… нет, это был не
сон! Ты чуть вошел, я вмиг узнала, Вся обомлела, запылала И в мыслях молвила:
вот он! Не правда ль? Я тебя слыхала: Ты говорил со мной в тиши, Когда я бедным
помогала Или молитвой услаждала Тоску волнуемой души? И в это самое мгновенье Не
ты ли, милое виденье, В прозрачной темноте мелькнул, Приникнул тихо к изголовью?
Не ты ль, с отрадой и любовью, Слова надежды мне шепнул? Кто ты, мой ангел ли
хранитель, Или коварный искуситель: Мои сомненья разреши. Быть может, это все
пустое, Обман неопытной души! И суждено совсем иное… Но так и быть! Судьбу мою
Отныне я тебе вручаю, Перед тобою слезы лью, Твоей защиты умоляю… Вообрази: я
здесь одна, Никто меня не понимает, Рассудок мой изнемогает, И молча гибнуть я
должна. Я жду тебя: единым взором Надежды сердца оживи Иль сон тяжелый перерви,
Увы, заслуженным укором! Кончаю! Страшно перечесть… Стыдом и страхом замираю… Но
мне порукой ваша честь, И смело ей себя вверяю…
82 changes: 82 additions & 0 deletions testdata/long-text.txt
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,82 @@
Я к вам пишу — чего же боле?
Что я могу еще сказать?
Теперь, я знаю, в вашей воле
Меня презреньем наказать.
Но вы, к моей несчастной доле
Хоть каплю жалости храня,
Вы не оставите меня.
Сначала я молчать хотела;
Поверьте: моего стыда
Вы не узнали б никогда,
Когда б надежду я имела
Хоть редко, хоть в неделю раз
В деревне нашей видеть вас,
Чтоб только слышать ваши речи,
Вам слово молвить, и потом
Все думать, думать об одном
И день и ночь до новой встречи.
Но, говорят, вы нелюдим;
В глуши, в деревне всё вам скучно,
А мы… ничем мы не блестим,
Хоть вам и рады простодушно.

Зачем вы посетили нас?
В глуши забытого селенья
Я никогда не знала б вас,
Не знала б горького мученья.
Души неопытной волненья
Смирив со временем (как знать?),
По сердцу я нашла бы друга,
Была бы верная супруга
И добродетельная мать.

Другой!.. Нет, никому на свете
Не отдала бы сердца я!
То в вышнем суждено совете…
То воля неба: я твоя;
Вся жизнь моя была залогом
Свиданья верного с тобой;
Я знаю, ты мне послан богом,
До гроба ты хранитель мой…
Ты в сновиденьях мне являлся,
Незримый, ты мне был уж мил,
Твой чудный взгляд меня томил,
В душе твой голос раздавался
Давно… нет, это был не сон!
Ты чуть вошел, я вмиг узнала,
Вся обомлела, запылала
И в мыслях молвила: вот он!
Не правда ль? Я тебя слыхала:
Ты говорил со мной в тиши,
Когда я бедным помогала
Или молитвой услаждала
Тоску волнуемой души?
И в это самое мгновенье
Не ты ли, милое виденье,
В прозрачной темноте мелькнул,
Приникнул тихо к изголовью?
Не ты ль, с отрадой и любовью,
Слова надежды мне шепнул?
Кто ты, мой ангел ли хранитель,
Или коварный искуситель:
Мои сомненья разреши.
Быть может, это все пустое,
Обман неопытной души!
И суждено совсем иное…
Но так и быть! Судьбу мою
Отныне я тебе вручаю,
Перед тобою слезы лью,
Твоей защиты умоляю…
Вообрази: я здесь одна,
Никто меня не понимает,
Рассудок мой изнемогает,
И молча гибнуть я должна.
Я жду тебя: единым взором
Надежды сердца оживи
Иль сон тяжелый перерви,
Увы, заслуженным укором!

Кончаю! Страшно перечесть…
Стыдом и страхом замираю…
Но мне порукой ваша честь,
И смело ей себя вверяю…
80 changes: 58 additions & 22 deletions wrap.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,21 +10,28 @@ package tablewriter
import (
"math"
"strings"
"unicode"

"github.com/mattn/go-runewidth"
)

var (
const (
nl = "\n"
sp = " "
)

const defaultPenalty = 1e5

// WrapString Wrap wraps s into a paragraph of lines of length lim, with minimal
// WrapString wraps s into a paragraph of lines of length lim, with minimal
// raggedness.
func WrapString(s string, lim int) ([]string, int) {
words := strings.Split(strings.Replace(s, nl, sp, -1), sp)
if s == sp {
return []string{sp}, lim
}
words := splitWords(s)
if len(words) == 0 {
return []string{""}, lim
}
var lines []string
max := 0
for _, v := range words {
Expand All @@ -39,6 +46,29 @@ func WrapString(s string, lim int) ([]string, int) {
return lines, lim
}

func splitWords(s string) []string {
words := make([]string, 0, len(s)/5)
var wordBegin int
wordPending := false
for i, c := range s {
if unicode.IsSpace(c) {
if wordPending {
words = append(words, s[wordBegin:i])
wordPending = false
}
continue
}
if !wordPending {
wordBegin = i
wordPending = true
}
}
if wordPending {
words = append(words, s[wordBegin:])
}
return words
}

// WrapWords is the low-level line-breaking algorithm, useful if you need more
// control over the details of the text wrapping process. For most uses,
// WrapString will be sufficient and more convenient.
Expand All @@ -52,35 +82,41 @@ func WrapString(s string, lim int) ([]string, int) {
// added to the error.
func WrapWords(words []string, spc, lim, pen int) [][]string {
n := len(words)

length := make([][]int, n)
if n == 0 {
return nil
}
lengths := make([]int, n)
for i := 0; i < n; i++ {
length[i] = make([]int, n)
length[i][i] = runewidth.StringWidth(words[i])
for j := i + 1; j < n; j++ {
length[i][j] = length[i][j-1] + spc + runewidth.StringWidth(words[j])
}
lengths[i] = runewidth.StringWidth(words[i])
}
nbrk := make([]int, n)
cost := make([]int, n)
for i := range cost {
cost[i] = math.MaxInt32
}
remainderLen := lengths[n-1]
for i := n - 1; i >= 0; i-- {
if length[i][n-1] <= lim {
if i < n-1 {
remainderLen += spc + lengths[i]
}
if remainderLen <= lim {
cost[i] = 0
nbrk[i] = n
} else {
for j := i + 1; j < n; j++ {
d := lim - length[i][j-1]
c := d*d + cost[j]
if length[i][j-1] > lim {
c += pen // too-long lines get a worse penalty
}
if c < cost[i] {
cost[i] = c
nbrk[i] = j
}
continue
}
phraseLen := lengths[i]
for j := i + 1; j < n; j++ {
if j > i+1 {
phraseLen += spc + lengths[j-1]
}
d := lim - phraseLen
c := d*d + cost[j]
if phraseLen > lim {
c += pen // too-long lines get a worse penalty
}
if c < cost[i] {
cost[i] = c
nbrk[i] = j
}
}
}
Expand Down
87 changes: 87 additions & 0 deletions wrap_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,9 @@
package tablewriter

import (
"os"
"reflect"
"runtime"
"strings"
"testing"

Expand Down Expand Up @@ -56,3 +59,87 @@ func TestDisplayWidth(t *testing.T) {
input = "\033[43;30m" + input + "\033[00m"
checkEqual(t, DisplayWidth(input), want)
}

// WrapString was extremely memory greedy, it performed insane number of
// allocations for what it was doing. See BenchmarkWrapString for details.
func TestWrapStringAllocation(t *testing.T) {
originalTextBytes, err := os.ReadFile("testdata/long-text.txt")
if err != nil {
t.Fatal(err)
}
originalText := string(originalTextBytes)

wantWrappedBytes, err := os.ReadFile("testdata/long-text-wrapped.txt")
if err != nil {
t.Fatal(err)
}
wantWrappedText := string(wantWrappedBytes)

var ms runtime.MemStats
runtime.ReadMemStats(&ms)
heapAllocBefore := int64(ms.HeapAlloc / 1024 / 1024)

// When
gotLines, gotLim := WrapString(originalText, 80)

// Then
wantLim := 80
if gotLim != wantLim {
t.Errorf("Invalid limit: want=%d, got=%d", wantLim, gotLim)
}

gotWrappedText := strings.Join(gotLines, "\n")
if gotWrappedText != wantWrappedText {
t.Errorf("Invalid lines: want=\n%s\n got=\n%s", wantWrappedText, gotWrappedText)
}

runtime.ReadMemStats(&ms)
heapAllocAfter := int64(ms.HeapAlloc / 1024 / 1024)
heapAllocDelta := heapAllocAfter - heapAllocBefore
if heapAllocDelta > 1 {
t.Fatalf("heap allocation should not be greater than 1Mb, got=%dMb", heapAllocDelta)
}
}

// Before optimization:
// BenchmarkWrapString-16 1 2490331031 ns/op 2535184104 B/op 50905550 allocs/op
// After optimization:
// BenchmarkWrapString-16 1652 658098 ns/op 230223 B/op 5176 allocs/op
func BenchmarkWrapString(b *testing.B) {
d, err := os.ReadFile("testdata/long-text.txt")
if err != nil {
b.Fatal(err)
}
for i := 0; i < b.N; i++ {
WrapString(string(d), 128)
}
}

func TestSplitWords(t *testing.T) {
for _, tt := range []struct {
in string
out []string
}{{
in: "",
out: []string{},
}, {
in: "a",
out: []string{"a"},
}, {
in: "a b",
out: []string{"a", "b"},
}, {
in: " a b ",
out: []string{"a", "b"},
}, {
in: "\r\na\t\t \r\t b\r\n ",
out: []string{"a", "b"},
}} {
t.Run(tt.in, func(t *testing.T) {
got := splitWords(tt.in)
if !reflect.DeepEqual(tt.out, got) {
t.Errorf("want=%s, got=%s", tt.out, got)
}
})
}
}

0 comments on commit df64c4b

Please sign in to comment.