h/sort-cli/main.go

package main

import (
	"bytes"
	"flag"
	"fmt"
	"golang.org/x/exp/maps"
	"io"
	"log"
	"os"
	"sort"
	"strings"
)

type Cfg struct {
	Key     int
	Numeric bool
	Reverse bool
	Unique  bool
}

func main() {
	fmt.Print(run())

	os.Exit(0)
}

func run() string {
	var c Cfg

	flag.IntVar(&c.Key, "k", 0, "sort via column")
	flag.BoolVar(&c.Numeric, "n", false, "compare according to string numerical value")
	flag.BoolVar(&c.Reverse, "r", false, "reverse the result of comparisons")
	flag.BoolVar(&c.Unique, "u", false, "output only the first of an equal run")

	flag.Parse()

	return doSort(&c, flag.Args())
}

func doSort(cfg *Cfg, sources []string) string {
	r := splitLines(load(sources), []byte("\n"), cfg.Key)

	if cfg.Unique {
		r = uniques(r)
	}

	rs := &Rows{s: r, cfg: cfg}
	sort.Sort(rs)

	return rs.String()
}

type Row struct {
	Line   []byte
	Column []byte
}

type Rows struct {
	s   []*Row
	cfg *Cfg
}

func (p *Rows) Len() int { return len(p.s) }
func (p *Rows) Less(i, j int) bool {
	s := p.s

	lr := s[i]
	rr := s[j]

	if p.cfg.Reverse {
		lr, rr = rr, lr
	}

	var l, r []rune

	if p.cfg.Key == 0 {
		l, r = bytes.Runes(lr.Line), bytes.Runes(rr.Line)
	} else {
		l, r = bytes.Runes(lr.Column), bytes.Runes(rr.Column)
	}

	ln := len(l)
	rn := len(r)

	for i := 0; i < min(ln, rn); i++ {
		if l[i] == r[i] {
			continue
		}

		return l[i] < r[i]
	}

	return ln < rn
}

func (p *Rows) Swap(i, j int) {
	s := p.s

	s[i], s[j] = s[j], s[i]
}

func (p *Rows) String() string {
	r := p.s

	var n int
	for _, s := range r {
		n += len(s.Line)
	}
	n += len(r) * len("\n")

	var sb strings.Builder
	sb.Grow(n)

	for _, c := range r {
		for _, rn := range c.Line {
			sb.WriteByte(rn)
		}

		sb.WriteString("\n")
	}

	return sb.String()
}

func load(sources []string) []byte {
	if len(sources) == 0 {
		return loadStdin()
	}

	inputs := make([][]byte, 0, len(sources))
	for _, path := range sources {
		if path == "-" {
			inputs = append(inputs, loadStdin())

			continue
		}

		inputs = append(inputs, loadFile(path))
	}

	var totalSize int
	for _, s := range inputs {
		totalSize += len(s)
	}

	var b bytes.Buffer
	b.Grow(totalSize)

	for _, c := range inputs {
		b.Write(c)
	}

	return b.Bytes()
}

func loadStdin() []byte {
	b, err := io.ReadAll(os.Stdin)
	if err != nil {
		log.Fatalf("can't read stdin: %e", err)
	}

	return b
}

func loadFile(path string) []byte {
	if _, err := os.Stat(path); err != nil {
		log.Fatalf("file not exists: %s", path)
	}

	content, err := os.ReadFile(path)
	if err != nil {
		log.Fatalf("file open file: %s", err)
	}

	return content
}

func splitLines(b, sp []byte, key int) []*Row {
	r := make([]*Row, 0, bytes.Count(b, sp))
	for _, b := range bytes.Split(b, sp) {
		var column []byte

		if key != 0 {
			bs := bytes.Split(b, []byte(" "))

			if len(bs) < key {
				continue // TODO is it error or not?
				// log.Fatalf("Column for key \"%d\" doesn't exists", cfg.Key)
			}

			column = bs[key-1]
		}

		r = append(r, &Row{Line: b, Column: column})
	}

	return r
}

func uniques(r []*Row) []*Row {
	m := make(map[string]*Row, len(r))

	for _, r := range r {
		r := r

		m[string(r.Line)] = r
	}

	return maps.Values(m)
}