From 5ff387bbe4ab90b587725808526e97a4f01f97a1 Mon Sep 17 00:00:00 2001
From: radare <pancake@nopcode.org>
Date: Sun, 1 Dec 2019 14:10:13 +0100
Subject: [PATCH] strings: split_nth and add tests

---
 vlib/builtin/js/string.v   |  4 --
 vlib/builtin/string.v      | 80 +++++++++++++-------------------------
 vlib/builtin/string_test.v | 34 +++++++++++++++-
 3 files changed, 60 insertions(+), 58 deletions(-)

diff --git a/vlib/builtin/js/string.v b/vlib/builtin/js/string.v
index e5853c9a0a..8f8a78717e 100644
--- a/vlib/builtin/js/string.v
+++ b/vlib/builtin/js/string.v
@@ -63,10 +63,6 @@ pub fn (s string) split(delim string) []string {
 	return s.split(delim)
 }
 
-pub fn (s string) split_single(delim byte) []string {
-	return s.split(delim.str())
-}
-
 pub fn (s string) split_into_lines() []string {
 	return s.split('\n')
 }
diff --git a/vlib/builtin/string.v b/vlib/builtin/string.v
index 515d53e194..4c181f3e3b 100644
--- a/vlib/builtin/string.v
+++ b/vlib/builtin/string.v
@@ -285,76 +285,50 @@ fn (s string) add(a string) string {
 }
 
 pub fn (s string) split(delim string) []string {
-	// println('string split delim="$delim" s="$s"')
+	return s.split_nth(delim, 0)
+}
+
+pub fn (s string) split_nth(delim string, nth int) []string {
 	mut res := []string
-	// if delim.len == 0 {
-	// 	res << s
-	// 	return res
-	// }
+	mut i := 0
 	if delim.len == 0 {
+		i = 1
 		for ch in s {
+			if nth > 0 && i >= nth {
+				res << s.substr(i, s.len)
+				break
+			}
 			res << ch.str()
+			i++
 		}
 		return res
 	}
-	if delim.len == 1 {
-		return s.split_single(delim[0])
-	}
-	mut i := 0
-	mut start := 0// - 1
-	for i < s.len {
-		// printiln(i)
-		mut a := s[i] == delim[0]
-		mut j := 1
-		for j < delim.len && a {
-			a = a && s[i + j] == delim[j]
+	mut start := 0
+	for i <= s.len {
+		mut is_delim := s[i] == delim[0]
+		mut j := 0
+		for is_delim && j < delim.len {
+			is_delim = is_delim && s[i + j] == delim[j]
 			j++
 		}
+		was_last := nth > 0 && res.len == nth
+		if was_last{break}
 		last := i == s.len - 1
-		if a || last {
-			if last {
+		if is_delim || last {
+			if !is_delim && last {
 				i++
 			}
 			mut val := s.substr(start, i)
-			// println('got it "$val" start=$start i=$i delim="$delim"')
-			if val.len > 0 {
-				// todo perf
-				// val now is '___VAL'. remove '___' from the start
-				if val.starts_with(delim) {
-					// println('!!')
-					val = val.right(delim.len)
-				}
-				res << val.trim_space()
+			if val.starts_with(delim) {
+				val = val.right(delim.len)
 			}
-			start = i
+			res << val
+			start = i + delim.len
 		}
 		i++
 	}
-	return res
-}
-
-pub fn (s string) split_single(delim byte) []string {
-	mut res := []string
-	if int(delim) == 0 {
-		res << s
-		return res
-	}
-	mut i := 0
-	mut start := 0
-	for i < s.len {
-		is_delim := s[i] == delim
-		last := i == s.len - 1
-		if is_delim || last {
-			if !is_delim && i == s.len - 1 {
-				i++
-			}
-			val := s.substr(start, i)
-			if val.len > 0 {
-				res << val
-			}
-			start = i + 1
-		}
-		i++
+	if s.ends_with (delim) && (nth < 1 || res.len < nth) {
+		res << ''
 	}
 	return res
 }
diff --git a/vlib/builtin/string_test.v b/vlib/builtin/string_test.v
index 4bfe4c50d9..3f50b6d64c 100644
--- a/vlib/builtin/string_test.v
+++ b/vlib/builtin/string_test.v
@@ -87,6 +87,37 @@ fn test_sort() {
 	assert vals[3] == 'arr'
 }
 
+fn test_split_nth() {
+	a := "1,2,3"
+	assert (a.split(',').len == 3)
+	assert (a.split_nth(',', -1).len == 3)
+	assert (a.split_nth(',', 0).len == 3)
+	assert (a.split_nth(',', 1).len == 1)
+	assert (a.split_nth(',', 2).len == 2)
+	assert (a.split_nth(',', 10).len == 3)
+	b := "1::2::3"
+	assert (b.split('::').len == 3)
+	assert (b.split_nth('::', -1).len == 3)
+	assert (b.split_nth('::', 0).len == 3)
+	assert (b.split_nth('::', 1).len == 1)
+	assert (b.split_nth('::', 2).len == 2)
+	assert (b.split_nth('::', 10).len == 3)
+	c := "ABCDEF"
+	assert (c.split('').len == 6)
+	assert (c.split_nth('', 3).len == 3)
+	assert (c.split_nth('BC', -1).len == 2)
+	d := ","
+	assert (d.split(',').len == 2)
+	assert (d.split_nth('', 3).len == 1)
+	assert (d.split_nth(',', -1).len == 2)
+	assert (d.split_nth(',', 3).len == 2)
+	e := ",,,0,,,,,a,,b,"
+	// assert (e.split(',,').len == 5)
+	// assert (e.split_nth(',,', 3).len == 2)
+	assert (e.split_nth(',', -1).len == 12)
+	assert (e.split_nth(',', 3).len == 3)
+}
+
 fn test_split() {
 	mut s := 'volt/twitch.v:34'
 	mut vals := s.split(':')
@@ -109,10 +140,11 @@ fn test_split() {
 	// /////////
 	s = 'lalala'
 	vals = s.split('a')
-	assert vals.len == 3
+	assert vals.len == 4
 	assert vals[0] == 'l'
 	assert vals[1] == 'l'
 	assert vals[2] == 'l'
+	assert vals[3] == ''
 	// /////////
 	s = 'awesome'
 	a := s.split('')