Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
gogo protobuf plugin that generates threadsafe "registry types" for message types that declare the registry MessageOption
package registry
import (
"github.com/gogo/protobuf/protoc-gen-gogo/generator"
"github.com/dan-compton/protobuf/danproto"
)
type plugin struct {
*generator.Generator
generator.PluginImports
messages []*generator.Descriptor
}
func NewRegistry() *plugin {
return &plugin{}
}
func (p *plugin) Name() string {
return "registry"
}
func (p *plugin) Init(g *generator.Generator) {
p.Generator = g
}
func (p *plugin) Generate(file *generator.FileDescriptor) {
p.PluginImports = generator.NewPluginImports(p.Generator)
p.messages = make([]*generator.Descriptor, 0)
syncPkg := p.NewImport("sync")
for _, message := range file.Messages() {
if !danproto.IsRegistry(file.FileDescriptorProto, message.DescriptorProto) {
continue
}
p.messages = append(p.messages, message)
// NOTE(dan) Not sure if this should be limited to 1 map and conform to a Registry interface.
baseTypeName := generator.CamelCaseSlice(message.TypeName())
ccTypeName := generator.CamelCaseSlice(message.TypeName())
ccTypeName += "Registry"
p.P(`type `, ccTypeName, ` struct {`)
p.P(syncPkg.Use(), `.RWMutex`)
p.In()
for _, field := range message.Field {
fieldname := p.GetFieldName(message, field)
goTyp, _ := p.GoType(message, field)
if p.IsMap(field) {
m := p.GoMapType(nil, field)
keyTyp, _ := p.GoType(message, m.KeyField)
valTyp, _ := p.GoType(message, m.ValueField)
p.P(fieldname, ` map[`, keyTyp, `]`, valTyp)
} else {
p.P(fieldname, ` `, goTyp)
}
}
p.Out()
p.P(`}`)
p.P(``)
// AsRegistry()
// Generate a new registry from the old type
p.P(`func (this *`, baseTypeName, `) AsRegistry`, `() *`, ccTypeName, `{`)
p.In()
p.P(`that := &`, ccTypeName, `{}`)
for _, field := range message.Field {
fieldname := p.GetFieldName(message, field)
p.P(`that.`, fieldname, ` = this.`, fieldname, ``)
}
p.P(`return that`)
p.Out()
p.P(`}`)
// NewTRegistry(*oldT)
p.P(`func New`, ccTypeName, ` (that *`, baseTypeName, `) *`, ccTypeName, ` {`)
p.In()
p.P(`this := &`, ccTypeName, `{}`)
for _, field := range message.Field {
fieldname := p.GetFieldName(message, field)
p.P(`this.`, fieldname, ` = that.`, fieldname, ``)
}
p.P(`return this`)
p.Out()
p.P(`}`)
for _, field := range message.Field {
p.In()
fieldname := p.GetFieldName(message, field)
if !p.IsMap(field) {
continue
} else {
m := p.GoMapType(nil, field)
keyTyp, _ := p.GoType(message, m.KeyField)
valTyp, _ := p.GoType(message, m.ValueField)
// CopyT() (map<keyT, valT)
p.P(`func (this *`, ccTypeName, `) Copy`, fieldname, `() map[`, keyTyp, `]`, valTyp, `{`)
p.In()
p.P(`t := make(map[`, keyTyp, `]`, valTyp, `)`)
p.P(`this.Lock()`)
p.P(`defer this.Unlock()`)
p.P(`for k,v := range this.`, fieldname, ` {`)
p.In()
p.P(`t[k] = v`)
p.Out()
p.P(`}`)
p.P(`return t`)
p.Out()
p.P(`}`)
// Keys()
p.P(`func (this *`, ccTypeName, `) Keys`, fieldname, `() []`, keyTyp, `{`)
p.In()
p.P(`var ks []`, keyTyp)
p.P(`this.RLock()`)
p.P(`defer this.RUnlock()`)
p.P(`for k, _ := range this.`, fieldname, ` {`)
p.In()
p.P(`ks = append(ks, k)`)
p.Out()
p.P(`}`)
p.P(`return ks`)
p.Out()
p.P(`}`)
// GetT(key keyT) (valT, ok)
p.P(`func (this *`, ccTypeName, `) Get`, fieldname, `(k `, keyTyp, `) (`, valTyp, `, bool) {`)
p.In()
p.P(`this.RLock()`)
p.P(`defer this.RUnlock()`)
p.P(`v, ok := this.`, fieldname, `[k]`)
p.P(`return v, ok`)
p.Out()
p.P(`}`)
p.P(``)
// SetT(key keyT, val valT)
p.P(`func (this *`, ccTypeName, `) Set`, fieldname, `(k `, ` `, keyTyp, `, v `, valTyp, `) {`)
p.In()
p.P(`this.Lock()`)
p.P(`defer this.Unlock()`)
p.P(`this.`, fieldname, `[k] = v`)
p.P(`return`)
p.Out()
p.P(`}`)
p.P(``)
}
p.Out()
p.P()
}
}
}
func init() {
generator.RegisterPlugin(NewRegistry())
}
/*
//Test code
import "testing"
var allUserPermissions = &Permissions{
Name: "user",
Permission: map[string]bool{
"a": true,
"b": true,
"c": true,
},
}
var allTeamPermissions = &Permissions{
Name: "team",
Permission: map[string]bool{
"aA": true,
"bB": true,
"cC": true,
},
}
var PR = &PermsRegistryRegistry{
Name: "permissions",
Registry: map[string]*Permissions{
allUserPermissions.Name: allUserPermissions,
allTeamPermissions.Name: allTeamPermissions,
},
}
func TestPermissionsRegistryTest(t *testing.T) {
r, ok := PR.GetRegistry("user")
t.Log(r, ok)
}
*/
/*
And some test protobuf
syntax = "proto3";
import "github.com/gogo/protobuf/gogoproto/gogo.proto";
import "github.com/dan-compton/protobuf/danproto/dan.proto";
package danprototest.regs;
message Permissions {
option (danproto.registry) = true;
option (gogoproto.sizer) = true;
string name = 1;
map<string, bool> permission = 2;
}
message PermsRegistry {
option (danproto.registry) = true;
string name = 1;
map<string, Permissions> registry = 2;
}
*/
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.