オーダーを保持するmap的コンテナ実装（go-generate向け）

orderedmap.go

/*
orderedmap package

Copyright (c) 2017 Noboru Irieda

This software is released under the MIT License.
http://opensource.org/licenses/mit-license.php
*/

package orderedmap

// github.com/cheekybits/genny 参照
//go:generate genny -in=$GOFILE -out=gen-$GOFILE gen "Type=Sample"

import (
	"github.com/cheekybits/genny/generic"
)

// Type ...
type Type generic.Type

// TypeItem ...
type TypeItem struct {
	index int
	value *Type
}

// TypeMap ...
type TypeMap struct {
	keys  []string
	items map[string]*TypeItem
}

// NewTypeMap ...
func NewTypeMap() *TypeMap {
	return &TypeMap{
		keys:  []string{},
		items: map[string]*TypeItem{},
	}
}

func (m *TypeMap) Len() int { return len(m.keys) }
func (m *TypeMap) Swap(i, j int) {
	iv := m.items[m.keys[i]]
	jv := m.items[m.keys[j]]
	jv.index, iv.index = iv.index, jv.index
	m.keys[i], m.keys[j] = m.keys[j], m.keys[i]
}
func (m *TypeMap) Less(i, j int) bool { return m.keys[i] < m.keys[j] }

// Get ...
func (m *TypeMap) Get(key string) *Type {
	item, ok := m.items[key]
	if ok {
		return item.value
	}
	return nil
}

// Set ...
func (m *TypeMap) Set(key string, value *Type) {
	m.items[key] = &TypeItem{len(m.keys), value}
	m.keys = append(m.keys, key)
}

// Del ...
func (m *TypeMap) Del(key string) {
	item, ok := m.items[key]
	if ok {
		m.keys = append(m.keys[:item.index], m.keys[item.index+1:]...)
		delete(m.items, key)
	}
}

// Iter ...
func (m *TypeMap) Iter(f func(key string, value *Type) bool) {
	for _, key := range m.keys {
		TypeItem, ok := m.items[key]
		if ok {
			if !f(key, TypeItem.value) {
				break
			}
		}
	}
}

orderedmap_test.go

package orderedmap

import (
	"fmt"
	"sort"
)

func ExampleSampleMap() {
	m := NewSampleMap()
	m.Set("hoge1", &Sample{})
	m.Set("moge1", &Sample{})
	m.Set("hoge2", &Sample{})
	m.Set("moge2", &Sample{})
	sort.Sort(m)
	m.Del("moge1")
	m.Iter(func(k string, v *Sample) bool {
		fmt.Println(k, v)
		return true
	})
	// output:
	// hoge1 &{}
	// hoge2 &{}
	// moge2 &{}
}

typedefinition.go

package orderedmap

type Sample struct{}

go generate orderedmap.goの結果

gen-orderedmap.go

// This file was automatically generated by genny.
// Any changes will be lost if this file is regenerated.
// see https://github.com/cheekybits/genny

package orderedmap

// github.com/cheekybits/genny 参照

// SampleItem ...
type SampleItem struct {
	index int
	value *Sample
}

// SampleMap ...
type SampleMap struct {
	keys  []string
	items map[string]*SampleItem
}

// NewSampleMap ...
func NewSampleMap() *SampleMap {
	return &SampleMap{
		keys:  []string{},
		items: map[string]*SampleItem{},
	}
}

func (m *SampleMap) Len() int { return len(m.keys) }
func (m *SampleMap) Swap(i, j int) {
	iv := m.items[m.keys[i]]
	jv := m.items[m.keys[j]]
	jv.index, iv.index = iv.index, jv.index
	m.keys[i], m.keys[j] = m.keys[j], m.keys[i]
}
func (m *SampleMap) Less(i, j int) bool { return m.keys[i] < m.keys[j] }

// Get ...
func (m *SampleMap) Get(key string) *Sample {
	item, ok := m.items[key]
	if ok {
		return item.value
	}
	return nil
}

// Set ...
func (m *SampleMap) Set(key string, value *Sample) {
	m.items[key] = &SampleItem{len(m.keys), value}
	m.keys = append(m.keys, key)
}

// Del ...
func (m *SampleMap) Del(key string) {
	item, ok := m.items[key]
	if ok {
		m.keys = append(m.keys[:item.index], m.keys[item.index+1:]...)
		delete(m.items, key)
	}
}

// Iter ...
func (m *SampleMap) Iter(f func(key string, value *Sample) bool) {
	for _, key := range m.keys {
		SampleItem, ok := m.items[key]
		if ok {
			if !f(key, SampleItem.value) {
				break
			}
		}
	}
}

万能のOrderedMapとか作ると結局動的タイプアサーションが必要になったりするので
go-generateで任意の構造体バリューを持てる実装を作ってみた。

型安全でメモリは浪費するけどSet順をキープしつつキーでソートも可能にしてみた。

typedefinition.goに任意の定義を宣言しておいて、
gennyに指定する型にその定義名を渡すと「定義名Map」コンテナ実装がジェネレートされる。

• コンカレントに参照する場合はsync.RWMutexを埋め込むといいよ。
• ソート機能いらんかったらLen/Swap/Lessを消そう。
• メモリ効率重視ならBTreeインデックスにしよう。
• MarshalJSON/UnmarshalJSONが必要なら生やすといいよ。
• キー型変えるのも必要に応じてどうぞ！