optimize regex usage and improve stringToSize function

mounts_linux.go

@@ -34,7 +34,7 @@ const (
 	// (6) optional fields: zero or more fields terminated by "-".
 	mountinfoOptionalFields = 6
 	// (7) separator between optional fields.
-	//mountinfoSeparator = 7
+	// mountinfoSeparator = 7
 	// (8) filesystem type: name of filesystem of the form.
 	mountinfoFsType = 8
 	// (9) mount source: filesystem specific information or "none".
@@ -43,6 +43,8 @@ const (
 	mountinfoSuperOptions = 10
 )
 
+var statRegex = regexp.MustCompile(`^/dev/mapper/(.*)-(.*)`)
+
 // Stat returns the mountpoint's stat information.
 func (m *Mount) Stat() unix.Statfs_t {
 	return m.Metadata.(unix.Statfs_t)
@@ -109,8 +111,7 @@ func mounts() ([]Mount, []string, error) {
 
 		// Resolve /dev/mapper/* device names.
 		if strings.HasPrefix(d.Device, "/dev/mapper/") {
-			re := regexp.MustCompile(`^/dev/mapper/(.*)-(.*)`)
-			match := re.FindAllStringSubmatch(d.Device, -1)
+			match := statRegex.FindAllStringSubmatch(d.Device, -1)
 			if len(match) > 0 && len(match[0]) == 3 {
 				d.Device = filepath.Join("/dev", match[0][1], match[0][2])
 			}

table.go

@@ -500,10 +500,25 @@ func sizeToString(size uint64) (str string) {
 	return
 }
 
+var (
+	sizeRegex  = regexp.MustCompile(`^(\d+)([KMGTPE]?)$`)
+	sizeShifts = map[byte]uint{
+		'K': 10,
+		'M': 20,
+		'G': 30,
+		'T': 40,
+		'P': 50,
+		'E': 60,
+	}
+)
+
 // stringToSize transforms an SI size into a number.
-func stringToSize(s string) (size uint64, err error) {
-	regex := regexp.MustCompile(`^(\d+)([KMGTPE]?)$`)
-	matches := regex.FindStringSubmatch(s)
+func stringToSize(s string) (uint64, error) {
+	if s == "" {
+		return 0, fmt.Errorf("empty string")
+	}
+
+	matches := sizeRegex.FindStringSubmatch(s)
 	if len(matches) == 0 {
 		return 0, fmt.Errorf("'%s' is not valid, must have integer with optional SI prefix", s)
 	}
@@ -512,29 +527,16 @@ func stringToSize(s string) (size uint64, err error) {
 	if err != nil {
 		return 0, err
 	}
-	if matches[2] != "" {
-		prefix := matches[2]
-		switch prefix {
-		case "K":
-			size = num << 10
-		case "M":
-			size = num << 20
-		case "G":
-			size = num << 30
-		case "T":
-			size = num << 40
-		case "P":
-			size = num << 50
-		case "E":
-			size = num << 60
-		default:
-			err = fmt.Errorf("prefix '%s' not allowed, valid prefixes are K, M, G, T, P, E", prefix)
-			return
-		}
-	} else {
-		size = num
+
+	if matches[2] == "" {
+		return num, nil
 	}
-	return
+
+	if shift, ok := sizeShifts[matches[2][0]]; ok {
+		return num << shift, nil
+	}
+
+	return 0, fmt.Errorf("prefix '%s' not allowed, valid prefixes are K, M, G, T, P, E", matches[2])
 }
 
 // stringToColumn converts a column name to its index.

table_test.go

@@ -0,0 +1,63 @@
+package main
+
+import "testing"
+
+func Test_stringToSize(t *testing.T) {
+	tests := []struct {
+		name    string
+		want    uint64
+		wantErr bool
+	}{
+		{"0", 0, false},
+		{"42", 42, false},
+
+		{"1K", 1 << 10, false},
+		{"2M", 2 << 20, false},
+		{"3G", 3 << 30, false},
+		{"4T", 4 << 40, false},
+		{"5P", 5 << 50, false},
+		{"6E", 6 << 60, false},
+
+		{"", 0, true},
+		{"abc", 0, true},
+		{"10Z", 0, true},
+		{"-5", 0, true},
+		{"18446744073709551615", ^uint64(0), false},
+
+		{" 10K", 0, true},
+		{"10K ", 0, true},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, err := stringToSize(tt.name)
+			if (err != nil) != tt.wantErr {
+				t.Errorf("stringToSize() error = %v, wantErr %v", err, tt.wantErr)
+				return
+			}
+			if got != tt.want {
+				t.Errorf("stringToSize() got = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}
+
+func BenchmarkStringToSize(b *testing.B) {
+	cases := []string{
+		"42",
+		"1K",
+		"128M",
+		"512G",
+		"2T",
+		"5P",
+		"6E",
+		"invalid",
+	}
+
+	for _, input := range cases {
+		b.Run(input, func(b *testing.B) {
+			for i := 0; i < b.N; i++ {
+				_, _ = stringToSize(input)
+			}
+		})
+	}
+}