ruspa/tree/tree_utils.go

package tree

import (
	"errors"
	"fmt"
	"math"
	"regexp"
	"strconv"
	"sync"
)

var ErrUnparsableSize = errors.New("unparsable size")
var matchSize = regexp.MustCompile("^(\\d+(\\.\\d+)?)([KMGTP]?B?)$")

type ByteSize uint

const (
	B ByteSize = 1 << (10 * iota)
	KB
	MB
	GB
	TB
	PB
)

func isLastLevel(node Node, maxLevel int) bool {
	return maxLevel-1 == node.Level()
}

func parseUnit(unit string) ByteSize {
	switch unit {
	case "B":
		return B
	case "KB":
		return KB
	case "MB":
		return MB
	case "GB":
		return GB
	case "TB":
		return TB
	case "PB":
		// I reckon this will never happen, but you never know
		return PB
	default:
		panic(ErrUnknownUnit)
	}
}

func parseUnitSafe(unit string) (ByteSize, error) {
	var returnVal ByteSize
	var err error
	defer func() {
		if r := recover(); r != nil {
			returnVal = 0
			err = ErrUnparsableSize
		}
	}()
	returnVal = parseUnit(unit)
	return returnVal, err
}

func fmtUnit(unit ByteSize) string {
	switch unit {
	case B:
		return "B"
	case KB:
		return "KB"
	case MB:
		return "MB"
	case GB:
		return "GB"
	case TB:
		return "TB"
	case PB:
		// I reckon this will never happen, but you never know
		return "PB"
	default:
		panic(ErrUnknownUnit)
	}
}

func fmtSize(size int64, unit ByteSize) string {
	var dimension float64
	dimension = float64(size) / float64(unit)
	return fmt.Sprintf("%.2f %s", dimension, fmtUnit(unit))
}

func ParseSize(size string) (int64, error) {
	result := matchSize.FindStringSubmatch(size)
	sizeNum := result[1]
	if sizeNum == "" {
		return 0, ErrUnparsableSize
	}
	unit, err := parseUnitSafe(result[3])
	if err != nil {
		return 0, err
	}
	num, err := strconv.ParseFloat(sizeNum, 64)
	if err != nil {
		return 0, err
	}

	return int64(math.Floor(num * float64(unit))), nil
}

func isIntoLimits(size, min, max int64) bool {
	if min <= size {
		if max == 0 {
			return true
		}
		if size <= max {
			return true
		}
	}
	return false
}

func merge(cs []<-chan int64) <-chan int64 {
	var wg sync.WaitGroup
	out := make(chan int64)

	// Start an output goroutine for each input channel in cs.  output
	// copies values from c to out until c is closed, then calls wg.Done.
	output := func(c <-chan int64) {
		for n := range c {
			out <- n
		}
		wg.Done()
	}
	wg.Add(len(cs))
	for _, c := range cs {
		go output(c)
	}

	// Start a goroutine to close out once all the output goroutines are
	// done.  This must start after the wg.Add call.
	go func() {
		wg.Wait()
		close(out)
	}()
	return out
}