From f1da30c25846bd64e227af7e777c0f8f0c029ca2 Mon Sep 17 00:00:00 2001 From: Maxim Vladimirskiy Date: Thu, 25 Aug 2022 15:15:06 +0300 Subject: [PATCH] Fix WrapWords excessive memory utilization --- table_test.go | 2 +- testdata/long-text-wrapped.txt | 27 +++++++++++ testdata/long-text.txt | 82 ++++++++++++++++++++++++++++++++ wrap.go | 80 ++++++++++++++++++++++--------- wrap_test.go | 87 ++++++++++++++++++++++++++++++++++ 5 files changed, 255 insertions(+), 23 deletions(-) create mode 100644 testdata/long-text-wrapped.txt create mode 100644 testdata/long-text.txt diff --git a/table_test.go b/table_test.go index ec3af2a..f100731 100644 --- a/table_test.go +++ b/table_test.go @@ -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 | diff --git a/testdata/long-text-wrapped.txt b/testdata/long-text-wrapped.txt new file mode 100644 index 0000000..7ac2e8c --- /dev/null +++ b/testdata/long-text-wrapped.txt @@ -0,0 +1,27 @@ +Я к вам пишу — чего же боле? Что я могу еще сказать? Теперь, я знаю, в вашей +воле Меня презреньем наказать. Но вы, к моей несчастной доле Хоть каплю жалости +храня, Вы не оставите меня. Сначала я молчать хотела; Поверьте: моего стыда Вы +не узнали б никогда, Когда б надежду я имела Хоть редко, хоть в неделю раз В +деревне нашей видеть вас, Чтоб только слышать ваши речи, Вам слово молвить, и +потом Все думать, думать об одном И день и ночь до новой встречи. Но, говорят, +вы нелюдим; В глуши, в деревне всё вам скучно, А мы… ничем мы не блестим, +Хоть вам и рады простодушно. Зачем вы посетили нас? В глуши забытого селенья +Я никогда не знала б вас, Не знала б горького мученья. Души неопытной волненья +Смирив со временем (как знать?), По сердцу я нашла бы друга, Была бы верная +супруга И добродетельная мать. Другой!.. Нет, никому на свете Не отдала бы +сердца я! То в вышнем суждено совете… То воля неба: я твоя; Вся жизнь моя была +залогом Свиданья верного с тобой; Я знаю, ты мне послан богом, До гроба ты +хранитель мой… Ты в сновиденьях мне являлся, Незримый, ты мне был уж мил, Твой +чудный взгляд меня томил, В душе твой голос раздавался Давно… нет, это был не +сон! Ты чуть вошел, я вмиг узнала, Вся обомлела, запылала И в мыслях молвила: +вот он! Не правда ль? Я тебя слыхала: Ты говорил со мной в тиши, Когда я бедным +помогала Или молитвой услаждала Тоску волнуемой души? И в это самое мгновенье Не +ты ли, милое виденье, В прозрачной темноте мелькнул, Приникнул тихо к изголовью? +Не ты ль, с отрадой и любовью, Слова надежды мне шепнул? Кто ты, мой ангел ли +хранитель, Или коварный искуситель: Мои сомненья разреши. Быть может, это все +пустое, Обман неопытной души! И суждено совсем иное… Но так и быть! Судьбу мою +Отныне я тебе вручаю, Перед тобою слезы лью, Твоей защиты умоляю… Вообрази: я +здесь одна, Никто меня не понимает, Рассудок мой изнемогает, И молча гибнуть я +должна. Я жду тебя: единым взором Надежды сердца оживи Иль сон тяжелый перерви, +Увы, заслуженным укором! Кончаю! Страшно перечесть… Стыдом и страхом замираю… Но +мне порукой ваша честь, И смело ей себя вверяю… \ No newline at end of file diff --git a/testdata/long-text.txt b/testdata/long-text.txt new file mode 100644 index 0000000..5a71cc3 --- /dev/null +++ b/testdata/long-text.txt @@ -0,0 +1,82 @@ + Я к вам пишу — чего же боле? + Что я могу еще сказать? + Теперь, я знаю, в вашей воле + Меня презреньем наказать. + Но вы, к моей несчастной доле + Хоть каплю жалости храня, + Вы не оставите меня. + Сначала я молчать хотела; + Поверьте: моего стыда + Вы не узнали б никогда, + Когда б надежду я имела + Хоть редко, хоть в неделю раз + В деревне нашей видеть вас, + Чтоб только слышать ваши речи, + Вам слово молвить, и потом + Все думать, думать об одном + И день и ночь до новой встречи. + Но, говорят, вы нелюдим; + В глуши, в деревне всё вам скучно, + А мы… ничем мы не блестим, + Хоть вам и рады простодушно. + + Зачем вы посетили нас? + В глуши забытого селенья + Я никогда не знала б вас, + Не знала б горького мученья. + Души неопытной волненья + Смирив со временем (как знать?), + По сердцу я нашла бы друга, + Была бы верная супруга + И добродетельная мать. + + Другой!.. Нет, никому на свете + Не отдала бы сердца я! + То в вышнем суждено совете… + То воля неба: я твоя; + Вся жизнь моя была залогом + Свиданья верного с тобой; + Я знаю, ты мне послан богом, + До гроба ты хранитель мой… + Ты в сновиденьях мне являлся, + Незримый, ты мне был уж мил, + Твой чудный взгляд меня томил, + В душе твой голос раздавался + Давно… нет, это был не сон! + Ты чуть вошел, я вмиг узнала, + Вся обомлела, запылала + И в мыслях молвила: вот он! + Не правда ль? Я тебя слыхала: + Ты говорил со мной в тиши, + Когда я бедным помогала + Или молитвой услаждала + Тоску волнуемой души? + И в это самое мгновенье + Не ты ли, милое виденье, + В прозрачной темноте мелькнул, + Приникнул тихо к изголовью? + Не ты ль, с отрадой и любовью, + Слова надежды мне шепнул? + Кто ты, мой ангел ли хранитель, + Или коварный искуситель: + Мои сомненья разреши. + Быть может, это все пустое, + Обман неопытной души! + И суждено совсем иное… + Но так и быть! Судьбу мою + Отныне я тебе вручаю, + Перед тобою слезы лью, + Твоей защиты умоляю… + Вообрази: я здесь одна, + Никто меня не понимает, + Рассудок мой изнемогает, + И молча гибнуть я должна. + Я жду тебя: единым взором + Надежды сердца оживи + Иль сон тяжелый перерви, + Увы, заслуженным укором! + + Кончаю! Страшно перечесть… + Стыдом и страхом замираю… + Но мне порукой ваша честь, + И смело ей себя вверяю… diff --git a/wrap.go b/wrap.go index a5a24da..4a2bc61 100644 --- a/wrap.go +++ b/wrap.go @@ -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 { @@ -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. @@ -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 } } } diff --git a/wrap_test.go b/wrap_test.go index a03f9fc..c2af5fa 100644 --- a/wrap_test.go +++ b/wrap_test.go @@ -8,6 +8,9 @@ package tablewriter import ( + "os" + "reflect" + "runtime" "strings" "testing" @@ -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) + } + }) + } +}