<commit_before>package main

import "fmt"
import "github.com/lucasb-eyer/go-colorful"
import "image"
import "image/draw"
import "image/png"
import "log"
import "math"
import "os"
import "text/template"
import "strings"

const LICENSE string =
`/* heatmap - High performance heatmap creation in C.
 *
 * The MIT License (MIT)
 *
 * Copyright (c) 2013 Lucas Beyer
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy of
 * this software and associated documentation files (the "Software"), to deal in
 * the Software without restriction, including without limitation the rights to
 * use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
 * the Software, and to permit persons to whom the Software is furnished to do so,
 * subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all
 * copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
 * FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
 * COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
 * IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
 * CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 */

`

const HEADER_TEMPLATE string =
`#ifndef _HEATMAP_COLORSCHEMES_{{ToUpper .}}_H
#define _HEATMAP_COLORSCHEMES_{{ToUpper .}}_H

#ifdef __cplusplus
extern "C" {
#endif

/* This one has only N discrete colors. */
extern const heatmap_colorscheme_t* heatmap_cs_{{.}}_discrete;
/* This is a very soft gradient along abovementioned discrete colors. */
extern const heatmap_colorscheme_t* heatmap_cs_{{.}}_soft;
/* This is a mix of the above two. Makes for a pretty result in many cases. */
extern const heatmap_colorscheme_t* heatmap_cs_{{.}}_mixed;
/* An exponential version of the default mix of the above two. */
/* Use this if your maximum is very "spiked". */
extern const heatmap_colorscheme_t* heatmap_cs_{{.}}_mixed_exp;

#ifdef __cplusplus
}
#endif

#endif /* _HEATMAP_COLORSCHEMES_{{ToUpper .}}_H */
`

// This table contains the "keypoints" of the colorgradient you want to generate.
// The position of each keypoint has to live in the range [0,1]
type GradientTableEntry struct{
    Col colorful.Color
    Pos float64
}
type GradientTable []GradientTableEntry

// This is the meat of the gradient computation. It returns a HCL-blend between
// the two colors around `t`.
// Note: It relies heavily on the fact that the gradient keypoints are sorted.
func (self GradientTable) GetInterpolatedColorFor(t float64) colorful.Color {
    for i := 0 ; i < len(self) - 1 ; i++ {
        c1 := self[i]
        c2 := self[i+1]
        if c1.Pos <= t && t <= c2.Pos {
            // We are in between c1 and c2. Go blend them!
            t := (t - c1.Pos)/(c2.Pos - c1.Pos)
            return c1.Col.BlendHcl(c2.Col, t).Clamped()
        }
    }

    // Nothing found? Means we're at (or past) the last gradient keypoint.
    return self[len(self)-1].Col
}

// This returns the color of the closest gradient keypoint.
// Note: This too relies on the fact that the gradient keypoints are sorted.
func (self GradientTable) GetColorFor(t float64) colorful.Color {
    for i := 0 ; i < len(self) - 1 ; i++ {
        if t < (self[i].Pos + self[i+1].Pos)*0.5 {
            return self[i].Col
        }
    }

    return self[len(self)-1].Col
}

// This is a very nice thing Golang forces you to do!
// It is necessary so that we can write out the literal of the colortable below.
func MustParseHex(s string) colorful.Color {
    c, err := colorful.Hex(s)
    if err != nil {
        log.Fatalf("Invalid color: %v: %v\n", s, err.Error())
    }
    return c
}

func savepng(img image.Image, fname string) {
    toimg, err := os.Create(fname)
    if err != nil {
        log.Printf("Error saving png: %v\n", err)
        return
    }
    defer toimg.Close()

    png.Encode(toimg, img)
}

func CreateColorschemes(keypoints GradientTable, name string, w, h int) {
    c_file, err := os.Create(name + ".c")
    if err != nil { log.Fatal(err) }
    defer func() {
        if err := c_file.Close(); err != nil {
            log.Fatal(err)
        }
    }()

    h_file, err := os.Create(name + ".h")
    if err != nil { log.Fatal(err) }
    defer func() {
        if err := h_file.Close(); err != nil {
            log.Fatal(err)
        }
    }()

    // Might as well already create the header, we know all we need by now.
    tmpl, err := template.New("test").Funcs(template.FuncMap{"ToUpper": strings.ToUpper}).Parse(LICENSE + HEADER_TEMPLATE)
    if err != nil { log.Panic(err) }
    if err = tmpl.Execute(h_file, name); err != nil { log.Panic(err) }

    img_discrete  := image.NewRGBA(image.Rect(0,0,w,len(keypoints)))
    img_soft      := image.NewRGBA(image.Rect(0,0,w,h))
    img_mixed     := image.NewRGBA(image.Rect(0,0,w,h))
    img_mixed_exp := image.NewRGBA(image.Rect(0,0,w,h))

    fmt.Fprint(c_file, LICENSE)
    fmt.Fprintln(c_file, "#ifdef __cplusplus")
    fmt.Fprintln(c_file, "extern \"C\" {")
    fmt.Fprintln(c_file, "#endif\n")
    fmt.Fprintln(c_file, `#include "heatmap.h"`)
    fmt.Fprintf(c_file, `#include "colorschemes/%v.h"`+"\n\n", name)

    // discrete
    fmt.Fprintln(c_file, "static const unsigned char discrete_data[] = {")
    fmt.Fprint(c_file, "    0, 0, 0, 0")
    for y, kp := range keypoints {
        r, g, b := kp.Col.RGB255()
        fmt.Fprintf(c_file, ", %v, %v, %v, 255", r, g, b)
        draw.Draw(img_discrete, image.Rect(0, y, w, y+1), &image.Uniform{kp.Col}, image.ZP, draw.Src)
    }
    fmt.Fprintln(c_file, "\n};")
    fmt.Fprintln(c_file, "static const heatmap_colorscheme_t discrete = { discrete_data, sizeof(discrete_data)/sizeof(discrete_data[0]/4) };")
    fmt.Fprintln(c_file, "const heatmap_colorscheme_t* heatmap_cs_" + name + "_discrete = &discrete;\n")
    savepng(img_discrete, name + "_discrete.png")

    // soft
    fmt.Fprintln(c_file, "static const unsigned char soft_data[] = {")
    fmt.Fprint(c_file, "    0, 0, 0, 0")
    for y := 0 ; y < h ; y++ {
        t := float64(y)/float64(h)
        a := uint8(math.Min(t*30.0, 1.0)*255.0)
        c := keypoints.GetInterpolatedColorFor(t)
        r, g, b := c.RGB255()
        fmt.Fprintf(c_file, ", %v, %v, %v, %v", r, g, b, a)
        draw.Draw(img_soft, image.Rect(0, y, w, y+1), &image.Uniform{c}, image.ZP, draw.Src)
    }
    fmt.Fprintln(c_file, "\n};")
    fmt.Fprintln(c_file, "static const heatmap_colorscheme_t soft = { soft_data, sizeof(soft_data)/sizeof(soft_data[0]/4) };")
    fmt.Fprintln(c_file, "const heatmap_colorscheme_t* heatmap_cs_" + name + "_soft = &soft;\n")
    savepng(img_soft, name + "_soft.png")

    // mixed
    fmt.Fprintln(c_file, "static const unsigned char mixed_data[] = {")
    fmt.Fprint(c_file, "    0, 0, 0, 0")
    for y := 0 ; y < h ; y++ {
        t := float64(y)/float64(h)
        a := uint8(math.Min(t*30.0, 1.0)*255.0)
        // Here, we actually want to create a gradient overlaid with a stairs-like gradient.
        // Thus c1 is the soft gradient color and c2 the closest keypoint's color.
        c1 := keypoints.GetInterpolatedColorFor(t)
        c2 := keypoints.GetColorFor(t)
        c := c1.BlendRgb(c2, 0.2)
        r, g, b := c.RGB255()
        fmt.Fprintf(c_file, ", %v, %v, %v, %v", r, g, b, a)
        draw.Draw(img_mixed, image.Rect(0, y, w, y+1), &image.Uniform{c}, image.ZP, draw.Src)
    }
    fmt.Fprintln(c_file, "\n};")
    fmt.Fprintln(c_file, "static const heatmap_colorscheme_t mixed = { mixed_data, sizeof(mixed_data)/sizeof(mixed_data[0]/4) };")
    fmt.Fprintln(c_file, "const heatmap_colorscheme_t* heatmap_cs_" + name + "_mixed = &mixed;\n")
    savepng(img_mixed, name + "_mixed.png")

    // exp
    fmt.Fprintln(c_file, "static const unsigned char mixed_exp_data[] = {")
    fmt.Fprint(c_file, "    0, 0, 0, 0")
    for y := 0 ; y < h ; y++ {
        t := float64(y)/float64(h)
        a := uint8(math.Min(t*100.0, 1.0)*255.0) // TODO: Change a
        t = 1.0 - t;
        t = 1.0 - t*t * t*t * t*t * t*t * t*t
        c1 := keypoints.GetInterpolatedColorFor(t)
        c2 := keypoints.GetColorFor(t)
        c := c1.BlendRgb(c2, 0.2)
        r, g, b := c.RGB255()
        fmt.Fprintf(c_file, ", %v, %v, %v, %v", r, g, b, a)
        draw.Draw(img_mixed_exp, image.Rect(0, y, w, y+1), &image.Uniform{c}, image.ZP, draw.Src)
    }
    fmt.Fprintln(c_file, "\n};")
    fmt.Fprintln(c_file, "static const heatmap_colorscheme_t mixed_exp = { mixed_exp_data, sizeof(mixed_exp_data)/sizeof(mixed_exp_data[0]/4) };")
    fmt.Fprintln(c_file, "const heatmap_colorscheme_t* heatmap_cs_" + name + "_mixed_exp = &mixed_exp;\n")
    savepng(img_mixed_exp, name + "_mixed_exp.png")

    fmt.Fprintln(c_file, "#ifdef __cplusplus")
    fmt.Fprintln(c_file, "}")
    fmt.Fprintln(c_file, "#endif")
}
<commit_msg>Colorscheme gen: Use template for C file too.<commit_after>package main

import "fmt"
import "github.com/lucasb-eyer/go-colorful"
import "image"
import "image/draw"
import "image/png"
import "log"
import "math"
import "os"
import "text/template"
import "strings"

const LICENSE string =
`/* heatmap - High performance heatmap creation in C.
 *
 * The MIT License (MIT)
 *
 * Copyright (c) 2013 Lucas Beyer
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy of
 * this software and associated documentation files (the "Software"), to deal in
 * the Software without restriction, including without limitation the rights to
 * use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
 * the Software, and to permit persons to whom the Software is furnished to do so,
 * subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all
 * copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
 * FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
 * COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
 * IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
 * CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 */

`

const HEADER_TEMPLATE string = LICENSE +
`#ifndef _HEATMAP_COLORSCHEMES_{{ToUpper .}}_H
#define _HEATMAP_COLORSCHEMES_{{ToUpper .}}_H

#ifdef __cplusplus
extern "C" {
#endif

/* This one has only N discrete colors. */
extern const heatmap_colorscheme_t* heatmap_cs_{{.}}_discrete;
/* This is a very soft gradient along abovementioned discrete colors. */
extern const heatmap_colorscheme_t* heatmap_cs_{{.}}_soft;
/* This is a mix of the above two. Makes for a pretty result in many cases. */
extern const heatmap_colorscheme_t* heatmap_cs_{{.}}_mixed;
/* An exponential version of the default mix of the above two. */
/* Use this if your maximum is very "spiked". */
extern const heatmap_colorscheme_t* heatmap_cs_{{.}}_mixed_exp;

#ifdef __cplusplus
}
#endif

#endif /* _HEATMAP_COLORSCHEMES_{{ToUpper .}}_H */
`

const SOURCE_TEMPLATE string = LICENSE +
`#ifdef __cplusplus
extern "C" {
#endif

#include "heatmap.h"
#include "colorschemes/{{.Name}}.h"

static const unsigned char discrete_data[] = {
    0, 0, 0, 0{{.Discrete}}
};
static const heatmap_colorscheme_t discrete = { discrete_data, sizeof(discrete_data)/sizeof(discrete_data[0]/4) };
const heatmap_colorscheme_t* heatmap_cs_{{.Name}}_discrete = &discrete;

static const unsigned char soft_data[] = {
    0, 0, 0, 0{{.Soft}}
};
static const heatmap_colorscheme_t soft = { soft_data, sizeof(soft_data)/sizeof(soft_data[0]/4) };
const heatmap_colorscheme_t* heatmap_cs_{{.Name}}_soft = &soft;

static const unsigned char mixed_data[] = {
    0, 0, 0, 0{{.Mixed}}
};
static const heatmap_colorscheme_t mixed = { mixed_data, sizeof(mixed_data)/sizeof(mixed_data[0]/4) };
const heatmap_colorscheme_t* heatmap_cs_{{.Name}}_mixed = &mixed;

static const unsigned char mixed_exp_data[] = {
    0, 0, 0, 0{{.MixedExp}}
};
static const heatmap_colorscheme_t mixed_exp = { mixed_exp_data, sizeof(mixed_exp_data)/sizeof(mixed_exp_data[0]/4) };
const heatmap_colorscheme_t* heatmap_cs_{{.Name}}_mixed_exp = &mixed_exp;

#ifdef __cplusplus
}
#endif
`

// This table contains the "keypoints" of the colorgradient you want to generate.
// The position of each keypoint has to live in the range [0,1]
type GradientTableEntry struct{
    Col colorful.Color
    Pos float64
}
type GradientTable []GradientTableEntry

// This is the meat of the gradient computation. It returns a HCL-blend between
// the two colors around `t`.
// Note: It relies heavily on the fact that the gradient keypoints are sorted.
func (self GradientTable) GetInterpolatedColorFor(t float64) colorful.Color {
    for i := 0 ; i < len(self) - 1 ; i++ {
        c1 := self[i]
        c2 := self[i+1]
        if c1.Pos <= t && t <= c2.Pos {
            // We are in between c1 and c2. Go blend them!
            t := (t - c1.Pos)/(c2.Pos - c1.Pos)
            return c1.Col.BlendHcl(c2.Col, t).Clamped()
        }
    }

    // Nothing found? Means we're at (or past) the last gradient keypoint.
    return self[len(self)-1].Col
}

// This returns the color of the closest gradient keypoint.
// Note: This too relies on the fact that the gradient keypoints are sorted.
func (self GradientTable) GetColorFor(t float64) colorful.Color {
    for i := 0 ; i < len(self) - 1 ; i++ {
        if t < (self[i].Pos + self[i+1].Pos)*0.5 {
            return self[i].Col
        }
    }

    return self[len(self)-1].Col
}

// This is a very nice thing Golang forces you to do!
// It is necessary so that we can write out the literal of the colortable below.
func MustParseHex(s string) colorful.Color {
    c, err := colorful.Hex(s)
    if err != nil {
        log.Fatalf("Invalid color: %v: %v\n", s, err.Error())
    }
    return c
}

func savepng(img image.Image, fname string) {
    toimg, err := os.Create(fname)
    if err != nil {
        log.Printf("Error saving png: %v\n", err)
        return
    }
    defer toimg.Close()

    png.Encode(toimg, img)
}

func CreateColorschemes(keypoints GradientTable, name string, w, h int) {
    c_file, err := os.Create(name + ".c")
    if err != nil { log.Fatal(err) }
    defer func() {
        if err := c_file.Close(); err != nil {
            log.Fatal(err)
        }
    }()

    h_file, err := os.Create(name + ".h")
    if err != nil { log.Fatal(err) }
    defer func() {
        if err := h_file.Close(); err != nil {
            log.Fatal(err)
        }
    }()

    // Might as well already create the header, we know all we need by now.
    tmpl, err := template.New("test").Funcs(template.FuncMap{"ToUpper": strings.ToUpper}).Parse(HEADER_TEMPLATE)
    if err != nil { log.Panic(err) }
    if err = tmpl.Execute(h_file, name); err != nil { log.Panic(err) }

    img_discrete  := image.NewRGBA(image.Rect(0,0,w,len(keypoints)))
    img_soft      := image.NewRGBA(image.Rect(0,0,w,h))
    img_mixed     := image.NewRGBA(image.Rect(0,0,w,h))
    img_mixed_exp := image.NewRGBA(image.Rect(0,0,w,h))
    discrete, soft, mixed, mixed_exp := "", "", "", ""

    // discrete
    for y, kp := range keypoints {
        r, g, b := kp.Col.RGB255()
        discrete += fmt.Sprintf(", %v, %v, %v, 255", r, g, b)
        draw.Draw(img_discrete, image.Rect(0, y, w, y+1), &image.Uniform{kp.Col}, image.ZP, draw.Src)
    }
    savepng(img_discrete, name + "_discrete.png")

    // soft
    for y := 0 ; y < h ; y++ {
        t := float64(y)/float64(h)
        a := uint8(math.Min(t*30.0, 1.0)*255.0)
        c := keypoints.GetInterpolatedColorFor(t)
        r, g, b := c.RGB255()
        soft += fmt.Sprintf(", %v, %v, %v, %v", r, g, b, a)
        draw.Draw(img_soft, image.Rect(0, y, w, y+1), &image.Uniform{c}, image.ZP, draw.Src)
    }
    savepng(img_soft, name + "_soft.png")

    // mixed
    for y := 0 ; y < h ; y++ {
        t := float64(y)/float64(h)
        a := uint8(math.Min(t*30.0, 1.0)*255.0)
        // Here, we actually want to create a gradient overlaid with a stairs-like gradient.
        // Thus c1 is the soft gradient color and c2 the closest keypoint's color.
        c1 := keypoints.GetInterpolatedColorFor(t)
        c2 := keypoints.GetColorFor(t)
        c := c1.BlendRgb(c2, 0.2)
        r, g, b := c.RGB255()
        mixed += fmt.Sprintf(", %v, %v, %v, %v", r, g, b, a)
        draw.Draw(img_mixed, image.Rect(0, y, w, y+1), &image.Uniform{c}, image.ZP, draw.Src)
    }
    savepng(img_mixed, name + "_mixed.png")

    // exp
    for y := 0 ; y < h ; y++ {
        t := float64(y)/float64(h)
        a := uint8(math.Min(t*100.0, 1.0)*255.0) // TODO: Change a
        t = 1.0 - t;
        t = 1.0 - t*t * t*t * t*t * t*t * t*t
        c1 := keypoints.GetInterpolatedColorFor(t)
        c2 := keypoints.GetColorFor(t)
        c := c1.BlendRgb(c2, 0.2)
        r, g, b := c.RGB255()
        mixed_exp += fmt.Sprintf(", %v, %v, %v, %v", r, g, b, a)
        draw.Draw(img_mixed_exp, image.Rect(0, y, w, y+1), &image.Uniform{c}, image.ZP, draw.Src)
    }
    savepng(img_mixed_exp, name + "_mixed_exp.png")

    // Write the C file now that we got all data.
    tmpl, err = template.New("test").Parse(SOURCE_TEMPLATE)
    if err != nil { log.Panic(err) }
    if err = tmpl.Execute(c_file, struct { Name, Discrete, Soft, Mixed, MixedExp string }{ name, discrete, soft, mixed, mixed_exp }); err != nil { log.Panic(err) }
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"log"
	"reflect"
	"testing"
)

type testCase struct {
	pl  []Partition
	ppl []Partition
	err error
	cfg *RebalanceConfig
}

func wrap(p []Partition) *PartitionList {
	return &PartitionList{
		Version:    1,
		Partitions: p,
	}
}

func TestBalancing(t *testing.T) {
	cfgLeader := DefaultRebalanceConfig()
	cfgLeader.AllowLeaderRebalancing = true

	cfg3Replicas := DefaultRebalanceConfig()
	cfg3Replicas.MinReplicasForRebalancing = 3

	cfg6Brokers := DefaultRebalanceConfig()
	cfg6Brokers.Brokers = []BrokerID{1, 2, 3, 4, 5, 6}

	tc := []testCase{
		testCase{
			pl: []Partition{
				Partition{Topic: "a", Partition: 1, Replicas: []BrokerID{1, 2, 3}, Weight: 1.0},
				Partition{Topic: "a", Partition: 2, Replicas: []BrokerID{1, 3, 2}, Weight: 1.0},
				Partition{Topic: "a", Partition: 3, Replicas: []BrokerID{1, 4, 5}, Weight: 1.0},
			},
			ppl: []Partition{
				Partition{Topic: "a", Partition: 1, Replicas: []BrokerID{4, 2, 3}, Weight: 1.0, NumReplicas: 3, Brokers: []BrokerID{1, 2, 3, 4, 5}},
			},
			cfg: &cfgLeader,
		},

		testCase{
			pl: []Partition{
				Partition{Topic: "a", Partition: 1, Replicas: []BrokerID{1, 2, 3}, Weight: 1.0},
				Partition{Topic: "a", Partition: 2, Replicas: []BrokerID{2, 1, 4}, Weight: 1.0},
				Partition{Topic: "a", Partition: 3, Replicas: []BrokerID{1, 2, 5}, Weight: 1.0},
			},
			ppl: []Partition{
				Partition{Topic: "a", Partition: 2, Replicas: []BrokerID{2, 3, 4}, Weight: 1.0, NumReplicas: 3, Brokers: []BrokerID{1, 2, 3, 4, 5}},
			},
		},
		testCase{
			pl: []Partition{
				Partition{Topic: "a", Partition: 1, Replicas: []BrokerID{1, 2, 3}, Weight: 1.0},
				Partition{Topic: "a", Partition: 2, Replicas: []BrokerID{2, 3, 4}, Weight: 1.0},
				Partition{Topic: "a", Partition: 3, Replicas: []BrokerID{1, 2, 5}, Weight: 1.0},
			},
			ppl: []Partition{
				Partition{Topic: "a", Partition: 1, Replicas: []BrokerID{1, 4, 3}, Weight: 1.0, NumReplicas: 3, Brokers: []BrokerID{1, 2, 3, 4, 5}},
			},
		},
		testCase{
			pl: []Partition{
				Partition{Topic: "a", Partition: 1, Replicas: []BrokerID{1, 4, 3}, Weight: 1.0},
				Partition{Topic: "a", Partition: 2, Replicas: []BrokerID{2, 3, 4}, Weight: 1.0},
				Partition{Topic: "a", Partition: 3, Replicas: []BrokerID{1, 2, 5}, Weight: 1.0},
			},
			ppl: []Partition{
				Partition{Topic: "a", Partition: 3, Replicas: []BrokerID{1, 3, 5}, Weight: 1.0, NumReplicas: 3, Brokers: []BrokerID{1, 2, 3, 4, 5}},
			},
		},

		testCase{
			pl: []Partition{
				Partition{Topic: "a", Partition: 1, Replicas: []BrokerID{1, 2}, Weight: 1.0},
				Partition{Topic: "a", Partition: 2, Replicas: []BrokerID{2, 3}, Weight: 1.0},
				Partition{Topic: "b", Partition: 1, Replicas: []BrokerID{4, 3, 2}, Weight: 1.0},
			},
			ppl: []Partition{
				Partition{Topic: "b", Partition: 1, Replicas: []BrokerID{4, 3, 1}, Weight: 1.0, NumReplicas: 3, Brokers: []BrokerID{1, 2, 3, 4}},
			},
			cfg: &cfg3Replicas,
		},

		testCase{
			pl: []Partition{
				Partition{Topic: "a", Partition: 1, Replicas: []BrokerID{1, 2, 3}, Weight: 1.0},
				Partition{Topic: "a", Partition: 2, Replicas: []BrokerID{1, 2, 3}, Weight: 1.0},
			},
			ppl: []Partition{
				Partition{Topic: "a", Partition: 1, Replicas: []BrokerID{1, 4, 3}, Weight: 1.0, NumReplicas: 3, Brokers: []BrokerID{1, 2, 3, 4, 5, 6}},
			},
			cfg: &cfg6Brokers,
		},
		testCase{
			pl: []Partition{
				Partition{Topic: "a", Partition: 1, Replicas: []BrokerID{1, 4, 3}, Weight: 1.0},
				Partition{Topic: "a", Partition: 2, Replicas: []BrokerID{1, 2, 3}, Weight: 1.0},
			},
			ppl: []Partition{
				Partition{Topic: "a", Partition: 1, Replicas: []BrokerID{1, 4, 5}, Weight: 1.0, NumReplicas: 3, Brokers: []BrokerID{1, 2, 3, 4, 5, 6}},
			},
			cfg: &cfg6Brokers,
		},
		testCase{
			pl: []Partition{
				Partition{Topic: "a", Partition: 1, Replicas: []BrokerID{1, 4, 5}, Weight: 1.0},
				Partition{Topic: "a", Partition: 2, Replicas: []BrokerID{1, 2, 3}, Weight: 1.0},
			},
			cfg: &cfg6Brokers,
		},
	}

	for _, c := range tc {
		log.SetOutput(&bytes.Buffer{})
		pl := wrap(c.pl)

		cfg := DefaultRebalanceConfig()
		if c.cfg != nil {
			cfg = *c.cfg
		}

		ppl, err := Balance(pl, cfg)

		if !reflect.DeepEqual(wrap(c.ppl), ppl) {
			t.Errorf("expected %v, got %v", wrap(c.ppl), ppl)
			t.Logf("pl %v", c.pl)
		}
		if c.err != nil && err == nil || c.err == nil && err != nil {
			t.Errorf("expected error %v, got %v", c.err, err)
		}
	}
}
<commit_msg>unit test for replica add/del: 86.4 cov<commit_after>package main

import (
	"bytes"
	"log"
	"reflect"
	"testing"
)

type testCase struct {
	pl  []Partition
	ppl []Partition
	err error
	cfg *RebalanceConfig
}

func wrap(p []Partition) *PartitionList {
	return &PartitionList{
		Version:    1,
		Partitions: p,
	}
}

func TestBalancing(t *testing.T) {
	cfgLeader := DefaultRebalanceConfig()
	cfgLeader.AllowLeaderRebalancing = true

	cfg3Replicas := DefaultRebalanceConfig()
	cfg3Replicas.MinReplicasForRebalancing = 3

	cfg6Brokers := DefaultRebalanceConfig()
	cfg6Brokers.Brokers = []BrokerID{1, 2, 3, 4, 5, 6}

	tc := []testCase{
		testCase{
			pl: []Partition{
				Partition{Topic: "a", Partition: 1, Replicas: []BrokerID{1, 2, 3}, Weight: 1.0},
				Partition{Topic: "a", Partition: 2, Replicas: []BrokerID{1, 3, 2}, Weight: 1.0},
				Partition{Topic: "a", Partition: 3, Replicas: []BrokerID{1, 4, 5}, Weight: 1.0},
			},
			ppl: []Partition{
				Partition{Topic: "a", Partition: 1, Replicas: []BrokerID{4, 2, 3}, Weight: 1.0, NumReplicas: 3, Brokers: []BrokerID{1, 2, 3, 4, 5}},
			},
			cfg: &cfgLeader,
		},

		testCase{
			pl: []Partition{
				Partition{Topic: "a", Partition: 1, Replicas: []BrokerID{1, 2, 3}, Weight: 1.0},
				Partition{Topic: "a", Partition: 2, Replicas: []BrokerID{2, 1, 4}, Weight: 1.0},
				Partition{Topic: "a", Partition: 3, Replicas: []BrokerID{1, 2, 5}, Weight: 1.0},
			},
			ppl: []Partition{
				Partition{Topic: "a", Partition: 2, Replicas: []BrokerID{2, 3, 4}, Weight: 1.0, NumReplicas: 3, Brokers: []BrokerID{1, 2, 3, 4, 5}},
			},
		},
		testCase{
			pl: []Partition{
				Partition{Topic: "a", Partition: 1, Replicas: []BrokerID{1, 2, 3}, Weight: 1.0},
				Partition{Topic: "a", Partition: 2, Replicas: []BrokerID{2, 3, 4}, Weight: 1.0},
				Partition{Topic: "a", Partition: 3, Replicas: []BrokerID{1, 2, 5}, Weight: 1.0},
			},
			ppl: []Partition{
				Partition{Topic: "a", Partition: 1, Replicas: []BrokerID{1, 4, 3}, Weight: 1.0, NumReplicas: 3, Brokers: []BrokerID{1, 2, 3, 4, 5}},
			},
		},
		testCase{
			pl: []Partition{
				Partition{Topic: "a", Partition: 1, Replicas: []BrokerID{1, 4, 3}, Weight: 1.0},
				Partition{Topic: "a", Partition: 2, Replicas: []BrokerID{2, 3, 4}, Weight: 1.0},
				Partition{Topic: "a", Partition: 3, Replicas: []BrokerID{1, 2, 5}, Weight: 1.0},
			},
			ppl: []Partition{
				Partition{Topic: "a", Partition: 3, Replicas: []BrokerID{1, 3, 5}, Weight: 1.0, NumReplicas: 3, Brokers: []BrokerID{1, 2, 3, 4, 5}},
			},
		},

		testCase{
			pl: []Partition{
				Partition{Topic: "a", Partition: 1, Replicas: []BrokerID{1, 2}, Weight: 1.0},
				Partition{Topic: "a", Partition: 2, Replicas: []BrokerID{2, 3}, Weight: 1.0},
				Partition{Topic: "b", Partition: 1, Replicas: []BrokerID{4, 3, 2}, Weight: 1.0},
			},
			ppl: []Partition{
				Partition{Topic: "b", Partition: 1, Replicas: []BrokerID{4, 3, 1}, Weight: 1.0, NumReplicas: 3, Brokers: []BrokerID{1, 2, 3, 4}},
			},
			cfg: &cfg3Replicas,
		},

		testCase{
			pl: []Partition{
				Partition{Topic: "a", Partition: 1, Replicas: []BrokerID{1, 2, 3}, Weight: 1.0},
				Partition{Topic: "a", Partition: 2, Replicas: []BrokerID{1, 2, 3}, Weight: 1.0},
			},
			ppl: []Partition{
				Partition{Topic: "a", Partition: 1, Replicas: []BrokerID{1, 4, 3}, Weight: 1.0, NumReplicas: 3, Brokers: []BrokerID{1, 2, 3, 4, 5, 6}},
			},
			cfg: &cfg6Brokers,
		},
		testCase{
			pl: []Partition{
				Partition{Topic: "a", Partition: 1, Replicas: []BrokerID{1, 4, 3}, Weight: 1.0},
				Partition{Topic: "a", Partition: 2, Replicas: []BrokerID{1, 2, 3}, Weight: 1.0},
			},
			ppl: []Partition{
				Partition{Topic: "a", Partition: 1, Replicas: []BrokerID{1, 4, 5}, Weight: 1.0, NumReplicas: 3, Brokers: []BrokerID{1, 2, 3, 4, 5, 6}},
			},
			cfg: &cfg6Brokers,
		},
		testCase{
			pl: []Partition{
				Partition{Topic: "a", Partition: 1, Replicas: []BrokerID{1, 4, 5}, Weight: 1.0},
				Partition{Topic: "a", Partition: 2, Replicas: []BrokerID{1, 2, 3}, Weight: 1.0},
			},
			cfg: &cfg6Brokers,
		},

		// remove extra replica
		testCase{
			pl: []Partition{
				Partition{Topic: "a", Partition: 1, Replicas: []BrokerID{1, 2, 3}, Weight: 1.0, NumReplicas: 2},
			},
			ppl: []Partition{
				Partition{Topic: "a", Partition: 1, Replicas: []BrokerID{1, 3}, Weight: 1.0, NumReplicas: 2, Brokers: []BrokerID{1, 2, 3}},
			},
		},

		// add missing replica
		testCase{
			pl: []Partition{
				Partition{Topic: "a", Partition: 1, Replicas: []BrokerID{1, 2}, Weight: 1.0, NumReplicas: 3, Brokers: []BrokerID{1, 2, 3}},
			},
			ppl: []Partition{
				Partition{Topic: "a", Partition: 1, Replicas: []BrokerID{1, 2, 3}, Weight: 1.0, NumReplicas: 3, Brokers: []BrokerID{1, 2, 3}},
			},
		},
	}

	for _, c := range tc {
		log.SetOutput(&bytes.Buffer{})
		pl := wrap(c.pl)

		cfg := DefaultRebalanceConfig()
		if c.cfg != nil {
			cfg = *c.cfg
		}

		ppl, err := Balance(pl, cfg)

		if !reflect.DeepEqual(wrap(c.ppl), ppl) {
			t.Errorf("expected %v, got %v", wrap(c.ppl), ppl)
			t.Logf("pl %v", c.pl)
		}
		if c.err != nil && err == nil || c.err == nil && err != nil {
			t.Errorf("expected error %v, got %v", c.err, err)
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/goamz/goamz/aws"
	"github.com/goamz/goamz/sqs"
	"gopkg.in/alecthomas/kingpin.v1"
	"log"
	"os"
)

var (
	app       = kingpin.New("dead-letter-requeue", "Requeues messages from a SQS dead-letter queue to the active one.")
	queueName = app.Arg("queue-name", "Name of the SQS queue (e.g. prod-mgmt-website-data-www100-jimdo-com).").Required().String()
)

func main() {
	var awsAccessKey = os.Getenv("AWS_ACCESS_KEY_ID")
	if len(awsAccessKey) == 0 {
		log.Fatalf("Please set environment variable AWS_ACCESS_KEY_ID")
	}
	var awsSecretKey = os.Getenv("AWS_SECRET_ACCESS_KEY")
	if len(awsSecretKey) == 0 {
		log.Fatalf("Please set environment variable AWS_SECRET_ACCESS_KEY")
	}

	kingpin.MustParse(app.Parse(os.Args[1:]))

	activeQueueName := *queueName

	var deadLetterQueueName = activeQueueName + "_dead_letter"

	var auth = aws.Auth{
		AccessKey: awsAccessKey,
		SecretKey: awsSecretKey,
	}

	conn := sqs.New(auth, aws.EUWest)

	deadLetterQueue, err := conn.GetQueue(deadLetterQueueName)
	if err != nil {
		log.Fatalf(err.Error())
		os.Exit(1)
	}

	activeQueue, err := conn.GetQueue(activeQueueName)
	if err != nil {
		log.Fatalf(err.Error())
		os.Exit(1)
	}

	log.Printf("Looking for messages to requeue.")
	for {
		resp, err := deadLetterQueue.ReceiveMessageWithParameters(
			map[string]string{
				"WaitTimeSeconds":     "20",
				"MaxNumberOfMessages": "10",
				"VisibilityTimeout":   "20"})
		if err != nil {
			log.Fatalf(err.Error())
			os.Exit(1)
		}

		messages := resp.Messages
		numberOfMessages := len(messages)
		if numberOfMessages == 0 {
			log.Printf("Requeuing messages done.")
			os.Exit(0)
		} else {
			log.Printf("Moving %v message(s)...", numberOfMessages)
		}

		_, err = activeQueue.SendMessageBatch(messages)
		if err != nil {
			log.Fatalf(err.Error())
			os.Exit(1)
		}

		_, err = deadLetterQueue.DeleteMessageBatch(messages)
		if err != nil {
			log.Fatalf(err.Error())
			os.Exit(1)
		}
	}
}
<commit_msg>Reorder imports<commit_after>package main

import (
	"log"
	"os"

	"github.com/goamz/goamz/aws"
	"github.com/goamz/goamz/sqs"
	"gopkg.in/alecthomas/kingpin.v1"
)

var (
	app       = kingpin.New("dead-letter-requeue", "Requeues messages from a SQS dead-letter queue to the active one.")
	queueName = app.Arg("queue-name", "Name of the SQS queue (e.g. prod-mgmt-website-data-www100-jimdo-com).").Required().String()
)

func main() {
	var awsAccessKey = os.Getenv("AWS_ACCESS_KEY_ID")
	if len(awsAccessKey) == 0 {
		log.Fatalf("Please set environment variable AWS_ACCESS_KEY_ID")
	}
	var awsSecretKey = os.Getenv("AWS_SECRET_ACCESS_KEY")
	if len(awsSecretKey) == 0 {
		log.Fatalf("Please set environment variable AWS_SECRET_ACCESS_KEY")
	}

	kingpin.MustParse(app.Parse(os.Args[1:]))

	activeQueueName := *queueName

	var deadLetterQueueName = activeQueueName + "_dead_letter"

	var auth = aws.Auth{
		AccessKey: awsAccessKey,
		SecretKey: awsSecretKey,
	}

	conn := sqs.New(auth, aws.EUWest)

	deadLetterQueue, err := conn.GetQueue(deadLetterQueueName)
	if err != nil {
		log.Fatalf(err.Error())
		os.Exit(1)
	}

	activeQueue, err := conn.GetQueue(activeQueueName)
	if err != nil {
		log.Fatalf(err.Error())
		os.Exit(1)
	}

	log.Printf("Looking for messages to requeue.")
	for {
		resp, err := deadLetterQueue.ReceiveMessageWithParameters(
			map[string]string{
				"WaitTimeSeconds":     "20",
				"MaxNumberOfMessages": "10",
				"VisibilityTimeout":   "20"})
		if err != nil {
			log.Fatalf(err.Error())
			os.Exit(1)
		}

		messages := resp.Messages
		numberOfMessages := len(messages)
		if numberOfMessages == 0 {
			log.Printf("Requeuing messages done.")
			os.Exit(0)
		} else {
			log.Printf("Moving %v message(s)...", numberOfMessages)
		}

		_, err = activeQueue.SendMessageBatch(messages)
		if err != nil {
			log.Fatalf(err.Error())
			os.Exit(1)
		}

		_, err = deadLetterQueue.DeleteMessageBatch(messages)
		if err != nil {
			log.Fatalf(err.Error())
			os.Exit(1)
		}
	}
}
<|endoftext|>


<commit_before>package funlz

import (
	"bytes"
	"compress/flate"
	"hash/crc32"
	"io"
	"io/ioutil"
	"log"
	"testing"
)

var _ = log.Printf

func compress(in []byte) []byte {
	var out bytes.Buffer
	c := NewWriter(&out)
	c.Write(in)
	c.Flush()
	return out.Bytes()
}

func compressNull(in []byte) {
	c := NewWriter(ioutil.Discard)
	c.Write(in)
	c.Flush()
}

func inflate(in []byte) []byte {
	var out bytes.Buffer
	c, _ := flate.NewWriter(&out, 1)
	c.Write(in)
	c.Flush()
	return out.Bytes()
}

func inflateNull(in []byte) {
	c, _ := flate.NewWriter(ioutil.Discard, 1)
	c.Write(in)
	c.Flush()
}

func decompress(in []byte) []byte {
	inb := bytes.NewBuffer(in)
	d := NewReader(inb)
	out, _ := ioutil.ReadAll(d)
	return out
}

func deflate(in []byte) []byte {
	inb := bytes.NewBuffer(in)
	d := flate.NewReader(inb)
	out, _ := ioutil.ReadAll(d)
	return out
}

type writeAndFlusher interface {
	io.Writer
	Flush() error
}

func compByPart(c writeAndFlusher, n int) {
	rnd := uint32(0)
	p := uint32(0)
	for i := 0; i < n; i++ {
		rnd = rnd*5 + 1
		c.Write(original[p : p+rnd%128])
		p += rnd % 128
		p = p % (1 << 18)
	}
}

type circBuffReader struct {
	b   []byte
	pos int
}

func (c *circBuffReader) Read(b []byte) (int, error) {
	bytes := 0
	for len(b) != 0 {
		n := copy(b, c.b[c.pos:])
		b = b[n:]
		c.pos += n
		if c.pos == len(c.b) {
			c.pos = 0
		}
		bytes += n
	}
	return bytes, nil
}

type circDecomp struct {
	mk func(io.Reader) io.Reader
	b  []byte
	r  io.Reader
}

func (c *circDecomp) Read(b []byte) (int, error) {
	if c.r == nil {
		c.r = c.mk(bytes.NewReader(c.b))
	}
	bytes, err := c.r.Read(b)
	if err == nil {
		c.r = nil
	}
	return bytes, nil
}

var buf [128]byte

func decompByPart(u io.Reader, n int) {
	rnd := uint32(0)
	for i := 0; i < n; i++ {
		rnd = rnd*5 + 1
		u.Read(buf[:rnd%128])
	}
}

func eq(a, b []byte) int {
	if len(a) != len(b) {
		return -2
	}
	for i, c := range a {
		if c != b[i] {
			log.Print("!eq ", c, b[i], a[i-1], b[i-1], a[i+1], b[i+1])
			return i
		}
	}
	return -1
}

var patterns [][2][]byte

func add(raw, comp string) {
	patterns = append(patterns, [2][]byte{[]byte(raw), []byte(comp)})
}

func init() {
	add("asdfasdf", "\x03asdf\x20\x03")
	add("aaaaaaaa", "\x00a\x50\x00")
	add("aaaaaaaab", "\x00a\x50\x00\x00b")
	add("baaaaaaaab", "\x01ba\x50\x00\x00b")
	if backref > 1 {
		add("baaaaacaaaaaaaaaaaaaaaaaaaaaaaaaaaab", "\x01ba\x20\x00\x00c\x30\x05\xf0\x00\x06\x00b")
		add("baaaaacaaaaaaaaaaaaaaaaaaaaaaaaaaab", "\x01ba\x20\x00\x00c\x30\x05\xf0\x00\x05\x00b")
	} else {
		add("baaaaacaaaaaaaaaaaaaaaaaaaaaaaaaaaab", "\x01ba\x20\x00\x00c\x20\x04\xf0\x00\x07\x00b")
		add("baaaaacaaaaaaaaaaaaaaaaaaaaaaaaaaab", "\x01ba\x20\x00\x00c\x20\x04\xf0\x00\x06\x00b")
	}
	add("This is a new era of my life with all good things", "\x1f\x11This is a new era of my life with all good things")
	add("This is a new era of my life with all good things. That is my new life.", "\x1f\x17This is a new era of my life with all good things. That\x20\x32\x01my\x30\x33\x20\x29\x00.")
}

var original, original1, compressed, compressed11111, flatted, flatted11111 []byte

func init() {
	original, _ = ioutil.ReadFile("GettingReal.html")
	original1, _ = ioutil.ReadFile("GettingReal.html")
	compressed = compress(original)
	compressed11111 = compress(original[:11111])
	flatted = inflate(original)
	flatted11111 = inflate(original[:11111])
}

var crc32c = crc32.MakeTable(crc32.Castagnoli)

func TestWriter(t *testing.T) {
	for _, p := range patterns {
		u, c := p[0], p[1]
		o := compress(u)
		if eq(c, o) != -1 {
			t.Errorf("not equal\n%#v\n%#v", c, o)
		}
	}
}

func TestReader(t *testing.T) {
	for _, p := range patterns {
		u, c := p[0], p[1]
		o := decompress(c)
		if eq(u, o) != -1 {
			t.Errorf("not equal\n%#v\n%#v", u, o)
		}
	}
}

func TestBigFile(t *testing.T) {
	log.Print("BigFile")
	decompressed := decompress(compressed)
	if p := eq(original, original1); p != -1 {
		t.Errorf("damage original at %d", p)
	}
	log.Printf("orig/comp %d/%d", len(original), len(compressed))
	if p := eq(original, decompressed); p != -1 {
		t.Errorf("big are not equal %d %d %d %d", len(original), len(compressed), len(decompressed), p)
	}
}

func TestHugeFile(t *testing.T) {
	var crc1, crc2 uint32
	var sz1, sz2 int
	fl := circBuffReader{b: original}
	rd, wr := io.Pipe()
	fin := make(chan struct{})
	comp := NewWriter(wr)
	decomp := NewReader(rd)
	log.Print("HugeFile")
	go func() {
		var buf [512]byte
		var last int
		for {
			n, _ := decomp.Read(buf[:])
			if n == 0 {
				break
			}
			crc2 = crc32.Update(crc2, crc32c, buf[:n])
			sz2 += n
			if sz2/1000000 > last {
				last = sz2 / 1000000
				log.Printf("hugefile: %d bytes", sz2)
			}
		}
		close(fin)
	}()
	var buf [512]byte
	for sz1 = 0; sz1 < 3*1<<29; sz1 += len(buf) {
		k, _ := fl.Read(buf[:])
		if k != 512 {
			panic("k!=512")
		}
		crc1 = crc32.Update(crc1, crc32c, buf[:])
		comp.Write(buf[:])
	}
	comp.Flush()
	wr.Close()
	<-fin
	if sz1 != sz2 {
		t.Errorf("sz1=%d sz2=%d", sz1, sz2)
	}
	if crc1 != crc2 {
		t.Errorf("crc32 mismatch")
	}
}

func BenchmarkCompressBig(b *testing.B) {
	for i := 0; i < b.N; i++ {
		compressNull(original)
	}
}

func BenchmarkCompressMedium(b *testing.B) {
	for i := 0; i < b.N; i++ {
		compressNull(original[:11111])
	}
}

func BenchmarkCompressByPart(b *testing.B) {
	c := NewWriter(ioutil.Discard)
	compByPart(c, b.N)
}

func BenchmarkDecompressBig(b *testing.B) {
	for i := 0; i < b.N; i++ {
		decompress(compressed)
	}
}

func BenchmarkDecompressMedium(b *testing.B) {
	for i := 0; i < b.N; i++ {
		decompress(compressed11111)
	}
}

func BenchmarkDecompressByPart(b *testing.B) {
	u := &circDecomp{mk: func(r io.Reader) io.Reader { return NewReader(r) }, b: compressed}
	decompByPart(u, b.N)
}

func BenchmarkFlateBig(b *testing.B) {
	for i := 0; i < b.N; i++ {
		inflateNull(original)
	}
}

func BenchmarkFlateMedium(b *testing.B) {
	for i := 0; i < b.N; i++ {
		inflateNull(original[:11111])
	}
}

func BenchmarkFlateByPart(b *testing.B) {
	c, _ := flate.NewWriter(ioutil.Discard, 1)
	compByPart(c, b.N)
}

func BenchmarkDeflateBig(b *testing.B) {
	for i := 0; i < b.N; i++ {
		deflate(flatted)
	}
}

func BenchmarkDeflateMedium(b *testing.B) {
	for i := 0; i < b.N; i++ {
		deflate(flatted11111)
	}
}

func BenchmarkDeflatePart(b *testing.B) {
	u := &circDecomp{mk: func(r io.Reader) io.Reader { return flate.NewReader(r) }, b: flatted}
	decompByPart(u, b.N)
}
<commit_msg>tests: fix for flush-mark and skip Huge file in short mode<commit_after>package funlz

import (
	"bytes"
	"compress/flate"
	"hash/crc32"
	"io"
	"io/ioutil"
	"log"
	"testing"
)

var _ = log.Printf

func compress(in []byte) []byte {
	var out bytes.Buffer
	c := NewWriter(&out)
	c.Write(in)
	c.Flush()
	return out.Bytes()
}

func compressNull(in []byte) {
	c := NewWriter(ioutil.Discard)
	c.Write(in)
	c.Flush()
}

func inflate(in []byte) []byte {
	var out bytes.Buffer
	c, _ := flate.NewWriter(&out, 1)
	c.Write(in)
	c.Flush()
	return out.Bytes()
}

func inflateNull(in []byte) {
	c, _ := flate.NewWriter(ioutil.Discard, 1)
	c.Write(in)
	c.Flush()
}

func decompress(in []byte) []byte {
	inb := bytes.NewBuffer(in)
	d := NewReader(inb)
	out, _ := ioutil.ReadAll(d)
	return out
}

func deflate(in []byte) []byte {
	inb := bytes.NewBuffer(in)
	d := flate.NewReader(inb)
	out, _ := ioutil.ReadAll(d)
	return out
}

type writeAndFlusher interface {
	io.Writer
	Flush() error
}

func compByPart(c writeAndFlusher, n int) {
	rnd := uint32(0)
	p := uint32(0)
	for i := 0; i < n; i++ {
		rnd = rnd*5 + 1
		c.Write(original[p : p+rnd%128])
		p += rnd % 128
		p = p % (1 << 18)
	}
}

type circBuffReader struct {
	b   []byte
	pos int
}

func (c *circBuffReader) Read(b []byte) (int, error) {
	bytes := 0
	for len(b) != 0 {
		n := copy(b, c.b[c.pos:])
		b = b[n:]
		c.pos += n
		if c.pos == len(c.b) {
			c.pos = 0
		}
		bytes += n
	}
	return bytes, nil
}

type circDecomp struct {
	mk func(io.Reader) io.Reader
	b  []byte
	r  io.Reader
}

func (c *circDecomp) Read(b []byte) (int, error) {
	if c.r == nil {
		c.r = c.mk(bytes.NewReader(c.b))
	}
	bytes, err := c.r.Read(b)
	if err == nil {
		c.r = nil
	}
	return bytes, nil
}

var buf [128]byte

func decompByPart(u io.Reader, n int) {
	rnd := uint32(0)
	for i := 0; i < n; i++ {
		rnd = rnd*5 + 1
		u.Read(buf[:rnd%128])
	}
}

func eq(a, b []byte) int {
	if len(a) != len(b) {
		return -2
	}
	for i, c := range a {
		if c != b[i] {
			log.Print("!eq ", c, b[i], a[i-1], b[i-1], a[i+1], b[i+1])
			return i
		}
	}
	return -1
}

var patterns [][2][]byte

func add(raw, comp string) {
	patterns = append(patterns, [2][]byte{[]byte(raw), []byte(comp)})
}

func init() {
	add("asdfasdf", "\x04asdf\x20\x03\x00")
	add("aaaaaaaa", "\x01a\x50\x00\x00")
	add("aaaaaaaab", "\x01a\x50\x00\x01b\x00")
	add("baaaaaaaab", "\x02ba\x50\x00\x01b\x00")
	if backref > 1 {
		add("baaaaacaaaaaaaaaaaaaaaaaaaaaaaaaaaab", "\x02ba\x20\x00\x01c\x30\x05\xf0\x00\x06\x01b\x00")
		add("baaaaacaaaaaaaaaaaaaaaaaaaaaaaaaaab", "\x02ba\x20\x00\x01c\x30\x05\xf0\x00\x05\x01b\x00")
	} else {
		add("baaaaacaaaaaaaaaaaaaaaaaaaaaaaaaaaab", "\x02ba\x20\x00\x01c\x20\x04\xf0\x00\x07\x01b\x00")
		add("baaaaacaaaaaaaaaaaaaaaaaaaaaaaaaaab", "\x02ba\x20\x00\x01c\x20\x04\xf0\x00\x06\x01b\x00")
	}
	add("This is a new era of my life with all good things", "\x1f\x12This is a new era of my life with all good things\x00")
	add("This is a new era of my life with all good things. That is my new life.", "\x1f\x18This is a new era of my life with all good things. That\x20\x32\x02my\x30\x33\x20\x29\x01.\x00")
}

var original, original1, compressed, compressed11111, flatted, flatted11111 []byte

func init() {
	original, _ = ioutil.ReadFile("GettingReal.html")
	original1, _ = ioutil.ReadFile("GettingReal.html")
	compressed = compress(original)
	compressed11111 = compress(original[:11111])
	flatted = inflate(original)
	flatted11111 = inflate(original[:11111])
}

var crc32c = crc32.MakeTable(crc32.Castagnoli)

func TestWriter(t *testing.T) {
	for _, p := range patterns {
		u, c := p[0], p[1]
		o := compress(u)
		if eq(c, o) != -1 {
			t.Errorf("not equal\n%#v\n%#v", c, o)
		}
	}
}

func TestReader(t *testing.T) {
	for _, p := range patterns {
		u, c := p[0], p[1]
		o := decompress(c)
		if eq(u, o) != -1 {
			t.Errorf("not equal\n%#v\n%#v", u, o)
		}
	}
}

func TestBigFile(t *testing.T) {
	log.Print("BigFile")
	decompressed := decompress(compressed)
	if p := eq(original, original1); p != -1 {
		t.Errorf("damage original at %d", p)
	}
	log.Printf("orig/comp %d/%d", len(original), len(compressed))
	if p := eq(original, decompressed); p != -1 {
		t.Errorf("big are not equal %d %d %d %d", len(original), len(compressed), len(decompressed), p)
	}
}

func TestHugeFile(t *testing.T) {
	if testing.Short() {
		t.Skip("skip in short mode")
	}
	var crc1, crc2 uint32
	var sz1, sz2 int
	fl := circBuffReader{b: original}
	rd, wr := io.Pipe()
	fin := make(chan struct{})
	comp := NewWriter(wr)
	decomp := NewReader(rd)
	log.Print("HugeFile")
	go func() {
		var buf [512]byte
		var last int
		for {
			n, _ := decomp.Read(buf[:])
			if n == 0 {
				break
			}
			crc2 = crc32.Update(crc2, crc32c, buf[:n])
			sz2 += n
			if sz2/1000000 > last {
				last = sz2 / 1000000
				log.Printf("hugefile: %d bytes", sz2)
			}
		}
		close(fin)
	}()
	var buf [512]byte
	for sz1 = 0; sz1 < 3*1<<29; sz1 += len(buf) {
		k, _ := fl.Read(buf[:])
		if k != 512 {
			panic("k!=512")
		}
		crc1 = crc32.Update(crc1, crc32c, buf[:])
		comp.Write(buf[:])
	}
	comp.Flush()
	wr.Close()
	<-fin
	if sz1 != sz2 {
		t.Errorf("sz1=%d sz2=%d", sz1, sz2)
	}
	if crc1 != crc2 {
		t.Errorf("crc32 mismatch")
	}
}

func BenchmarkCompressBig(b *testing.B) {
	for i := 0; i < b.N; i++ {
		compressNull(original)
	}
}

func BenchmarkCompressMedium(b *testing.B) {
	for i := 0; i < b.N; i++ {
		compressNull(original[:11111])
	}
}

func BenchmarkCompressByPart(b *testing.B) {
	c := NewWriter(ioutil.Discard)
	compByPart(c, b.N)
}

func BenchmarkDecompressBig(b *testing.B) {
	for i := 0; i < b.N; i++ {
		decompress(compressed)
	}
}

func BenchmarkDecompressMedium(b *testing.B) {
	for i := 0; i < b.N; i++ {
		decompress(compressed11111)
	}
}

func BenchmarkDecompressByPart(b *testing.B) {
	u := &circDecomp{mk: func(r io.Reader) io.Reader { return NewReader(r) }, b: compressed}
	decompByPart(u, b.N)
}

func BenchmarkFlateBig(b *testing.B) {
	for i := 0; i < b.N; i++ {
		inflateNull(original)
	}
}

func BenchmarkFlateMedium(b *testing.B) {
	for i := 0; i < b.N; i++ {
		inflateNull(original[:11111])
	}
}

func BenchmarkFlateByPart(b *testing.B) {
	c, _ := flate.NewWriter(ioutil.Discard, 1)
	compByPart(c, b.N)
}

func BenchmarkDeflateBig(b *testing.B) {
	for i := 0; i < b.N; i++ {
		deflate(flatted)
	}
}

func BenchmarkDeflateMedium(b *testing.B) {
	for i := 0; i < b.N; i++ {
		deflate(flatted11111)
	}
}

func BenchmarkDeflatePart(b *testing.B) {
	u := &circDecomp{mk: func(r io.Reader) io.Reader { return flate.NewReader(r) }, b: flatted}
	decompByPart(u, b.N)
}
<|endoftext|>


<commit_before>// Copyright 2015 Google Inc. All Rights Reserved.
// Author: jacobsa@google.com (Aaron Jacobs)

package gcs

import (
	"errors"
	"io"
	"net/http"
	"unicode/utf8"

	"golang.org/x/net/context"
	"google.golang.org/cloud"
	"google.golang.org/cloud/storage"
)

// Bucket represents a GCS bucket, pre-bound with a bucket name and necessary
// authorization information.
//
// Each method that may block accepts a context object that is used for
// deadlines and cancellation. Users need not package authorization information
// into the context object (using cloud.WithContext or similar).
type Bucket interface {
	Name() string

	// List the objects in the bucket that meet the criteria defined by the
	// query, returning a result object that contains the results and potentially
	// a cursor for retrieving the next portion of the larger set of results.
	ListObjects(ctx context.Context, query *storage.Query) (*storage.Objects, error)

	// Create a reader for the contents of the object with the given name. The
	// caller must arrange for the reader to be closed when it is no longer
	// needed.
	NewReader(ctx context.Context, objectName string) (io.ReadCloser, error)

	// Return an ObjectWriter that can be used to create or overwrite an object
	// with the given attributes. attrs.Name must be specified. Otherwise, nil-
	// and zero-valud attributes are ignored.
	//
	// Object names must:
	//
	// *  be non-empty.
	// *  be no longer than 1024 bytes.
	// *  be valid UTF-8.
	// *  not contain the code point U+000A (line feed).
	// *  not contain the code point U+000D (carriage return).
	//
	// See here for authoritative documentation:
	//     https://cloud.google.com/storage/docs/bucket-naming#objectnames
	NewWriter(ctx context.Context, attrs *storage.ObjectAttrs) (ObjectWriter, error)

	// TODO(jacobsa): Comments.
	UpdateObject(
		ctx context.Context,
		newAttrs *storage.ObjectAttrs) (*storage.Object, error)

	// Delete the object with the given name.
	DeleteObject(ctx context.Context, name string) error
}

type bucket struct {
	projID string
	client *http.Client
	name   string
}

func (b *bucket) Name() string {
	return b.name
}

func (b *bucket) ListObjects(ctx context.Context, query *storage.Query) (*storage.Objects, error) {
	authContext := cloud.WithContext(ctx, b.projID, b.client)
	return storage.ListObjects(authContext, b.name, query)
}

func (b *bucket) NewReader(ctx context.Context, objectName string) (io.ReadCloser, error) {
	authContext := cloud.WithContext(ctx, b.projID, b.client)
	return storage.NewReader(authContext, b.name, objectName)
}

func (b *bucket) NewWriter(
	ctx context.Context,
	attrs *storage.ObjectAttrs) (ObjectWriter, error) {
	authContext := cloud.WithContext(ctx, b.projID, b.client)

	// As of 2015-02, the wrapped storage package doesn't check this for us,
	// causing silently transformed names:
	//     https://github.com/GoogleCloudPlatform/gcloud-golang/issues/111
	if !utf8.ValidString(attrs.Name) {
		return nil, errors.New("Invalid object name: not valid UTF-8")
	}

	// Create and initialize the wrapped writer.
	wrapped := storage.NewWriter(authContext, b.name, attrs.Name)
	wrapped.ObjectAttrs = *attrs

	w := &objectWriter{
		wrapped: wrapped,
	}

	return w, nil
}

func (b *bucket) UpdateObject(
	ctx context.Context,
	newAttrs *storage.ObjectAttrs) (*storage.Object, error) {
	authContext := cloud.WithContext(ctx, b.projID, b.client)
	return storage.UpdateAttrs(authContext, b.Name(), newAttrs.Name, *newAttrs)
}

func (b *bucket) DeleteObject(ctx context.Context, name string) error {
	authContext := cloud.WithContext(ctx, b.projID, b.client)
	return storage.DeleteObject(authContext, b.name, name)
}
<commit_msg>Added documentation for UpdateObject.<commit_after>// Copyright 2015 Google Inc. All Rights Reserved.
// Author: jacobsa@google.com (Aaron Jacobs)

package gcs

import (
	"errors"
	"io"
	"net/http"
	"unicode/utf8"

	"golang.org/x/net/context"
	"google.golang.org/cloud"
	"google.golang.org/cloud/storage"
)

// Bucket represents a GCS bucket, pre-bound with a bucket name and necessary
// authorization information.
//
// Each method that may block accepts a context object that is used for
// deadlines and cancellation. Users need not package authorization information
// into the context object (using cloud.WithContext or similar).
type Bucket interface {
	Name() string

	// List the objects in the bucket that meet the criteria defined by the
	// query, returning a result object that contains the results and potentially
	// a cursor for retrieving the next portion of the larger set of results.
	ListObjects(ctx context.Context, query *storage.Query) (*storage.Objects, error)

	// Create a reader for the contents of the object with the given name. The
	// caller must arrange for the reader to be closed when it is no longer
	// needed.
	NewReader(ctx context.Context, objectName string) (io.ReadCloser, error)

	// Return an ObjectWriter that can be used to create or overwrite an object
	// with the given attributes. attrs.Name must be specified. Otherwise, nil-
	// and zero-valud attributes are ignored.
	//
	// Object names must:
	//
	// *  be non-empty.
	// *  be no longer than 1024 bytes.
	// *  be valid UTF-8.
	// *  not contain the code point U+000A (line feed).
	// *  not contain the code point U+000D (carriage return).
	//
	// See here for authoritative documentation:
	//     https://cloud.google.com/storage/docs/bucket-naming#objectnames
	NewWriter(ctx context.Context, attrs *storage.ObjectAttrs) (ObjectWriter, error)

	// Update the object specified by newAttrs.Name, patching using the non-zero
	// fields of newAttrs.
	UpdateObject(
		ctx context.Context,
		newAttrs *storage.ObjectAttrs) (*storage.Object, error)

	// Delete the object with the given name.
	DeleteObject(ctx context.Context, name string) error
}

type bucket struct {
	projID string
	client *http.Client
	name   string
}

func (b *bucket) Name() string {
	return b.name
}

func (b *bucket) ListObjects(ctx context.Context, query *storage.Query) (*storage.Objects, error) {
	authContext := cloud.WithContext(ctx, b.projID, b.client)
	return storage.ListObjects(authContext, b.name, query)
}

func (b *bucket) NewReader(ctx context.Context, objectName string) (io.ReadCloser, error) {
	authContext := cloud.WithContext(ctx, b.projID, b.client)
	return storage.NewReader(authContext, b.name, objectName)
}

func (b *bucket) NewWriter(
	ctx context.Context,
	attrs *storage.ObjectAttrs) (ObjectWriter, error) {
	authContext := cloud.WithContext(ctx, b.projID, b.client)

	// As of 2015-02, the wrapped storage package doesn't check this for us,
	// causing silently transformed names:
	//     https://github.com/GoogleCloudPlatform/gcloud-golang/issues/111
	if !utf8.ValidString(attrs.Name) {
		return nil, errors.New("Invalid object name: not valid UTF-8")
	}

	// Create and initialize the wrapped writer.
	wrapped := storage.NewWriter(authContext, b.name, attrs.Name)
	wrapped.ObjectAttrs = *attrs

	w := &objectWriter{
		wrapped: wrapped,
	}

	return w, nil
}

func (b *bucket) UpdateObject(
	ctx context.Context,
	newAttrs *storage.ObjectAttrs) (*storage.Object, error) {
	authContext := cloud.WithContext(ctx, b.projID, b.client)
	return storage.UpdateAttrs(authContext, b.Name(), newAttrs.Name, *newAttrs)
}

func (b *bucket) DeleteObject(ctx context.Context, name string) error {
	authContext := cloud.WithContext(ctx, b.projID, b.client)
	return storage.DeleteObject(authContext, b.name, name)
}
<|endoftext|>


<commit_before>package decoder

import (
	"bytes"
	"encoding/xml"
	"io"
)

const TERMINATOR = "nonexisting_node_this_is_a_hack"

type InnerXMLBuffer []byte

func (self *InnerXMLBuffer) ReadByte() (byte, error) {
	var b byte
	var err error
	if len(*self) > 0 {
		b = (*self)[0]
		*self = append((*self)[:0], (*self)[1:]...)
	} else {
		err = io.EOF
	}

	return b, err
}

func (self *InnerXMLBuffer) Read(b []byte) (int, error) {
	if len(*self) == 0 {
		return 0, io.EOF
	}

	n := copy(b, *self)
	*self = (*self)[n:]
	return n, nil
}

type InnerDecoder struct {
	*xml.Decoder
	*InnerXMLBuffer
}

func NewInnerDecoder() *InnerDecoder {
	buffer := new(InnerXMLBuffer)
	return &InnerDecoder{
		Decoder:        xml.NewDecoder(buffer),
		InnerXMLBuffer: buffer,
	}
}

func (self *InnerDecoder) PutXML(b []byte) {
	*self.InnerXMLBuffer = bytes.TrimSpace(b)
	// We need this fake tag to trick xml.Decoder and exit the parsing loop once inner XML buffer is empty.
	// We check if buffer is empty after every tag parsed. Without this tag we can have "chardata" before
	// outer closing tag and xml.Decoder.Token will never return till EOF,
	// but after EOF it won't read any more data from buffer even if we provide any
	*self.InnerXMLBuffer = append(*self.InnerXMLBuffer, "<"+TERMINATOR+"/>"...)
}

func (self *InnerDecoder) IsEmpty() bool {
	return len(*self.InnerXMLBuffer) == 0
}
<commit_msg>Trying some optimizations<commit_after>package decoder

import (
	"bytes"
	"encoding/xml"
	"io"
)

const TERMINATOR = "nonexisting_node_this_is_a_hack"

type InnerXMLBuffer []byte

func (self *InnerXMLBuffer) ReadByte() (byte, error) {
	var b byte
	var err error
	if len(*self) > 0 {
		b = (*self)[0]
		*self = (*self)[1:]
	} else {
		err = io.EOF
	}

	return b, err
}

func (self *InnerXMLBuffer) Read(b []byte) (int, error) {
	if len(*self) == 0 {
		return 0, io.EOF
	}

	n := copy(b, *self)
	*self = (*self)[n:]
	return n, nil
}

type InnerDecoder struct {
	*xml.Decoder
	*InnerXMLBuffer
}

func NewInnerDecoder() *InnerDecoder {
	buffer := new(InnerXMLBuffer)
	return &InnerDecoder{
		Decoder:        xml.NewDecoder(buffer),
		InnerXMLBuffer: buffer,
	}
}

func (self *InnerDecoder) PutXML(b []byte) {
	*self.InnerXMLBuffer = bytes.TrimSpace(b)
	// We need this fake tag to trick xml.Decoder and exit the parsing loop once inner XML buffer is empty.
	// We check if buffer is empty after every tag parsed. Without this tag we can have "chardata" before
	// outer closing tag and xml.Decoder.Token will never return till EOF,
	// but after EOF it won't read any more data from buffer even if we provide any
	*self.InnerXMLBuffer = append(*self.InnerXMLBuffer, "<"+TERMINATOR+"/>"...)
}

func (self *InnerDecoder) IsEmpty() bool {
	return len(*self.InnerXMLBuffer) == 0
}
<|endoftext|>


<commit_before>// +build 386 amd64 arm

package native

// Printer type. Don't touch the internals.
type Printer struct {
	w []byte
}

// Create a new printer with empty output.
func Out() *Printer {
	return &Printer{[]byte{}}
}

// Create a new printer with output prefixed with the given byte slice.
func Outwith(b []byte) *Printer {
	return &Printer{b}
}

// Create a new printer with an empty slice with the capacity given below.
func OutCap(initialcap int) *Printer {
	return &Printer{make([]byte, 0, initialcap)}
}

// Output a byte.
func (p *Printer) Byte(d byte) *Printer {
	p.w = append(p.w, d)
	return p
}

// Output 2 native endian bytes.
func (p *Printer) U16(d uint16) *Printer {
	p.w = append(p.w, byte(d), byte(d>>8))
	return p
}

// Output 4 native endian bytes.
func (p *Printer) U32(d uint32) *Printer {
	p.w = append(p.w, byte(d), byte(d>>8), byte(d>>16), byte(d>>24))
	return p
}

// Output 4 native endian bytes.
func (p *Printer) U64(d uint64) *Printer {
	p.w = append(p.w, byte(d), byte(d>>8), byte(d>>16), byte(d>>24), byte(d>>32), byte(d>>40), byte(d>>48), byte(d>>56))
	return p
}

var z16 = make([]byte, 16)

// Align to boundary
func (p *Printer) Align(n int) *Printer {
	r := len(p.w) % n
	if r == 0 {
		return p
	}
	r = n - r
	for r > 0 {
		cur := r
		if cur > 16 {
			cur = 16
		}
		p.w = append(p.w, z16[:cur]...)
		r -= cur
	}

	return p
}

// Skip (zero-fill) some bytes.
func (p *Printer) Skip(n int) *Printer {
	for n > 0 {
		cur := n
		if cur > 16 {
			cur = 16
		}
		p.w = append(p.w, z16[:cur]...)
		n -= cur
	}

	return p
}

// Output a raw byte slice with no length prefix.
func (p *Printer) Bytes(d []byte) *Printer {
	p.w = append(p.w, d...)
	return p
}

// Output a raw string with no length prefix.
func (p *Printer) String(d string) *Printer {
	p.w = append(p.w, []byte(d)...)
	return p
}

// Output a string with a 4 byte native endian length prefix and no trailing null.
func (p *Printer) U32String(d string) *Printer {
	return p.U32(uint32(len(d))).String(d)
}

// Output bytes with a 4 byte native endian length prefix and no trailing null.
func (p *Printer) U32Bytes(d []byte) *Printer {
	return p.U32(uint32(len(d))).Bytes(d)
}

// Output a string with a 2 byte native endian length prefix and no trailing null.
func (p *Printer) U16String(d string) *Printer {
	if len(d) > 0xffff {
		panic("binprinter: string too long")
	}
	return p.U16(uint16(len(d))).String(d)
}

// Output a string with a 1 byte native endian length prefix and no trailing null.
func (p *Printer) U8String(d string) *Printer {
	if len(d) > 0xff {
		panic("binprinter: string too long")
	}
	return p.Byte(byte(len(d))).String(d)
}

// Output a string terminated by a null-byte
func (p *Printer) String0(d string) *Printer {
	return p.String(d).Byte(0)
}

// Get the output as a byte slice.
func (p *Printer) Out() []byte {
	return p.w
}

// Start counting bytes for the length field in question.
func (p *Printer) LenStart(l *Len) *Printer {
	l.start = len(p.w)
	return p
}

// Add a 16 bit field at the current location that will be filled with the length.
func (p *Printer) LenU16(l *Len) *Printer {
	l.ls = append(l.ls, ls{uint32(len(p.w)), 2})
	return p.U16(0)
}

// Call LenDone for all the arguments
func (p *Printer) LensDone(ls ...*Len) *Printer {
	for _, l := range ls {
		p.LenDone(l)
	}
	return p
}

// Fill fields associated with this length with the current offset.
func (p *Printer) LenDone(l *Len) *Printer {
	plen := len(p.w) - l.start
	for _, ls := range l.ls {
		switch ls.size {
		case 2:
			PutU16(p.w[ls.offset:], uint16(plen))
		}
	}
	return p
}

// Type for handling length fields.
type Len struct {
	ls    []ls
	start int
}

type ls struct {
	offset uint32
	size   uint32
}
<commit_msg>Add LenU32<commit_after>// +build 386 amd64 arm

package native

// Printer type. Don't touch the internals.
type Printer struct {
	w []byte
}

// Create a new printer with empty output.
func Out() *Printer {
	return &Printer{[]byte{}}
}

// Create a new printer with output prefixed with the given byte slice.
func Outwith(b []byte) *Printer {
	return &Printer{b}
}

// Create a new printer with an empty slice with the capacity given below.
func OutCap(initialcap int) *Printer {
	return &Printer{make([]byte, 0, initialcap)}
}

// Output a byte.
func (p *Printer) Byte(d byte) *Printer {
	p.w = append(p.w, d)
	return p
}

// Output 2 native endian bytes.
func (p *Printer) U16(d uint16) *Printer {
	p.w = append(p.w, byte(d), byte(d>>8))
	return p
}

// Output 4 native endian bytes.
func (p *Printer) U32(d uint32) *Printer {
	p.w = append(p.w, byte(d), byte(d>>8), byte(d>>16), byte(d>>24))
	return p
}

// Output 4 native endian bytes.
func (p *Printer) U64(d uint64) *Printer {
	p.w = append(p.w, byte(d), byte(d>>8), byte(d>>16), byte(d>>24), byte(d>>32), byte(d>>40), byte(d>>48), byte(d>>56))
	return p
}

var z16 = make([]byte, 16)

// Align to boundary
func (p *Printer) Align(n int) *Printer {
	r := len(p.w) % n
	if r == 0 {
		return p
	}
	r = n - r
	for r > 0 {
		cur := r
		if cur > 16 {
			cur = 16
		}
		p.w = append(p.w, z16[:cur]...)
		r -= cur
	}

	return p
}

// Skip (zero-fill) some bytes.
func (p *Printer) Skip(n int) *Printer {
	for n > 0 {
		cur := n
		if cur > 16 {
			cur = 16
		}
		p.w = append(p.w, z16[:cur]...)
		n -= cur
	}

	return p
}

// Output a raw byte slice with no length prefix.
func (p *Printer) Bytes(d []byte) *Printer {
	p.w = append(p.w, d...)
	return p
}

// Output a raw string with no length prefix.
func (p *Printer) String(d string) *Printer {
	p.w = append(p.w, []byte(d)...)
	return p
}

// Output a string with a 4 byte native endian length prefix and no trailing null.
func (p *Printer) U32String(d string) *Printer {
	return p.U32(uint32(len(d))).String(d)
}

// Output bytes with a 4 byte native endian length prefix and no trailing null.
func (p *Printer) U32Bytes(d []byte) *Printer {
	return p.U32(uint32(len(d))).Bytes(d)
}

// Output a string with a 2 byte native endian length prefix and no trailing null.
func (p *Printer) U16String(d string) *Printer {
	if len(d) > 0xffff {
		panic("binprinter: string too long")
	}
	return p.U16(uint16(len(d))).String(d)
}

// Output a string with a 1 byte native endian length prefix and no trailing null.
func (p *Printer) U8String(d string) *Printer {
	if len(d) > 0xff {
		panic("binprinter: string too long")
	}
	return p.Byte(byte(len(d))).String(d)
}

// Output a string terminated by a null-byte
func (p *Printer) String0(d string) *Printer {
	return p.String(d).Byte(0)
}

// Get the output as a byte slice.
func (p *Printer) Out() []byte {
	return p.w
}

// Start counting bytes for the length field in question.
func (p *Printer) LenStart(l *Len) *Printer {
	l.start = len(p.w)
	return p
}

// Add a 16 bit field at the current location that will be filled with the length.
func (p *Printer) LenU16(l *Len) *Printer {
	l.ls = append(l.ls, ls{uint32(len(p.w)), 2})
	return p.U16(0)
}

// Add a 32 bit field at the current location that will be filled with the length.
func (p *Printer) LenU16(l *Len) *Printer {
	l.ls = append(l.ls, ls{uint32(len(p.w)), 4})
	return p.U16(0)
}

// Call LenDone for all the arguments
func (p *Printer) LensDone(ls ...*Len) *Printer {
	for _, l := range ls {
		p.LenDone(l)
	}
	return p
}

// Fill fields associated with this length with the current offset.
func (p *Printer) LenDone(l *Len) *Printer {
	plen := len(p.w) - l.start
	for _, ls := range l.ls {
		switch ls.size {
		case 2:
			PutU16(p.w[ls.offset:], uint16(plen))
		case 4:
			PutU32(p.w[ls.offset:], uint32(plen))
		}
	}
	return p
}

// Type for handling length fields.
type Len struct {
	ls    []ls
	start int
}

type ls struct {
	offset uint32
	size   uint32
}
<|endoftext|>


<commit_before>package bitcoin

import (
	"errors"
	"fmt"

	"github.com/bfix/gospel/math"
)

// Curve is the elliptic curve used for Bitcoin (Secp256k1)
type Curve struct {
	// P is the generator of the underlying field "F_p"
	// = 2^256 - 2^32 - 2^9 - 2^8 - 2^7 - 2^6 - 2^4 - 1
	P *math.Int
	// Gx is the x-coord of the base point
	Gx *math.Int
	// Gy is the y-coord of the base point
	Gy *math.Int
	// N is the order of G
	N *math.Int
	// curve parameter (=7)
	B *math.Int
}

var (
	// Curve is the reference to a curve instance (meant as singleton)
	c = &Curve{
		P:  math.NewIntFromHex("FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFEFFFFFC2F"),
		Gx: math.NewIntFromHex("79BE667EF9DCBBAC55A06295CE870B07029BFCDB2DCE28D959F2815B16F81798"),
		Gy: math.NewIntFromHex("483ADA7726A3C4655DA4FBFC0E1108A8FD17B448A68554199C47D08FFB10D4B8"),
		N:  math.NewIntFromHex("FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFEBAAEDCE6AF48A03BBFD25E8CD0364141"),
		B:  math.SEVEN,
	}
)

// GetCurve returns the singleton curve instance
func GetCurve() *Curve {
	return c
}

// GetBasePoint returns the base Point of the curve
func GetBasePoint() *Point {
	return NewPoint(c.Gx, c.Gy)
}

// MultBase multiplies the base Point of the curve with a scalar value k
func MultBase(k *math.Int) *Point {
	return GetBasePoint().Mult(k)
}

// Solve the curve equation for given x-ccordinate (returns positive y)
func Solve(x *math.Int) (*math.Int, bool) {
	// compute 'y = +√(x³ + 7)'
	y2 := pAdd(pCub(x), math.SEVEN)
	y, err := math.SqrtModP(y2, c.P)
	return y, err == nil
}

// SignY of a point (y-coordinate) on the Bitcoin curve:
// +1 if above x-axis, -1 if below.
func SignY(p *Point) (int, error) {
	y, ok := Solve(p.X())
	if ok {
		// positive solution?
		if y.Equals(p.Y()) {
			return 1, nil
		}
		// negative solution?
		if y.Equals(c.P.Sub(p.Y())) {
			return -1, nil
		}
	}
	return 0, errors.New("Point not on curve")
}

// Inf is the point at "infinity"
var Inf = NewPoint(math.ZERO, math.ZERO)

// Point (x,y) on the curve
type Point struct { // exported Point type
	x, y *math.Int // coordinate values
}

// NewPoint instaniates a new Point
func NewPoint(a, b *math.Int) *Point {
	return &Point{x: a, y: b}
}

// NewPointFromBytes reconstructs a Point from binary representation.
func NewPointFromBytes(b []byte) (p *Point, compr bool, err error) {
	p = NewPoint(math.ZERO, math.ZERO)
	err = nil
	compr = true
	switch b[0] {
	case 0:
	case 4:
		p.x = math.NewIntFromBytes(b[1:33])
		p.y = math.NewIntFromBytes(b[33:])
		compr = false
	case 3:
		p.x = math.NewIntFromBytes(b[1:])
		p.y, err = computeY(p.x, 1)
		if err != nil {
			return
		}
	case 2:
		p.x = math.NewIntFromBytes(b[1:])
		p.y, err = computeY(p.x, 0)
		if err != nil {
			return
		}
	default:
		err = errors.New("Invalid binary Point representation")
	}
	return
}

// X returns the x-coordinate of a point.
func (p *Point) X() *math.Int {
	return p.x
}

// Y returns the y-coordinate of a point.
func (p *Point) Y() *math.Int {
	return p.y
}

// Neg returns -P for the point P.
func (p *Point) Neg() *Point {
	return NewPoint(p.x, p.y.Neg())
}

// String returns a human-readable representation of a point.
func (p *Point) String() string {
	return fmt.Sprintf("(%v,%v)", p.x, p.y)
}

// Equals checks if two Points are equal
func (p *Point) Equals(q *Point) bool {
	return p.x.Cmp(q.x) == 0 && p.y.Cmp(q.y) == 0
}

// IsInf checks if a Point is at infinity
func (p *Point) IsInf() bool {
	return p.x.Cmp(math.ZERO) == 0 && p.y.Cmp(math.ZERO) == 0
}

// IsOnCurve checks if a Point (x,y) is on the curve
func (p *Point) IsOnCurve() bool {
	// y² = x³ + 7
	y2 := pSqr(p.y)
	x3 := pCub(p.x)
	return y2.Cmp(pAdd(x3, c.B)) == 0
}

// Add two Points on the curve
func (p *Point) Add(q *Point) *Point {
	if p.Equals(q) {
		return p.Double()
	}
	if p.Equals(Inf) {
		return q
	}
	if q.Equals(Inf) {
		return p
	}
	_p1 := newJacPoint(p.x, p.y, math.ONE)
	_p2 := newJacPoint(q.x, q.y, math.ONE)
	return _p1.add(_p2).conv()
}

// Double a Point on the curve
func (p *Point) Double() *Point {
	if p.Equals(Inf) {
		return Inf
	}
	return newJacPoint(p.x, p.y, math.ONE).double().conv()
}

// Mult multiplies a Point on the curve with a scalar value k using
// a Montgomery multiplication approach
func (p *Point) Mult(k *math.Int) *Point {
	return newJacPoint(p.x, p.y, math.ONE).mult(k).conv()
}

// Bytes returns a byte representation of Point (compressed or uncompressed).
func (p *Point) Bytes(compressed bool) []byte {
	if p.Equals(Inf) {
		return []byte{0}
	}
	var res []byte
	if compressed {
		rc := byte(2)
		if p.y.Bit(0) == 1 {
			rc = 3
		}
		res = append(res, rc)
		res = append(res, coordAsBytes(p.x)...)
	} else {
		res = append(res, 4)
		res = append(res, coordAsBytes(p.x)...)
		res = append(res, coordAsBytes(p.y)...)
	}
	return res
}

// JacPoint is a point on the curve that is represented internally in
// Jacobian coordinates (X,Y,Z) with "x = X/Z^2" and "y = Y/Z^3". See:
// [http://www.hyperelliptic.org/EFD/g1p/auto-shortw-jacobian-0.html]
type jacPoint struct { // internal Point type
	x, y, z *math.Int // using Jacobian coordinates
}

// NewJacPoint instaniates a new Point
func newJacPoint(a, b, c *math.Int) *jacPoint {
	return &jacPoint{x: a, y: b, z: c}
}

// Point at infinity
var jacInf = newJacPoint(Inf.x, Inf.y, math.ONE)

// check if a Point is at infinity
func (p *jacPoint) isInf() bool {
	return p.x.Equals(math.ZERO) && p.y.Equals(math.ZERO)
}

// convert internal Point to external representation
func (p *jacPoint) conv() *Point {
	zi := pInv(p.z)
	x := pMul(p.x, pSqr(zi))
	y := pMul(p.y, pCub(zi))
	return NewPoint(x, y)
}

// add two jacPoints on the curve
// [http://www.hyperelliptic.org/EFD/g1p/data/shortw/jacobian-0/addJacition/addJac-2007-bl]
func (p *jacPoint) add(q *jacPoint) *jacPoint {
	if p.isInf() {
		return q
	}
	if q.isInf() {
		return p
	}
	z1z1 := pSqr(p.z)
	z2z2 := pSqr(q.z)
	u1 := pMul(p.x, z2z2)
	u2 := pMul(q.x, z1z1)
	s1 := pMul(pMul(p.y, q.z), z2z2)
	s2 := pMul(pMul(q.y, p.z), z1z1)
	h := pSub(u2, u1)
	i := pSqr(pMul(math.TWO, h))
	j := pMul(h, i)
	r := pMul(math.TWO, pSub(s2, s1))
	v := pMul(u1, i)
	w := pAdd(p.z, q.z)
	x := pSub(pSub(pSqr(r), j), pMul(math.TWO, v))
	y := pSub(pMul(r, pSub(v, x)), pMul(math.TWO, pMul(s1, j)))
	z := pMul(pSub(pSub(pSqr(w), z1z1), z2z2), h)
	return newJacPoint(x, y, z)
}

// double a Point on the curve
// [http://www.hyperelliptic.org/EFD/g1p/data/shortw/jacobian-0/doubling/dbl-2009-alnr]
func (p *jacPoint) double() *jacPoint {
	if p.isInf() {
		return p
	}
	a := pSqr(p.x)
	b := pSqr(p.y)
	zz := pSqr(p.z)
	c := pSqr(b)
	d := pMul(math.TWO, pSub(pSub(pSqr(p.x.Add(b)), a), c))
	e := pMul(math.THREE, a)
	f := pSqr(e)
	x := pSub(f, pMul(math.TWO, d))
	y := pSub(pMul(e, pSub(d, x)), pMul(math.EIGHT, c))
	z := pSub(pSub(pSqr(pAdd(p.y, p.z)), b), zz)
	return newJacPoint(x, y, z)
}

// Multiply a Point on the curve with a scalar value k using
// a Montgomery multiplication algorithm
func (p *jacPoint) mult(k *math.Int) *jacPoint {
	if p.isInf() {
		return p
	}
	if k.Cmp(math.ZERO) == 0 {
		return jacInf
	}
	r := jacInf
	for _, val := range k.Bytes() {
		for pos := 0; pos < 8; pos++ {
			r = r.double()
			if val&0x80 == 0x80 {
				r = p.add(r)
			}
			val <<= 1
		}
	}
	return r
}

// helper: convert coordinate to byte array of correct length
func coordAsBytes(v *math.Int) []byte {
	bv := v.Bytes()
	plen := 32 - len(bv)
	if plen == 0 {
		return bv
	}
	b := make([]byte, plen)
	return append(b, bv...)
}

// helper: reconstruct y-coordinate of Point
func computeY(x *math.Int, m uint) (y *math.Int, err error) {
	y = math.ZERO
	err = nil
	y2 := pAdd(pCub(x), c.B)
	y, err = math.SqrtModP(y2, c.P)
	if err == nil {
		if y.Bit(0) != m {
			y = c.P.Sub(y)
		}
	}
	return
}

// modulus
func nMod(a *math.Int) *math.Int {
	return a.Mod(c.N)
}

func pInv(a *math.Int) *math.Int {
	return a.ModInverse(c.P)
}

func pMul(a, b *math.Int) *math.Int {
	return a.Mul(b).Mod(c.P)
}

func pSqr(a *math.Int) *math.Int {
	return pMul(a, a)
}

func pCub(a *math.Int) *math.Int {
	return pMul(pSqr(a), a)
}

func pSub(a, b *math.Int) *math.Int {
	x := a.Sub(b)
	if x.Sign() == -1 {
		x = x.Add(c.P)
	}
	return x
}

func pAdd(a, b *math.Int) *math.Int {
	return a.Add(b).Mod(c.P)
}

func nInv(a *math.Int) *math.Int {
	return a.ModInverse(c.N)
}

func nMul(a, b *math.Int) *math.Int {
	return a.Mul(b).Mod(c.N)
}

func nRnd(a *math.Int) *math.Int {
	return math.NewIntRndRange(a, c.N)
}
<commit_msg>Fixed conversion of Jacobian point to affine point.<commit_after>package bitcoin

import (
	"errors"
	"fmt"

	"github.com/bfix/gospel/math"
)

// Curve is the elliptic curve used for Bitcoin (Secp256k1)
type Curve struct {
	// P is the generator of the underlying field "F_p"
	// = 2^256 - 2^32 - 2^9 - 2^8 - 2^7 - 2^6 - 2^4 - 1
	P *math.Int
	// Gx is the x-coord of the base point
	Gx *math.Int
	// Gy is the y-coord of the base point
	Gy *math.Int
	// N is the order of G
	N *math.Int
	// curve parameter (=7)
	B *math.Int
}

var (
	// Curve is the reference to a curve instance (meant as singleton)
	c = &Curve{
		P:  math.NewIntFromHex("FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFEFFFFFC2F"),
		Gx: math.NewIntFromHex("79BE667EF9DCBBAC55A06295CE870B07029BFCDB2DCE28D959F2815B16F81798"),
		Gy: math.NewIntFromHex("483ADA7726A3C4655DA4FBFC0E1108A8FD17B448A68554199C47D08FFB10D4B8"),
		N:  math.NewIntFromHex("FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFEBAAEDCE6AF48A03BBFD25E8CD0364141"),
		B:  math.SEVEN,
	}
)

// GetCurve returns the singleton curve instance
func GetCurve() *Curve {
	return c
}

// GetBasePoint returns the base Point of the curve
func GetBasePoint() *Point {
	return NewPoint(c.Gx, c.Gy)
}

// MultBase multiplies the base Point of the curve with a scalar value k
func MultBase(k *math.Int) *Point {
	return GetBasePoint().Mult(k)
}

// Solve the curve equation for given x-ccordinate (returns positive y)
func Solve(x *math.Int) (*math.Int, bool) {
	// compute 'y = +√(x³ + 7)'
	y2 := pAdd(pCub(x), math.SEVEN)
	y, err := math.SqrtModP(y2, c.P)
	return y, err == nil
}

// SignY of a point (y-coordinate) on the Bitcoin curve:
// +1 if above x-axis, -1 if below.
func SignY(p *Point) (int, error) {
	y, ok := Solve(p.X())
	if ok {
		// positive solution?
		if y.Equals(p.Y()) {
			return 1, nil
		}
		// negative solution?
		if y.Equals(c.P.Sub(p.Y())) {
			return -1, nil
		}
	}
	return 0, errors.New("Point not on curve")
}

// Inf is the point at "infinity"
var Inf = NewPoint(math.ZERO, math.ZERO)

// Point (x,y) on the curve
type Point struct { // exported Point type
	x, y *math.Int // coordinate values
}

// NewPoint instaniates a new Point
func NewPoint(a, b *math.Int) *Point {
	return &Point{x: a, y: b}
}

// NewPointFromBytes reconstructs a Point from binary representation.
func NewPointFromBytes(b []byte) (p *Point, compr bool, err error) {
	p = NewPoint(math.ZERO, math.ZERO)
	err = nil
	compr = true
	switch b[0] {
	case 0:
	case 4:
		p.x = math.NewIntFromBytes(b[1:33])
		p.y = math.NewIntFromBytes(b[33:])
		compr = false
	case 3:
		p.x = math.NewIntFromBytes(b[1:])
		p.y, err = computeY(p.x, 1)
		if err != nil {
			return
		}
	case 2:
		p.x = math.NewIntFromBytes(b[1:])
		p.y, err = computeY(p.x, 0)
		if err != nil {
			return
		}
	default:
		err = errors.New("Invalid binary Point representation")
	}
	return
}

// X returns the x-coordinate of a point.
func (p *Point) X() *math.Int {
	return p.x
}

// Y returns the y-coordinate of a point.
func (p *Point) Y() *math.Int {
	return p.y
}

// Neg returns -P for the point P.
func (p *Point) Neg() *Point {
	return NewPoint(p.x, p.y.Neg())
}

// String returns a human-readable representation of a point.
func (p *Point) String() string {
	return fmt.Sprintf("(%v,%v)", p.x, p.y)
}

// Equals checks if two Points are equal
func (p *Point) Equals(q *Point) bool {
	return p.x.Cmp(q.x) == 0 && p.y.Cmp(q.y) == 0
}

// IsInf checks if a Point is at infinity
func (p *Point) IsInf() bool {
	return p.x.Cmp(math.ZERO) == 0 && p.y.Cmp(math.ZERO) == 0
}

// IsOnCurve checks if a Point (x,y) is on the curve
func (p *Point) IsOnCurve() bool {
	// y² = x³ + 7
	y2 := pSqr(p.y)
	x3 := pCub(p.x)
	return y2.Cmp(pAdd(x3, c.B)) == 0
}

// Add two Points on the curve
func (p *Point) Add(q *Point) *Point {
	if p.Equals(q) {
		return p.Double()
	}
	if p.Equals(Inf) {
		return q
	}
	if q.Equals(Inf) {
		return p
	}
	_p1 := newJacPoint(p.x, p.y, math.ONE)
	_p2 := newJacPoint(q.x, q.y, math.ONE)
	return _p1.add(_p2).conv()
}

// Double a Point on the curve
func (p *Point) Double() *Point {
	if p.Equals(Inf) {
		return Inf
	}
	return newJacPoint(p.x, p.y, math.ONE).double().conv()
}

// Mult multiplies a Point on the curve with a scalar value k using
// a Montgomery multiplication approach
func (p *Point) Mult(k *math.Int) *Point {
	return newJacPoint(p.x, p.y, math.ONE).mult(k).conv()
}

// Bytes returns a byte representation of Point (compressed or uncompressed).
func (p *Point) Bytes(compressed bool) []byte {
	if p.Equals(Inf) {
		return []byte{0}
	}
	var res []byte
	if compressed {
		rc := byte(2)
		if p.y.Bit(0) == 1 {
			rc = 3
		}
		res = append(res, rc)
		res = append(res, coordAsBytes(p.x)...)
	} else {
		res = append(res, 4)
		res = append(res, coordAsBytes(p.x)...)
		res = append(res, coordAsBytes(p.y)...)
	}
	return res
}

// JacPoint is a point on the curve that is represented internally in
// Jacobian coordinates (X,Y,Z) with "x = X/Z^2" and "y = Y/Z^3". See:
// [http://www.hyperelliptic.org/EFD/g1p/auto-shortw-jacobian-0.html]
type jacPoint struct { // internal Point type
	x, y, z *math.Int // using Jacobian coordinates
}

// NewJacPoint instaniates a new Point
func newJacPoint(a, b, c *math.Int) *jacPoint {
	return &jacPoint{x: a, y: b, z: c}
}

// Point at infinity
var jacInf = newJacPoint(Inf.x, Inf.y, math.ONE)

// check if a Point is at infinity
func (p *jacPoint) isInf() bool {
	return p.x.Equals(math.ZERO) && p.y.Equals(math.ZERO)
}

// convert internal Point to external representation
func (p *jacPoint) conv() *Point {
	if p.z.Equals(math.ZERO) {
		return NewPoint(math.ZERO, math.ZERO)
	}
	zi := pInv(p.z)
	x := pMul(p.x, pSqr(zi))
	y := pMul(p.y, pCub(zi))
	return NewPoint(x, y)
}

// add two jacPoints on the curve
// [http://www.hyperelliptic.org/EFD/g1p/data/shortw/jacobian-0/addJacition/addJac-2007-bl]
func (p *jacPoint) add(q *jacPoint) *jacPoint {
	if p.isInf() {
		return q
	}
	if q.isInf() {
		return p
	}
	z1z1 := pSqr(p.z)
	z2z2 := pSqr(q.z)
	u1 := pMul(p.x, z2z2)
	u2 := pMul(q.x, z1z1)
	s1 := pMul(pMul(p.y, q.z), z2z2)
	s2 := pMul(pMul(q.y, p.z), z1z1)
	h := pSub(u2, u1)
	i := pSqr(pMul(math.TWO, h))
	j := pMul(h, i)
	r := pMul(math.TWO, pSub(s2, s1))
	v := pMul(u1, i)
	w := pAdd(p.z, q.z)
	x := pSub(pSub(pSqr(r), j), pMul(math.TWO, v))
	y := pSub(pMul(r, pSub(v, x)), pMul(math.TWO, pMul(s1, j)))
	z := pMul(pSub(pSub(pSqr(w), z1z1), z2z2), h)
	return newJacPoint(x, y, z)
}

// double a Point on the curve
// [http://www.hyperelliptic.org/EFD/g1p/data/shortw/jacobian-0/doubling/dbl-2009-alnr]
func (p *jacPoint) double() *jacPoint {
	if p.isInf() {
		return p
	}
	a := pSqr(p.x)
	b := pSqr(p.y)
	zz := pSqr(p.z)
	c := pSqr(b)
	d := pMul(math.TWO, pSub(pSub(pSqr(p.x.Add(b)), a), c))
	e := pMul(math.THREE, a)
	f := pSqr(e)
	x := pSub(f, pMul(math.TWO, d))
	y := pSub(pMul(e, pSub(d, x)), pMul(math.EIGHT, c))
	z := pSub(pSub(pSqr(pAdd(p.y, p.z)), b), zz)
	return newJacPoint(x, y, z)
}

// Multiply a Point on the curve with a scalar value k using
// a Montgomery multiplication algorithm
func (p *jacPoint) mult(k *math.Int) *jacPoint {
	if p.isInf() {
		return p
	}
	if k.Cmp(math.ZERO) == 0 {
		return jacInf
	}
	r := jacInf
	for _, val := range k.Bytes() {
		for pos := 0; pos < 8; pos++ {
			r = r.double()
			if val&0x80 == 0x80 {
				r = p.add(r)
			}
			val <<= 1
		}
	}
	return r
}

// helper: convert coordinate to byte array of correct length
func coordAsBytes(v *math.Int) []byte {
	bv := v.Bytes()
	plen := 32 - len(bv)
	if plen == 0 {
		return bv
	}
	b := make([]byte, plen)
	return append(b, bv...)
}

// helper: reconstruct y-coordinate of Point
func computeY(x *math.Int, m uint) (y *math.Int, err error) {
	y = math.ZERO
	err = nil
	y2 := pAdd(pCub(x), c.B)
	y, err = math.SqrtModP(y2, c.P)
	if err == nil {
		if y.Bit(0) != m {
			y = c.P.Sub(y)
		}
	}
	return
}

// modulus
func nMod(a *math.Int) *math.Int {
	return a.Mod(c.N)
}

func pInv(a *math.Int) *math.Int {
	return a.ModInverse(c.P)
}

func pMul(a, b *math.Int) *math.Int {
	return a.Mul(b).Mod(c.P)
}

func pSqr(a *math.Int) *math.Int {
	return pMul(a, a)
}

func pCub(a *math.Int) *math.Int {
	return pMul(pSqr(a), a)
}

func pSub(a, b *math.Int) *math.Int {
	x := a.Sub(b)
	if x.Sign() == -1 {
		x = x.Add(c.P)
	}
	return x
}

func pAdd(a, b *math.Int) *math.Int {
	return a.Add(b).Mod(c.P)
}

func nInv(a *math.Int) *math.Int {
	return a.ModInverse(c.N)
}

func nMul(a, b *math.Int) *math.Int {
	return a.Mul(b).Mod(c.N)
}

func nRnd(a *math.Int) *math.Int {
	return math.NewIntRndRange(a, c.N)
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/md5"
	"crypto/sha1"
	"crypto/sha256"
	"encoding/json"
	"fmt"
	"hash"
	"io"
	"io/ioutil"
	"mime"
	"os"
	"strconv"
	"strings"

	"github.com/progrium/go-basher"
)

var Version string

func assert(err error) {
	if err != nil {
		println("!!", err.Error())
		os.Exit(2)
	}
}

func fatal(msg string) {
	println("!!", msg)
	os.Exit(2)
}

func UploadUrl(args []string) {
	bytes, err := ioutil.ReadAll(os.Stdin)
	assert(err)
	var release map[string]interface{}
	assert(json.Unmarshal(bytes, &release))
	url, ok := release["upload_url"].(string)
	if !ok {
		os.Exit(2)
	}
	url = strings.Replace(url, "{", "", 1)
	url = strings.Replace(url, "}", "", 1)
	fmt.Println(url)
}

func ReleaseIdFromTagname(args []string) {
	tagname := args[0]
	bytes, err := ioutil.ReadAll(os.Stdin)
	assert(err)
	var releases []map[string]interface{}
	assert(json.Unmarshal(bytes, &releases))
	for _, release := range releases {
		if release["tag_name"].(string) == tagname {
			fmt.Println(strconv.Itoa(int(release["id"].(float64))))
			return
		}
	}
	os.Exit(2)
}

func MimeType(args []string) {
	filename := args[0]
	ext := filename[strings.LastIndex(filename, "."):]
	mime.AddExtensionType(".gz", "application/gzip")
	mime.AddExtensionType(".tgz", "application/gzip")
	mime.AddExtensionType(".tar", "application/tar")
	mime.AddExtensionType(".zip", "application/zip")
	mimetype := mime.TypeByExtension(ext)
	if mimetype != "" {
		fmt.Println(mimetype)
	} else {
		fmt.Println("application/octet-stream")
	}
}

func Checksum(args []string) {
	if len(args) < 1 {
		fatal("No algorithm specified")
	}
	var h hash.Hash
	switch args[0] {
	case "md5":
		h = md5.New()
	case "sha1":
		h = sha1.New()
	case "sha256":
		h = sha256.New()
	default:
		fatal("Algorithm '" + args[0] + "' is unsupported")
	}
	io.Copy(h, os.Stdin)
	fmt.Printf("%x\n", h.Sum(nil))
}

func main() {
	os.Setenv("VERSION", Version)
	basher.Application(map[string]func([]string){
		"upload-url":              UploadUrl,
		"release-id-from-tagname": ReleaseIdFromTagname,
		"mimetype":                MimeType,
		"checksum":                Checksum,
	}, []string{
		"bash/gh-release.bash",
	}, Asset, true)
}
<commit_msg>fix case where file has no extension and caused panic: runtime error: slice bounds out of range<commit_after>package main

import (
	"crypto/md5"
	"crypto/sha1"
	"crypto/sha256"
	"encoding/json"
	"fmt"
	"hash"
	"io"
	"io/ioutil"
	"mime"
	"os"
	"path/filepath"
	"strconv"
	"strings"

	"github.com/progrium/go-basher"
)

var Version string

func assert(err error) {
	if err != nil {
		println("!!", err.Error())
		os.Exit(2)
	}
}

func fatal(msg string) {
	println("!!", msg)
	os.Exit(2)
}

func UploadUrl(args []string) {
	bytes, err := ioutil.ReadAll(os.Stdin)
	assert(err)
	var release map[string]interface{}
	assert(json.Unmarshal(bytes, &release))
	url, ok := release["upload_url"].(string)
	if !ok {
		os.Exit(2)
	}
	url = strings.Replace(url, "{", "", 1)
	url = strings.Replace(url, "}", "", 1)
	fmt.Println(url)
}

func ReleaseIdFromTagname(args []string) {
	tagname := args[0]
	bytes, err := ioutil.ReadAll(os.Stdin)
	assert(err)
	var releases []map[string]interface{}
	assert(json.Unmarshal(bytes, &releases))
	for _, release := range releases {
		if release["tag_name"].(string) == tagname {
			fmt.Println(strconv.Itoa(int(release["id"].(float64))))
			return
		}
	}
	os.Exit(2)
}

func MimeType(args []string) {
	filename := args[0]
	ext := filepath.Ext(filename)
	mime.AddExtensionType(".gz", "application/gzip")
	mime.AddExtensionType(".tgz", "application/gzip")
	mime.AddExtensionType(".tar", "application/tar")
	mime.AddExtensionType(".zip", "application/zip")
	mimetype := mime.TypeByExtension(ext)
	if mimetype != "" {
		fmt.Println(mimetype)
	} else {
		fmt.Println("application/octet-stream")
	}
}

func Checksum(args []string) {
	if len(args) < 1 {
		fatal("No algorithm specified")
	}
	var h hash.Hash
	switch args[0] {
	case "md5":
		h = md5.New()
	case "sha1":
		h = sha1.New()
	case "sha256":
		h = sha256.New()
	default:
		fatal("Algorithm '" + args[0] + "' is unsupported")
	}
	io.Copy(h, os.Stdin)
	fmt.Printf("%x\n", h.Sum(nil))
}

func main() {
	os.Setenv("VERSION", Version)
	basher.Application(map[string]func([]string){
		"upload-url":              UploadUrl,
		"release-id-from-tagname": ReleaseIdFromTagname,
		"mimetype":                MimeType,
		"checksum":                Checksum,
	}, []string{
		"bash/gh-release.bash",
	}, Asset, true)
}
<|endoftext|>


<commit_before>/*
Copyright 2018 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package main

import (
	"errors"
	"flag"
	"io/ioutil"
	"os"
	"path/filepath"

	"github.com/sirupsen/logrus"
	_ "k8s.io/client-go/plugin/pkg/client/auth/gcp" // support gcp users in .kube/config
	"k8s.io/test-infra/prow/config"
	prowflagutil "k8s.io/test-infra/prow/flagutil"
	"k8s.io/test-infra/prow/github"
	_ "k8s.io/test-infra/prow/hook"
	"k8s.io/test-infra/prow/logrusutil"
	"k8s.io/test-infra/prow/plugins"
	"k8s.io/test-infra/prow/plugins/updateconfig"
)

type options struct {
	sourcePath string

	configPath    string
	jobConfigPath string
	pluginConfig  string

	dryRun     bool
	kubernetes prowflagutil.KubernetesOptions
}

func gatherOptions() options {
	o := options{}
	fs := flag.NewFlagSet(os.Args[0], flag.ExitOnError)

	fs.StringVar(&o.sourcePath, "source-path", "", "Path to root of source directory to use for config updates.")

	fs.StringVar(&o.configPath, "config-path", "/etc/config/config.yaml", "Path to config.yaml.")
	fs.StringVar(&o.jobConfigPath, "job-config-path", "", "Path to prow job configs.")
	fs.StringVar(&o.pluginConfig, "plugin-config", "/etc/plugins/plugins.yaml", "Path to plugin config file.")

	fs.BoolVar(&o.dryRun, "dry-run", true, "Whether or not to make mutating API calls to GitHub.")
	o.kubernetes.AddFlags(fs)

	fs.Parse(os.Args[1:])
	return o
}

func (o *options) Validate() error {
	if o.sourcePath == "" {
		return errors.New("--source-path must be provided")
	}

	if err := o.kubernetes.Validate(o.dryRun); err != nil {
		return err
	}

	return nil
}

func main() {
	logrusutil.ComponentInit("config-bootstrapper")

	o := gatherOptions()
	if err := o.Validate(); err != nil {
		logrus.WithError(err).Fatal("Invalid options")
	}

	configAgent := &config.Agent{}
	if err := configAgent.Start(o.configPath, o.jobConfigPath); err != nil {
		logrus.WithError(err).Fatal("Error starting config agent.")
	}

	pluginAgent := &plugins.ConfigAgent{}
	if err := pluginAgent.Start(o.pluginConfig); err != nil {
		logrus.WithError(err).Fatal("Error starting plugin configuration agent.")
	}

	client, err := o.kubernetes.InfrastructureClusterClient(o.dryRun)
	if err != nil {
		logrus.WithError(err).Fatal("Error getting Kubernetes client.")
	}

	// act like the whole repo just got committed
	var changes []github.PullRequestChange
	filepath.Walk(o.sourcePath, func(path string, info os.FileInfo, err error) error {
		if info.IsDir() {
			return nil
		}
		// we know path will be below sourcePath, but we can't
		// communicate that to the filepath module. We can ignore
		// this error as we can be certain it won't occur
		if relPath, err := filepath.Rel(o.sourcePath, path); err == nil {
			changes = append(changes, github.PullRequestChange{
				Filename: relPath,
				Status:   github.PullRequestFileAdded,
			})
			logrus.Infof("added to mock change: %s", relPath)
		} else {
			logrus.WithError(err).Warn("unexpected error determining relative path to file")
		}
		return nil
	})

	for cm, data := range updateconfig.FilterChanges(pluginAgent.Config().ConfigUpdater, changes, logrus.NewEntry(logrus.StandardLogger())) {
		if cm.Namespace == "" {
			cm.Namespace = configAgent.Config().ProwJobNamespace
		}
		logger := logrus.WithFields(logrus.Fields{"configmap": map[string]string{"name": cm.Name, "namespace": cm.Namespace}})
		if err := updateconfig.Update(&osFileGetter{root: o.sourcePath}, client.CoreV1().ConfigMaps(cm.Namespace), cm.Name, cm.Namespace, data, nil, logger); err != nil {
			logger.WithError(err).Error("failed to update config on cluster")
		} else {
			logger.Info("Successfully processed configmap")
		}
	}
}

type osFileGetter struct {
	root string
}

func (g *osFileGetter) GetFile(filename string) ([]byte, error) {
	return ioutil.ReadFile(filepath.Join(g.root, filename))
}
<commit_msg>config-bootstrapper: propagate failures to return code<commit_after>/*
Copyright 2018 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package main

import (
	"errors"
	"flag"
	"io/ioutil"
	"os"
	"path/filepath"

	"github.com/sirupsen/logrus"
	_ "k8s.io/client-go/plugin/pkg/client/auth/gcp" // support gcp users in .kube/config
	"k8s.io/test-infra/prow/config"
	prowflagutil "k8s.io/test-infra/prow/flagutil"
	"k8s.io/test-infra/prow/github"
	_ "k8s.io/test-infra/prow/hook"
	"k8s.io/test-infra/prow/logrusutil"
	"k8s.io/test-infra/prow/plugins"
	"k8s.io/test-infra/prow/plugins/updateconfig"
)

type options struct {
	sourcePath string

	configPath    string
	jobConfigPath string
	pluginConfig  string

	dryRun     bool
	kubernetes prowflagutil.KubernetesOptions
}

func gatherOptions() options {
	o := options{}
	fs := flag.NewFlagSet(os.Args[0], flag.ExitOnError)

	fs.StringVar(&o.sourcePath, "source-path", "", "Path to root of source directory to use for config updates.")

	fs.StringVar(&o.configPath, "config-path", "/etc/config/config.yaml", "Path to config.yaml.")
	fs.StringVar(&o.jobConfigPath, "job-config-path", "", "Path to prow job configs.")
	fs.StringVar(&o.pluginConfig, "plugin-config", "/etc/plugins/plugins.yaml", "Path to plugin config file.")

	fs.BoolVar(&o.dryRun, "dry-run", true, "Whether or not to make mutating API calls to GitHub.")
	o.kubernetes.AddFlags(fs)

	fs.Parse(os.Args[1:])
	return o
}

func (o *options) Validate() error {
	if o.sourcePath == "" {
		return errors.New("--source-path must be provided")
	}

	if err := o.kubernetes.Validate(o.dryRun); err != nil {
		return err
	}

	return nil
}

func main() {
	var errors int
	logrusutil.ComponentInit("config-bootstrapper")

	o := gatherOptions()
	if err := o.Validate(); err != nil {
		logrus.WithError(err).Fatal("Invalid options")
	}

	configAgent := &config.Agent{}
	if err := configAgent.Start(o.configPath, o.jobConfigPath); err != nil {
		logrus.WithError(err).Fatal("Error starting config agent.")
	}

	pluginAgent := &plugins.ConfigAgent{}
	if err := pluginAgent.Start(o.pluginConfig); err != nil {
		logrus.WithError(err).Fatal("Error starting plugin configuration agent.")
	}

	client, err := o.kubernetes.InfrastructureClusterClient(o.dryRun)
	if err != nil {
		logrus.WithError(err).Fatal("Error getting Kubernetes client.")
	}

	// act like the whole repo just got committed
	var changes []github.PullRequestChange
	filepath.Walk(o.sourcePath, func(path string, info os.FileInfo, err error) error {
		if info.IsDir() {
			return nil
		}
		// we know path will be below sourcePath, but we can't
		// communicate that to the filepath module. We can ignore
		// this error as we can be certain it won't occur
		if relPath, err := filepath.Rel(o.sourcePath, path); err == nil {
			changes = append(changes, github.PullRequestChange{
				Filename: relPath,
				Status:   github.PullRequestFileAdded,
			})
			logrus.Infof("added to mock change: %s", relPath)
		} else {
			logrus.WithError(err).Warn("unexpected error determining relative path to file")
			errors++
		}
		return nil
	})

	for cm, data := range updateconfig.FilterChanges(pluginAgent.Config().ConfigUpdater, changes, logrus.NewEntry(logrus.StandardLogger())) {
		if cm.Namespace == "" {
			cm.Namespace = configAgent.Config().ProwJobNamespace
		}
		logger := logrus.WithFields(logrus.Fields{"configmap": map[string]string{"name": cm.Name, "namespace": cm.Namespace}})
		if err := updateconfig.Update(&osFileGetter{root: o.sourcePath}, client.CoreV1().ConfigMaps(cm.Namespace), cm.Name, cm.Namespace, data, nil, logger); err != nil {
			logger.WithError(err).Error("failed to update config on cluster")
			errors++
		} else {
			logger.Info("Successfully processed configmap")
		}
	}

	if errors > 0 {
		logrus.WithField("fail-count", errors).Fatalf("errors occurred during update")
	}
}

type osFileGetter struct {
	root string
}

func (g *osFileGetter) GetFile(filename string) ([]byte, error) {
	return ioutil.ReadFile(filepath.Join(g.root, filename))
}
<|endoftext|>