golang.go

package generator

import (
	"fmt"
	"io"
	"sync"
)

// WriteGolangHeader writes the file header for a generated Go source file to
// the specified io.Writer.
func WriteGolangHeader(out io.Writer, src, pkg string, imports ...string) error {
	if _, err := fmt.Fprintln(out, "// Code generated by go generate."); err != nil {
		return err
	}
	fmt.Fprintln(out, "// source:", src)
	fmt.Fprintln(out, "// DO NOT EDIT!")
	fmt.Fprintln(out)

	fmt.Fprintln(out, "package", pkg)
	_, err := fmt.Fprintln(out) // might be final line if len(imports) == 0

	if len(imports) > 1 {
		fmt.Fprintln(out, "import (")
		for _, i := range imports {
			fmt.Fprintf(out, "\t%q", i)
			fmt.Fprintln(out)
		}
		fmt.Fprintln(out, ")")
		_, err = fmt.Fprintln(out)
	} else if len(imports) > 0 {
		fmt.Fprintf(out, "import %q", imports[0])
		fmt.Fprintln(out)
		_, err = fmt.Fprintln(out)
	}

	return err
}

// GolangCommentWriter returns an io.WriteCloser to wrap an existing
// io.Writer.  Data written will be formatted as Golang comments, including
// word wrapping.  (If word wrapping is not desired, maxLineLen can be set to
// zero.)
//
// indentLevel is the number of tab stops by which to indent the output.
//
// The caller *must* call Close() on the returned object when done writing, in
// order to flush any buffered text.  Doing so will not close the underlying
// io.Writer.
func GolangCommentWriter(out io.Writer, indentLevel, maxLineLen int) io.WriteCloser {
	return &golangCommentWriter{
		out:         out,
		indentLevel: indentLevel,
		maxLineLen:  maxLineLen,
	}
}

// golangCommentWriter is the package-private implementation returned by
// GolangCommentWriter().
type golangCommentWriter struct {
	out         io.Writer  // where to send formatted output
	indentLevel int        // number of leading tabs
	maxLineLen  int        // maximum characters per line
	sp          []byte     // whitespace yet to be written
	word        []byte     // word yet to be written
	col         int        // column (0 indexed)
	expectCol   int        // column if sp+word is written
	needCopy    bool       // set if Write() returns without flush
	mutex       sync.Mutex // Write() and Close() are not thread-safe
}

// writeNewline outputs a newline, and performs related housekeeping.
func (w *golangCommentWriter) writeNewline() error {
	w.col = 0
	w.expectCol = 0
	_, err := w.out.Write([]byte{'\n'})
	return err
}

// writeLineLeader outputs indentation and start-of-comment token.
func (w *golangCommentWriter) writeLineLeader() error {
	for i := 0; i < w.indentLevel; i++ {
		if _, err := w.out.Write([]byte{'\t'}); err != nil {
			return err
		}
		w.col += 8
	}

	if _, err := w.out.Write([]byte{'/', '/'}); err != nil {
		return err
	}
	w.col += 2

	return nil
}

// flushWord outputs the buffered whitespace and word, and resets the buffers.
func (w *golangCommentWriter) flushWord() error {
	if w.col == 0 {
		if err := w.writeLineLeader(); err != nil {
			return err
		}

		if len(w.sp) == 0 && len(w.word) > 0 {
			if _, err := w.out.Write([]byte{' '}); err != nil {
				return err
			}
			w.col++
		}
	}

	for i := 0; i < len(w.sp); i++ {
		switch w.sp[i] {
		case ' ':
			w.col++
		case '\t':
			w.col += 8 - w.col%8
		}

		if _, err := w.out.Write([]byte{w.sp[i]}); err != nil {
			return err
		}
	}

	if _, err := w.out.Write(w.word); err != nil {
		return err
	}

	w.col += len(w.word)
	w.expectCol = w.col

	w.sp = []byte{}
	w.word = []byte{}
	w.needCopy = false

	return nil
}

// extendOrAppend returns a subslice from src, unless needCopy is set, in
// which case it copies from src to dest.
func (w *golangCommentWriter) extendOrAppend(dest, src []byte, i, j int) []byte {
	if w.needCopy {
		return append(dest, src[j])
	}
	return src[i : j+1]
}

func (w *golangCommentWriter) Write(b []byte) (int, error) {
	w.mutex.Lock()
	defer w.mutex.Unlock()

	for i, j := 0, 0; j < len(b); j++ {
		switch b[j] {
		case ' ':
			if len(w.word) > 0 {
				if err := w.flushWord(); err != nil {
					return i, err
				}
				i = j
			}
			w.sp = w.extendOrAppend(w.sp, b, i, j)
			w.expectCol++

		case '\t':
			if len(w.word) > 0 {
				if err := w.flushWord(); err != nil {
					return i, err
				}
				i = j
			}
			w.sp = w.extendOrAppend(w.sp, b, i, j)
			w.expectCol += 8 - w.expectCol%8

		case '\n':
			if len(w.word) > 0 {
				if err := w.flushWord(); err != nil {
					return i, err
				}
			} else if w.col == 0 {
				// blank line
				w.sp = []byte{} // discard line's trailing whitespace
				if err := w.writeLineLeader(); err != nil {
					return i, err
				}
			}
			if err := w.writeNewline(); err != nil {
				return i, err
			}
			i = j + 1

		default:
			w.word = w.extendOrAppend(w.word, b, i+len(w.sp), j)
			w.expectCol++
		}

		if w.maxLineLen > 0 && w.expectCol >= w.maxLineLen && w.col != 0 {
			// wrap line
			if err := w.writeNewline(); err != nil {
				return i, err
			}
		}
	}

	if !w.needCopy && (len(w.sp) > 0 || len(w.word) > 0) {
		// need to make a copy of unflushed data, since b may not be
		// around when we want to flush it later
		w.sp = append([]byte{}, w.sp...)
		w.word = append([]byte{}, w.word...)
		w.needCopy = true
	}

	return len(b), nil
}

func (w *golangCommentWriter) Close() error {
	w.mutex.Lock()
	defer w.mutex.Unlock()

	var err error
	if len(w.word) > 0 {
		err = w.flushWord()
	} // else ignore trailing whitespace
	if err == nil && w.col != 0 {
		err = w.writeNewline()
	}
	return err
}