github.com/jxskiss/gopkg@v0.17.3/set/template.go

// +build ignore

// This program generates set implementation for concrete types.
// It can be invoked by running `go generate`.
package main

import (
	"log"
	"os"
	"text/template"
	"time"
)

type Type struct {
	Type      string
	SetType   string
	ZeroValue string
}

var targetTypes = []Type{
	{"int", "Int", "0"},
	{"int32", "Int32", "0"},
	{"uint32", "Uint32", "0"},
	{"int64", "Int64", "0"},
	{"uint64", "Uint64", "0"},
	{"string", "String", `""`},
}

func main() {
	timestamp := func() func() string {
		now := time.Now().Format(time.RFC3339)
		return func() string { return now }
	}()

	tmpl := template.Must(template.New("").
		Funcs(template.FuncMap{"Timestamp": timestamp}).
		Parse(codetmpl))

	for _, t := range targetTypes {
		func() {
			f, err := os.Create(t.Type + ".go")
			if err != nil {
				log.Fatal(err)
			}
			defer f.Close()
			if err = tmpl.Execute(f, t); err != nil {
				log.Fatal(err)
			}
		}()
	}
}

var codetmpl = `// Code generated by go generate at {{ Timestamp }}; DO NOT EDIT.

package set

import "encoding/json"

// {{ .SetType }} is {{ .Type }} set collection.
// The zero value of {{ .SetType }} is an empty instance ready to use. A zero {{ .SetType }}
// value shall not be copied, or it may result incorrect behavior.
type {{ .SetType }} struct {
	m map[{{ .Type }}]struct{}
}

// New{{ .SetType }} creates {{ .SetType }} instance.
func New{{ .SetType }}(vals ...{{ .Type }}) {{ .SetType }} {
	size := max(len(vals), minSize)
	set := {{ .SetType }}{
		m: make(map[{{ .Type }}]struct{}, size),
	}
	set.Add(vals...)
	return set
}

// New{{ .SetType }}WithSize creates {{ .SetType }} instance with given initial size.
func New{{ .SetType }}WithSize(size int) {{ .SetType }} {
	set := {{ .SetType }}{
		m: make(map[{{ .Type }}]struct{}, size),
	}
	return set
}

// Size returns the size of set.
func (s {{ .SetType }}) Size() int { return len(s.m) }

// Add adds values into the set.
func (s *{{ .SetType }}) Add(vals ...{{ .Type }}) {
	if s.m == nil {
		size := max(len(vals), minSize)
		s.m = make(map[{{ .Type }}]struct{}, size)
	}
	for idx := range vals {
		s.m[vals[idx]] = struct{}{}
	}
}

// Del deletes values from the set.
func (s *{{ .SetType }}) Del(vals ...{{ .Type }}) {
	for idx := range vals {
		delete(s.m, vals[idx])
	}
}

// Pop pops an element from the set, in no particular order.
func (s *{{ .SetType }}) Pop() {{ .Type }} {
	for val := range s.m {
		delete(s.m, val)
		return val
	}
	return {{ .ZeroValue }}
}

// Iterate iterates the set in no particular order and call the given function
// for each set element.
func (s {{ .SetType }}) Iterate(fn func({{ .Type }})) {
	for val := range s.m {
		fn(val)
	}
}

// Contains returns true if the set contains all the values.
func (s {{ .SetType }}) Contains(vals ...{{ .Type }}) bool {
	if len(vals) == 0 {
		return false
	}
	for _, v := range vals {
		if _, ok := s.m[v]; !ok {
			return false
		}
	}
	return true
}

// ContainsAny returns true if the set contains any of the values.
func (s {{ .SetType }}) ContainsAny(vals ...{{ .Type }}) bool {
	for _, v := range vals {
		if _, ok := s.m[v]; ok {
			return true
		}
	}
	return false
}

// Diff returns new {{ .SetType }} about the values which other set doesn't contain.
func (s {{ .SetType }}) Diff(other {{ .SetType }}) {{ .SetType }} {
	res := New{{ .SetType }}WithSize(s.Size())

	for val := range s.m {
		if _, ok := other.m[val]; !ok {
			res.m[val] = struct{}{}
		}
	}
	return res
}

// DiffSlice is similar to Diff, but takes a slice as parameter.
func (s {{ .SetType }}) DiffSlice(other []{{ .Type }}) {{ .SetType }} {
	if len(s.m) > len(other) {
		tmp := New{{ .SetType }}WithSize(len(other))
		dup := 0
		for _, val := range other {
			if _, ok := s.m[val]; ok {
				dup++
			}
			tmp.m[val] = struct{}{}
		}
		res := New{{ .SetType }}WithSize(max(s.Size()-dup, 0))
		for val := range s.m {
			if _, ok := tmp.m[val]; !ok {
				res.m[val] = struct{}{}
			}
		}
		return res
	} else {
		res := New{{ .SetType }}WithSize(s.Size())
		for val := range s.m {
			res.m[val] = struct{}{}
		}
		for _, val := range other {
			if _, ok := res.m[val]; ok {
				delete(res.m, val)
			}
		}
		return res
	}
}

// FilterInclude returns a new slice which contains values that present in
// the provided slice and also present in the {{ .SetType }} set.
func (s {{ .SetType }}) FilterInclude(slice []{{ .Type }}) []{{ .Type }} {
	res := make([]{{ .Type }}, 0, min(s.Size(), len(slice)))
	for _, val := range slice {
		if _, ok := s.m[val]; ok {
			res = append(res, val)
		}
	}
	return res
}

// FilterExclude returns a new slice which contains values that present in
// the provided slice but don't present in the {{ .SetType }} set.
func (s {{ .SetType }}) FilterExclude(slice []{{ .Type }}) []{{ .Type }} {
	res := make([]{{ .Type }}, 0, len(slice))
	for _, val := range slice {
		if _, ok := s.m[val]; !ok {
			res = append(res, val)
		}
	}
	return res
}

// Intersect returns new {{ .SetType }} about values which other set also contains.
func (s {{ .SetType }}) Intersect(other {{ .SetType }}) {{ .SetType }} {
	res := New{{ .SetType }}WithSize(min(s.Size(), other.Size()))

	// loop over the smaller set
	if len(s.m) <= len(other.m) {
		for val := range s.m {
			if _, ok := other.m[val]; ok {
				res.m[val] = struct{}{}
			}
		}
	} else {
		for val := range other.m {
			if _, ok := s.m[val]; ok {
				res.m[val] = struct{}{}
			}
		}
	}
	return res
}

// IntersectSlice is similar to Intersect, but takes a slice as parameter.
func (s {{ .SetType }}) IntersectSlice(other []{{ .Type }}) {{ .SetType }} {
	res := New{{ .SetType }}WithSize(min(s.Size(), len(other)))

	for _, val := range other {
		if _, ok := s.m[val]; ok {
			res.m[val] = struct{}{}
		}
	}
	return res
}

// Union returns new {{ .SetType }} about values either in the set or the other set.
func (s {{ .SetType }}) Union(other {{ .SetType }}) {{ .SetType }} {
	res := New{{ .SetType }}WithSize(s.Size() + other.Size())

	for val := range s.m {
		res.m[val] = struct{}{}
	}
	for val := range other.m {
		res.m[val] = struct{}{}
	}
	return res
}

// UnionSlice is similar to Union, but takes a slice as parameter.
func (s {{ .SetType }}) UnionSlice(other []{{ .Type }}) {{ .SetType }} {
	res := New{{ .SetType }}WithSize(s.Size() + len(other))

	for val := range s.m {
		res.m[val] = struct{}{}
	}
	for _, val := range other {
		res.m[val] = struct{}{}
	}
	return res
}

// Slice converts set into {{ .Type }} slice.
func (s {{ .SetType }}) Slice() []{{ .Type }} {
	res := make([]{{ .Type }}, 0, len(s.m))

	for val := range s.m {
		res = append(res, val)
	}
	return res
}

// Map converts set into map[{{ .Type }}]bool.
func (s {{ .SetType }}) Map() map[{{ .Type }}]bool {
	res := make(map[{{ .Type }}]bool, len(s.m))

	for val := range s.m {
		res[val] = true
	}
	return res
}

// MarshalJSON implements json.Marshaler interface, the set will be
// marshaled as an {{ .Type }} array.
func (s {{ .SetType }}) MarshalJSON() ([]byte, error) {
	res := s.Slice()
	return json.Marshal(res)
}

// UnmarshalJSON implements json.Unmarshaler interface, it will unmarshal
// an {{ .Type }} array to the set.
func (s *{{ .SetType }}) UnmarshalJSON(b []byte) error {
	vals := make([]{{ .Type }}, 0)
	err := json.Unmarshal(b, &vals)
	if err == nil {
		s.Add(vals...)
	}
	return err
}

// MarshalYAML implements yaml.Marshaler interface of the yaml package,
// the set will be marshaled as an {{ .Type }} array.
func (s {{ .SetType }}) MarshalYAML() (interface{}, error) {
	res := s.Slice()
	return res, nil
}

// UnmarshalYAML implements yaml.Unmarshaler interface of the yaml package,
// it will unmarshal an {{ .Type }} array to the set.
func (s *{{ .SetType }}) UnmarshalYAML(unmarshal func(interface{}) error) error {
	vals := make([]{{ .Type }}, 0)
	err := unmarshal(&vals)
	if err == nil {
		s.Add(vals...)
	}
	return err
}
`