generator.go

package hero

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strings"
	"sync"
)

var formatMap = map[string]string{
	String:    "%s",
	Interface: "fmt.Sprintf(\"%%v\", %s)",
	Int:       "hero.FormatInt(int64(%s), buffer)",
	Uint:      "hero.FormatUint(uint64(%s), buffer)",
	Float:     "hero.FormatFloat(float64(%s), buffer)",
	Bool:      "hero.FormatBool(%s, buffer)",
}

func writeToFile(path string, buffer *bytes.Buffer) {
	err := ioutil.WriteFile(path, buffer.Bytes(), os.ModePerm)
	if err != nil {
		panic(err)
	}
}

func genAbsPath(path string) string {
	if !filepath.IsAbs(path) {
		path, _ = filepath.Abs(path)
	}
	return path
}

func checkError(err error) {
	if err != nil {
		log.Fatal(err)
	}
}

// gen generates code to buffer.
func gen(n *node, buffer *bytes.Buffer) {
	for _, child := range n.children {
		switch child.t {
		case TypeCode:
			buffer.Write(child.chunk.Bytes())
		case TypeHTML:
			buffer.WriteString(fmt.Sprintf(
				"buffer.WriteString(`%s`)",
				child.chunk.String(),
			))
		case TypeRawValue, TypeEscapedValue:
			var format string

			switch child.subtype {
			case Int, Uint, Float, Bool, String, Interface:
				format = formatMap[child.subtype]
				if child.subtype != String &&
					child.subtype != Interface {
					goto WriteFormat
				}
			case Bytes:
				if child.t == TypeRawValue {
					format = "buffer.Write(%s)"
					goto WriteFormat
				}
				format = "*(*string)(unsafe.Pointer(&(%s)))"
			default:
				log.Fatal("unknown value type: " + child.subtype)
			}

			if child.t == TypeEscapedValue {
				format = fmt.Sprintf(
					"hero.EscapeHTML(%s, buffer)", format,
				)
			} else {
				format = fmt.Sprintf(
					"buffer.WriteString(%s)", format,
				)
			}

		WriteFormat:
			buffer.WriteString(fmt.Sprintf(format, child.chunk.String()))
		case TypeBlock, TypeInclude:
			gen(child, buffer)
		default:
			continue
		}

		buffer.WriteByte(BreakLine)
	}
}

// Generate generates go code from source to test. pkgName represents the
// package name of the generated code.
func Generate(source, dest, pkgName string) {
	defer cleanGlobal()

	source, dest = genAbsPath(source), genAbsPath(dest)

	stat, err := os.Stat(source)
	checkError(err)

	fmt.Println("Parsing...")
	if stat.IsDir() {
		parseDir(source)
	} else {
		source, file := filepath.Split(source)
		parseFile(source, file)
	}

	stat, err = os.Stat(dest)
	if os.IsNotExist(err) {
		if err = os.MkdirAll(dest, os.ModePerm); err != nil {
			log.Fatal(err)
		}
	} else if !stat.IsDir() {
		log.Fatal(dest + " is not dir")
	} else if err != nil {
		log.Fatal(err)
	}

	fmt.Println("Generating...")

	var wg sync.WaitGroup
	for path, n := range parsedNodes {
		wg.Add(1)

		fileName := filepath.Join(dest, fmt.Sprintf(
			"%s.go",
			strings.Join(strings.Split(path[len(source)+1:], "/"), "_"),
		))

		go func(n *node, source, fileName string) {
			defer wg.Done()

			buffer := bytes.NewBufferString(`
				// Code generated by hero.
			`)
			buffer.WriteString(fmt.Sprintf("// source: %s", source))
			buffer.WriteString(`
				// DO NOT EDIT!
			`)
			buffer.WriteString(fmt.Sprintf("package %s\n", pkgName))
			buffer.WriteString(`
				import "html"
				import "unsafe"

				import "github.com/shiyanhui/hero"
			`)

			imports := n.childrenByType(TypeImport)
			for _, item := range imports {
				buffer.Write(item.chunk.Bytes())
			}

			definitions := n.childrenByType(TypeDefinition)
			if len(definitions) == 0 {
				writeToFile(fileName, buffer)
				return
			}

			buffer.Write(definitions[0].chunk.Bytes())
			buffer.WriteString(`{
			`)
			gen(n, buffer)
			buffer.WriteString(`
			}`)

			writeToFile(fileName, buffer)
		}(n, path, fileName)
	}
	wg.Wait()

	fmt.Println("Executing goimports...")
	execCommand("goimports -w " + dest)

	fmt.Println("Executing go vet...")
	execCommand("go tool vet -v " + dest)
}