<commit_before>package quotes

import (
	// "fmt"
	// "html/template"
	// "net/http"

	"github.com/GoogleCloudPlatform/go-endpoints/endpoints"
	"golang.org/x/net/context"
	// "google.golang.org/appengine"
	"google.golang.org/appengine/datastore"
)

type Quote struct {
	UID     *datastore.Key `json:"uid" datastore:"-"`
	Author  string         `json:"author"`
	Message string         `json:"message"`
}

type QuotesAPI struct{}

type Quotes struct {
	Quotes []Quote `json:"quotes"`
}

type AddRequest struct {
	Author  string
	Message string
}

func (QuotesAPI) Add(c context.Context, r *AddRequest) (*Quote, error) {
	k := datastore.NewIncompleteKey(c, "Quote", quoteKey(c))

	t := &Quote{Author: r.Author, Message: r.Message}

	k, err := datastore.Put(c, k, t)
	if err != nil {
		return nil, err
	}
	t.UID = k
	return t, nil
}

func (QuotesAPI) List(c context.Context) (*Quotes, error) {
	quotes := []Quote{}
	// If we omitted the .Ancestor from this query there would be
	// a slight chance that Quote that had just been written would not
	// show up in a query.
	keys, err := datastore.NewQuery("Quote").Ancestor(quoteKey(c)).GetAll(c, &quotes)
	if err != nil {
		return nil, err
	}

	for i, k := range keys {
		quotes[i].UID = k
	}
	return &Quotes{quotes}, nil
}

func init() {
	// register the quotes API with cloud endpoints.
	api, err := endpoints.RegisterService(QuotesAPI{}, "quotesService", "v1", "Quotes API", true)
	if err != nil {
		panic(err)
	}

	// adapt the name, method, and path for each method.
	info := api.MethodByName("List").Info()
	info.Name, info.HTTPMethod, info.Path = "getQuotes", "GET", "quotesService"

	info = api.MethodByName("Add").Info()
	info.Name, info.HTTPMethod, info.Path = "addQuote", "POST", "quotesService"

	// start handling cloud endpoint requests.
	endpoints.HandleHTTP()
}

func quoteKey(c context.Context) *datastore.Key {
	return datastore.NewKey(c, "Quote", "default_quote", 0, nil)
}

// func root(w http.ResponseWriter, r *http.Request) {
// 	fmt.Fprint(w, rootForm)
// }

// const rootForm = `
//
//
//
Enter Author's quote

// 			<form action="/quote" method="post" accept-charset="utf-8">
// 				<input type="text" name="author" value="Author's Name: " id="author"/>
// 				<input type="text" name="quote" value="Write quote..." id="quote"/>
// 				<input type="submit" value="Submit Quote"/>
// 			</form>
//
//
// `

// var quoteTemplate = template.Must(template.New("quote").Parse(quoteTemplateHTML))

// func quote(w http.ResponseWriter, r *http.Request) {
// 	c := appengine.NewContext(r)
// 	q1 := Quote{
// 		Author:  r.FormValue("author"),
// 		Message: r.FormValue("quote"),
// 	}

// 	key, err := datastore.Put(c, datastore.NewIncompleteKey(c, "Quote", quoteKey(c)), &q1)
// 	if err != nil {
// 		http.Error(w, err.Error(), http.StatusInternalServerError)
// 		return
// 	}

// 	var q2 Quote
// 	err1 := datastore.Get(c, key, &q2)
// 	if err1 != nil {
// 		http.Error(w, err1.Error(), http.StatusInternalServerError)
// 		return
// 	}

// 	err2 := quoteTemplate.Execute(w, q2.Message)
// 	if err2 != nil {
// 		http.Error(w, err2.Error(), http.StatusInternalServerError)
// 		return
// 	}
// }

// const quoteTemplateHTML = `
//
//
//
You wrote:

//
{{html .}}

//
//
// `

// func fetch(w http.ResponseWriter, r *http.Request) {
// 	c := appengine.NewContext(r)
// 	q := datastore.NewQuery("Quote").Ancestor(quoteKey(c)).Limit(10)

// 	quotes := make([]Quote, 0, 10)
// 	_, err := q.GetAll(c, &quotes)
// 	if err != nil {
// 		http.Error(w, err.Error(), http.StatusInternalServerError)
// 		return
// 	}

// 	err1 := fetchTemplate.Execute(w, quotes)
// 	if err1 != nil {
// 		http.Error(w, err1.Error(), http.StatusInternalServerError)
// 	}
// }

// var fetchTemplate = template.Must(template.New("fetch").Parse(`
//
//
//     <title>All the Quotes</title>
//
//
//
All the Quotes

//     {{range .}}
//       {{with .Author}}
//
{{.}} wrote:

//       {{else}}
//
An anonymous person wrote:

//       {{end}}
//
{{.Message}}

//     {{end}}
//
//
// `))
<commit_msg>refactoring<commit_after>package quotes

import (
	"github.com/GoogleCloudPlatform/go-endpoints/endpoints"
	"golang.org/x/net/context"
	"google.golang.org/appengine/datastore"
)

type Quote struct {
	UID     *datastore.Key `json:"uid" datastore:"-"`
	Author  string         `json:"author"`
	Message string         `json:"message"`
}

type QuotesAPI struct{}

type Quotes struct {
	Quotes []Quote `json:"quotes"`
}

type AddRequest struct {
	Author  string
	Message string
}

func (QuotesAPI) Add(c context.Context, r *AddRequest) (*Quote, error) {
	k := datastore.NewIncompleteKey(c, "Quote", quoteKey(c))

	t := &Quote{Author: r.Author, Message: r.Message}

	k, err := datastore.Put(c, k, t)
	if err != nil {
		return nil, err
	}
	t.UID = k
	return t, nil
}

func (QuotesAPI) List(c context.Context) (*Quotes, error) {
	quotes := []Quote{}
	// If we omitted the .Ancestor from this query there would be
	// a slight chance that Quote that had just been written would not
	// show up in a query.
	keys, err := datastore.NewQuery("Quote").Ancestor(quoteKey(c)).GetAll(c, &quotes)
	if err != nil {
		return nil, err
	}

	for i, k := range keys {
		quotes[i].UID = k
	}
	return &Quotes{quotes}, nil
}

func quoteKey(c context.Context) *datastore.Key {
	return datastore.NewKey(c, "Quote", "default_quote", 0, nil)
}

func init() {
	// register the quotes API with cloud endpoints.
	api, err := endpoints.RegisterService(QuotesAPI{}, "quotesService", "v1", "Quotes API", true)
	if err != nil {
		panic(err)
	}

	// adapt the name, method, and path for each method.
	info := api.MethodByName("List").Info()
	info.Name, info.HTTPMethod, info.Path = "getQuotes", "GET", "quotesService"

	info = api.MethodByName("Add").Info()
	info.Name, info.HTTPMethod, info.Path = "addQuote", "POST", "quotesService"

	// start handling cloud endpoint requests.
	endpoints.HandleHTTP()
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"fmt"
	"math"
	"strconv"

	"github.com/mjibson/go-dsp/spectral"
	"github.com/wayneashleyberry/terminal-dimensions"
)

// Rasterizer fills one or more buffers with discreet audio samples
type Rasterizer func([][]float64)

// RenderMono rasterizes a TFunc into a mono Portaudio channel
func RenderMono(tfunc TFunc) Rasterizer {
	var t uint64
	return func(out [][]float64) {
		for i := range out[0] {
			var tc = Timecode(t + uint64(i))
			out[0][i] = float64(tfunc(tc))
		}
		t = t + uint64(len(out[0]))
	}
}

// RenderStereo rasterizes two TFuncs into a stereo Portaudio channel
func RenderStereo(left TFunc, right TFunc) Rasterizer {
	var t uint64
	return func(out [][]float64) {
		for i := range out[0] {
			var tc = Timecode(t + uint64(i))
			out[0][i] = float64(left(tc))
			out[1][i] = float64(right(tc))
		}
		t = t + uint64(len(out[0]))
	}
}

// Visualize renders a Rasterizer's power spectral density to the screen
func Visualize(sampleRate Frequency, colormap Colormap, inner Rasterizer) Rasterizer {
	termWidth := termWidth()

	segmentLength := termWidth*2 - 2 // FFT results in segmentLength/2+1 output samples

	// stores the rendered output
	lineQueueDepth := 8
	lineQueue := make(chan []float64, lineQueueDepth)
	go printer(lineQueue, segmentLength, colormap)

	options := newOptions(segmentLength)

	return func(out [][]float64) {
		// adjust PSD options if they are too agressive for the input data
		if segmentLength > len(out[0]) {
			options = newOptions(len(out[0]))
		}

		// compute raster values from the inner rasterizer
		inner(out)

		// divide the rasterized sound into [segmentLength] sample segments
		var segments [][]float64
		for _, chanBuf := range out {
			chanSegs := spectral.Segment(chanBuf, options.NFFT, options.Noverlap)
			segments = append(segments, chanSegs...)
		}

		// subtract the mean of each segment from each segment's sample
		// https://www.mathworks.com/matlabcentral/answers/267658-why-am-i-getting-huge-values-on-low-frequencies-of-my-psd#comment_342452
		for i := range segments {
			segment := segments[i]
			var avg float64
			for j := range segment {
				avg += segment[j]
			}
			avg /= float64(len(segment))
			for j := range segment {
				segment[j] -= avg
			}
		}

		// compute the averge PSD for all segments
		psd := make([]float64, options.Pad/2+1) // output of FFT is len(time domain)/2+1
		var freqs []float64
		for _, seg := range segments {
			var pxx []float64
			// compute the PSD for each segment
			pxx, freqs = spectral.Pwelch(seg, float64(sampleRate), options)
			for i := range pxx {
				psd[i] += pxx[i]
			}
			// println()
		}

		// normalize the PSD
		for i := range psd {
			// divide by number of segments to get average
			psd[i] /= float64(len(segments))
			// linearize the PSD
			psd[i] *= math.Sqrt(freqs[i])
		}

		// skip line rendering if there are too many lines waiting to be rendered
		if len(lineQueue) < lineQueueDepth {
			lineQueue <- psd
		}
	}
}

func newOptions(segmentLength int) *spectral.PwelchOptions {
	overlap := segmentLength / 2
	options := new(spectral.PwelchOptions)
	options.Scale_off = true
	options.Noverlap = overlap
	options.NFFT = segmentLength
	options.Pad = segmentLength
	return options
}

// Renders printed lines from a queue
func printer(lineQueue chan []float64, segmentLength int, colormap Colormap) {
	outLine := bytes.NewBuffer(make([]byte, 30*segmentLength, 30*segmentLength))
	for {
		psd := <-lineQueue
		outLine.Reset()
		for i := range psd {
			r, g, b := colormap(psd[i])
			outLine.WriteString("\x1b[48;2;")
			outLine.WriteString(strconv.Itoa(int(r)))
			outLine.WriteString(";")
			outLine.WriteString(strconv.Itoa(int(g)))
			outLine.WriteString(";")
			outLine.WriteString(strconv.Itoa(int(b)))
			outLine.WriteString("m \x1b[0m")
		}
		fmt.Println(outLine.String())
	}
}

// get the width of the terminal - always even
func termWidth() int {
	w, e := terminaldimensions.Width()
	if e != nil {
		return 80
	}

	// make w even
	if w%2 != 0 {
		w--
	}
	return int(w)
}
<commit_msg>Render PSD on a logarithmic scale<commit_after>package main

import (
	"bytes"
	"fmt"
	"math"
	"strconv"

	"github.com/mjibson/go-dsp/spectral"
	"github.com/wayneashleyberry/terminal-dimensions"
)

// Rasterizer fills one or more buffers with discreet audio samples
type Rasterizer func([][]float64)

// RenderMono rasterizes a TFunc into a mono Portaudio channel
func RenderMono(tfunc TFunc) Rasterizer {
	var t uint64
	return func(out [][]float64) {
		for i := range out[0] {
			var tc = Timecode(t + uint64(i))
			out[0][i] = float64(tfunc(tc))
		}
		t = t + uint64(len(out[0]))
	}
}

// RenderStereo rasterizes two TFuncs into a stereo Portaudio channel
func RenderStereo(left TFunc, right TFunc) Rasterizer {
	var t uint64
	return func(out [][]float64) {
		for i := range out[0] {
			var tc = Timecode(t + uint64(i))
			out[0][i] = float64(left(tc))
			out[1][i] = float64(right(tc))
		}
		t = t + uint64(len(out[0]))
	}
}

// Visualize renders a Rasterizer's power spectral density to the screen
func Visualize(sampleRate Frequency, colormap Colormap, inner Rasterizer) Rasterizer {
	termWidth := termWidth()

	segmentLength := termWidth*2 - 2 // FFT results in segmentLength/2+1 output samples

	// stores the rendered output
	lineQueueDepth := 8
	lineQueue := make(chan []float64, lineQueueDepth)
	go printer(lineQueue, segmentLength, colormap)

	options := newOptions(segmentLength)

	return func(out [][]float64) {
		// adjust PSD options if they are too agressive for the input data
		if segmentLength > len(out[0]) {
			options = newOptions(len(out[0]))
		}

		// compute raster values from the inner rasterizer
		inner(out)

		// divide the rasterized sound into [segmentLength] sample segments
		var segments [][]float64
		for _, chanBuf := range out {
			chanSegs := spectral.Segment(chanBuf, options.NFFT, options.Noverlap)
			segments = append(segments, chanSegs...)
		}

		// subtract the mean of each segment from each segment's sample
		// https://www.mathworks.com/matlabcentral/answers/267658-why-am-i-getting-huge-values-on-low-frequencies-of-my-psd#comment_342452
		for i := range segments {
			segment := segments[i]
			var avg float64
			for j := range segment {
				avg += segment[j]
			}
			avg /= float64(len(segment))
			for j := range segment {
				segment[j] -= avg
			}
		}

		// compute the averge PSD for all segments
		psd := make([]float64, options.Pad/2+1) // output of FFT is len(time domain)/2+1
		var freqs []float64
		for _, seg := range segments {
			var pxx []float64
			// compute the PSD for each segment
			pxx, freqs = spectral.Pwelch(seg, float64(sampleRate), options)
			for i := range pxx {
				psd[i] += pxx[i]
			}
			// println()
		}

		// normalize the PSD
		for i := range psd {
			// divide by number of segments to get average
			psd[i] /= float64(len(segments))
			// linearize the PSD. power is expressed in relationship to the square root of the frequency
			psd[i] *= math.Sqrt(freqs[i])
			// plot data intensity on a logarithmic scale between zero and 1
			psd[i] = math.Log10(psd[i] + 1)
		}

		// skip line rendering if there are too many lines waiting to be rendered
		if len(lineQueue) < lineQueueDepth {
			lineQueue <- psd
		}
	}
}

func newOptions(segmentLength int) *spectral.PwelchOptions {
	overlap := segmentLength / 2
	options := new(spectral.PwelchOptions)
	options.Scale_off = true
	options.Noverlap = overlap
	options.NFFT = segmentLength
	options.Pad = segmentLength
	return options
}

// Renders printed lines from a queue
func printer(lineQueue chan []float64, segmentLength int, colormap Colormap) {
	outLine := bytes.NewBuffer(make([]byte, 30*segmentLength, 30*segmentLength))
	for {
		psd := <-lineQueue
		outLine.Reset()
		for i := range psd {
			r, g, b := colormap(psd[i])
			outLine.WriteString("\x1b[48;2;")
			outLine.WriteString(strconv.Itoa(int(r)))
			outLine.WriteString(";")
			outLine.WriteString(strconv.Itoa(int(g)))
			outLine.WriteString(";")
			outLine.WriteString(strconv.Itoa(int(b)))
			outLine.WriteString("m \x1b[0m")
		}
		fmt.Println(outLine.String())
	}
}

// get the width of the terminal - always even
func termWidth() int {
	w, e := terminaldimensions.Width()
	if e != nil {
		return 80
	}

	// make w even
	if w%2 != 0 {
		w--
	}
	return int(w)
}
<|endoftext|>


<commit_before>package testing2

import (
	"fmt"
	"reflect"
	"testing"
)

// NoCheck means don't check this value
var NoCheck = struct{}{}

// NonNil means this value should not be nil
var NonNil = struct{}{}

type TestCase struct {
	state   bool
	expects [][]interface{}
	args    [][]reflect.Value
}

func Tests() *TestCase {
	return &TestCase{
		state: wantExpect,
	}
}

func Expect(expect ...interface{}) *TestCase {
	return Tests().Expect(expect...)
}

const (
	wantArg    = true
	wantExpect = !wantArg
)

func (test *TestCase) checkState(state bool) {
	if test.state != state {
		panic("don't call same function continully")
	}

	test.state = !state
}

func (test *TestCase) Expect(expect ...interface{}) *TestCase {
	test.checkState(wantExpect)

	test.expects = append(test.expects, expect)
	return test
}

func (test *TestCase) True() *TestCase {
	return test.Expect(true)
}

func (test *TestCase) False() *TestCase {
	return test.Expect(false)
}

func (test *TestCase) Nil() *TestCase {
	return test.Expect(nil)
}

func (test *TestCase) Arg(args ...interface{}) *TestCase {
	test.checkState(wantArg)

	argVals := make([]reflect.Value, len(args))
	for i := range args {
		argVals[i] = reflect.ValueOf(args[i])
	}

	test.args = append(test.args, argVals)

	return test
}

func (test *TestCase) Run(t testing.TB, fn interface{}) {
	for index := range test.expects {
		var args = test.args[index]
		var expects = test.expects[index]

		results := reflect.ValueOf(fn).Call(args)

		for i, r := range results {
			if expect := expects[i]; expect != NoCheck {
				indexs := fmt.Sprintf("(%d:%d)", index+1, i+1)
				if expect == nil {
					indexNil(t, 1, indexs, "nil", r.Interface())
				} else if expect == NonNil {
					indexNil(t, 1, indexs, "not nil", "nil")
				} else {
					indexDeepEq(t, 1, indexs, expect, r.Interface())
				}
			}
		}
	}
}

func indexNil(t testing.TB, skip int, index string, expect, got interface{}) {
	if !isNil(got) {
		indexErrorInfo(t, skip+1, index, expect, got, false)
	}
}

func indexNNil(t testing.TB, skip int, index string, expect, got interface{}) {
	if isNil(got) {
		indexErrorInfo(t, skip+1, index, expect, got, false)
	}
}

// indexEq assert expect and got is equal, else print error message
func indexDeepEq(t testing.TB, skip int, index string, expect, got interface{}) {
	if !reflect.DeepEqual(expect, got) {
		indexErrorInfo(t, skip+1, index, expect, got, true)
	}
}
<commit_msg>testing2: support multiple funcion calls<commit_after>package testing2

import (
	"fmt"
	"reflect"
	"testing"
)

// NoCheck means don't check this value
var NoCheck = struct{}{}

// NonNil means this value should not be nil
var NonNil = struct{}{}

type TestCase struct {
	state   bool
	expects [][]interface{}
	args    [][]reflect.Value
}

func Tests() *TestCase {
	return &TestCase{
		state: wantExpect,
	}
}

func Expect(expect ...interface{}) *TestCase {
	return Tests().Expect(expect...)
}

const (
	wantArg    = true
	wantExpect = !wantArg
)

func (test *TestCase) checkState(state bool) {
	if test.state != state {
		panic("don't call same function continully")
	}

	test.state = !state
}

func (test *TestCase) Expect(expect ...interface{}) *TestCase {
	test.checkState(wantExpect)

	test.expects = append(test.expects, expect)
	return test
}

func (test *TestCase) True() *TestCase {
	return test.Expect(true)
}

func (test *TestCase) False() *TestCase {
	return test.Expect(false)
}

func (test *TestCase) Nil() *TestCase {
	return test.Expect(nil)
}

func (test *TestCase) Arg(args ...interface{}) *TestCase {
	test.checkState(wantArg)

	argVals := make([]reflect.Value, len(args))
	for i := range args {
		argVals[i] = reflect.ValueOf(args[i])
	}

	test.args = append(test.args, argVals)

	return test
}

func (test *TestCase) Run(t testing.TB, fn ...interface{}) {
	for index := range test.expects {
		var expects = test.expects[index]

		var results = test.args[index]
		for _, f := range fn {
			results = reflect.ValueOf(f).Call(results)
		}

		for i, r := range results {
			if expect := expects[i]; expect != NoCheck {
				indexs := fmt.Sprintf("(%d:%d)", index+1, i+1)
				if expect == nil {
					indexNil(t, 1, indexs, "nil", r.Interface())
				} else if expect == NonNil {
					indexNil(t, 1, indexs, "not nil", "nil")
				} else {
					indexDeepEq(t, 1, indexs, expect, r.Interface())
				}
			}
		}
	}
}

func indexNil(t testing.TB, skip int, index string, expect, got interface{}) {
	if !isNil(got) {
		indexErrorInfo(t, skip+1, index, expect, got, false)
	}
}

func indexNNil(t testing.TB, skip int, index string, expect, got interface{}) {
	if isNil(got) {
		indexErrorInfo(t, skip+1, index, expect, got, false)
	}
}

// indexEq assert expect and got is equal, else print error message
func indexDeepEq(t testing.TB, skip int, index string, expect, got interface{}) {
	if !reflect.DeepEqual(expect, got) {
		indexErrorInfo(t, skip+1, index, expect, got, true)
	}
}
<|endoftext|>


<commit_before>// Copyright (C) 2014 Daniel Harrison

package hfile

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"log"
	"os"
	"sort"

	"github.com/edsrzf/mmap-go"
	"github.com/golang/snappy"
)

type Reader struct {
	mmap         mmap.MMap
	majorVersion uint32
	minorVersion uint32

	header Header
	index  []Block

	debug bool
}

type Header struct {
	index int

	fileInfoOffset             uint64
	dataIndexOffset            uint64
	dataIndexCount             uint32
	metaIndexOffset            uint64
	metaIndexCount             uint32
	totalUncompressedDataBytes uint64
	entryCount                 uint32
	compressionCodec           uint32
}

type Block struct {
	offset        uint64
	size          uint32
	firstKeyBytes []byte
}

func NewReaderFromConfig(cfg *CollectionConfig, debug bool) (*Reader, error) {
	return NewReader(cfg.Name, cfg.Path, cfg.Mlock, debug)
}

func NewReader(name, path string, lock, debug bool) (*Reader, error) {
	f, err := os.OpenFile(path, os.O_RDONLY, 0)

	if err != nil {
		return nil, err
	}

	hfile := new(Reader)
	hfile.debug = debug

	hfile.mmap, err = mmap.Map(f, mmap.RDONLY, 0)

	if err != nil {
		return nil, err
	}

	fi, err := f.Stat()
	if err != nil {
		return nil, err
	}

	if lock {
		mb := 1024.0 * 1024.0
		log.Printf("[Reader.NewReader] locking %s (%.02fmb)...\n", name, float64(fi.Size())/mb)
		if err = hfile.mmap.Lock(); err != nil {
			log.Printf("[Reader.NewReader] error locking %s: %s\n", name, err.Error())
			return nil, err
		}
		log.Printf("[Reader.NewReader] locked %s.\n", name)
	} else if debug {
		log.Printf("[Reader.NewReader] Not locking %s...\n", name)
	}

	v := binary.BigEndian.Uint32(hfile.mmap[len(hfile.mmap)-4:])
	hfile.majorVersion = v & 0x00ffffff
	hfile.minorVersion = v >> 24

	hfile.header, err = hfile.newHeader(hfile.mmap)
	if err != nil {
		return hfile, err
	}
	err = hfile.loadIndex(hfile.mmap)
	if err != nil {
		return hfile, err
	}

	return hfile, nil
}

func (r *Reader) PrintDebugInfo(out io.Writer) {
	fmt.Fprintln(out, "entries: ", r.header.entryCount)
	fmt.Fprintln(out, "blocks: ", len(r.index))
	for i, blk := range r.index {
		fmt.Fprintf(out, "\t#%d: %s (%v)\n", i, blk.firstKeyBytes, blk.firstKeyBytes)
	}
}

func (r *Reader) newHeader(mmap mmap.MMap) (Header, error) {
	header := Header{}

	if r.majorVersion != 1 || r.minorVersion != 0 {
		return header, fmt.Errorf("wrong version: %d.%d", r.majorVersion, r.minorVersion)
	}

	header.index = len(mmap) - 60
	buf := bytes.NewReader(mmap[header.index:])

	headerMagic := make([]byte, 8)
	buf.Read(headerMagic)
	if bytes.Compare(headerMagic, TrailerMagic) != 0 {
		return header, errors.New("bad header magic")
	}

	binary.Read(buf, binary.BigEndian, &header.fileInfoOffset)
	binary.Read(buf, binary.BigEndian, &header.dataIndexOffset)
	binary.Read(buf, binary.BigEndian, &header.dataIndexCount)
	binary.Read(buf, binary.BigEndian, &header.metaIndexOffset)
	binary.Read(buf, binary.BigEndian, &header.metaIndexCount)
	binary.Read(buf, binary.BigEndian, &header.totalUncompressedDataBytes)
	binary.Read(buf, binary.BigEndian, &header.entryCount)
	binary.Read(buf, binary.BigEndian, &header.compressionCodec)
	return header, nil
}

func (r *Reader) loadIndex(mmap mmap.MMap) error {

	dataIndexEnd := r.header.metaIndexOffset
	if r.header.metaIndexOffset == 0 {
		dataIndexEnd = uint64(r.header.index)
	}
	buf := bytes.NewReader(mmap[r.header.dataIndexOffset:dataIndexEnd])

	dataIndexMagic := make([]byte, 8)
	buf.Read(dataIndexMagic)
	if bytes.Compare(dataIndexMagic, IndexMagic) != 0 {
		return errors.New("bad data index magic")
	}

	for buf.Len() > 0 {
		dataBlock := Block{}

		binary.Read(buf, binary.BigEndian, &dataBlock.offset)
		binary.Read(buf, binary.BigEndian, &dataBlock.size)

		firstKeyLen, _ := binary.ReadUvarint(buf)
		dataBlock.firstKeyBytes = make([]byte, firstKeyLen)
		buf.Read(dataBlock.firstKeyBytes)

		r.index = append(r.index, dataBlock)
	}

	return nil
}

func After(a, b []byte) bool {
	return bytes.Compare(a, b) > 0
}

func (b *Block) IsAfter(key []byte) bool {
	return After(b.firstKeyBytes, key)
}

func (r *Reader) FindBlock(from int, key []byte) int {
	remaining := len(r.index) - from - 1
	if r.debug {
		log.Printf("[Reader.findBlock] cur %d, remaining %d\n", from, remaining)
	}

	if remaining <= 0 {
		if r.debug {
			log.Println("[Reader.findBlock] last block")
		}
		return from // s.cur is the last block, so it is only choice.
	}

	if r.index[from+1].IsAfter(key) {
		if r.debug {
			log.Println("[Reader.findBlock] next block is past key")
		}
		return from
	}

	offset := sort.Search(remaining, func(i int) bool {
		return r.index[from+i+1].IsAfter(key)
	})

	return from + offset
}

func (r *Reader) GetBlock(i int) (*bytes.Reader, error) {
	var buf *bytes.Reader

	block := r.index[i]

	switch r.header.compressionCodec {
	case CompressionNone:
		buf = bytes.NewReader(r.mmap[block.offset : block.offset+uint64(block.size)])
	case CompressionSnappy:
		uncompressedByteSize := binary.BigEndian.Uint32(r.mmap[block.offset : block.offset+4])
		if uncompressedByteSize != block.size {
			return nil, errors.New("mismatched uncompressed block size")
		}
		compressedByteSize := binary.BigEndian.Uint32(r.mmap[block.offset+4 : block.offset+8])
		compressedBytes := r.mmap[block.offset+8 : block.offset+8+uint64(compressedByteSize)]
		uncompressedBytes, err := snappy.Decode(nil, compressedBytes)
		if err != nil {
			return nil, err
		}
		buf = bytes.NewReader(uncompressedBytes)
	default:
		return nil, errors.New("Unsupported compression codec " + string(r.header.compressionCodec))
	}

	dataBlockMagic := make([]byte, 8)
	buf.Read(dataBlockMagic)
	if bytes.Compare(dataBlockMagic, DataMagic) != 0 {
		return nil, errors.New("bad data block magic")
	}

	return buf, nil
}
<commit_msg>optionally print start keys in hex<commit_after>// Copyright (C) 2014 Daniel Harrison

package hfile

import (
	"bytes"
	"encoding/binary"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"log"
	"os"
	"sort"

	"github.com/edsrzf/mmap-go"
	"github.com/golang/snappy"
)

type Reader struct {
	mmap         mmap.MMap
	majorVersion uint32
	minorVersion uint32

	header Header
	index  []Block

	debug bool
}

type Header struct {
	index int

	fileInfoOffset             uint64
	dataIndexOffset            uint64
	dataIndexCount             uint32
	metaIndexOffset            uint64
	metaIndexCount             uint32
	totalUncompressedDataBytes uint64
	entryCount                 uint32
	compressionCodec           uint32
}

type Block struct {
	offset        uint64
	size          uint32
	firstKeyBytes []byte
}

func NewReaderFromConfig(cfg *CollectionConfig, debug bool) (*Reader, error) {
	return NewReader(cfg.Name, cfg.Path, cfg.Mlock, debug)
}

func NewReader(name, path string, lock, debug bool) (*Reader, error) {
	f, err := os.OpenFile(path, os.O_RDONLY, 0)

	if err != nil {
		return nil, err
	}

	hfile := new(Reader)
	hfile.debug = debug

	hfile.mmap, err = mmap.Map(f, mmap.RDONLY, 0)

	if err != nil {
		return nil, err
	}

	fi, err := f.Stat()
	if err != nil {
		return nil, err
	}

	if lock {
		mb := 1024.0 * 1024.0
		log.Printf("[Reader.NewReader] locking %s (%.02fmb)...\n", name, float64(fi.Size())/mb)
		if err = hfile.mmap.Lock(); err != nil {
			log.Printf("[Reader.NewReader] error locking %s: %s\n", name, err.Error())
			return nil, err
		}
		log.Printf("[Reader.NewReader] locked %s.\n", name)
	} else if debug {
		log.Printf("[Reader.NewReader] Not locking %s...\n", name)
	}

	v := binary.BigEndian.Uint32(hfile.mmap[len(hfile.mmap)-4:])
	hfile.majorVersion = v & 0x00ffffff
	hfile.minorVersion = v >> 24

	hfile.header, err = hfile.newHeader(hfile.mmap)
	if err != nil {
		return hfile, err
	}
	err = hfile.loadIndex(hfile.mmap)
	if err != nil {
		return hfile, err
	}

	return hfile, nil
}

func (r *Reader) PrintDebugInfo(out io.Writer, includeStartKeys bool) {
	fmt.Fprintln(out, "entries: ", r.header.entryCount)
	fmt.Fprintf(out, "compressed: %v (%d)\n", r.header.compressionCodec != CompressionNone, r.header.compressionCodec)
	fmt.Fprintln(out, "blocks: ", len(r.index))
	if includeStartKeys {
		for i, blk := range r.index {
			fmt.Fprintf(out, "\t#%d: %s\n", i, hex.EncodeToString(blk.firstKeyBytes))
		}
	}
}

func (r *Reader) newHeader(mmap mmap.MMap) (Header, error) {
	header := Header{}

	if r.majorVersion != 1 || r.minorVersion != 0 {
		return header, fmt.Errorf("wrong version: %d.%d", r.majorVersion, r.minorVersion)
	}

	header.index = len(mmap) - 60
	buf := bytes.NewReader(mmap[header.index:])

	headerMagic := make([]byte, 8)
	buf.Read(headerMagic)
	if bytes.Compare(headerMagic, TrailerMagic) != 0 {
		return header, errors.New("bad header magic")
	}

	binary.Read(buf, binary.BigEndian, &header.fileInfoOffset)
	binary.Read(buf, binary.BigEndian, &header.dataIndexOffset)
	binary.Read(buf, binary.BigEndian, &header.dataIndexCount)
	binary.Read(buf, binary.BigEndian, &header.metaIndexOffset)
	binary.Read(buf, binary.BigEndian, &header.metaIndexCount)
	binary.Read(buf, binary.BigEndian, &header.totalUncompressedDataBytes)
	binary.Read(buf, binary.BigEndian, &header.entryCount)
	binary.Read(buf, binary.BigEndian, &header.compressionCodec)
	return header, nil
}

func (r *Reader) loadIndex(mmap mmap.MMap) error {

	dataIndexEnd := r.header.metaIndexOffset
	if r.header.metaIndexOffset == 0 {
		dataIndexEnd = uint64(r.header.index)
	}
	buf := bytes.NewReader(mmap[r.header.dataIndexOffset:dataIndexEnd])

	dataIndexMagic := make([]byte, 8)
	buf.Read(dataIndexMagic)
	if bytes.Compare(dataIndexMagic, IndexMagic) != 0 {
		return errors.New("bad data index magic")
	}

	for buf.Len() > 0 {
		dataBlock := Block{}

		binary.Read(buf, binary.BigEndian, &dataBlock.offset)
		binary.Read(buf, binary.BigEndian, &dataBlock.size)

		firstKeyLen, _ := binary.ReadUvarint(buf)
		dataBlock.firstKeyBytes = make([]byte, firstKeyLen)
		buf.Read(dataBlock.firstKeyBytes)

		r.index = append(r.index, dataBlock)
	}

	return nil
}

func After(a, b []byte) bool {
	return bytes.Compare(a, b) > 0
}

func (b *Block) IsAfter(key []byte) bool {
	return After(b.firstKeyBytes, key)
}

func (r *Reader) FindBlock(from int, key []byte) int {
	remaining := len(r.index) - from - 1
	if r.debug {
		log.Printf("[Reader.findBlock] cur %d, remaining %d\n", from, remaining)
	}

	if remaining <= 0 {
		if r.debug {
			log.Println("[Reader.findBlock] last block")
		}
		return from // s.cur is the last block, so it is only choice.
	}

	if r.index[from+1].IsAfter(key) {
		if r.debug {
			log.Println("[Reader.findBlock] next block is past key")
		}
		return from
	}

	offset := sort.Search(remaining, func(i int) bool {
		return r.index[from+i+1].IsAfter(key)
	})

	return from + offset
}

func (r *Reader) GetBlock(i int) (*bytes.Reader, error) {
	var buf *bytes.Reader

	block := r.index[i]

	switch r.header.compressionCodec {
	case CompressionNone:
		buf = bytes.NewReader(r.mmap[block.offset : block.offset+uint64(block.size)])
	case CompressionSnappy:
		uncompressedByteSize := binary.BigEndian.Uint32(r.mmap[block.offset : block.offset+4])
		if uncompressedByteSize != block.size {
			return nil, errors.New("mismatched uncompressed block size")
		}
		compressedByteSize := binary.BigEndian.Uint32(r.mmap[block.offset+4 : block.offset+8])
		compressedBytes := r.mmap[block.offset+8 : block.offset+8+uint64(compressedByteSize)]
		uncompressedBytes, err := snappy.Decode(nil, compressedBytes)
		if err != nil {
			return nil, err
		}
		buf = bytes.NewReader(uncompressedBytes)
	default:
		return nil, errors.New("Unsupported compression codec " + string(r.header.compressionCodec))
	}

	dataBlockMagic := make([]byte, 8)
	buf.Read(dataBlockMagic)
	if bytes.Compare(dataBlockMagic, DataMagic) != 0 {
		return nil, errors.New("bad data block magic")
	}

	return buf, nil
}
<|endoftext|>


<commit_before>package iso9660

import (
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"os"
	"path/filepath"

	"github.com/c4milo/gotoolkit"
)

var (
	// ErrInvalidImage is returned when an attempt to unpack the image Primary Volume Descriptor failed or
	// when the end of the image was reached without finding a primary volume descriptor
	ErrInvalidImage = func(err error) error { return fmt.Errorf("invalid-iso9660-image: %s", err) }
	// ErrCorruptedImage is returned when a seek operation, on the image, failed.
	ErrCorruptedImage = func(err error) error { return fmt.Errorf("corrupted-image: %s", err) }
)

type imageReader struct {
	io.ReadSeeker
}

func (ir *imageReader) ReadAt(p []byte, off int64) (n int, err error) {
	// 0 means io.SeekStart. Not using constant for backwards compatibility.
	if _, err = ir.Seek(off, 0); err == nil {
		n, err = io.ReadFull(ir, p)
	}
	return
}

// Reader defines the state of the ISO9660 image reader. It needs to be instantiated
// from its constructor.
type Reader struct {
	// File descriptor to the opened ISO image
	image *imageReader
	// Copy of unencoded Primary Volume Descriptor
	pvd PrimaryVolume
	// Queue used to walk through file system iteratively
	queue gotoolkit.Queue
	// Current sector
	sector uint32
	// Current bytes read from current sector.
	read uint32
}

// NewReader creates a new ISO 9660 image reader.
func NewReader(rs io.ReadSeeker) (*Reader, error) {
	// Starts reading from image data area
	sector := dataAreaSector
	// Iterates over volume descriptors until it finds the primary volume descriptor
	// or an error condition.
	for {
		offset, err := rs.Seek(int64(sector*sectorSize), os.SEEK_SET)
		if err != nil {
			return nil, ErrCorruptedImage(err)
		}

		var volDesc VolumeDescriptor
		if err := binary.Read(rs, binary.BigEndian, &volDesc); err != nil {
			return nil, ErrCorruptedImage(err)
		}

		if volDesc.Type == primaryVol {
			// backs up to the beginning of the sector again in order to unpack
			// the entire primary volume descriptor more easily.
			if _, err := rs.Seek(offset, os.SEEK_SET); err != nil {
				return nil, ErrCorruptedImage(err)
			}

			reader := new(Reader)
			reader.image = &imageReader{rs}
			reader.queue = new(gotoolkit.SliceQueue)

			if err := reader.unpackPVD(); err != nil {
				return nil, ErrCorruptedImage(err)
			}

			return reader, nil
		}

		if volDesc.Type == volSetTerminator {
			return nil, ErrInvalidImage(errors.New("Volume Set Terminator reached. A Primary Volume Descriptor was not found."))
		}
		sector++
	}
}

// Skip skips the given number of directory records.
func (r *Reader) Skip(n int) error {
	var drecord File
	var len byte
	var err error
	for i := 0; i < n; i++ {
		if len, err = r.unpackDRecord(&drecord); err != nil {
			return err
		}
		r.read += uint32(len)
	}
	return nil
}

// Next moves onto the next directory record present in the image.
// It does not use the Path Table since the goal is to read everything
// from the ISO image.
func (r *Reader) Next() (os.FileInfo, error) {
	if r.queue.IsEmpty() {
		return nil, io.EOF
	}

	// We only dequeue the directory when it does not contain more children
	// or when it is empty and there is no children to iterate over.
	item, err := r.queue.Peek()
	if err != nil {
		panic(err)
	}

	f := item.(File)
	if r.sector == 0 {
		r.sector = f.ExtentLocationBE
		_, err := r.image.Seek(int64(r.sector*sectorSize), os.SEEK_SET)
		if err != nil {
			return nil, ErrCorruptedImage(err)
		}

		// Skips . and .. directories
		if err = r.Skip(2); err != nil {
			return nil, ErrCorruptedImage(err)
		}
	}

	var drecord File
	var len byte
	if (r.read % sectorSize) == 0 {
		r.sector++
		_, err := r.image.Seek(int64(r.sector*sectorSize), os.SEEK_SET)
		if err != nil {
			return nil, ErrCorruptedImage(err)
		}
	}

	if len, err = r.unpackDRecord(&drecord); err != nil && err != io.EOF {
		return nil, ErrCorruptedImage(err)
	}
	r.read += uint32(len)

	if err == io.EOF {
		// directory record is empty, sector space wasted, move onto next sector.
		rsize := (sectorSize - (r.read % sectorSize))
		buf := make([]byte, rsize)
		if err := binary.Read(r.image, binary.BigEndian, buf); err != nil {
			return nil, ErrCorruptedImage(err)
		}
		r.read += rsize
	}

	// If there is no more entries in the current directory, dequeue it.
	if r.read >= f.ExtentLengthBE {
		r.read = 0
		r.sector = 0
		r.queue.Dequeue()
	}

	// End of directory listing, drecord is empty so we don't bother
	// to return it and keep iterating to look for the next actual
	// directory or file.
	if drecord.fileID == "" {
		return r.Next()
	}

	parent := f.Name()
	if parent == "\x00" {
		parent = "/"
	}
	drecord.fileID = filepath.Join(parent, drecord.fileID)

	if drecord.IsDir() {
		r.queue.Enqueue(drecord)
	} else {
		drecord.image = r.image
	}

	return &drecord, nil
}

// unpackDRecord unpacks directory record bits into Go's struct
func (r *Reader) unpackDRecord(f *File) (byte, error) {
	// Gets the directory record length
	var len byte
	if err := binary.Read(r.image, binary.BigEndian, &len); err != nil {
		return len, ErrCorruptedImage(err)
	}

	if len == 0 {
		return len + 1, io.EOF
	}

	// Reads directory record into Go struct
	var drecord DirectoryRecord
	if err := binary.Read(r.image, binary.BigEndian, &drecord); err != nil {
		return len, ErrCorruptedImage(err)
	}

	f.DirectoryRecord = drecord
	// Gets the name
	name := make([]byte, drecord.FileIDLength)
	if err := binary.Read(r.image, binary.BigEndian, name); err != nil {
		return len, ErrCorruptedImage(err)
	}
	f.fileID = string(name)

	// Padding field as per section 9.1.12 in ECMA-119
	if (drecord.FileIDLength % 2) == 0 {
		var zero byte
		if err := binary.Read(r.image, binary.BigEndian, &zero); err != nil {
			return len, ErrCorruptedImage(err)
		}
	}

	// System use field as per section 9.1.13 in ECMA-119
	// Directory record has 34 bytes in addition to the name's
	// variable length and the padding field mentioned in section 9.1.12
	totalLen := 34 + drecord.FileIDLength - (drecord.FileIDLength % 2)
	sysUseLen := int64(len - totalLen)
	if sysUseLen > 0 {
		sysData := make([]byte, sysUseLen)
		if err := binary.Read(r.image, binary.BigEndian, sysData); err != nil {
			return len, ErrCorruptedImage(err)
		}
	}
	return len, nil
}

// unpackPVD unpacks Primary Volume Descriptor in three phases. This is
// because the root directory record is a variable-length record and Go's binary
// package doesn't support unpacking variable-length structs easily.
func (r *Reader) unpackPVD() error {
	// Unpack first half
	var pvd1 PrimaryVolumePart1
	if err := binary.Read(r.image, binary.BigEndian, &pvd1); err != nil {
		return ErrCorruptedImage(err)
	}
	r.pvd.PrimaryVolumePart1 = pvd1

	// Unpack root directory record
	var drecord File
	if _, err := r.unpackDRecord(&drecord); err != nil {
		return ErrCorruptedImage(err)
	}
	r.pvd.DirectoryRecord = drecord.DirectoryRecord
	r.queue.Enqueue(drecord)

	// Unpack second half
	var pvd2 PrimaryVolumePart2
	if err := binary.Read(r.image, binary.BigEndian, &pvd2); err != nil {
		return ErrCorruptedImage(err)
	}
	r.pvd.PrimaryVolumePart2 = pvd2

	return nil
}
<commit_msg>Follow ReadAt specification as much as we can<commit_after>package iso9660

import (
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"os"
	"path/filepath"

	"github.com/c4milo/gotoolkit"
)

var (
	// ErrInvalidImage is returned when an attempt to unpack the image Primary Volume Descriptor failed or
	// when the end of the image was reached without finding a primary volume descriptor
	ErrInvalidImage = func(err error) error { return fmt.Errorf("invalid-iso9660-image: %s", err) }
	// ErrCorruptedImage is returned when a seek operation, on the image, failed.
	ErrCorruptedImage = func(err error) error { return fmt.Errorf("corrupted-image: %s", err) }
)

type imageReader struct {
	io.ReadSeeker
}

func (ir *imageReader) ReadAt(p []byte, off int64) (n int, err error) {
	// If the underlying ReadSeeker happen to implement ReadAt, let's use it.
	if readerAt, ok := ir.ReadSeeker.(io.ReaderAt); ok {
		return readerAt.ReadAt(p, off)
	}
	// If not we are going to emulate it with Seeking and Reading.
	// ReadAt should not affect the Seeker's position,
	// thus we need to get the current position first, and Seek back to it after we're done.
	// 1 means io.SeekCurrent. Not using constant for backwards compatibility.
	// 0 means io.SeekStart.
	if curPos, err := ir.Seek(0, 1); err == nil {
		// We can only Seek back later if we could get the current position
		defer ir.Seek(curPos, 0)
	}
	if _, err = ir.Seek(off, 0); err == nil {
		// ReadAt should fill the buffer if there are no errors in the stream, so let's use ReadFull.
		n, err = io.ReadFull(ir, p)
		if err == io.ErrUnexpectedEOF {
			// ReadAt should return EOF instead of ErrUnexpectedEOF
			err = io.EOF
		}
	}
	return
}

// Reader defines the state of the ISO9660 image reader. It needs to be instantiated
// from its constructor.
type Reader struct {
	// File descriptor to the opened ISO image
	image *imageReader
	// Copy of unencoded Primary Volume Descriptor
	pvd PrimaryVolume
	// Queue used to walk through file system iteratively
	queue gotoolkit.Queue
	// Current sector
	sector uint32
	// Current bytes read from current sector.
	read uint32
}

// NewReader creates a new ISO 9660 image reader.
func NewReader(rs io.ReadSeeker) (*Reader, error) {
	// Starts reading from image data area
	sector := dataAreaSector
	// Iterates over volume descriptors until it finds the primary volume descriptor
	// or an error condition.
	for {
		offset, err := rs.Seek(int64(sector*sectorSize), os.SEEK_SET)
		if err != nil {
			return nil, ErrCorruptedImage(err)
		}

		var volDesc VolumeDescriptor
		if err := binary.Read(rs, binary.BigEndian, &volDesc); err != nil {
			return nil, ErrCorruptedImage(err)
		}

		if volDesc.Type == primaryVol {
			// backs up to the beginning of the sector again in order to unpack
			// the entire primary volume descriptor more easily.
			if _, err := rs.Seek(offset, os.SEEK_SET); err != nil {
				return nil, ErrCorruptedImage(err)
			}

			reader := new(Reader)
			reader.image = &imageReader{rs}
			reader.queue = new(gotoolkit.SliceQueue)

			if err := reader.unpackPVD(); err != nil {
				return nil, ErrCorruptedImage(err)
			}

			return reader, nil
		}

		if volDesc.Type == volSetTerminator {
			return nil, ErrInvalidImage(errors.New("Volume Set Terminator reached. A Primary Volume Descriptor was not found."))
		}
		sector++
	}
}

// Skip skips the given number of directory records.
func (r *Reader) Skip(n int) error {
	var drecord File
	var len byte
	var err error
	for i := 0; i < n; i++ {
		if len, err = r.unpackDRecord(&drecord); err != nil {
			return err
		}
		r.read += uint32(len)
	}
	return nil
}

// Next moves onto the next directory record present in the image.
// It does not use the Path Table since the goal is to read everything
// from the ISO image.
func (r *Reader) Next() (os.FileInfo, error) {
	if r.queue.IsEmpty() {
		return nil, io.EOF
	}

	// We only dequeue the directory when it does not contain more children
	// or when it is empty and there is no children to iterate over.
	item, err := r.queue.Peek()
	if err != nil {
		panic(err)
	}

	f := item.(File)
	if r.sector == 0 {
		r.sector = f.ExtentLocationBE
		_, err := r.image.Seek(int64(r.sector*sectorSize), os.SEEK_SET)
		if err != nil {
			return nil, ErrCorruptedImage(err)
		}

		// Skips . and .. directories
		if err = r.Skip(2); err != nil {
			return nil, ErrCorruptedImage(err)
		}
	}

	var drecord File
	var len byte
	if (r.read % sectorSize) == 0 {
		r.sector++
		_, err := r.image.Seek(int64(r.sector*sectorSize), os.SEEK_SET)
		if err != nil {
			return nil, ErrCorruptedImage(err)
		}
	}

	if len, err = r.unpackDRecord(&drecord); err != nil && err != io.EOF {
		return nil, ErrCorruptedImage(err)
	}
	r.read += uint32(len)

	if err == io.EOF {
		// directory record is empty, sector space wasted, move onto next sector.
		rsize := (sectorSize - (r.read % sectorSize))
		buf := make([]byte, rsize)
		if err := binary.Read(r.image, binary.BigEndian, buf); err != nil {
			return nil, ErrCorruptedImage(err)
		}
		r.read += rsize
	}

	// If there is no more entries in the current directory, dequeue it.
	if r.read >= f.ExtentLengthBE {
		r.read = 0
		r.sector = 0
		r.queue.Dequeue()
	}

	// End of directory listing, drecord is empty so we don't bother
	// to return it and keep iterating to look for the next actual
	// directory or file.
	if drecord.fileID == "" {
		return r.Next()
	}

	parent := f.Name()
	if parent == "\x00" {
		parent = "/"
	}
	drecord.fileID = filepath.Join(parent, drecord.fileID)

	if drecord.IsDir() {
		r.queue.Enqueue(drecord)
	} else {
		drecord.image = r.image
	}

	return &drecord, nil
}

// unpackDRecord unpacks directory record bits into Go's struct
func (r *Reader) unpackDRecord(f *File) (byte, error) {
	// Gets the directory record length
	var len byte
	if err := binary.Read(r.image, binary.BigEndian, &len); err != nil {
		return len, ErrCorruptedImage(err)
	}

	if len == 0 {
		return len + 1, io.EOF
	}

	// Reads directory record into Go struct
	var drecord DirectoryRecord
	if err := binary.Read(r.image, binary.BigEndian, &drecord); err != nil {
		return len, ErrCorruptedImage(err)
	}

	f.DirectoryRecord = drecord
	// Gets the name
	name := make([]byte, drecord.FileIDLength)
	if err := binary.Read(r.image, binary.BigEndian, name); err != nil {
		return len, ErrCorruptedImage(err)
	}
	f.fileID = string(name)

	// Padding field as per section 9.1.12 in ECMA-119
	if (drecord.FileIDLength % 2) == 0 {
		var zero byte
		if err := binary.Read(r.image, binary.BigEndian, &zero); err != nil {
			return len, ErrCorruptedImage(err)
		}
	}

	// System use field as per section 9.1.13 in ECMA-119
	// Directory record has 34 bytes in addition to the name's
	// variable length and the padding field mentioned in section 9.1.12
	totalLen := 34 + drecord.FileIDLength - (drecord.FileIDLength % 2)
	sysUseLen := int64(len - totalLen)
	if sysUseLen > 0 {
		sysData := make([]byte, sysUseLen)
		if err := binary.Read(r.image, binary.BigEndian, sysData); err != nil {
			return len, ErrCorruptedImage(err)
		}
	}
	return len, nil
}

// unpackPVD unpacks Primary Volume Descriptor in three phases. This is
// because the root directory record is a variable-length record and Go's binary
// package doesn't support unpacking variable-length structs easily.
func (r *Reader) unpackPVD() error {
	// Unpack first half
	var pvd1 PrimaryVolumePart1
	if err := binary.Read(r.image, binary.BigEndian, &pvd1); err != nil {
		return ErrCorruptedImage(err)
	}
	r.pvd.PrimaryVolumePart1 = pvd1

	// Unpack root directory record
	var drecord File
	if _, err := r.unpackDRecord(&drecord); err != nil {
		return ErrCorruptedImage(err)
	}
	r.pvd.DirectoryRecord = drecord.DirectoryRecord
	r.queue.Enqueue(drecord)

	// Unpack second half
	var pvd2 PrimaryVolumePart2
	if err := binary.Read(r.image, binary.BigEndian, &pvd2); err != nil {
		return ErrCorruptedImage(err)
	}
	r.pvd.PrimaryVolumePart2 = pvd2

	return nil
}
<|endoftext|>


<commit_before>// Copyright 2014 Jens Rantil. All rights reserved.  Use of this source code is
// governed by a BSD-style license that can be found in the LICENSE file.

package csv

import (
	"bufio"
	"bytes"
	"io"
	"unicode/utf8"
)

// A Reader reads records from a CSV-encoded file.
//
// Can be created by calling either NewReader or using NewDialectReader.
type Reader struct {
	opts   Dialect
	r      *bufio.Reader
	tmpBuf bytes.Buffer
}

// Creates a reader that conforms to RFC 4180 and behaves identical as a
// encoding/csv.Reader.
//
// See `Default*` constants for default dialect used.
func NewReader(r io.Reader) *Reader {
	opts := Dialect{}
	opts.setDefaults()
	return NewDialectReader(r, opts)
}

// Create a custom CSV reader.
func NewDialectReader(r io.Reader, opts Dialect) *Reader {
	opts.setDefaults()
	return &Reader{
		opts: opts,
		r:    bufio.NewReader(r),
	}
}

// ReadAll reads all the remaining records from r. Each record is a slice of
// fields. A successful call returns err == nil, not err == EOF. Because
// ReadAll is defined to read until EOF, it does not treat end of file as an
// error to be reported.
func (r *Reader) ReadAll() ([][]string, error) {
	allRows := make([][]string, 0, 1)
	for {
		fields, err := r.Read()
		if err == io.EOF {
			return allRows, nil
		}
		if err != nil {
			return nil, err
		}
		allRows = append(allRows, fields)
	}

	// Required by Go 1.0 to compile. Unreachable code.
	return allRows, nil
}

// Read reads one record from r. The record is a slice of strings with each
// string representing one field.
func (r *Reader) Read() ([]string, error) {
	// TODO: Possible optimization; store the maximum number of columns for
	// faster preallocation.
	record := make([]string, 0, 2)

	for {
		field, err := r.readField()
		record = append(record, field)
		if err != nil {
			return record, err
		}

		if nextIsLineTerminator, _ := r.nextIsLineTerminator(); nextIsLineTerminator {
			// Skipping so that next read call is good to go.
			err = r.skipLineTerminator()
			// Error is not expected since it should be in the Unreader buffer, but
			// might as well return it just in case.
			return record, err
		}
		nextIsDelimiter, err := r.nextIsDelimiter()
		if !nextIsDelimiter {
			// Herein lies the devil!
			return record, err
		} else {
			r.skipDelimiter()
		}
	}

	// Required by Go 1.0 to compile. Unreachable code.
	return record, nil
}

func (r *Reader) readField() (string, error) {
	char, _, err := r.r.ReadRune()
	if err != nil {
		return "", err
	}

	// Let the next individual reader functions handle this.
	r.r.UnreadRune()

	if char == r.opts.QuoteChar {
		return r.readQuotedField()
	}
	return r.readUnquotedField()
}

func (r *Reader) nextIsLineTerminator() (bool, error) {
	return r.nextIsString(r.opts.LineTerminator)
}

func (r *Reader) nextIsDelimiter() (bool, error) {
	return r.nextIsString(string(r.opts.Delimiter))
}

func (r *Reader) nextIsString(s string) (bool, error) {
	bs := []byte(s)
	n := len(bs)
	nextBytes, err := r.r.Peek(n)
	return bytes.Equal(nextBytes, bs), err
}

func (r *Reader) skipLineTerminator() error {
	for _ = range r.opts.LineTerminator {
		_, _, err := r.r.ReadRune()
		if err != nil {
			return err
		}
	}
	return nil
}

func (r *Reader) skipDelimiter() error {
	_, _, err := r.r.ReadRune()
	return err
}

func (r *Reader) readQuotedField() (string, error) {
	char, _, err := r.r.ReadRune()
	if err != nil {
		return "", err
	}
	if char != r.opts.QuoteChar {
		panic("Expected first character to be quote character.")
	}

	s := &r.tmpBuf
	defer r.tmpBuf.Reset() // TODO: Not using defer here is faster.
	for {
		char, _, err := r.r.ReadRune()
		if err != nil {
			return s.String(), err
		}
		if char != r.opts.QuoteChar {
			s.WriteRune(char)
		} else {
			switch r.opts.DoubleQuote {
			case DoDoubleQuote:
				char, _, err = r.r.ReadRune()
				if err != nil {
					return s.String(), err
				}
				if char == r.opts.QuoteChar {
					s.WriteRune(char)
				} else {
					r.r.UnreadRune()
					return s.String(), nil
				}
			case NoDoubleQuote:
				if s.Len() == 0 {
					return s.String(), nil
				}
				lastRune, size := utf8.DecodeLastRuneInString(s.String())
				if lastRune == utf8.RuneError && size == 1 {
					panic("Field contained malformed rune.")
				}
				if lastRune == r.opts.EscapeChar {
					// Replace previous escape character.
					s.Truncate(s.Len() - utf8.RuneLen(char))
					s.WriteRune(char)
				} else {
					return s.String(), nil
				}
			default:
				panic("Unrecognized double quote mode.")
			}
		}
	}

	// Required by Go 1.0 to compile. Unreachable code.
	return s.String(), nil
}

func (r *Reader) readUnquotedField() (string, error) {
	// TODO: Use bytes.Buffer
	s := &r.tmpBuf
	defer r.tmpBuf.Reset() // TODO: Not using defer here is faster.
	for {
		char, _, err := r.r.ReadRune()
		if err != nil || char == r.opts.Delimiter {
			// TODO Can a non quoted string be escaped? In that case, it should be
			// handled here. Should probably have a look at how Python's csv module
			// is handling this.

			// Putting it back for the outer loop to read separators. This makes more
			// compatible with readQuotedField().
			r.r.UnreadRune()

			return s.String(), err
		} else {
			s.WriteRune(char)
		}
		if ok, _ := r.nextIsLineTerminator(); ok {
			return s.String(), nil
		}
	}

	// Required by Go 1.0 to compile. Unreachable code.
	return s.String(), nil
}
<commit_msg>perf(reader): remove unnecessary string2byte conversions<commit_after>// Copyright 2014 Jens Rantil. All rights reserved.  Use of this source code is
// governed by a BSD-style license that can be found in the LICENSE file.

package csv

import (
	"bufio"
	"bytes"
	"io"
	"unicode/utf8"
)

// A Reader reads records from a CSV-encoded file.
//
// Can be created by calling either NewReader or using NewDialectReader.
type Reader struct {
	opts                    Dialect
	r                       *bufio.Reader
	tmpBuf                  bytes.Buffer
	optimizedDelimiter      []byte
	optimizedLineTerminator []byte
}

// Creates a reader that conforms to RFC 4180 and behaves identical as a
// encoding/csv.Reader.
//
// See `Default*` constants for default dialect used.
func NewReader(r io.Reader) *Reader {
	opts := Dialect{}
	opts.setDefaults()
	return NewDialectReader(r, opts)
}

// Create a custom CSV reader.
func NewDialectReader(r io.Reader, opts Dialect) *Reader {
	opts.setDefaults()
	return &Reader{
		opts:                    opts,
		r:                       bufio.NewReader(r),
		optimizedDelimiter:      []byte(string(opts.Delimiter)),
		optimizedLineTerminator: []byte(opts.LineTerminator),
	}
}

// ReadAll reads all the remaining records from r. Each record is a slice of
// fields. A successful call returns err == nil, not err == EOF. Because
// ReadAll is defined to read until EOF, it does not treat end of file as an
// error to be reported.
func (r *Reader) ReadAll() ([][]string, error) {
	allRows := make([][]string, 0, 1)
	for {
		fields, err := r.Read()
		if err == io.EOF {
			return allRows, nil
		}
		if err != nil {
			return nil, err
		}
		allRows = append(allRows, fields)
	}

	// Required by Go 1.0 to compile. Unreachable code.
	return allRows, nil
}

// Read reads one record from r. The record is a slice of strings with each
// string representing one field.
func (r *Reader) Read() ([]string, error) {
	// TODO: Possible optimization; store the maximum number of columns for
	// faster preallocation.
	record := make([]string, 0, 2)

	for {
		field, err := r.readField()
		record = append(record, field)
		if err != nil {
			return record, err
		}

		if nextIsLineTerminator, _ := r.nextIsLineTerminator(); nextIsLineTerminator {
			// Skipping so that next read call is good to go.
			err = r.skipLineTerminator()
			// Error is not expected since it should be in the Unreader buffer, but
			// might as well return it just in case.
			return record, err
		}
		nextIsDelimiter, err := r.nextIsDelimiter()
		if !nextIsDelimiter {
			// Herein lies the devil!
			return record, err
		} else {
			r.skipDelimiter()
		}
	}

	// Required by Go 1.0 to compile. Unreachable code.
	return record, nil
}

func (r *Reader) readField() (string, error) {
	char, _, err := r.r.ReadRune()
	if err != nil {
		return "", err
	}

	// Let the next individual reader functions handle this.
	r.r.UnreadRune()

	if char == r.opts.QuoteChar {
		return r.readQuotedField()
	}
	return r.readUnquotedField()
}

func (r *Reader) nextIsLineTerminator() (bool, error) {
	return r.nextIsBytes(r.optimizedLineTerminator)
}

func (r *Reader) nextIsDelimiter() (bool, error) {
	return r.nextIsBytes(r.optimizedDelimiter)
}

func (r *Reader) nextIsBytes(bs []byte) (bool, error) {
	n := len(bs)
	nextBytes, err := r.r.Peek(n)
	return bytes.Equal(nextBytes, bs), err
}

func (r *Reader) skipLineTerminator() error {
	for _ = range r.opts.LineTerminator {
		_, _, err := r.r.ReadRune()
		if err != nil {
			return err
		}
	}
	return nil
}

func (r *Reader) skipDelimiter() error {
	_, _, err := r.r.ReadRune()
	return err
}

func (r *Reader) readQuotedField() (string, error) {
	char, _, err := r.r.ReadRune()
	if err != nil {
		return "", err
	}
	if char != r.opts.QuoteChar {
		panic("Expected first character to be quote character.")
	}

	s := &r.tmpBuf
	defer r.tmpBuf.Reset() // TODO: Not using defer here is faster.
	for {
		char, _, err := r.r.ReadRune()
		if err != nil {
			return s.String(), err
		}
		if char != r.opts.QuoteChar {
			s.WriteRune(char)
		} else {
			switch r.opts.DoubleQuote {
			case DoDoubleQuote:
				char, _, err = r.r.ReadRune()
				if err != nil {
					return s.String(), err
				}
				if char == r.opts.QuoteChar {
					s.WriteRune(char)
				} else {
					r.r.UnreadRune()
					return s.String(), nil
				}
			case NoDoubleQuote:
				if s.Len() == 0 {
					return s.String(), nil
				}
				lastRune, size := utf8.DecodeLastRuneInString(s.String())
				if lastRune == utf8.RuneError && size == 1 {
					panic("Field contained malformed rune.")
				}
				if lastRune == r.opts.EscapeChar {
					// Replace previous escape character.
					s.Truncate(s.Len() - utf8.RuneLen(char))
					s.WriteRune(char)
				} else {
					return s.String(), nil
				}
			default:
				panic("Unrecognized double quote mode.")
			}
		}
	}

	// Required by Go 1.0 to compile. Unreachable code.
	return s.String(), nil
}

func (r *Reader) readUnquotedField() (string, error) {
	// TODO: Use bytes.Buffer
	s := &r.tmpBuf
	defer r.tmpBuf.Reset() // TODO: Not using defer here is faster.
	for {
		char, _, err := r.r.ReadRune()
		if err != nil || char == r.opts.Delimiter {
			// TODO Can a non quoted string be escaped? In that case, it should be
			// handled here. Should probably have a look at how Python's csv module
			// is handling this.

			// Putting it back for the outer loop to read separators. This makes more
			// compatible with readQuotedField().
			r.r.UnreadRune()

			return s.String(), err
		} else {
			s.WriteRune(char)
		}
		if ok, _ := r.nextIsLineTerminator(); ok {
			return s.String(), nil
		}
	}

	// Required by Go 1.0 to compile. Unreachable code.
	return s.String(), nil
}
<|endoftext|>


<commit_before>// Copyright (c) 2017, Daniel Martí <mvdan@mvdan.cc>
// See LICENSE for licensing information

package main

import (
	"bytes"
	"fmt"
	"go/ast"
	"go/importer"
	"go/parser"
	"go/printer"
	"go/token"
	"go/types"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"text/template"
)

const mainFile = "goreduce_main.go"

var (
	mainTmpl = template.Must(template.New("main").Parse(`package main

func main() {
	{{ . }}()
}
`))
	rawPrinter = printer.Config{Mode: printer.RawFormat}
)

type reducer struct {
	dir     string
	logOut  io.Writer
	matchRe *regexp.Regexp

	fset     *token.FileSet
	origFset *token.FileSet
	pkg      *ast.Package
	files    []*ast.File
	file     *ast.File
	funcDecl *ast.FuncDecl
	origMain *ast.FuncDecl

	tconf types.Config
	info  *types.Info

	useIdents map[types.Object][]*ast.Ident

	outBin  string
	goArgs  []string
	dstFile *os.File
	dstBuf  *bytes.Buffer

	didChange bool

	tried map[string]bool

	walker
}

var errNoReduction = fmt.Errorf("could not reduce program")

func reduce(dir, funcName, match string, logOut io.Writer, bflags ...string) error {
	r := &reducer{
		dir:    dir,
		logOut: logOut,
		tried:  make(map[string]bool),
		dstBuf: bytes.NewBuffer(nil),
	}
	tdir, err := ioutil.TempDir("", "goreduce")
	if err != nil {
		return err
	}
	defer os.RemoveAll(tdir)
	r.tconf.Importer = importer.Default()
	if r.matchRe, err = regexp.Compile(match); err != nil {
		return err
	}
	r.fset = token.NewFileSet()
	pkgs, err := parser.ParseDir(r.fset, r.dir, nil, parser.ParseComments)
	if err != nil {
		return err
	}
	if len(pkgs) != 1 {
		return fmt.Errorf("expected 1 package, got %d", len(pkgs))
	}
	for _, pkg := range pkgs {
		r.pkg = pkg
	}
	r.origFset = token.NewFileSet()
	parser.ParseDir(r.origFset, r.dir, nil, 0)
	tfnames := make([]string, 0, len(r.pkg.Files)+1)
	for fpath, file := range r.pkg.Files {
		if r.funcDecl == nil {
			if fd := findFunc(file, funcName); fd != nil {
				r.file = file
				r.funcDecl = fd
			}
		}
		r.files = append(r.files, file)
		tfname := filepath.Join(tdir, filepath.Base(fpath))
		f, err := os.Create(tfname)
		if err != nil {
			return err
		}
		if fd := delFunc(file, "main"); fd != nil && file == r.file {
			r.origMain = fd
		}
		file.Name.Name = "main"
		if err := rawPrinter.Fprint(f, r.fset, file); err != nil {
			return err
		}
		if file == r.file {
			r.dstFile = f
			defer r.dstFile.Close()
		} else if err := f.Close(); err != nil {
			return err
		}
		tfnames = append(tfnames, tfname)
	}
	if r.funcDecl == nil {
		return fmt.Errorf("top-level func %s does not exist", funcName)
	}
	mfname := filepath.Join(tdir, mainFile)
	mf, err := os.Create(mfname)
	if err != nil {
		return err
	}
	if err := mainTmpl.Execute(mf, funcName); err != nil {
		return err
	}
	if err := mf.Close(); err != nil {
		return err
	}
	tfnames = append(tfnames, mfname)
	r.outBin = filepath.Join(tdir, "bin")
	r.goArgs = []string{"build", "-o", r.outBin}
	r.goArgs = append(r.goArgs, buildFlags...)
	r.goArgs = append(r.goArgs, bflags...)
	r.goArgs = append(r.goArgs, tfnames...)
	// Check that the output matches before we apply any changes
	if err := r.checkRun(); err != nil {
		return err
	}
	if anyChanges := r.reduceLoop(); !anyChanges {
		return errNoReduction
	}
	fname := r.fset.Position(r.file.Pos()).Filename
	f, err := os.Create(fname)
	if err != nil {
		return err
	}
	r.file.Name.Name = r.pkg.Name
	if r.origMain != nil {
		r.file.Decls = append(r.file.Decls, r.origMain)
	}
	if err := printer.Fprint(f, r.fset, r.file); err != nil {
		return err
	}
	return f.Close()
}

func (r *reducer) logChange(node ast.Node, format string, a ...interface{}) {
	if *verbose {
		pos := r.origFset.Position(node.Pos())
		fmt.Fprintf(r.logOut, "%s:%d: %s\n", pos.Filename, pos.Line,
			fmt.Sprintf(format, a...))
	}
}

func (r *reducer) checkRun() error {
	out, err := r.buildAndRun()
	if err != nil {
		return err
	}
	if !r.matchRe.Match(out) {
		return fmt.Errorf("error does not match:\n%s", string(out))
	}
	return nil
}

func (r *reducer) okChange() bool {
	if r.didChange {
		return false
	}
	if _, err := r.tconf.Check(r.dir, r.fset, r.files, nil); err != nil {
		if terr, ok := err.(types.Error); ok && terr.Soft {
			println("unexpected go/types soft error: " + terr.Msg)
		}
	}
	r.dstBuf.Reset()
	if err := rawPrinter.Fprint(r.dstBuf, r.fset, r.file); err != nil {
		return false
	}
	newSrc := r.dstBuf.String()
	if r.tried[newSrc] {
		return false
	}
	r.tried[newSrc] = true
	if err := r.dstFile.Truncate(0); err != nil {
		return false
	}
	if _, err := r.dstFile.Seek(0, io.SeekStart); err != nil {
		return false
	}
	if _, err := r.dstFile.Write(r.dstBuf.Bytes()); err != nil {
		return false
	}
	if err := r.checkRun(); err != nil {
		return false
	}
	// Reduction worked
	r.didChange = true
	return true
}

func (r *reducer) reduceLoop() (anyChanges bool) {
	r.info = &types.Info{
		Defs: make(map[*ast.Ident]types.Object),
		Uses: make(map[*ast.Ident]types.Object),
	}
	for {
		if _, err := r.tconf.Check(r.dir, r.fset, r.files, r.info); err != nil {
			panic("types.Check should not error here")
		}
		r.fillUses()
		r.didChange = false
		r.walk(r.file, r.reduceNode)
		if !r.didChange {
			return
		}
		anyChanges = true
	}
}

func (r *reducer) fillUses() {
	r.useIdents = make(map[types.Object][]*ast.Ident)
	for id, obj := range r.info.Uses {
		if pkg := obj.Pkg(); pkg == nil || pkg.Name() != "main" {
			// builtin or declared outside of our pkg
			continue
		}
		r.useIdents[obj] = append(r.useIdents[obj], id)
	}
}

func findFunc(file *ast.File, name string) *ast.FuncDecl {
	for _, decl := range file.Decls {
		fd, ok := decl.(*ast.FuncDecl)
		if ok && fd.Name.Name == name {
			return fd
		}
	}
	return nil
}

func delFunc(file *ast.File, name string) *ast.FuncDecl {
	for i, decl := range file.Decls {
		fd, ok := decl.(*ast.FuncDecl)
		if ok && fd.Name.Name == name {
			file.Decls = append(file.Decls[:i], file.Decls[i+1:]...)
			return fd
		}
	}
	return nil
}

func (r *reducer) buildAndRun() ([]byte, error) {
	cmd := exec.Command("go", r.goArgs...)
	if out, err := cmd.CombinedOutput(); err != nil {
		if _, ok := err.(*exec.ExitError); ok {
			return out, nil
		}
		return nil, err
	}
	if out, err := exec.Command(r.outBin).CombinedOutput(); err != nil {
		if _, ok := err.(*exec.ExitError); ok {
			return out, nil
		}
		return nil, err
	}
	return nil, fmt.Errorf("expected an error to occur")
}
<commit_msg>Actually bail if type-checking fails<commit_after>// Copyright (c) 2017, Daniel Martí <mvdan@mvdan.cc>
// See LICENSE for licensing information

package main

import (
	"bytes"
	"fmt"
	"go/ast"
	"go/importer"
	"go/parser"
	"go/printer"
	"go/token"
	"go/types"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"text/template"
)

const mainFile = "goreduce_main.go"

var (
	mainTmpl = template.Must(template.New("main").Parse(`package main

func main() {
	{{ . }}()
}
`))
	rawPrinter = printer.Config{Mode: printer.RawFormat}
)

type reducer struct {
	dir     string
	logOut  io.Writer
	matchRe *regexp.Regexp

	fset     *token.FileSet
	origFset *token.FileSet
	pkg      *ast.Package
	files    []*ast.File
	file     *ast.File
	funcDecl *ast.FuncDecl
	origMain *ast.FuncDecl

	tconf types.Config
	info  *types.Info

	useIdents map[types.Object][]*ast.Ident

	outBin  string
	goArgs  []string
	dstFile *os.File
	dstBuf  *bytes.Buffer

	didChange bool

	tried map[string]bool

	walker
}

var errNoReduction = fmt.Errorf("could not reduce program")

func reduce(dir, funcName, match string, logOut io.Writer, bflags ...string) error {
	r := &reducer{
		dir:    dir,
		logOut: logOut,
		tried:  make(map[string]bool),
		dstBuf: bytes.NewBuffer(nil),
	}
	tdir, err := ioutil.TempDir("", "goreduce")
	if err != nil {
		return err
	}
	defer os.RemoveAll(tdir)
	r.tconf.Importer = importer.Default()
	if r.matchRe, err = regexp.Compile(match); err != nil {
		return err
	}
	r.fset = token.NewFileSet()
	pkgs, err := parser.ParseDir(r.fset, r.dir, nil, parser.ParseComments)
	if err != nil {
		return err
	}
	if len(pkgs) != 1 {
		return fmt.Errorf("expected 1 package, got %d", len(pkgs))
	}
	for _, pkg := range pkgs {
		r.pkg = pkg
	}
	r.origFset = token.NewFileSet()
	parser.ParseDir(r.origFset, r.dir, nil, 0)
	tfnames := make([]string, 0, len(r.pkg.Files)+1)
	for fpath, file := range r.pkg.Files {
		if r.funcDecl == nil {
			if fd := findFunc(file, funcName); fd != nil {
				r.file = file
				r.funcDecl = fd
			}
		}
		r.files = append(r.files, file)
		tfname := filepath.Join(tdir, filepath.Base(fpath))
		f, err := os.Create(tfname)
		if err != nil {
			return err
		}
		if fd := delFunc(file, "main"); fd != nil && file == r.file {
			r.origMain = fd
		}
		file.Name.Name = "main"
		if err := rawPrinter.Fprint(f, r.fset, file); err != nil {
			return err
		}
		if file == r.file {
			r.dstFile = f
			defer r.dstFile.Close()
		} else if err := f.Close(); err != nil {
			return err
		}
		tfnames = append(tfnames, tfname)
	}
	if r.funcDecl == nil {
		return fmt.Errorf("top-level func %s does not exist", funcName)
	}
	mfname := filepath.Join(tdir, mainFile)
	mf, err := os.Create(mfname)
	if err != nil {
		return err
	}
	if err := mainTmpl.Execute(mf, funcName); err != nil {
		return err
	}
	if err := mf.Close(); err != nil {
		return err
	}
	tfnames = append(tfnames, mfname)
	r.outBin = filepath.Join(tdir, "bin")
	r.goArgs = []string{"build", "-o", r.outBin}
	r.goArgs = append(r.goArgs, buildFlags...)
	r.goArgs = append(r.goArgs, bflags...)
	r.goArgs = append(r.goArgs, tfnames...)
	// Check that the output matches before we apply any changes
	if err := r.checkRun(); err != nil {
		return err
	}
	if anyChanges := r.reduceLoop(); !anyChanges {
		return errNoReduction
	}
	fname := r.fset.Position(r.file.Pos()).Filename
	f, err := os.Create(fname)
	if err != nil {
		return err
	}
	r.file.Name.Name = r.pkg.Name
	if r.origMain != nil {
		r.file.Decls = append(r.file.Decls, r.origMain)
	}
	if err := printer.Fprint(f, r.fset, r.file); err != nil {
		return err
	}
	return f.Close()
}

func (r *reducer) logChange(node ast.Node, format string, a ...interface{}) {
	if *verbose {
		pos := r.origFset.Position(node.Pos())
		fmt.Fprintf(r.logOut, "%s:%d: %s\n", pos.Filename, pos.Line,
			fmt.Sprintf(format, a...))
	}
}

func (r *reducer) checkRun() error {
	out, err := r.buildAndRun()
	if err != nil {
		return err
	}
	if !r.matchRe.Match(out) {
		return fmt.Errorf("error does not match:\n%s", string(out))
	}
	return nil
}

func (r *reducer) okChange() bool {
	if r.didChange {
		return false
	}
	if _, err := r.tconf.Check(r.dir, r.fset, r.files, nil); err != nil {
		if terr, ok := err.(types.Error); ok && terr.Soft {
			println("unexpected go/types soft error: " + terr.Msg)
		}
		return false
	}
	r.dstBuf.Reset()
	if err := rawPrinter.Fprint(r.dstBuf, r.fset, r.file); err != nil {
		return false
	}
	newSrc := r.dstBuf.String()
	if r.tried[newSrc] {
		return false
	}
	r.tried[newSrc] = true
	if err := r.dstFile.Truncate(0); err != nil {
		return false
	}
	if _, err := r.dstFile.Seek(0, io.SeekStart); err != nil {
		return false
	}
	if _, err := r.dstFile.Write(r.dstBuf.Bytes()); err != nil {
		return false
	}
	if err := r.checkRun(); err != nil {
		return false
	}
	// Reduction worked
	r.didChange = true
	return true
}

func (r *reducer) reduceLoop() (anyChanges bool) {
	r.info = &types.Info{
		Defs: make(map[*ast.Ident]types.Object),
		Uses: make(map[*ast.Ident]types.Object),
	}
	for {
		if _, err := r.tconf.Check(r.dir, r.fset, r.files, r.info); err != nil {
			panic("types.Check should not error here")
		}
		r.fillUses()
		r.didChange = false
		r.walk(r.file, r.reduceNode)
		if !r.didChange {
			return
		}
		anyChanges = true
	}
}

func (r *reducer) fillUses() {
	r.useIdents = make(map[types.Object][]*ast.Ident)
	for id, obj := range r.info.Uses {
		if pkg := obj.Pkg(); pkg == nil || pkg.Name() != "main" {
			// builtin or declared outside of our pkg
			continue
		}
		r.useIdents[obj] = append(r.useIdents[obj], id)
	}
}

func findFunc(file *ast.File, name string) *ast.FuncDecl {
	for _, decl := range file.Decls {
		fd, ok := decl.(*ast.FuncDecl)
		if ok && fd.Name.Name == name {
			return fd
		}
	}
	return nil
}

func delFunc(file *ast.File, name string) *ast.FuncDecl {
	for i, decl := range file.Decls {
		fd, ok := decl.(*ast.FuncDecl)
		if ok && fd.Name.Name == name {
			file.Decls = append(file.Decls[:i], file.Decls[i+1:]...)
			return fd
		}
	}
	return nil
}

func (r *reducer) buildAndRun() ([]byte, error) {
	cmd := exec.Command("go", r.goArgs...)
	if out, err := cmd.CombinedOutput(); err != nil {
		if _, ok := err.(*exec.ExitError); ok {
			return out, nil
		}
		return nil, err
	}
	if out, err := exec.Command(r.outBin).CombinedOutput(); err != nil {
		if _, ok := err.(*exec.ExitError); ok {
			return out, nil
		}
		return nil, err
	}
	return nil, fmt.Errorf("expected an error to occur")
}
<|endoftext|>


<commit_before>// Copyright (c) 2017, Daniel Martí <mvdan@mvdan.cc>
// See LICENSE for licensing information

package main

import (
	"bytes"
	"fmt"
	"go/ast"
	"go/importer"
	"go/parser"
	"go/printer"
	"go/token"
	"go/types"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"text/template"
)

const mainFile = "goreduce_main.go"

var (
	mainTmpl = template.Must(template.New("main").Parse(`package main

func main() {
	{{ . }}()
}
`))
	rawPrinter = printer.Config{Mode: printer.RawFormat}
)

type reducer struct {
	dir     string
	logOut  io.Writer
	matchRe *regexp.Regexp

	fset     *token.FileSet
	origFset *token.FileSet
	pkg      *ast.Package
	files    []*ast.File
	file     *ast.File
	funcDecl *ast.FuncDecl
	origMain *ast.FuncDecl

	tconf types.Config
	info  *types.Info

	numUses map[types.Object]int

	outBin  string
	goArgs  []string
	dstFile *os.File
	dstBuf  *bytes.Buffer

	didChange bool

	tried map[string]bool

	walker
}

var errNoReduction = fmt.Errorf("could not reduce program")

func reduce(dir, funcName, match string, logOut io.Writer, bflags ...string) error {
	r := &reducer{
		dir:    dir,
		logOut: logOut,
		tried:  make(map[string]bool),
		dstBuf: bytes.NewBuffer(nil),
	}
	tdir, err := ioutil.TempDir("", "goreduce")
	if err != nil {
		return err
	}
	defer os.RemoveAll(tdir)
	r.tconf.Importer = importer.Default()
	if r.matchRe, err = regexp.Compile(match); err != nil {
		return err
	}
	r.fset = token.NewFileSet()
	pkgs, err := parser.ParseDir(r.fset, r.dir, nil, 0)
	if err != nil {
		return err
	}
	if len(pkgs) != 1 {
		return fmt.Errorf("expected 1 package, got %d", len(pkgs))
	}
	r.origFset = token.NewFileSet()
	parser.ParseDir(r.origFset, r.dir, nil, 0)
	for _, pkg := range pkgs {
		r.pkg = pkg
	}
	for _, file := range r.pkg.Files {
		r.files = append(r.files, file)
	}
	r.file, r.funcDecl = findFunc(r.files, funcName)
	if r.file == nil {
		return fmt.Errorf("top-level func %s does not exist", funcName)
	}
	tfnames := make([]string, 0, len(r.files)+1)
	for _, file := range r.files {
		fname := r.fset.Position(file.Pos()).Filename
		tfname := filepath.Join(tdir, filepath.Base(fname))
		f, err := os.Create(tfname)
		if err != nil {
			return nil
		}
		if fd := delFunc(file, "main"); fd != nil && file == r.file {
			r.origMain = fd
		}
		file.Name.Name = "main"
		if err := rawPrinter.Fprint(f, r.fset, file); err != nil {
			return err
		}
		if file == r.file {
			r.dstFile = f
			defer r.dstFile.Close()
		} else if err := f.Close(); err != nil {
			return err
		}
		tfnames = append(tfnames, tfname)
	}
	mfname := filepath.Join(tdir, mainFile)
	mf, err := os.Create(mfname)
	if err != nil {
		return err
	}
	if err := mainTmpl.Execute(mf, funcName); err != nil {
		return err
	}
	if err := mf.Close(); err != nil {
		return err
	}
	tfnames = append(tfnames, mfname)
	r.outBin = filepath.Join(tdir, "bin")
	r.goArgs = []string{"build", "-o", r.outBin}
	r.goArgs = append(r.goArgs, buildFlags...)
	r.goArgs = append(r.goArgs, bflags...)
	r.goArgs = append(r.goArgs, tfnames...)
	// Check that the output matches before we apply any changes
	if err := r.checkRun(); err != nil {
		return err
	}
	if anyChanges := r.reduceLoop(); !anyChanges {
		return errNoReduction
	}
	fname := r.fset.Position(r.file.Pos()).Filename
	f, err := os.Create(fname)
	if err != nil {
		return err
	}
	r.file.Name.Name = r.pkg.Name
	if r.origMain != nil {
		r.file.Decls = append(r.file.Decls, r.origMain)
	}
	if err := printer.Fprint(f, r.fset, r.file); err != nil {
		return err
	}
	return f.Close()
}

func (r *reducer) logChange(node ast.Node, format string, a ...interface{}) {
	if *verbose {
		pos := r.origFset.Position(node.Pos())
		fmt.Fprintf(r.logOut, "%s:%d: %s\n", pos.Filename, pos.Line,
			fmt.Sprintf(format, a...))
	}
}

func (r *reducer) checkRun() error {
	out, err := r.buildAndRun()
	if err != nil {
		return err
	}
	if !r.matchRe.Match(out) {
		return fmt.Errorf("error does not match:\n%s", string(out))
	}
	return nil
}

// panicOnSoftErr will make the tool panic if we get stuck because of a
// soft type-checking error that we should have avoided. Leaving this on
// would slow the tool down a bit.
const panicOnSoftErr = false

func (r *reducer) okChange() bool {
	if r.didChange {
		return false
	}
	if panicOnSoftErr {
		if _, err := r.tconf.Check(r.dir, r.fset, r.files, nil); err != nil {
			if terr, ok := err.(types.Error); ok && terr.Soft {
				panic("unexpected go/types soft error: " + terr.Msg)
			}
		}
	}
	r.dstBuf.Reset()
	if err := rawPrinter.Fprint(r.dstBuf, r.fset, r.file); err != nil {
		return false
	}
	if s := r.dstBuf.String(); r.tried[s] {
		return false
	} else {
		r.tried[s] = true
	}
	if err := r.dstFile.Truncate(0); err != nil {
		return false
	}
	if _, err := r.dstFile.Seek(0, io.SeekStart); err != nil {
		return false
	}
	if _, err := r.dstFile.Write(r.dstBuf.Bytes()); err != nil {
		return false
	}
	if err := r.checkRun(); err != nil {
		return false
	}
	// Reduction worked
	r.didChange = true
	return true
}

func (r *reducer) reduceLoop() (anyChanges bool) {
	r.info = &types.Info{
		Defs: make(map[*ast.Ident]types.Object),
		Uses: make(map[*ast.Ident]types.Object),
	}
	for {
		if _, err := r.tconf.Check(r.dir, r.fset, r.files, r.info); err != nil {
			panic("types.Check should not error here")
		}
		r.fillUses()
		r.didChange = false
		r.walk(r.file, r.reduceNode)
		if !r.didChange {
			return
		}
		anyChanges = true
	}
}

func (r *reducer) fillUses() {
	r.numUses = make(map[types.Object]int)
	for _, obj := range r.info.Uses {
		if pkg := obj.Pkg(); pkg == nil || pkg.Name() != "main" {
			// builtin or declared outside of our pkg
			continue
		}
		switch obj.(type) {
		case *types.PkgName, *types.Var:
			r.numUses[obj]++
		}
	}
}

func findFunc(files []*ast.File, name string) (*ast.File, *ast.FuncDecl) {
	for _, file := range files {
		for _, decl := range file.Decls {
			fd, ok := decl.(*ast.FuncDecl)
			if ok && fd.Name.Name == name {
				return file, fd
			}
		}
	}
	return nil, nil
}

func delFunc(file *ast.File, name string) *ast.FuncDecl {
	for i, decl := range file.Decls {
		fd, ok := decl.(*ast.FuncDecl)
		if ok && fd.Name.Name == name {
			file.Decls = append(file.Decls[:i], file.Decls[i+1:]...)
			return fd
		}
	}
	return nil
}

func (r *reducer) buildAndRun() ([]byte, error) {
	cmd := exec.Command("go", r.goArgs...)
	if out, err := cmd.CombinedOutput(); err != nil {
		if _, ok := err.(*exec.ExitError); ok {
			return out, nil
		}
		return nil, err
	}
	if out, err := exec.Command(r.outBin).CombinedOutput(); err != nil {
		if _, ok := err.(*exec.ExitError); ok {
			return out, nil
		}
		return nil, err
	}
	return nil, fmt.Errorf("expected an error to occur")
}
<commit_msg>Rework package files logic<commit_after>// Copyright (c) 2017, Daniel Martí <mvdan@mvdan.cc>
// See LICENSE for licensing information

package main

import (
	"bytes"
	"fmt"
	"go/ast"
	"go/importer"
	"go/parser"
	"go/printer"
	"go/token"
	"go/types"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"text/template"
)

const mainFile = "goreduce_main.go"

var (
	mainTmpl = template.Must(template.New("main").Parse(`package main

func main() {
	{{ . }}()
}
`))
	rawPrinter = printer.Config{Mode: printer.RawFormat}
)

type reducer struct {
	dir     string
	logOut  io.Writer
	matchRe *regexp.Regexp

	fset     *token.FileSet
	origFset *token.FileSet
	pkg      *ast.Package
	files    []*ast.File
	file     *ast.File
	funcDecl *ast.FuncDecl
	origMain *ast.FuncDecl

	tconf types.Config
	info  *types.Info

	numUses map[types.Object]int

	outBin  string
	goArgs  []string
	dstFile *os.File
	dstBuf  *bytes.Buffer

	didChange bool

	tried map[string]bool

	walker
}

var errNoReduction = fmt.Errorf("could not reduce program")

func reduce(dir, funcName, match string, logOut io.Writer, bflags ...string) error {
	r := &reducer{
		dir:    dir,
		logOut: logOut,
		tried:  make(map[string]bool),
		dstBuf: bytes.NewBuffer(nil),
	}
	tdir, err := ioutil.TempDir("", "goreduce")
	if err != nil {
		return err
	}
	defer os.RemoveAll(tdir)
	r.tconf.Importer = importer.Default()
	if r.matchRe, err = regexp.Compile(match); err != nil {
		return err
	}
	r.fset = token.NewFileSet()
	pkgs, err := parser.ParseDir(r.fset, r.dir, nil, 0)
	if err != nil {
		return err
	}
	if len(pkgs) != 1 {
		return fmt.Errorf("expected 1 package, got %d", len(pkgs))
	}
	for _, pkg := range pkgs {
		r.pkg = pkg
	}
	r.origFset = token.NewFileSet()
	parser.ParseDir(r.origFset, r.dir, nil, 0)
	tfnames := make([]string, 0, len(r.pkg.Files)+1)
	for fpath, file := range r.pkg.Files {
		if r.funcDecl == nil {
			if fd := findFunc(file, funcName); fd != nil {
				r.file = file
				r.funcDecl = fd
			}
		}
		r.files = append(r.files, file)
		tfname := filepath.Join(tdir, filepath.Base(fpath))
		f, err := os.Create(tfname)
		if err != nil {
			return nil
		}
		if fd := delFunc(file, "main"); fd != nil && file == r.file {
			r.origMain = fd
		}
		file.Name.Name = "main"
		if err := rawPrinter.Fprint(f, r.fset, file); err != nil {
			return err
		}
		if file == r.file {
			r.dstFile = f
			defer r.dstFile.Close()
		} else if err := f.Close(); err != nil {
			return err
		}
		tfnames = append(tfnames, tfname)
	}
	if r.funcDecl == nil {
		return fmt.Errorf("top-level func %s does not exist", funcName)
	}
	mfname := filepath.Join(tdir, mainFile)
	mf, err := os.Create(mfname)
	if err != nil {
		return err
	}
	if err := mainTmpl.Execute(mf, funcName); err != nil {
		return err
	}
	if err := mf.Close(); err != nil {
		return err
	}
	tfnames = append(tfnames, mfname)
	r.outBin = filepath.Join(tdir, "bin")
	r.goArgs = []string{"build", "-o", r.outBin}
	r.goArgs = append(r.goArgs, buildFlags...)
	r.goArgs = append(r.goArgs, bflags...)
	r.goArgs = append(r.goArgs, tfnames...)
	// Check that the output matches before we apply any changes
	if err := r.checkRun(); err != nil {
		return err
	}
	if anyChanges := r.reduceLoop(); !anyChanges {
		return errNoReduction
	}
	fname := r.fset.Position(r.file.Pos()).Filename
	f, err := os.Create(fname)
	if err != nil {
		return err
	}
	r.file.Name.Name = r.pkg.Name
	if r.origMain != nil {
		r.file.Decls = append(r.file.Decls, r.origMain)
	}
	if err := printer.Fprint(f, r.fset, r.file); err != nil {
		return err
	}
	return f.Close()
}

func (r *reducer) logChange(node ast.Node, format string, a ...interface{}) {
	if *verbose {
		pos := r.origFset.Position(node.Pos())
		fmt.Fprintf(r.logOut, "%s:%d: %s\n", pos.Filename, pos.Line,
			fmt.Sprintf(format, a...))
	}
}

func (r *reducer) checkRun() error {
	out, err := r.buildAndRun()
	if err != nil {
		return err
	}
	if !r.matchRe.Match(out) {
		return fmt.Errorf("error does not match:\n%s", string(out))
	}
	return nil
}

// panicOnSoftErr will make the tool panic if we get stuck because of a
// soft type-checking error that we should have avoided. Leaving this on
// would slow the tool down a bit.
const panicOnSoftErr = false

func (r *reducer) okChange() bool {
	if r.didChange {
		return false
	}
	if panicOnSoftErr {
		if _, err := r.tconf.Check(r.dir, r.fset, r.files, nil); err != nil {
			if terr, ok := err.(types.Error); ok && terr.Soft {
				panic("unexpected go/types soft error: " + terr.Msg)
			}
		}
	}
	r.dstBuf.Reset()
	if err := rawPrinter.Fprint(r.dstBuf, r.fset, r.file); err != nil {
		return false
	}
	newSrc := r.dstBuf.String()
	if r.tried[newSrc] {
		return false
	}
	r.tried[newSrc] = true
	if err := r.dstFile.Truncate(0); err != nil {
		return false
	}
	if _, err := r.dstFile.Seek(0, io.SeekStart); err != nil {
		return false
	}
	if _, err := r.dstFile.Write(r.dstBuf.Bytes()); err != nil {
		return false
	}
	if err := r.checkRun(); err != nil {
		return false
	}
	// Reduction worked
	r.didChange = true
	return true
}

func (r *reducer) reduceLoop() (anyChanges bool) {
	r.info = &types.Info{
		Defs: make(map[*ast.Ident]types.Object),
		Uses: make(map[*ast.Ident]types.Object),
	}
	for {
		if _, err := r.tconf.Check(r.dir, r.fset, r.files, r.info); err != nil {
			panic("types.Check should not error here")
		}
		r.fillUses()
		r.didChange = false
		r.walk(r.file, r.reduceNode)
		if !r.didChange {
			return
		}
		anyChanges = true
	}
}

func (r *reducer) fillUses() {
	r.numUses = make(map[types.Object]int)
	for _, obj := range r.info.Uses {
		if pkg := obj.Pkg(); pkg == nil || pkg.Name() != "main" {
			// builtin or declared outside of our pkg
			continue
		}
		switch obj.(type) {
		case *types.PkgName, *types.Var:
			r.numUses[obj]++
		}
	}
}

func findFunc(file *ast.File, name string) *ast.FuncDecl {
	for _, decl := range file.Decls {
		fd, ok := decl.(*ast.FuncDecl)
		if ok && fd.Name.Name == name {
			return fd
		}
	}
	return nil
}

func delFunc(file *ast.File, name string) *ast.FuncDecl {
	for i, decl := range file.Decls {
		fd, ok := decl.(*ast.FuncDecl)
		if ok && fd.Name.Name == name {
			file.Decls = append(file.Decls[:i], file.Decls[i+1:]...)
			return fd
		}
	}
	return nil
}

func (r *reducer) buildAndRun() ([]byte, error) {
	cmd := exec.Command("go", r.goArgs...)
	if out, err := cmd.CombinedOutput(); err != nil {
		if _, ok := err.(*exec.ExitError); ok {
			return out, nil
		}
		return nil, err
	}
	if out, err := exec.Command(r.outBin).CombinedOutput(); err != nil {
		if _, ok := err.(*exec.ExitError); ok {
			return out, nil
		}
		return nil, err
	}
	return nil, fmt.Errorf("expected an error to occur")
}
<|endoftext|>