gen.go

// Copyright 2019 the u-root Authors. 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 (
	"encoding/json"
	"fmt"
	"go/ast"
	"go/parser"
	"go/token"
	"log"
	"os"
	"path/filepath"
	"regexp"
	"strings"
)

const mainStr = "main"

func uncomment(s string) string {
	s = strings.TrimSpace(s)
	if len(s) == 0 {
		return ""
	}
	if strings.HasPrefix(s, "/*") {
		return strings.TrimSpace(s[2 : len(s)-2])
	}
	// comment starts with '//'
	s = s[2:]
	s = strings.Replace(s, "\n// ", "\n", -1)
	s = strings.Replace(s, "\n//", "\n", -1)
	return strings.TrimSpace(s)
}

func extractMan(name string) (string, error) {
	fset := token.NewFileSet()
	src, err := parser.ParseFile(fset, name, nil, parser.ParseComments)
	if err != nil {
		return "", err
	}

	if src.Name.Name != mainStr {
		return "", fmt.Errorf("not a main package")
	}

	hasMainFunc := false
	for _, decl := range src.Decls {
		f, ok := decl.(*ast.FuncDecl)
		hasMainFunc = hasMainFunc || ok && f.Name.Name == mainStr
	}
	if !hasMainFunc {
		return "", fmt.Errorf("file doesn't contain func main()")
	}
	if len(src.Comments) == 0 {
		return "", fmt.Errorf("file doesn't contain comments")
	}
	var man string
	// First comment group is expected to be copyright notice.
	for _, cg := range src.Comments[1:] {
		for _, comment := range cg.List {
			if comment.Slash > src.Name.Pos() {
				break
			}
			man += comment.Text + "\n"
		}
	}
	return uncomment(man), nil
}

func walk(mans map[string]string, root string) (err error) {
	re, err := regexp.Compile("_.*\\.go$")
	if err != nil {
		return fmt.Errorf("error compiling regexp: %v", err)
	}
	err = filepath.Walk(root, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		if info.IsDir() {
			return nil
		}
		dir := filepath.Dir(path)
		// Depth 1.
		if filepath.Dir(dir) != root {
			return nil
		}
		name := info.Name()
		if !strings.HasSuffix(name, ".go") {
			return nil
		}
		if re.MatchString(name) {
			return nil
		}
		cmd := filepath.Base(dir)
		man, err := extractMan(path)
		if err == nil && len(man) > 0 {
			mans[cmd] = man
		}
		return nil
	})
	if err != nil {
		return fmt.Errorf("error walking the path %q: %v", root, err)
	}
	return nil
}

// writeFile saves man data to a go file in JSON format.
// JSON is used to simplify the process of escaping
// special characters.
func writeFile(fname string, mans map[string]string) error {
	dir := filepath.Dir(fname)
	if _, err := os.Stat(dir); os.IsNotExist(err) {
		if err := os.Mkdir(dir, 0775); err != nil {
			return err
		}
	}
	f, err := os.Create(fname)
	if err != nil {
		return err
	}
	defer f.Close()
	b, err := json.Marshal(mans)
	if err != nil {
		return err
	}
	if _, err := fmt.Fprintln(f, "// Code generated by man/gen/gen.go. DO NOT EDIT."); err != nil {
		return err
	}
	if _, err := fmt.Fprintln(f, "package data"); err != nil {
		return err
	}
	if _, err := fmt.Fprintln(f); err != nil {
		return err
	}
	_, err = fmt.Fprintf(f, "var Data = %q\n", b)
	return err
}

func main() {
	mans := make(map[string]string)
	l := len(os.Args)
	for _, root := range os.Args[1 : l-1] {
		if err := walk(mans, root); err != nil {
			log.Fatal(err)
		}
	}
	dest := os.Args[l-1]
	if err := writeFile(dest, mans); err != nil {
		log.Fatal(err)
	}
}