go
// Copyright 2022 <mzh.scnu@qq.com>. All rights reserved. 
 // Use of this source code is governed by a BSD-style 
 // license that can be found in the LICENSE file. 
  
 package main 
  
 import ( 
         "flag" 
         "fmt" 
         "strings" 
  
         "google.golang.org/protobuf/compiler/protogen" 
 ) 
  
 func main() { 
         g := rpc{} 
  
         var flags flag.FlagSet 
         genCli := flags.Bool("gen-cli", false, "--tinyrpc_out=gen-cli=true:.") 
  
         protogen.Options{ 
                 ParamFunc: flags.Set, 
         }.Run(func(p *protogen.Plugin) error { 
                 err := g.Generate(p) 
                 if err != nil { 
                         return err 
                 } 
                 if *genCli { 
                         return generateCli(p) 
                 } 
                 return nil 
         }) 
 } 
  
 type rpc struct{} 
  
 // Generate generate service code 
 func (md *rpc) Generate(plugin *protogen.Plugin) error { 
         for _, f := range plugin.Files { 
                 if len(f.Services) == 0 { 
                         continue 
                 } 
                 fileName := f.GeneratedFilenamePrefix + ".svr.go" 
                 t := plugin.NewGeneratedFile(fileName, f.GoImportPath) 
                 t.P("// Code generated by protoc-gen-tinyrpc.") 
                 t.P() 
                 pkg := fmt.Sprintf("package %s", f.GoPackageName) 
                 t.P(pkg) 
                 t.P() 
                 for _, s := range f.Services { 
                         serviceCode := fmt.Sprintf(`%stype %s struct{}`, 
                                 getComments(s.Comments), s.Desc.Name()) 
                         t.P(serviceCode) 
                         t.P() 
                         for _, m := range s.Methods { 
                                 funcCode := fmt.Sprintf(`%sfunc(this *%s) %s(args *%s,reply *%s)error{ 
                                         // define your service ... 
                                         return nil 
                                 } 
                                 `, getComments(m.Comments), s.Desc.Name(), 
                                         m.Desc.Name(), m.Input.Desc.Name(), m.Output.Desc.Name()) 
                                 t.P(funcCode) 
                         } 
                 } 
         } 
         return nil 
 } 
  
 func generateCli(plugin *protogen.Plugin) error { 
         for _, f := range plugin.Files { 
                 if len(f.Services) == 0 { 
                         continue 
                 } 
                 fileName := f.GeneratedFilenamePrefix + ".cli.go" 
                 t := plugin.NewGeneratedFile(fileName, f.GoImportPath) 
                 t.P("// Code generated by protoc-gen-tinyrpc. DO NOT EDIT.") 
                 t.P() 
                 pkg := fmt.Sprintf("package %s", f.GoPackageName) 
                 t.P(pkg) 
                 t.P() 
                 importCode := `import ( 
                 "github.com/zehuamama/tinyrpc" 
                 "net" 
                 )` 
                 t.P(importCode) 
                 t.P() 
                 for _, s := range f.Services { 
                         clientCode := fmt.Sprintf(`%stype %sClient struct{ 
                                 client *tinyrpc.Client 
                         } 
                         `, 
                                 getComments(s.Comments), s.Desc.Name()) 
                         t.P(clientCode) 
                         t.P() 
                         newClientCode := fmt.Sprintf(`func New%sClient(conn net.Conn) %sClient { 
                                 return %sClient{client: tinyrpc.NewClient(conn)} 
                         } 
                         `, 
                                 s.Desc.Name(), s.Desc.Name(), s.Desc.Name()) 
                         t.P(newClientCode) 
                         t.P() 
                         for _, m := range s.Methods { 
                                 funcCode := fmt.Sprintf(`%sfunc(this *%sClient) %s(args *%s,reply *%s)error{ 
                                         return this.client.Call("%s.%s", args, reply) 
                                 } 
                                 `, getComments(m.Comments), s.Desc.Name(), 
                                         m.Desc.Name(), m.Input.Desc.Name(), m.Output.Desc.Name(), s.Desc.Name(), m.Desc.Name()) 
                                 t.P(funcCode) 
                         } 
                 } 
         } 
         return nil 
 } 
  
 // getComments get comment details 
 func getComments(comments protogen.CommentSet) string { 
         c := make([]string, 0) 
         c = append(c, strings.Split(string(comments.Leading), "\n")...) 
         c = append(c, strings.Split(string(comments.Trailing), "\n")...) 
  
         res := "" 
         for _, comment := range c { 
                 if strings.TrimSpace(comment) == "" { 
                         continue 
                 } 
                 res += "//" + comment + "\n" 
         } 
         return res 
 }