1// Copyright (C) 2019 Storj Labs, Inc.
2// See LICENSE for copying information.
3
4package memory
5
6import (
7	"errors"
8	"fmt"
9	"strconv"
10	"strings"
11)
12
13// base 2 and base 10 sizes.
14const (
15	B Size = 1 << (10 * iota)
16	KiB
17	MiB
18	GiB
19	TiB
20	PiB
21	EiB
22
23	KB Size = 1e3
24	MB Size = 1e6
25	GB Size = 1e9
26	TB Size = 1e12
27	PB Size = 1e15
28	EB Size = 1e18
29)
30
31// Size implements flag.Value for collecting memory size in bytes.
32type Size int64
33
34// Int returns bytes size as int.
35func (size Size) Int() int { return int(size) }
36
37// Int32 returns bytes size as int32.
38func (size Size) Int32() int32 { return int32(size) }
39
40// Int64 returns bytes size as int64.
41func (size Size) Int64() int64 { return int64(size) }
42
43// Float64 returns bytes size as float64.
44func (size Size) Float64() float64 { return float64(size) }
45
46// KiB returns size in kibibytes.
47func (size Size) KiB() float64 { return size.Float64() / KiB.Float64() }
48
49// MiB returns size in mebibytes.
50func (size Size) MiB() float64 { return size.Float64() / MiB.Float64() }
51
52// GiB returns size in gibibytes.
53func (size Size) GiB() float64 { return size.Float64() / GiB.Float64() }
54
55// TiB returns size in tebibytes.
56func (size Size) TiB() float64 { return size.Float64() / TiB.Float64() }
57
58// PiB returns size in pebibytes.
59func (size Size) PiB() float64 { return size.Float64() / PiB.Float64() }
60
61// EiB returns size in exbibytes.
62func (size Size) EiB() float64 { return size.Float64() / EiB.Float64() }
63
64// KB returns size in kilobytes.
65func (size Size) KB() float64 { return size.Float64() / KB.Float64() }
66
67// MB returns size in megabytes.
68func (size Size) MB() float64 { return size.Float64() / MB.Float64() }
69
70// GB returns size in gigabytes.
71func (size Size) GB() float64 { return size.Float64() / GB.Float64() }
72
73// TB returns size in terabytes.
74func (size Size) TB() float64 { return size.Float64() / TB.Float64() }
75
76// PB returns size in petabytes.
77func (size Size) PB() float64 { return size.Float64() / PB.Float64() }
78
79// EB returns size in exabytes.
80func (size Size) EB() float64 { return size.Float64() / EB.Float64() }
81
82// String converts size to a string using base-2 prefixes, unless the number
83// appears to be in base 10.
84func (size Size) String() string {
85	if countZeros(int64(size), 1000) > countZeros(int64(size), 1024) {
86		return size.Base10String()
87	}
88	return size.Base2String()
89}
90
91// countZeros considers a number num in base base. It counts zeros of that
92// number in that base from least significant to most, stopping when a non-zero
93// value is hit.
94func countZeros(num, base int64) (count int) {
95	for num != 0 && num%base == 0 {
96		num /= base
97		count++
98	}
99	return count
100}
101
102// Base2String converts size to a string using base-2 prefixes.
103func (size Size) Base2String() string {
104	if size == 0 {
105		return "0 B"
106	}
107
108	switch {
109	case abs(size) >= EiB*2/3:
110		return fmt.Sprintf("%.1f EiB", size.EiB())
111	case abs(size) >= PiB*2/3:
112		return fmt.Sprintf("%.1f PiB", size.PiB())
113	case abs(size) >= TiB*2/3:
114		return fmt.Sprintf("%.1f TiB", size.TiB())
115	case abs(size) >= GiB*2/3:
116		return fmt.Sprintf("%.1f GiB", size.GiB())
117	case abs(size) >= MiB*2/3:
118		return fmt.Sprintf("%.1f MiB", size.MiB())
119	case abs(size) >= KiB*2/3:
120		return fmt.Sprintf("%.1f KiB", size.KiB())
121	}
122
123	return strconv.FormatInt(size.Int64(), 10) + " B"
124}
125
126// Base10String converts size to a string using base-10 prefixes.
127func (size Size) Base10String() string {
128	if size == 0 {
129		return "0 B"
130	}
131
132	switch {
133	case abs(size) >= EB*2/3:
134		return fmt.Sprintf("%.2f EB", size.EB())
135	case abs(size) >= PB*2/3:
136		return fmt.Sprintf("%.2f PB", size.PB())
137	case abs(size) >= TB*2/3:
138		return fmt.Sprintf("%.2f TB", size.TB())
139	case abs(size) >= GB*2/3:
140		return fmt.Sprintf("%.2f GB", size.GB())
141	case abs(size) >= MB*2/3:
142		return fmt.Sprintf("%.2f MB", size.MB())
143	case abs(size) >= KB*2/3:
144		return fmt.Sprintf("%.2f KB", size.KB())
145	}
146
147	return strconv.FormatInt(size.Int64(), 10) + " B"
148}
149
150func abs(size Size) Size {
151	if size > 0 {
152		return size
153	}
154	return -size
155}
156
157func isLetter(b byte) bool {
158	return ('a' <= b && b <= 'z') || ('A' <= b && b <= 'Z')
159}
160
161// Set updates value from string.
162func (size *Size) Set(s string) error {
163	if s == "" {
164		return errors.New("empty size")
165	}
166
167	p := len(s)
168	for isLetter(s[p-1]) {
169		p--
170
171		if p < 0 {
172			return errors.New("p out of bounds")
173		}
174	}
175
176	value, suffix := s[:p], s[p:]
177	suffix = strings.ToUpper(suffix)
178	if suffix == "" || suffix[len(suffix)-1] != 'B' {
179		suffix += "B"
180	}
181
182	value = strings.TrimSpace(value)
183	v, err := strconv.ParseFloat(value, 64)
184	if err != nil {
185		return err
186	}
187
188	switch suffix {
189	case "EB":
190		*size = Size(v * EB.Float64())
191	case "EIB":
192		*size = Size(v * EiB.Float64())
193	case "PB":
194		*size = Size(v * PB.Float64())
195	case "PIB":
196		*size = Size(v * PiB.Float64())
197	case "TB":
198		*size = Size(v * TB.Float64())
199	case "TIB":
200		*size = Size(v * TiB.Float64())
201	case "GB":
202		*size = Size(v * GB.Float64())
203	case "GIB":
204		*size = Size(v * GiB.Float64())
205	case "MB":
206		*size = Size(v * MB.Float64())
207	case "MIB":
208		*size = Size(v * MiB.Float64())
209	case "KB":
210		*size = Size(v * KB.Float64())
211	case "KIB":
212		*size = Size(v * KiB.Float64())
213	case "B", "":
214		*size = Size(v)
215	default:
216		return fmt.Errorf("unknown suffix %q", suffix)
217	}
218
219	return nil
220}
221
222// Type implements pflag.Value.
223func (Size) Type() string { return "memory.Size" }
224
225// MarshalText returns size as a string.
226func (size Size) MarshalText() (string, error) {
227	return size.String(), nil
228}
229
230// UnmarshalText parses text as a string.
231func (size *Size) UnmarshalText(text []byte) error {
232	return size.Set(string(text))
233}
234
235// MarshalJSON returns size as a json string.
236func (size Size) MarshalJSON() ([]byte, error) {
237	return []byte(strconv.Quote(size.String())), nil
238}
239
240// UnmarshalJSON parses text from a json string.
241func (size *Size) UnmarshalJSON(text []byte) error {
242	unquoted, err := strconv.Unquote(string(text))
243	if err != nil {
244		return err
245	}
246	return size.Set(unquoted)
247}
248