main.go

package main

import (
	"bytes"
	"flag"
	"fmt"
	"go/ast"
	"go/importer"
	"go/parser"
	"go/token"
	"go/types"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strings"
)

var (
	typeNames = flag.String("type", "", "comma-separated list of type names; must be set")
)

func usage() {
	fmt.Fprintf(os.Stderr, "Usage of %s:\n", os.Args[0])
	flag.PrintDefaults()
}

func main() {
	flag.Usage = usage
	flag.Parse()
	log.SetPrefix("rpctocli: ")
	log.SetFlags(0)

	g := NewGenerator()
	dir, err := os.Getwd()
	if err != nil {
		log.Fatalf("Failed to get working directory: %v", err)
	}

	err = g.Parse(dir)
	if err != nil {
		log.Fatalf("Failed to parse and type-check files in directory '%s': %v", dir, err)
	}

	err = g.Extract()
	if err != nil {
		log.Fatalf("Failed to extract rpc info: %v", err)
	}

	err = g.Generate()
	if err != nil {
		log.Fatalf("Failed to generate code: %v", err)
	}

	output, err := g.Format()
	if err != nil {
		log.Fatalf("Failed to format code: %v", err)
	}

	//@todo how determine file name
	fpath := filepath.Join(g.pkg.dir, fmt.Sprintf("%s_cmd.go", g.pkg.name))
	err = ioutil.WriteFile(fpath, output, 0644)
	if err != nil {
		log.Fatalf("Failed to write file: %v", err)
	}

}

//A Package the generator is building a cli for
type Package struct {
	dir   string
	name  string
	files []*File
	defs  map[*ast.Ident]types.Object
}

func (pkg *Package) check(fs *token.FileSet, astFiles []*ast.File) error {
	pkg.defs = make(map[*ast.Ident]types.Object)
	config := types.Config{Importer: importer.Default(), FakeImportC: true}
	info := &types.Info{
		Defs: pkg.defs,
	}

	_, err := config.Check(pkg.dir, fs, astFiles, info)
	if err != nil {
		return fmt.Errorf("Failed to type check package: %v", err)
	}

	return nil
}

//A File in the package
type File struct {
	pkg  *Package
	file *ast.File
}

//The Generator itself
type Generator struct {
	buff     *bytes.Buffer
	pkg      *Package
	services map[string]*RPCService
}

//A RPCMethod is an in-memory representation of method on a rpc service
type RPCMethod struct {
	input  types.Type //can (pointer of) builtin or named type
	output types.Type //this is the underlying elem of the pointer param
	recv   *types.Named
	sig    *types.Signature
}

//A RPCService is an in-memory representation of a parsed rpc service
type RPCService struct {
	methods map[string]*RPCMethod
}

//NewGenerator initializes a generator
func NewGenerator() *Generator {
	return &Generator{
		buff:     bytes.NewBuffer(nil),
		services: map[string]*RPCService{},
	}
}

//Format the code and return resulting bytes
func (g *Generator) Format() ([]byte, error) {

	//@todo implement
	return g.buff.Bytes(), nil
}

//Generate go code and bufferin memory
func (g *Generator) Generate() error {

	imports := []string{
		"log",
		"net/rpc",
		"github.com/codegangsta/cli",
	}

	//headers&imports
	fmt.Fprintf(g.buff, "// Code generated by \"rpctocli %s\"; DO NOT EDIT\n", strings.Join(os.Args[1:], " "))
	fmt.Fprintf(g.buff, "package %s\n\n", g.pkg.name)
	fmt.Fprintf(g.buff, "import(\n")
	for _, imp := range imports {
		fmt.Fprintf(g.buff, "\t\"%s\"\n", imp)
	}
	fmt.Fprintf(g.buff, ")\n\n")

	//generate command
	for sname, s := range g.services {
		fmt.Fprintf(g.buff, "var %sCommand = cli.Command{", sname)
		fmt.Fprintf(g.buff, `Name: "%s",`, sname)
		fmt.Fprintf(g.buff, `Subcommands: []cli.Command{`)
		for mname := range s.methods {
			fmt.Fprintf(g.buff, "%sSubCommand, ", mname)
		}
		fmt.Fprintf(g.buff, `},`)
		fmt.Fprintf(g.buff, "}\n\n")

		for mname, m := range s.methods {

			argInit, flags, err := g.generateFlagsAndArgs(m)
			if err != nil {
				return err
			}

			fmt.Fprintf(g.buff, "var %sSubCommand = cli.Command{", mname)
			fmt.Fprintf(g.buff, `Name: "%s",`, mname)
			fmt.Fprintf(g.buff, `Flags: []cli.Flag{%s},`, flags)
			fmt.Fprintf(g.buff, `Action: func(ctx *cli.Context) {`)

			//@todo allow custom remote addr
			//@todo generate command factory methods instead
			//@todo provide custom logger possibility
			//@todo custom post action (what to do with reply pointer)

			replyTypeName := m.output.String()
			if named, ok := m.output.(*types.Named); ok {
				replyTypeName = named.Obj().Name()
			}

			fmt.Fprintf(g.buff, `

        client, err := rpc.DialHTTP("tcp", "127.0.0.1:1234")
      	if err != nil {
      		log.Fatal("dialing:", err)
      	}

      	%s
      	var reply %s
      	err = client.Call("%s.%s", args, &reply)
      	if err != nil {
      		log.Fatal("%s error:", err)
      	}

        log.Println(reply)

      `, argInit, replyTypeName, sname, mname, sname)

			fmt.Fprintf(g.buff, `},`)
			fmt.Fprintf(g.buff, "}\n\n")
		}
	}

	//expose a list of all generated commands
	fmt.Fprintf(g.buff, "var GeneratedCommands = []cli.Command{")
	for sname := range g.services {
		fmt.Fprintf(g.buff, "%sCommand, ", sname)
	}
	fmt.Fprintf(g.buff, "}")

	return nil
}

func (g *Generator) generateFlagsAndArgs(method *RPCMethod) (string, string, error) {
	arginit := ""
	flagInit := ""

	argtype := method.input
	expectPointer := false
	if argpt, ok := argtype.(*types.Pointer); ok {
		expectPointer = true
		argtype = argpt.Elem()
	}

	//arg field to cli option type
	options := map[string]string{}

	var typeName string
	if argnamed, ok := argtype.(*types.Named); ok {
		typeName = argnamed.Obj().Name()

		//if the underlying type is a struct, create options from exported fields
		if argstruct, ok := argnamed.Underlying().(*types.Struct); ok {

			for i := 0; i < argstruct.NumFields(); i++ {
				if !argstruct.Field(i).Exported() {
					continue
				}

				//we only do basic field types for now
				f := argstruct.Field(i)
				if basicfield, ok := f.Type().(*types.Basic); ok {
					options[f.Name()] = strings.Title(basicfield.String())
				} else {
					// @todo show warning?
					// log.Printf("warning: RPC arg struct type '%s' has an exported field '%s' that is not a basic type (int, string etc), skipping...", typeName, f.Name())
				}
			}

		} else if argbasic, ok := argnamed.Underlying().(*types.Basic); ok {
			_ = argbasic
			//how can know how to typecast from strings (which args always are)
			//ctx.Args().Get(0)

			// log.Printf("ABC: %s", argbasic.String())
			//named type is not a struct, perhaps basic type?
			//in any case, @todo use as arg
		}

	} else if argbasic, ok := argtype.(*types.Basic); ok {
		typeName = argbasic.Name()
	} else {
		log.Fatalf("Unexpected input type: %T", argtype)
	}

	if expectPointer {
		arginit = fmt.Sprintf("args := new(%s)\n", typeName)
	} else {
		arginit = fmt.Sprintf("var args %s\n", typeName)
	}

	for fname, flagtype := range options {
		arginit += fmt.Sprintf(`args.%s = ctx.%s("%s")`+"\n", fname, flagtype, fname)
		flagInit += fmt.Sprintf(`cli.%sFlag{Name: "%s"},`, flagtype, fname)
	}

	return arginit, flagInit, nil
}

//Extract rpc services and methods by analysing type definitions
func (g *Generator) Extract() error {
	rpcmethods := map[string]*RPCMethod{}

	for _, obj := range g.pkg.defs {
		if obj == nil {
			continue
		}

	DEFSWITCH:
		switch t := obj.Type().(type) {

		//as per "net/rpc" are we looking for methods with the following:
		// [x] the fn is a method (has receiver)
		// [x] the method's type is exported.
		// [x] the method is exported.
		// [x] the method has two arguments...
		// [x] both exported (or builtin) types.
		// [x] the method's second argument is a pointer.
		// [x] the method has return type error.
		case *types.Signature:

			//needs to be a method (have a receiver) and be exported
			if t.Recv() == nil || !obj.Exported() {
				break
			}

			//receiver must be named and exported
			var recvt types.Type
			if recvpt, ok := t.Recv().Type().(*types.Pointer); ok {
				recvt = recvpt.Elem()
			} else {
				recvt = t.Recv().Type()
			}

			recv, ok := recvt.(*types.Named)
			if ok {
				if !recv.Obj().Exported() {
					break
				}
			}

			//method must have two params
			if t.Params().Len() != 2 {
				break
			}

			//all param types must be exported or builtin
			for i := 0; i < t.Params().Len(); i++ {
				var paramt types.Type
				if parampt, ok := t.Params().At(i).Type().(*types.Pointer); ok {
					paramt = parampt.Elem()
				} else {

					//second arg must be a pointer
					if i == 1 {
						break DEFSWITCH
					}

					paramt = t.Params().At(i).Type()
				}

				//if param type is Named, it must be exported, else it must be buildin
				if paramnamedt, ok := paramt.(*types.Named); ok {
					if !paramnamedt.Obj().Exported() {
						break DEFSWITCH
					}

				} else if strings.Contains(paramt.String(), ".") {
					break DEFSWITCH
				}
			}

			//must have one result: error
			if t.Results().Len() != 1 || t.Results().At(0).Type().String() != "error" {
				break
			}

			rpcmethods[obj.Name()] = &RPCMethod{
				input:  t.Params().At(0).Type(),
				output: t.Params().At(1).Type().(*types.Pointer).Elem(), //this is checked above
				recv:   recv,
				sig:    t,
			}
		}
	}

	for n, rpcm := range rpcmethods {
		rpcs, ok := g.services[rpcm.recv.Obj().Name()]
		if !ok {
			rpcs = &RPCService{
				methods: map[string]*RPCMethod{},
			}

			g.services[rpcm.recv.Obj().Name()] = rpcs
		}

		rpcs.methods[n] = rpcm
	}

	log.Printf("%+v", g.services["Arith"].methods)

	return nil
}

//Parse all .go files in a directory
func (g *Generator) Parse(dir string) error {
	g.pkg = &Package{}

	fis, err := ioutil.ReadDir(dir)
	if err != nil {
		return fmt.Errorf("Failed to read the directory '%s': %v", dir, err)
	}

	astFiles := []*ast.File{}
	fset := token.NewFileSet()
	for _, fi := range fis {
		if !strings.HasSuffix(fi.Name(), ".go") {
			continue
		}

		astFile, err := parser.ParseFile(fset, fi.Name(), nil, parser.ParseComments)
		if err != nil {
			return fmt.Errorf("Failed to parse file '%s': %v", fi.Name(), err)
		}

		astFiles = append(astFiles, astFile)
		g.pkg.files = append(g.pkg.files, &File{
			pkg:  g.pkg,
			file: astFile,
		})
	}

	if len(g.pkg.files) == 0 {
		log.Fatalf("%s: no buildable Go files", dir)
	}

	g.pkg.dir = dir
	g.pkg.name = g.pkg.files[0].file.Name.Name
	return g.pkg.check(fset, astFiles)
}