<commit_before>package chess

import (
	"strconv"
	"strings"
)

// A Board represents a chess board and its relationship between squares and pieces.
type Board struct {
	bbWhiteKing   bitboard
	bbWhiteQueen  bitboard
	bbWhiteRook   bitboard
	bbWhiteBishop bitboard
	bbWhiteKnight bitboard
	bbWhitePawn   bitboard
	bbBlackKing   bitboard
	bbBlackQueen  bitboard
	bbBlackRook   bitboard
	bbBlackBishop bitboard
	bbBlackKnight bitboard
	bbBlackPawn   bitboard
	whiteSqs      bitboard
	blackSqs      bitboard
	emptySqs      bitboard
	whiteKingSq   Square
	blackKingSq   Square
}

// SquareMap returns a mapping of squares to pieces.  A square is only added to the map if it is occupied.
func (b *Board) SquareMap() map[Square]Piece {
	m := map[Square]Piece{}
	for sq := 0; sq < numOfSquaresInBoard; sq++ {
		p := b.piece(Square(sq))
		if p != NoPiece {
			m[Square(sq)] = p
		}
	}
	return m
}

// Draw returns visual representation of the board useful for debugging.
func (b *Board) Draw() string {
	s := "\n A B C D E F G H\n"
	for r := 7; r >= 0; r-- {
		s += Rank(r).String()
		for f := 0; f < numOfSquaresInRow; f++ {
			p := b.piece(getSquare(File(f), Rank(r)))
			if p == NoPiece {
				s += "-"
			} else {
				s += p.String()
			}
			s += " "
		}
		s += "\n"
	}
	return s
}

// String implements the fmt.Stringer interface and returns
// a string in the FEN board format: rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR
func (b *Board) String() string {
	fen := ""
	for r := 7; r >= 0; r-- {
		for f := 0; f < numOfSquaresInRow; f++ {
			sq := getSquare(File(f), Rank(r))
			p := b.piece(sq)
			if p != NoPiece {
				fen += p.getFENChar()
			} else {
				fen += "1"
			}
		}
		if r != 0 {
			fen += "/"
		}
	}
	for i := 8; i > 1; i-- {
		repeatStr := strings.Repeat("1", i)
		countStr := strconv.Itoa(i)
		fen = strings.Replace(fen, repeatStr, countStr, -1)
	}
	return fen
}

func newBoard(m map[Square]Piece) *Board {
	b := &Board{}
	for _, p1 := range allPieces {
		bm := map[Square]bool{}
		for sq, p2 := range m {
			if p1 == p2 {
				bm[sq] = true
			}
		}
		bb := newBitboard(bm)
		b.setBBForPiece(p1, bb)
	}
	b.calcConvienceBBs()
	return b
}

func (b *Board) update(m *Move) {
	p1 := b.piece(m.s1)
	s1BB := bbForSquare(m.s1)
	s2BB := bbForSquare(m.s2)

	// move s1 piece to s2
	for _, p := range allPieces {
		bb := b.bbForPiece(p)
		// remove what was at s2
		b.setBBForPiece(p, bb & ^s2BB)
		// move what was at s1 to s2
		if bb.Occupied(m.s1) {
			bb = b.bbForPiece(p)
			b.setBBForPiece(p, (bb & ^s1BB)|s2BB)
		}
	}
	// check promotion
	if m.promo != NoPieceType {
		newPiece := getPiece(m.promo, p1.Color())
		// remove pawn
		bbPawn := b.bbForPiece(p1)
		b.setBBForPiece(p1, bbPawn & ^s2BB)
		// add promo piece
		bbPromo := b.bbForPiece(newPiece)
		b.setBBForPiece(newPiece, bbPromo|s2BB)
	}
	// remove captured en passant piece
	if m.HasTag(EnPassant) {
		if p1.Color() == White {
			b.bbBlackPawn = ^(bbForSquare(m.s2) << 8) & b.bbBlackPawn
		} else {
			b.bbWhitePawn = ^(bbForSquare(m.s2) >> 8) & b.bbWhitePawn
		}
	}
	// move rook for castle
	if p1.Color() == White && m.HasTag(KingSideCastle) {
		b.bbWhiteRook = (b.bbWhiteRook & ^bbForSquare(H1) | bbForSquare(F1))
	} else if p1.Color() == White && m.HasTag(QueenSideCastle) {
		b.bbWhiteRook = (b.bbWhiteRook & ^bbForSquare(A1)) | bbForSquare(D1)
	} else if p1.Color() == Black && m.HasTag(KingSideCastle) {
		b.bbBlackRook = (b.bbBlackRook & ^bbForSquare(H8) | bbForSquare(F8))
	} else if p1.Color() == Black && m.HasTag(QueenSideCastle) {
		b.bbBlackRook = (b.bbBlackRook & ^bbForSquare(A8)) | bbForSquare(D8)
	}
	b.calcConvienceBBs()
}

func (b *Board) calcConvienceBBs() {
	whiteSqs := b.bbWhiteKing | b.bbWhiteQueen | b.bbWhiteRook | b.bbWhiteBishop | b.bbWhiteKnight | b.bbWhitePawn
	blackSqs := b.bbBlackKing | b.bbBlackQueen | b.bbBlackRook | b.bbBlackBishop | b.bbBlackKnight | b.bbBlackPawn
	emptySqs := ^(whiteSqs | blackSqs)
	b.whiteSqs = whiteSqs
	b.blackSqs = blackSqs
	b.emptySqs = emptySqs
	b.whiteKingSq = NoSquare
	b.blackKingSq = NoSquare

	for sq := 0; sq < numOfSquaresInBoard; sq++ {
		sqr := Square(sq)
		if b.bbWhiteKing.Occupied(sqr) {
			b.whiteKingSq = sqr
		} else if b.bbBlackKing.Occupied(sqr) {
			b.blackKingSq = sqr
		}
	}
}

func (b *Board) copy() *Board {
	return &Board{
		whiteSqs:      b.whiteSqs,
		blackSqs:      b.blackSqs,
		emptySqs:      b.emptySqs,
		whiteKingSq:   b.whiteKingSq,
		blackKingSq:   b.blackKingSq,
		bbWhiteKing:   b.bbWhiteKing,
		bbWhiteQueen:  b.bbWhiteQueen,
		bbWhiteRook:   b.bbWhiteRook,
		bbWhiteBishop: b.bbWhiteBishop,
		bbWhiteKnight: b.bbWhiteKnight,
		bbWhitePawn:   b.bbWhitePawn,
		bbBlackKing:   b.bbBlackKing,
		bbBlackQueen:  b.bbBlackQueen,
		bbBlackRook:   b.bbBlackRook,
		bbBlackBishop: b.bbBlackBishop,
		bbBlackKnight: b.bbBlackKnight,
		bbBlackPawn:   b.bbBlackPawn,
	}
}

func (b *Board) isOccupied(sq Square) bool {
	return !b.emptySqs.Occupied(sq)
}

func (b *Board) piece(sq Square) Piece {
	for _, p := range allPieces {
		bb := b.bbForPiece(p)
		if bb.Occupied(sq) {
			return p
		}
	}
	return NoPiece
}

func (b *Board) hasSufficientMaterial() bool {
	// queen, rook, or pawn exist
	if (b.bbWhiteQueen | b.bbWhiteRook | b.bbWhitePawn |
		b.bbBlackQueen | b.bbBlackRook | b.bbBlackPawn) > 0 {
		return true
	}
	// if king is missing then it is a test
	if b.bbWhiteKing == 0 || b.bbBlackKing == 0 {
		return true
	}
	count := map[PieceType]int{}
	pieceMap := b.SquareMap()
	for _, p := range pieceMap {
		count[p.Type()]++
	}
	// 	king versus king
	if count[Bishop] == 0 && count[Knight] == 0 {
		return false
	}
	// king and bishop versus king
	if count[Bishop] == 1 && count[Knight] == 0 {
		return false
	}
	// king and knight versus king
	if count[Bishop] == 0 && count[Knight] == 1 {
		return false
	}
	// king and bishop(s) versus king and bishop(s) with the bishops on the same colour.
	if count[Knight] == 0 {
		whiteCount := 0
		blackCount := 0
		for sq, p := range pieceMap {
			if p.Type() == Bishop {
				switch sq.color() {
				case White:
					whiteCount++
				case Black:
					blackCount++
				}
			}
		}
		if whiteCount == 0 || blackCount == 0 {
			return false
		}
	}
	return true
}

func (b *Board) bbForPiece(p Piece) bitboard {
	switch p {
	case WhiteKing:
		return b.bbWhiteKing
	case WhiteQueen:
		return b.bbWhiteQueen
	case WhiteRook:
		return b.bbWhiteRook
	case WhiteBishop:
		return b.bbWhiteBishop
	case WhiteKnight:
		return b.bbWhiteKnight
	case WhitePawn:
		return b.bbWhitePawn
	case BlackKing:
		return b.bbBlackKing
	case BlackQueen:
		return b.bbBlackQueen
	case BlackRook:
		return b.bbBlackRook
	case BlackBishop:
		return b.bbBlackBishop
	case BlackKnight:
		return b.bbBlackKnight
	case BlackPawn:
		return b.bbBlackPawn
	}
	return bitboard(0)
}

func (b *Board) setBBForPiece(p Piece, bb bitboard) {
	switch p {
	case WhiteKing:
		b.bbWhiteKing = bb
	case WhiteQueen:
		b.bbWhiteQueen = bb
	case WhiteRook:
		b.bbWhiteRook = bb
	case WhiteBishop:
		b.bbWhiteBishop = bb
	case WhiteKnight:
		b.bbWhiteKnight = bb
	case WhitePawn:
		b.bbWhitePawn = bb
	case BlackKing:
		b.bbBlackKing = bb
	case BlackQueen:
		b.bbBlackQueen = bb
	case BlackRook:
		b.bbBlackRook = bb
	case BlackBishop:
		b.bbBlackBishop = bb
	case BlackKnight:
		b.bbBlackKnight = bb
	case BlackPawn:
		b.bbBlackPawn = bb
	default:
		panic("HERE")
	}
}
<commit_msg>optimized board update<commit_after>package chess

import (
	"strconv"
	"strings"
)

// A Board represents a chess board and its relationship between squares and pieces.
type Board struct {
	bbWhiteKing   bitboard
	bbWhiteQueen  bitboard
	bbWhiteRook   bitboard
	bbWhiteBishop bitboard
	bbWhiteKnight bitboard
	bbWhitePawn   bitboard
	bbBlackKing   bitboard
	bbBlackQueen  bitboard
	bbBlackRook   bitboard
	bbBlackBishop bitboard
	bbBlackKnight bitboard
	bbBlackPawn   bitboard
	whiteSqs      bitboard
	blackSqs      bitboard
	emptySqs      bitboard
	whiteKingSq   Square
	blackKingSq   Square
}

// SquareMap returns a mapping of squares to pieces.  A square is only added to the map if it is occupied.
func (b *Board) SquareMap() map[Square]Piece {
	m := map[Square]Piece{}
	for sq := 0; sq < numOfSquaresInBoard; sq++ {
		p := b.piece(Square(sq))
		if p != NoPiece {
			m[Square(sq)] = p
		}
	}
	return m
}

// Draw returns visual representation of the board useful for debugging.
func (b *Board) Draw() string {
	s := "\n A B C D E F G H\n"
	for r := 7; r >= 0; r-- {
		s += Rank(r).String()
		for f := 0; f < numOfSquaresInRow; f++ {
			p := b.piece(getSquare(File(f), Rank(r)))
			if p == NoPiece {
				s += "-"
			} else {
				s += p.String()
			}
			s += " "
		}
		s += "\n"
	}
	return s
}

// String implements the fmt.Stringer interface and returns
// a string in the FEN board format: rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR
func (b *Board) String() string {
	fen := ""
	for r := 7; r >= 0; r-- {
		for f := 0; f < numOfSquaresInRow; f++ {
			sq := getSquare(File(f), Rank(r))
			p := b.piece(sq)
			if p != NoPiece {
				fen += p.getFENChar()
			} else {
				fen += "1"
			}
		}
		if r != 0 {
			fen += "/"
		}
	}
	for i := 8; i > 1; i-- {
		repeatStr := strings.Repeat("1", i)
		countStr := strconv.Itoa(i)
		fen = strings.Replace(fen, repeatStr, countStr, -1)
	}
	return fen
}

func newBoard(m map[Square]Piece) *Board {
	b := &Board{}
	for _, p1 := range allPieces {
		bm := map[Square]bool{}
		for sq, p2 := range m {
			if p1 == p2 {
				bm[sq] = true
			}
		}
		bb := newBitboard(bm)
		b.setBBForPiece(p1, bb)
	}
	b.calcConvienceBBs(nil)
	return b
}

func (b *Board) update(m *Move) {
	p1 := b.piece(m.s1)
	s1BB := bbForSquare(m.s1)
	s2BB := bbForSquare(m.s2)

	// move s1 piece to s2
	for _, p := range allPieces {
		bb := b.bbForPiece(p)
		// remove what was at s2
		b.setBBForPiece(p, bb & ^s2BB)
		// move what was at s1 to s2
		if bb.Occupied(m.s1) {
			bb = b.bbForPiece(p)
			b.setBBForPiece(p, (bb & ^s1BB)|s2BB)
		}
	}
	// check promotion
	if m.promo != NoPieceType {
		newPiece := getPiece(m.promo, p1.Color())
		// remove pawn
		bbPawn := b.bbForPiece(p1)
		b.setBBForPiece(p1, bbPawn & ^s2BB)
		// add promo piece
		bbPromo := b.bbForPiece(newPiece)
		b.setBBForPiece(newPiece, bbPromo|s2BB)
	}
	// remove captured en passant piece
	if m.HasTag(EnPassant) {
		if p1.Color() == White {
			b.bbBlackPawn = ^(bbForSquare(m.s2) << 8) & b.bbBlackPawn
		} else {
			b.bbWhitePawn = ^(bbForSquare(m.s2) >> 8) & b.bbWhitePawn
		}
	}
	// move rook for castle
	if p1.Color() == White && m.HasTag(KingSideCastle) {
		b.bbWhiteRook = (b.bbWhiteRook & ^bbForSquare(H1) | bbForSquare(F1))
	} else if p1.Color() == White && m.HasTag(QueenSideCastle) {
		b.bbWhiteRook = (b.bbWhiteRook & ^bbForSquare(A1)) | bbForSquare(D1)
	} else if p1.Color() == Black && m.HasTag(KingSideCastle) {
		b.bbBlackRook = (b.bbBlackRook & ^bbForSquare(H8) | bbForSquare(F8))
	} else if p1.Color() == Black && m.HasTag(QueenSideCastle) {
		b.bbBlackRook = (b.bbBlackRook & ^bbForSquare(A8)) | bbForSquare(D8)
	}
	b.calcConvienceBBs(m)
}

func (b *Board) calcConvienceBBs(m *Move) {
	whiteSqs := b.bbWhiteKing | b.bbWhiteQueen | b.bbWhiteRook | b.bbWhiteBishop | b.bbWhiteKnight | b.bbWhitePawn
	blackSqs := b.bbBlackKing | b.bbBlackQueen | b.bbBlackRook | b.bbBlackBishop | b.bbBlackKnight | b.bbBlackPawn
	emptySqs := ^(whiteSqs | blackSqs)
	b.whiteSqs = whiteSqs
	b.blackSqs = blackSqs
	b.emptySqs = emptySqs
	if m == nil {
		b.whiteKingSq = NoSquare
		b.blackKingSq = NoSquare

		for sq := 0; sq < numOfSquaresInBoard; sq++ {
			sqr := Square(sq)
			if b.bbWhiteKing.Occupied(sqr) {
				b.whiteKingSq = sqr
			} else if b.bbBlackKing.Occupied(sqr) {
				b.blackKingSq = sqr
			}
		}
	} else if m.s1 == b.whiteKingSq {
		b.whiteKingSq = m.s2
	} else if m.s1 == b.blackKingSq {
		b.blackKingSq = m.s2
	}
}

func (b *Board) copy() *Board {
	return &Board{
		whiteSqs:      b.whiteSqs,
		blackSqs:      b.blackSqs,
		emptySqs:      b.emptySqs,
		whiteKingSq:   b.whiteKingSq,
		blackKingSq:   b.blackKingSq,
		bbWhiteKing:   b.bbWhiteKing,
		bbWhiteQueen:  b.bbWhiteQueen,
		bbWhiteRook:   b.bbWhiteRook,
		bbWhiteBishop: b.bbWhiteBishop,
		bbWhiteKnight: b.bbWhiteKnight,
		bbWhitePawn:   b.bbWhitePawn,
		bbBlackKing:   b.bbBlackKing,
		bbBlackQueen:  b.bbBlackQueen,
		bbBlackRook:   b.bbBlackRook,
		bbBlackBishop: b.bbBlackBishop,
		bbBlackKnight: b.bbBlackKnight,
		bbBlackPawn:   b.bbBlackPawn,
	}
}

func (b *Board) isOccupied(sq Square) bool {
	return !b.emptySqs.Occupied(sq)
}

func (b *Board) piece(sq Square) Piece {
	for _, p := range allPieces {
		bb := b.bbForPiece(p)
		if bb.Occupied(sq) {
			return p
		}
	}
	return NoPiece
}

func (b *Board) hasSufficientMaterial() bool {
	// queen, rook, or pawn exist
	if (b.bbWhiteQueen | b.bbWhiteRook | b.bbWhitePawn |
		b.bbBlackQueen | b.bbBlackRook | b.bbBlackPawn) > 0 {
		return true
	}
	// if king is missing then it is a test
	if b.bbWhiteKing == 0 || b.bbBlackKing == 0 {
		return true
	}
	count := map[PieceType]int{}
	pieceMap := b.SquareMap()
	for _, p := range pieceMap {
		count[p.Type()]++
	}
	// 	king versus king
	if count[Bishop] == 0 && count[Knight] == 0 {
		return false
	}
	// king and bishop versus king
	if count[Bishop] == 1 && count[Knight] == 0 {
		return false
	}
	// king and knight versus king
	if count[Bishop] == 0 && count[Knight] == 1 {
		return false
	}
	// king and bishop(s) versus king and bishop(s) with the bishops on the same colour.
	if count[Knight] == 0 {
		whiteCount := 0
		blackCount := 0
		for sq, p := range pieceMap {
			if p.Type() == Bishop {
				switch sq.color() {
				case White:
					whiteCount++
				case Black:
					blackCount++
				}
			}
		}
		if whiteCount == 0 || blackCount == 0 {
			return false
		}
	}
	return true
}

func (b *Board) bbForPiece(p Piece) bitboard {
	switch p {
	case WhiteKing:
		return b.bbWhiteKing
	case WhiteQueen:
		return b.bbWhiteQueen
	case WhiteRook:
		return b.bbWhiteRook
	case WhiteBishop:
		return b.bbWhiteBishop
	case WhiteKnight:
		return b.bbWhiteKnight
	case WhitePawn:
		return b.bbWhitePawn
	case BlackKing:
		return b.bbBlackKing
	case BlackQueen:
		return b.bbBlackQueen
	case BlackRook:
		return b.bbBlackRook
	case BlackBishop:
		return b.bbBlackBishop
	case BlackKnight:
		return b.bbBlackKnight
	case BlackPawn:
		return b.bbBlackPawn
	}
	return bitboard(0)
}

func (b *Board) setBBForPiece(p Piece, bb bitboard) {
	switch p {
	case WhiteKing:
		b.bbWhiteKing = bb
	case WhiteQueen:
		b.bbWhiteQueen = bb
	case WhiteRook:
		b.bbWhiteRook = bb
	case WhiteBishop:
		b.bbWhiteBishop = bb
	case WhiteKnight:
		b.bbWhiteKnight = bb
	case WhitePawn:
		b.bbWhitePawn = bb
	case BlackKing:
		b.bbBlackKing = bb
	case BlackQueen:
		b.bbBlackQueen = bb
	case BlackRook:
		b.bbBlackRook = bb
	case BlackBishop:
		b.bbBlackBishop = bb
	case BlackKnight:
		b.bbBlackKnight = bb
	case BlackPawn:
		b.bbBlackPawn = bb
	default:
		panic("HERE")
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2017 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 model

import (
	"fmt"
	"time"
)

const (
	// OOMBumpUpRatio specifies how much memory will be added after observing OOM.
	OOMBumpUpRatio float64 = 1.2
	// OOMMinBumpUp specifies minimal increase of memeory after observing OOM.
	OOMMinBumpUp float64 = 100 * 1024 * 1024 // 100MB
)

// ContainerUsageSample is a measure of resource usage of a container over some
// interval.
type ContainerUsageSample struct {
	// Start of the measurement interval.
	MeasureStart time.Time
	// Average CPU usage in cores or memory usage in bytes.
	Usage ResourceAmount
	// CPU or memory request at the time of measurment.
	Request ResourceAmount
	// Which resource is this sample for.
	Resource ResourceName
}

// ContainerState stores information about a single container instance.
// Each ContainerState has a pointer to the aggregation that is used for
// aggregating its usage samples.
// It holds the recent history of CPU and memory utilization.
//   Note: samples are added to intervals based on their start timestamps.
type ContainerState struct {
	// Current request.
	Request Resources
	// Start of the latest CPU usage sample that was aggregated.
	LastCPUSampleStart time.Time
	// Max memory usage observed in the current aggregation interval.
	MemoryPeak ResourceAmount
	// End time of the current memory aggregation interval (not inclusive).
	WindowEnd time.Time
	// Start of the latest memory usage sample that was aggregated.
	lastMemorySampleStart time.Time
	// Aggregation to add usage samples to.
	aggregator ContainerStateAggregator
}

// NewContainerState returns a new ContainerState.
func NewContainerState(request Resources, aggregator ContainerStateAggregator) *ContainerState {
	return &ContainerState{
		Request:               request,
		LastCPUSampleStart:    time.Time{},
		WindowEnd:             time.Time{},
		lastMemorySampleStart: time.Time{},
		aggregator:            aggregator,
	}
}

func (sample *ContainerUsageSample) isValid(expectedResource ResourceName) bool {
	return sample.Usage >= 0 && sample.Resource == expectedResource
}

func (container *ContainerState) addCPUSample(sample *ContainerUsageSample) bool {
	// Order should not matter for the histogram, other than deduplication.
	if !sample.isValid(ResourceCPU) || !sample.MeasureStart.After(container.LastCPUSampleStart) {
		return false // Discard invalid, duplicate or out-of-order samples.
	}
	container.aggregator.AddSample(sample)
	container.LastCPUSampleStart = sample.MeasureStart
	return true
}

func (container *ContainerState) addMemorySample(sample *ContainerUsageSample) bool {
	ts := sample.MeasureStart
	if !sample.isValid(ResourceMemory) || ts.Before(container.lastMemorySampleStart) {
		return false // Discard invalid or outdated samples.
	}
	container.lastMemorySampleStart = ts
	if container.WindowEnd.IsZero() { // This is the first sample.
		container.WindowEnd = ts
	}

	// Each container aggregates one peak per aggregation interval. If the timestamp of the
	// current sample is earlier than the end of the current interval (WindowEnd) and is larger
	// than the current peak, the peak is updated in the aggregation by subtracting the old value
	// and adding the new value.
	addNewPeak := false
	if ts.Before(container.WindowEnd) {
		if container.MemoryPeak != 0 && sample.Usage > container.MemoryPeak {
			// Remove the old peak.
			oldPeak := ContainerUsageSample{
				MeasureStart: container.WindowEnd,
				Usage:        container.MemoryPeak,
				Request:      sample.Request,
				Resource:     ResourceMemory,
			}
			container.aggregator.SubtractSample(&oldPeak)
			addNewPeak = true
		}
	} else {
		// Shift the memory aggregation window to the next interval.
		shift := truncate(ts.Sub(container.WindowEnd), MemoryAggregationInterval) + MemoryAggregationInterval
		container.WindowEnd = container.WindowEnd.Add(shift)
		addNewPeak = true
	}
	if addNewPeak {
		newPeak := ContainerUsageSample{
			MeasureStart: container.WindowEnd,
			Usage:        sample.Usage,
			Request:      sample.Request,
			Resource:     ResourceMemory,
		}
		container.aggregator.AddSample(&newPeak)
		container.MemoryPeak = sample.Usage
	}
	return true
}

// RecordOOM adds info regarding OOM event in the model as an artificial memory sample.
func (container *ContainerState) RecordOOM(timestamp time.Time, requestedMemory ResourceAmount) error {
	// Discard old OOM
	if timestamp.Before(container.WindowEnd.Add(-1 * MemoryAggregationInterval)) {
		return fmt.Errorf("OOM event will be discarded - it is too old (%v)", timestamp)
	}
	// Get max of the request and the recent memory peak.
	memoryUsed := ResourceAmountMax(requestedMemory, container.MemoryPeak)
	memoryNeeded := ResourceAmountMax(memoryUsed+MemoryAmountFromBytes(OOMMinBumpUp),
		ScaleResource(memoryUsed, OOMBumpUpRatio))

	oomMemorySample := ContainerUsageSample{
		MeasureStart: timestamp,
		Usage:        memoryNeeded,
		Resource:     ResourceMemory,
	}
	if !container.addMemorySample(&oomMemorySample) {
		return fmt.Errorf("Adding OOM sample failed")
	}
	return nil
}

// AddSample adds a usage sample to the given ContainerState. Requires samples
// for a single resource to be passed in chronological order (i.e. in order of
// growing MeasureStart). Invalid samples (out of order or measure out of legal
// range) are discarded. Returns true if the sample was aggregated, false if it
// was discarded.
// Note: usage samples don't hold their end timestamp / duration. They are
// implicitly assumed to be disjoint when aggregating.
func (container *ContainerState) AddSample(sample *ContainerUsageSample) bool {
	switch sample.Resource {
	case ResourceCPU:
		return container.addCPUSample(sample)
	case ResourceMemory:
		return container.addMemorySample(sample)
	default:
		return false
	}
}

// Truncate returns the result of rounding d toward zero to a multiple of m.
// If m <= 0, Truncate returns d unchanged.
// This helper function is introduced to support older implementations of the
// time package that don't provide Duration.Truncate function.
func truncate(d, m time.Duration) time.Duration {
	if m <= 0 {
		return d
	}
	return d - d%m
}
<commit_msg>Update container.go<commit_after>/*
Copyright 2017 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 model

import (
	"fmt"
	"time"
)

const (
	// OOMBumpUpRatio specifies how much memory will be added after observing OOM.
	OOMBumpUpRatio float64 = 1.2
	// OOMMinBumpUp specifies minimal increase of memory after observing OOM.
	OOMMinBumpUp float64 = 100 * 1024 * 1024 // 100MB
)

// ContainerUsageSample is a measure of resource usage of a container over some
// interval.
type ContainerUsageSample struct {
	// Start of the measurement interval.
	MeasureStart time.Time
	// Average CPU usage in cores or memory usage in bytes.
	Usage ResourceAmount
	// CPU or memory request at the time of measurment.
	Request ResourceAmount
	// Which resource is this sample for.
	Resource ResourceName
}

// ContainerState stores information about a single container instance.
// Each ContainerState has a pointer to the aggregation that is used for
// aggregating its usage samples.
// It holds the recent history of CPU and memory utilization.
//   Note: samples are added to intervals based on their start timestamps.
type ContainerState struct {
	// Current request.
	Request Resources
	// Start of the latest CPU usage sample that was aggregated.
	LastCPUSampleStart time.Time
	// Max memory usage observed in the current aggregation interval.
	MemoryPeak ResourceAmount
	// End time of the current memory aggregation interval (not inclusive).
	WindowEnd time.Time
	// Start of the latest memory usage sample that was aggregated.
	lastMemorySampleStart time.Time
	// Aggregation to add usage samples to.
	aggregator ContainerStateAggregator
}

// NewContainerState returns a new ContainerState.
func NewContainerState(request Resources, aggregator ContainerStateAggregator) *ContainerState {
	return &ContainerState{
		Request:               request,
		LastCPUSampleStart:    time.Time{},
		WindowEnd:             time.Time{},
		lastMemorySampleStart: time.Time{},
		aggregator:            aggregator,
	}
}

func (sample *ContainerUsageSample) isValid(expectedResource ResourceName) bool {
	return sample.Usage >= 0 && sample.Resource == expectedResource
}

func (container *ContainerState) addCPUSample(sample *ContainerUsageSample) bool {
	// Order should not matter for the histogram, other than deduplication.
	if !sample.isValid(ResourceCPU) || !sample.MeasureStart.After(container.LastCPUSampleStart) {
		return false // Discard invalid, duplicate or out-of-order samples.
	}
	container.aggregator.AddSample(sample)
	container.LastCPUSampleStart = sample.MeasureStart
	return true
}

func (container *ContainerState) addMemorySample(sample *ContainerUsageSample) bool {
	ts := sample.MeasureStart
	if !sample.isValid(ResourceMemory) || ts.Before(container.lastMemorySampleStart) {
		return false // Discard invalid or outdated samples.
	}
	container.lastMemorySampleStart = ts
	if container.WindowEnd.IsZero() { // This is the first sample.
		container.WindowEnd = ts
	}

	// Each container aggregates one peak per aggregation interval. If the timestamp of the
	// current sample is earlier than the end of the current interval (WindowEnd) and is larger
	// than the current peak, the peak is updated in the aggregation by subtracting the old value
	// and adding the new value.
	addNewPeak := false
	if ts.Before(container.WindowEnd) {
		if container.MemoryPeak != 0 && sample.Usage > container.MemoryPeak {
			// Remove the old peak.
			oldPeak := ContainerUsageSample{
				MeasureStart: container.WindowEnd,
				Usage:        container.MemoryPeak,
				Request:      sample.Request,
				Resource:     ResourceMemory,
			}
			container.aggregator.SubtractSample(&oldPeak)
			addNewPeak = true
		}
	} else {
		// Shift the memory aggregation window to the next interval.
		shift := truncate(ts.Sub(container.WindowEnd), MemoryAggregationInterval) + MemoryAggregationInterval
		container.WindowEnd = container.WindowEnd.Add(shift)
		addNewPeak = true
	}
	if addNewPeak {
		newPeak := ContainerUsageSample{
			MeasureStart: container.WindowEnd,
			Usage:        sample.Usage,
			Request:      sample.Request,
			Resource:     ResourceMemory,
		}
		container.aggregator.AddSample(&newPeak)
		container.MemoryPeak = sample.Usage
	}
	return true
}

// RecordOOM adds info regarding OOM event in the model as an artificial memory sample.
func (container *ContainerState) RecordOOM(timestamp time.Time, requestedMemory ResourceAmount) error {
	// Discard old OOM
	if timestamp.Before(container.WindowEnd.Add(-1 * MemoryAggregationInterval)) {
		return fmt.Errorf("OOM event will be discarded - it is too old (%v)", timestamp)
	}
	// Get max of the request and the recent memory peak.
	memoryUsed := ResourceAmountMax(requestedMemory, container.MemoryPeak)
	memoryNeeded := ResourceAmountMax(memoryUsed+MemoryAmountFromBytes(OOMMinBumpUp),
		ScaleResource(memoryUsed, OOMBumpUpRatio))

	oomMemorySample := ContainerUsageSample{
		MeasureStart: timestamp,
		Usage:        memoryNeeded,
		Resource:     ResourceMemory,
	}
	if !container.addMemorySample(&oomMemorySample) {
		return fmt.Errorf("Adding OOM sample failed")
	}
	return nil
}

// AddSample adds a usage sample to the given ContainerState. Requires samples
// for a single resource to be passed in chronological order (i.e. in order of
// growing MeasureStart). Invalid samples (out of order or measure out of legal
// range) are discarded. Returns true if the sample was aggregated, false if it
// was discarded.
// Note: usage samples don't hold their end timestamp / duration. They are
// implicitly assumed to be disjoint when aggregating.
func (container *ContainerState) AddSample(sample *ContainerUsageSample) bool {
	switch sample.Resource {
	case ResourceCPU:
		return container.addCPUSample(sample)
	case ResourceMemory:
		return container.addMemorySample(sample)
	default:
		return false
	}
}

// Truncate returns the result of rounding d toward zero to a multiple of m.
// If m <= 0, Truncate returns d unchanged.
// This helper function is introduced to support older implementations of the
// time package that don't provide Duration.Truncate function.
func truncate(d, m time.Duration) time.Duration {
	if m <= 0 {
		return d
	}
	return d - d%m
}
<|endoftext|>


<commit_before>// +build !lambdabinary

package sparta

import (
	"fmt"
	"log"
	"os"
	"path/filepath"
	"time"

	"github.com/pkg/errors"
	"github.com/rs/zerolog"
	"github.com/spf13/cobra"
	validator "gopkg.in/go-playground/validator.v9"
)

func platformLogSysInfo(lambdaFunc string, logger *zerolog.Logger) {
	// NOP
}

// RegisterCodePipelineEnvironment is part of a CodePipeline deployment
// and defines the environments available for deployment. Environments
// are defined the `environmentName`. The values defined in the
// environmentVariables are made available to each service as
// environment variables. The environment key will be transformed into
// a configuration file for a CodePipeline CloudFormation action:
// TemplateConfiguration: !Sub "TemplateSource::${environmentName}".
func RegisterCodePipelineEnvironment(environmentName string,
	environmentVariables map[string]string) error {
	if _, exists := codePipelineEnvironments[environmentName]; exists {
		return errors.Errorf("Environment (%s) has already been defined", environmentName)
	}
	codePipelineEnvironments[environmentName] = environmentVariables
	return nil
}

// Main defines the primary handler for transforming an application into a Sparta package.  The
// serviceName is used to uniquely identify your service within a region and will
// be used for subsequent updates.  For provisioning, ensure that you've
// properly configured AWS credentials for the golang SDK.
// See http://docs.aws.amazon.com/sdk-for-go/api/aws/defaults.html#DefaultChainCredentials-constant
// for more information.
func Main(serviceName string, serviceDescription string, lambdaAWSInfos []*LambdaAWSInfo, api APIGateway, site *S3Site) error {
	return MainEx(serviceName,
		serviceDescription,
		lambdaAWSInfos,
		api,
		site,
		nil,
		false)
}

// MainEx provides an "extended" Main that supports customizing the standard Sparta
// workflow via the `workflowHooks` parameter.
func MainEx(serviceName string,
	serviceDescription string,
	lambdaAWSInfos []*LambdaAWSInfo,
	api APIGateway,
	site *S3Site,
	workflowHooks *WorkflowHooks,
	useCGO bool) error {

	//////////////////////////////////////////////////////////////////////////////
	// cmdRoot defines the root, non-executable command
	CommandLineOptions.Root.Short = fmt.Sprintf("%s - Sparta v.%s powered AWS Lambda Microservice",
		serviceName,
		SpartaVersion)
	CommandLineOptions.Root.Long = serviceDescription
	CommandLineOptions.Root.PersistentPreRunE = func(cmd *cobra.Command, args []string) error {

		// Save the ServiceName in case a custom command wants it
		OptionsGlobal.ServiceName = serviceName
		OptionsGlobal.ServiceDescription = serviceDescription

		validateErr := validate.Struct(OptionsGlobal)
		if nil != validateErr {
			return validateErr
		}

		// Format?
		// Running in AWS?
		disableColors := OptionsGlobal.DisableColors || isRunningInAWS()
		logger, loggerErr := NewLoggerForOutput(OptionsGlobal.LogLevel,
			OptionsGlobal.LogFormat,
			disableColors)
		if nil != loggerErr {
			return loggerErr
		}

		// This is a NOP, but makes megacheck happy b/c it doesn't know about
		// build flags
		platformLogSysInfo("", logger)
		OptionsGlobal.Logger = logger
		welcomeMessage := fmt.Sprintf("Service: %s", serviceName)

		// Header information...
		displayPrettyHeader(headerDivider, disableColors, logger)

		// Metadata about the build...
		logger.Info().
			Str("Option", cmd.Name()).
			Str("LinkFlags", OptionsGlobal.LinkerFlags).
			Str("UTC", time.Now().UTC().Format(time.RFC3339)).
			Msg(welcomeMessage)
		logger.Info().Msg(headerDivider)
		return nil
	}

	//////////////////////////////////////////////////////////////////////////////
	// Version
	CommandLineOptions.Root.AddCommand(CommandLineOptions.Version)

	//////////////////////////////////////////////////////////////////////////////
	// Build
	CommandLineOptions.Build.PreRunE = func(cmd *cobra.Command, args []string) error {
		validateErr := validate.Struct(optionsBuild)

		OptionsGlobal.Logger.Debug().
			Interface("ValidateErr", validateErr).
			Interface("OptionsProvision", optionsProvision).
			Msg("Build validation results")
		return validateErr
	}

	if nil == CommandLineOptions.Build.RunE {
		CommandLineOptions.Build.RunE = func(cmd *cobra.Command, args []string) error {
			buildID, buildIDErr := computeBuildID(optionsProvision.BuildID, OptionsGlobal.Logger)
			if nil != buildIDErr {
				return buildIDErr
			}

			// Save the BuildID
			StampedBuildID = buildID

			// Ok, for this we're going some way to tell the Build Command
			// where to write the output...I suppose we could just use a TeeWriter...
			sanitizedServiceName := sanitizedName(serviceName)
			templateName := fmt.Sprintf("%s-cftemplate.json", sanitizedServiceName)

			templateFilePath := filepath.Join(optionsProvision.OutputDir, templateName)
			templateFile, templateFileErr := os.Create(templateFilePath)
			if templateFileErr != nil {
				return templateFileErr
			}
			defer templateFile.Close()
			return Build(OptionsGlobal.Noop,
				serviceName,
				serviceDescription,
				lambdaAWSInfos,
				api,
				site,
				useCGO,
				buildID,
				optionsBuild.OutputDir,
				OptionsGlobal.BuildTags,
				OptionsGlobal.LinkerFlags,
				templateFile,
				workflowHooks,
				OptionsGlobal.Logger)
		}
	}
	CommandLineOptions.Root.AddCommand(CommandLineOptions.Build)

	//////////////////////////////////////////////////////////////////////////////
	// Provision
	CommandLineOptions.Provision.PreRunE = func(cmd *cobra.Command, args []string) error {
		validateErr := validate.Struct(optionsProvision)

		OptionsGlobal.Logger.Debug().
			Interface("validateErr", validateErr).
			Interface("optionsProvision", optionsProvision).
			Msg("Provision validation results")
		return validateErr
	}

	if nil == CommandLineOptions.Provision.RunE {
		CommandLineOptions.Provision.RunE = func(cmd *cobra.Command, args []string) error {
			buildID, buildIDErr := computeBuildID(optionsProvision.BuildID, OptionsGlobal.Logger)
			if nil != buildIDErr {
				return buildIDErr
			}
			StampedBuildID = buildID

			templateFile, templateFileErr := templateOutputFile(optionsProvision.OutputDir,
				serviceName)
			if templateFileErr != nil {
				return templateFileErr
			}

			// TODO: Build, then Provision
			buildErr := Build(OptionsGlobal.Noop,
				serviceName,
				serviceDescription,
				lambdaAWSInfos,
				api,
				site,
				useCGO,
				buildID,
				optionsProvision.OutputDir,
				OptionsGlobal.BuildTags,
				OptionsGlobal.LinkerFlags,
				templateFile,
				workflowHooks,
				OptionsGlobal.Logger)
			templateFile.Close()

			if buildErr != nil {
				return buildErr
			}
			// So for this, we need to take command
			// line params and turn them into a map...
			parseErr := optionsProvision.parseParams()
			if parseErr != nil {
				return parseErr
			}
			OptionsGlobal.Logger.Debug().
				Interface("params", optionsProvision.stackParams).
				Msg("ParseParams")

			// We don't need to walk the params because we
			// put values in the Metadata block for them all...
			// Save the BuildID

			return Provision(true || OptionsGlobal.Noop,
				templateFile.Name(),
				optionsProvision.stackParams,
				optionsProvision.stackTags,
				optionsProvision.InPlace,
				optionsProvision.PipelineTrigger,
				OptionsGlobal.Logger)
		}
	}
	CommandLineOptions.Root.AddCommand(CommandLineOptions.Provision)

	//////////////////////////////////////////////////////////////////////////////
	// Delete
	CommandLineOptions.Delete.RunE = func(cmd *cobra.Command, args []string) error {
		return Delete(serviceName, OptionsGlobal.Logger)
	}

	CommandLineOptions.Root.AddCommand(CommandLineOptions.Delete)

	//////////////////////////////////////////////////////////////////////////////
	// Execute
	if nil == CommandLineOptions.Execute.RunE {
		CommandLineOptions.Execute.RunE = func(cmd *cobra.Command, args []string) error {
			// Ensure the discovery service is initialized
			initializeDiscovery(OptionsGlobal.Logger)

			return Execute(serviceName,
				lambdaAWSInfos,
				OptionsGlobal.Logger)
		}
	}
	CommandLineOptions.Root.AddCommand(CommandLineOptions.Execute)

	//////////////////////////////////////////////////////////////////////////////
	// Describe
	if nil == CommandLineOptions.Describe.RunE {
		CommandLineOptions.Describe.RunE = func(cmd *cobra.Command, args []string) error {
			validateErr := validate.Struct(optionsDescribe)
			if nil != validateErr {
				return errors.Wrapf(validateErr, "Failed to validate `describe` options")
			}
			fileWriter, fileWriterErr := os.Create(optionsDescribe.OutputFile)
			if fileWriterErr != nil {
				return fileWriterErr
			}
			defer func() {
				closeErr := fileWriter.Close()
				if closeErr != nil {
					OptionsGlobal.Logger.Warn().
						Err(closeErr).
						Msg("Failed to close describe output writer")
				}
			}()

			describeErr := Describe(serviceName,
				serviceDescription,
				lambdaAWSInfos,
				api,
				site,
				optionsDescribe.S3Bucket,
				OptionsGlobal.BuildTags,
				OptionsGlobal.LinkerFlags,
				fileWriter,
				workflowHooks,
				OptionsGlobal.Logger)

			if describeErr == nil {
				describeErr = fileWriter.Sync()
			}
			return describeErr
		}
	}
	CommandLineOptions.Root.AddCommand(CommandLineOptions.Describe)

	//////////////////////////////////////////////////////////////////////////////
	// Explore
	if nil == CommandLineOptions.Explore.RunE {
		CommandLineOptions.Explore.RunE = func(cmd *cobra.Command, args []string) error {
			validateErr := validate.Struct(optionsExplore)
			if nil != validateErr {
				return validateErr
			}

			return ExploreWithInputFilter(serviceName,
				serviceDescription,
				lambdaAWSInfos,
				api,
				site,
				optionsExplore.InputExtensions,
				optionsDescribe.S3Bucket,
				OptionsGlobal.BuildTags,
				OptionsGlobal.LinkerFlags,
				OptionsGlobal.Logger)
		}
	}
	CommandLineOptions.Root.AddCommand(CommandLineOptions.Explore)

	//////////////////////////////////////////////////////////////////////////////
	// Profile
	if nil == CommandLineOptions.Profile.RunE {
		CommandLineOptions.Profile.RunE = func(cmd *cobra.Command, args []string) error {
			validateErr := validate.Struct(optionsProfile)
			if nil != validateErr {
				return validateErr
			}
			return Profile(serviceName,
				serviceDescription,
				optionsProfile.S3Bucket,
				optionsProfile.Port,
				OptionsGlobal.Logger)
		}
	}
	CommandLineOptions.Root.AddCommand(CommandLineOptions.Profile)

	//////////////////////////////////////////////////////////////////////////////
	// Status
	if nil == CommandLineOptions.Status.RunE {
		CommandLineOptions.Status.RunE = func(cmd *cobra.Command, args []string) error {
			validateErr := validate.Struct(optionsStatus)
			if nil != validateErr {
				return validateErr
			}
			return Status(serviceName,
				serviceDescription,
				optionsStatus.Redact,
				OptionsGlobal.Logger)
		}
	}
	CommandLineOptions.Root.AddCommand(CommandLineOptions.Status)

	// Run it!
	executedCmd, executeErr := CommandLineOptions.Root.ExecuteC()
	if executeErr != nil {
		if OptionsGlobal.Logger == nil {
			newLogger, newLoggerErr := NewLogger(zerolog.InfoLevel.String())
			if newLoggerErr != nil {
				fmt.Printf("Failed to create new logger: %v", newLoggerErr)
				zLogger := zerolog.New(os.Stderr).With().Timestamp().Logger()
				newLogger = &zLogger
			}
			OptionsGlobal.Logger = newLogger
		}
		if OptionsGlobal.Logger != nil {
			validationErr, validationErrOk := executeErr.(validator.ValidationErrors)
			if validationErrOk {
				for _, eachError := range validationErr {
					OptionsGlobal.Logger.Error().
						Interface("Error", eachError).
						Msg("Validation error")
				}
				// Only show the usage if there were input validation errors
				if executedCmd != nil {
					usageErr := executedCmd.Usage()
					if usageErr != nil {
						OptionsGlobal.Logger.Error().Err(usageErr).Msg("Usage error")
					}
				}
			} else {
				OptionsGlobal.Logger.Error().Err(executeErr).Msg("Failed to execute command")
			}
		} else {
			log.Printf("ERROR: %s", executeErr)
		}
	}

	// Cleanup, if for some reason the caller wants to re-execute later...
	CommandLineOptions.Root.PersistentPreRunE = nil
	return executeErr
}
<commit_msg>Include failstop header usage<commit_after>// +build !lambdabinary

package sparta

import (
	"fmt"
	"log"
	"os"
	"path/filepath"
	"time"

	"github.com/pkg/errors"
	"github.com/rs/zerolog"
	"github.com/spf13/cobra"
	validator "gopkg.in/go-playground/validator.v9"
)

func platformLogSysInfo(lambdaFunc string, logger *zerolog.Logger) {
	// NOP
}

// RegisterCodePipelineEnvironment is part of a CodePipeline deployment
// and defines the environments available for deployment. Environments
// are defined the `environmentName`. The values defined in the
// environmentVariables are made available to each service as
// environment variables. The environment key will be transformed into
// a configuration file for a CodePipeline CloudFormation action:
// TemplateConfiguration: !Sub "TemplateSource::${environmentName}".
func RegisterCodePipelineEnvironment(environmentName string,
	environmentVariables map[string]string) error {
	if _, exists := codePipelineEnvironments[environmentName]; exists {
		return errors.Errorf("Environment (%s) has already been defined", environmentName)
	}
	codePipelineEnvironments[environmentName] = environmentVariables
	return nil
}

// Main defines the primary handler for transforming an application into a Sparta package.  The
// serviceName is used to uniquely identify your service within a region and will
// be used for subsequent updates.  For provisioning, ensure that you've
// properly configured AWS credentials for the golang SDK.
// See http://docs.aws.amazon.com/sdk-for-go/api/aws/defaults.html#DefaultChainCredentials-constant
// for more information.
func Main(serviceName string, serviceDescription string, lambdaAWSInfos []*LambdaAWSInfo, api APIGateway, site *S3Site) error {
	return MainEx(serviceName,
		serviceDescription,
		lambdaAWSInfos,
		api,
		site,
		nil,
		false)
}

// MainEx provides an "extended" Main that supports customizing the standard Sparta
// workflow via the `workflowHooks` parameter.
func MainEx(serviceName string,
	serviceDescription string,
	lambdaAWSInfos []*LambdaAWSInfo,
	api APIGateway,
	site *S3Site,
	workflowHooks *WorkflowHooks,
	useCGO bool) error {

	//////////////////////////////////////////////////////////////////////////////
	// cmdRoot defines the root, non-executable command
	CommandLineOptions.Root.Short = fmt.Sprintf("%s - Sparta v.%s powered AWS Lambda Microservice",
		serviceName,
		SpartaVersion)
	CommandLineOptions.Root.Long = serviceDescription
	CommandLineOptions.Root.PersistentPreRunE = func(cmd *cobra.Command, args []string) error {

		// Save the ServiceName in case a custom command wants it
		OptionsGlobal.ServiceName = serviceName
		OptionsGlobal.ServiceDescription = serviceDescription

		validateErr := validate.Struct(OptionsGlobal)
		if nil != validateErr {
			return validateErr
		}

		// Format?
		// Running in AWS?
		disableColors := OptionsGlobal.DisableColors || isRunningInAWS()
		logger, loggerErr := NewLoggerForOutput(OptionsGlobal.LogLevel,
			OptionsGlobal.LogFormat,
			disableColors)
		if nil != loggerErr {
			return loggerErr
		}

		// This is a NOP, but makes megacheck happy b/c it doesn't know about
		// build flags
		platformLogSysInfo("", logger)
		OptionsGlobal.Logger = logger
		welcomeMessage := fmt.Sprintf("Service: %s", serviceName)

		// Header information...
		displayPrettyHeader(headerDivider, disableColors, logger)

		// Metadata about the build...
		logger.Info().
			Str("Option", cmd.Name()).
			Str("LinkFlags", OptionsGlobal.LinkerFlags).
			Str("UTC", time.Now().UTC().Format(time.RFC3339)).
			Msg(welcomeMessage)
		logger.Info().Msg(headerDivider)
		return nil
	}

	//////////////////////////////////////////////////////////////////////////////
	// Version
	CommandLineOptions.Root.AddCommand(CommandLineOptions.Version)

	//////////////////////////////////////////////////////////////////////////////
	// Build
	CommandLineOptions.Build.PreRunE = func(cmd *cobra.Command, args []string) error {
		validateErr := validate.Struct(optionsBuild)

		OptionsGlobal.Logger.Debug().
			Interface("ValidateErr", validateErr).
			Interface("OptionsProvision", optionsProvision).
			Msg("Build validation results")
		return validateErr
	}

	if nil == CommandLineOptions.Build.RunE {
		CommandLineOptions.Build.RunE = func(cmd *cobra.Command, args []string) error {
			buildID, buildIDErr := computeBuildID(optionsProvision.BuildID, OptionsGlobal.Logger)
			if nil != buildIDErr {
				return buildIDErr
			}

			// Save the BuildID
			StampedBuildID = buildID

			// Ok, for this we're going some way to tell the Build Command
			// where to write the output...I suppose we could just use a TeeWriter...
			sanitizedServiceName := sanitizedName(serviceName)
			templateName := fmt.Sprintf("%s-cftemplate.json", sanitizedServiceName)

			templateFilePath := filepath.Join(optionsProvision.OutputDir, templateName)
			templateFile, templateFileErr := os.Create(templateFilePath)
			if templateFileErr != nil {
				return templateFileErr
			}
			defer templateFile.Close()
			return Build(OptionsGlobal.Noop,
				serviceName,
				serviceDescription,
				lambdaAWSInfos,
				api,
				site,
				useCGO,
				buildID,
				optionsBuild.OutputDir,
				OptionsGlobal.BuildTags,
				OptionsGlobal.LinkerFlags,
				templateFile,
				workflowHooks,
				OptionsGlobal.Logger)
		}
	}
	CommandLineOptions.Root.AddCommand(CommandLineOptions.Build)

	//////////////////////////////////////////////////////////////////////////////
	// Provision
	CommandLineOptions.Provision.PreRunE = func(cmd *cobra.Command, args []string) error {
		validateErr := validate.Struct(optionsProvision)

		OptionsGlobal.Logger.Debug().
			Interface("validateErr", validateErr).
			Interface("optionsProvision", optionsProvision).
			Msg("Provision validation results")
		return validateErr
	}

	if nil == CommandLineOptions.Provision.RunE {
		CommandLineOptions.Provision.RunE = func(cmd *cobra.Command, args []string) error {
			buildID, buildIDErr := computeBuildID(optionsProvision.BuildID, OptionsGlobal.Logger)
			if nil != buildIDErr {
				return buildIDErr
			}
			StampedBuildID = buildID

			templateFile, templateFileErr := templateOutputFile(optionsProvision.OutputDir,
				serviceName)
			if templateFileErr != nil {
				return templateFileErr
			}

			// TODO: Build, then Provision
			buildErr := Build(OptionsGlobal.Noop,
				serviceName,
				serviceDescription,
				lambdaAWSInfos,
				api,
				site,
				useCGO,
				buildID,
				optionsProvision.OutputDir,
				OptionsGlobal.BuildTags,
				OptionsGlobal.LinkerFlags,
				templateFile,
				workflowHooks,
				OptionsGlobal.Logger)
			templateFile.Close()

			if buildErr != nil {
				return buildErr
			}
			// So for this, we need to take command
			// line params and turn them into a map...
			parseErr := optionsProvision.parseParams()
			if parseErr != nil {
				return parseErr
			}
			OptionsGlobal.Logger.Debug().
				Interface("params", optionsProvision.stackParams).
				Msg("ParseParams")

			// We don't need to walk the params because we
			// put values in the Metadata block for them all...
			// Save the BuildID

			return Provision(true || OptionsGlobal.Noop,
				templateFile.Name(),
				optionsProvision.stackParams,
				optionsProvision.stackTags,
				optionsProvision.InPlace,
				optionsProvision.PipelineTrigger,
				OptionsGlobal.Logger)
		}
	}
	CommandLineOptions.Root.AddCommand(CommandLineOptions.Provision)

	//////////////////////////////////////////////////////////////////////////////
	// Delete
	CommandLineOptions.Delete.RunE = func(cmd *cobra.Command, args []string) error {
		return Delete(serviceName, OptionsGlobal.Logger)
	}

	CommandLineOptions.Root.AddCommand(CommandLineOptions.Delete)

	//////////////////////////////////////////////////////////////////////////////
	// Execute
	if nil == CommandLineOptions.Execute.RunE {
		CommandLineOptions.Execute.RunE = func(cmd *cobra.Command, args []string) error {
			// Ensure the discovery service is initialized
			initializeDiscovery(OptionsGlobal.Logger)

			return Execute(serviceName,
				lambdaAWSInfos,
				OptionsGlobal.Logger)
		}
	}
	CommandLineOptions.Root.AddCommand(CommandLineOptions.Execute)

	//////////////////////////////////////////////////////////////////////////////
	// Describe
	if nil == CommandLineOptions.Describe.RunE {
		CommandLineOptions.Describe.RunE = func(cmd *cobra.Command, args []string) error {
			validateErr := validate.Struct(optionsDescribe)
			if nil != validateErr {
				return errors.Wrapf(validateErr, "Failed to validate `describe` options")
			}
			fileWriter, fileWriterErr := os.Create(optionsDescribe.OutputFile)
			if fileWriterErr != nil {
				return fileWriterErr
			}
			defer func() {
				closeErr := fileWriter.Close()
				if closeErr != nil {
					OptionsGlobal.Logger.Warn().
						Err(closeErr).
						Msg("Failed to close describe output writer")
				}
			}()

			describeErr := Describe(serviceName,
				serviceDescription,
				lambdaAWSInfos,
				api,
				site,
				optionsDescribe.S3Bucket,
				OptionsGlobal.BuildTags,
				OptionsGlobal.LinkerFlags,
				fileWriter,
				workflowHooks,
				OptionsGlobal.Logger)

			if describeErr == nil {
				describeErr = fileWriter.Sync()
			}
			return describeErr
		}
	}
	CommandLineOptions.Root.AddCommand(CommandLineOptions.Describe)

	//////////////////////////////////////////////////////////////////////////////
	// Explore
	if nil == CommandLineOptions.Explore.RunE {
		CommandLineOptions.Explore.RunE = func(cmd *cobra.Command, args []string) error {
			validateErr := validate.Struct(optionsExplore)
			if nil != validateErr {
				return validateErr
			}

			return ExploreWithInputFilter(serviceName,
				serviceDescription,
				lambdaAWSInfos,
				api,
				site,
				optionsExplore.InputExtensions,
				optionsDescribe.S3Bucket,
				OptionsGlobal.BuildTags,
				OptionsGlobal.LinkerFlags,
				OptionsGlobal.Logger)
		}
	}
	CommandLineOptions.Root.AddCommand(CommandLineOptions.Explore)

	//////////////////////////////////////////////////////////////////////////////
	// Profile
	if nil == CommandLineOptions.Profile.RunE {
		CommandLineOptions.Profile.RunE = func(cmd *cobra.Command, args []string) error {
			validateErr := validate.Struct(optionsProfile)
			if nil != validateErr {
				return validateErr
			}
			return Profile(serviceName,
				serviceDescription,
				optionsProfile.S3Bucket,
				optionsProfile.Port,
				OptionsGlobal.Logger)
		}
	}
	CommandLineOptions.Root.AddCommand(CommandLineOptions.Profile)

	//////////////////////////////////////////////////////////////////////////////
	// Status
	if nil == CommandLineOptions.Status.RunE {
		CommandLineOptions.Status.RunE = func(cmd *cobra.Command, args []string) error {
			validateErr := validate.Struct(optionsStatus)
			if nil != validateErr {
				return validateErr
			}
			return Status(serviceName,
				serviceDescription,
				optionsStatus.Redact,
				OptionsGlobal.Logger)
		}
	}
	CommandLineOptions.Root.AddCommand(CommandLineOptions.Status)

	// Run it!
	executedCmd, executeErr := CommandLineOptions.Root.ExecuteC()
	if executeErr != nil {
		if OptionsGlobal.Logger == nil {
			// Use a default console logger
			newLogger, newLoggerErr := NewLoggerForOutput(zerolog.InfoLevel.String(),
				"text",
				isRunningInAWS())
			if newLoggerErr != nil {
				fmt.Printf("Failed to create new logger: %v", newLoggerErr)
				zLogger := zerolog.New(os.Stderr).With().Timestamp().Logger()
				newLogger = &zLogger
			}
			OptionsGlobal.Logger = newLogger
		}
		if OptionsGlobal.Logger != nil {
			validationErr, validationErrOk := executeErr.(validator.ValidationErrors)
			if validationErrOk {
				for _, eachError := range validationErr {
					OptionsGlobal.Logger.Error().
						Interface("Error", eachError).
						Msg("Validation error")
				}
				// Only show the usage if there were input validation errors
				if executedCmd != nil {
					usageErr := executedCmd.Usage()
					if usageErr != nil {
						OptionsGlobal.Logger.Error().Err(usageErr).Msg("Usage error")
					}
				}
			} else {
				displayPrettyHeader(headerDivider, isRunningInAWS(), OptionsGlobal.Logger)
				OptionsGlobal.Logger.Error().Err(executeErr).Msg("Failed to execute command")
			}
		} else {
			log.Printf("ERROR: %s", executeErr)
		}
	}

	// Cleanup, if for some reason the caller wants to re-execute later...
	CommandLineOptions.Root.PersistentPreRunE = nil
	return executeErr
}
<|endoftext|>


<commit_before>package httputilmore

import (
	"bytes"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"strconv"
	"strings"

	"github.com/grokify/gotilla/encoding/jsonutil"
	"github.com/pkg/errors"
)

// GetWriteFile gets the conents of a URL and stores the body in
// the desired filename location.
func GetWriteFile(client *http.Client, url, filename string) error {
	resp, err := client.Get(url)
	if err != nil {
		return errors.Wrap(err, "httputilmore.GetStoreURL.client.Get()")
	}
	defer resp.Body.Close()
	dir, file := filepath.Split(filename)
	if len(strings.TrimSpace(dir)) > 0 {
		os.Chdir(dir)
	}
	f, err := os.Create(file)
	if err != nil {
		return errors.Wrap(err, "httputilmore.GetStoreURL.os.Create()")
	}
	defer f.Close()
	_, err = io.Copy(f, resp.Body)
	if err != nil {
		err = errors.Wrap(err, "httputilmore.GetStoreURL.io.Copy()")
	}
	return err
}

// GetWriteFile performs a HTTP GET request and saves the response body
// to the file path specified. It reeads the entire file into memory
// which is not ideal for large files.
func GetWriteFileSimple(url string, filename string, perm os.FileMode) ([]byte, error) {
	_, bytes, err := GetResponseAndBytes(url)
	if err != nil {
		return bytes, err
	}
	err = ioutil.WriteFile(filename, bytes, perm)
	return bytes, err
}

func GetJsonSimple(requrl string, header http.Header, data interface{}) (*http.Response, error) {
	req, err := http.NewRequest(http.MethodGet, requrl, nil)
	if err != nil {
		return nil, err
	}
	req.Header = header
	client := &http.Client{}
	resp, err := client.Do(req)
	if err != nil {
		return nil, err
	}
	_, err = jsonutil.UnmarshalIoReader(resp.Body, data)
	return resp, err
}

func PostJsonBytes(client *http.Client, requrl string, headers map[string]string, bodyBytes []byte) (*http.Response, error) {
	req, err := http.NewRequest(http.MethodPost, requrl, bytes.NewBuffer(bodyBytes))
	if err != nil {
		return nil, err
	}
	for k, v := range headers {
		k = strings.TrimSpace(k)
		kMatch := strings.ToLower(k)
		if kMatch == strings.ToLower(HeaderContentType) {
			continue
		}
		req.Header.Set(k, v)
	}
	req.Header.Set(HeaderContentType, ContentTypeAppJsonUtf8)
	if client == nil {
		client = &http.Client{}
	}
	return client.Do(req)
}

func PostJsonMarshal(client *http.Client, requrl string, headers map[string]string, body interface{}) (*http.Response, error) {
	bodyBytes, err := json.Marshal(body)
	if err != nil {
		return nil, err
	}
	return PostJsonMarshal(client, requrl, headers, bodyBytes)
}

// PostJsonSimple performs a HTTP POST request converting a body interface{} to
// JSON and adding the appropriate JSON Content-Type header.
func PostJsonSimple(requrl string, body interface{}) (*http.Response, error) {
	return PostJsonMarshal(nil, requrl, map[string]string{}, body)
}

// GetResponseAndBytes retreives a URL and returns the response body
// as a byte array in addition to the *http.Response.
func GetResponseAndBytes(url string) (*http.Response, []byte, error) {
	resp, err := http.Get(url)
	if err != nil {
		return resp, []byte{}, err
	}
	bytes, err := ioutil.ReadAll(resp.Body)
	return resp, bytes, err
}

func SendWwwFormUrlEncodedSimple(method, urlStr string, data url.Values) (*http.Response, error) {
	req, err := http.NewRequest(
		method,
		urlStr,
		strings.NewReader(data.Encode())) // URL-encoded payload
	if err != nil {
		return nil, err
	}
	req.Header.Add(HeaderContentType, ContentTypeAppFormUrlEncoded)
	req.Header.Add(HeaderContentLength, strconv.Itoa(len(data.Encode())))
	client := &http.Client{}
	return client.Do(req)
}
<commit_msg>loop bugfix<commit_after>package httputilmore

import (
	"bytes"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"strconv"
	"strings"

	"github.com/grokify/gotilla/encoding/jsonutil"
	"github.com/pkg/errors"
)

// GetWriteFile gets the conents of a URL and stores the body in
// the desired filename location.
func GetWriteFile(client *http.Client, url, filename string) error {
	resp, err := client.Get(url)
	if err != nil {
		return errors.Wrap(err, "httputilmore.GetStoreURL.client.Get()")
	}
	defer resp.Body.Close()
	dir, file := filepath.Split(filename)
	if len(strings.TrimSpace(dir)) > 0 {
		os.Chdir(dir)
	}
	f, err := os.Create(file)
	if err != nil {
		return errors.Wrap(err, "httputilmore.GetStoreURL.os.Create()")
	}
	defer f.Close()
	_, err = io.Copy(f, resp.Body)
	if err != nil {
		err = errors.Wrap(err, "httputilmore.GetStoreURL.io.Copy()")
	}
	return err
}

// GetWriteFile performs a HTTP GET request and saves the response body
// to the file path specified. It reeads the entire file into memory
// which is not ideal for large files.
func GetWriteFileSimple(url string, filename string, perm os.FileMode) ([]byte, error) {
	_, bytes, err := GetResponseAndBytes(url)
	if err != nil {
		return bytes, err
	}
	err = ioutil.WriteFile(filename, bytes, perm)
	return bytes, err
}

func GetJsonSimple(requrl string, header http.Header, data interface{}) (*http.Response, error) {
	req, err := http.NewRequest(http.MethodGet, requrl, nil)
	if err != nil {
		return nil, err
	}
	req.Header = header
	client := &http.Client{}
	resp, err := client.Do(req)
	if err != nil {
		return nil, err
	}
	_, err = jsonutil.UnmarshalIoReader(resp.Body, data)
	return resp, err
}

func PostJsonBytes(client *http.Client, requrl string, headers map[string]string, bodyBytes []byte) (*http.Response, error) {
	req, err := http.NewRequest(http.MethodPost, requrl, bytes.NewBuffer(bodyBytes))
	if err != nil {
		return nil, err
	}
	for k, v := range headers {
		k = strings.TrimSpace(k)
		kMatch := strings.ToLower(k)
		if kMatch == strings.ToLower(HeaderContentType) {
			continue
		}
		req.Header.Set(k, v)
	}
	req.Header.Set(HeaderContentType, ContentTypeAppJsonUtf8)
	if client == nil {
		client = &http.Client{}
	}
	return client.Do(req)
}

func PostJsonMarshal(client *http.Client, requrl string, headers map[string]string, body interface{}) (*http.Response, error) {
	bodyBytes, err := json.Marshal(body)
	if err != nil {
		return nil, err
	}
	return PostJsonBytes(client, requrl, headers, bodyBytes)
}

// PostJsonSimple performs a HTTP POST request converting a body interface{} to
// JSON and adding the appropriate JSON Content-Type header.
func PostJsonSimple(requrl string, body interface{}) (*http.Response, error) {
	return PostJsonMarshal(nil, requrl, map[string]string{}, body)
}

// GetResponseAndBytes retreives a URL and returns the response body
// as a byte array in addition to the *http.Response.
func GetResponseAndBytes(url string) (*http.Response, []byte, error) {
	resp, err := http.Get(url)
	if err != nil {
		return resp, []byte{}, err
	}
	bytes, err := ioutil.ReadAll(resp.Body)
	return resp, bytes, err
}

func SendWwwFormUrlEncodedSimple(method, urlStr string, data url.Values) (*http.Response, error) {
	req, err := http.NewRequest(
		method,
		urlStr,
		strings.NewReader(data.Encode())) // URL-encoded payload
	if err != nil {
		return nil, err
	}
	req.Header.Add(HeaderContentType, ContentTypeAppFormUrlEncoded)
	req.Header.Add(HeaderContentLength, strconv.Itoa(len(data.Encode())))
	client := &http.Client{}
	return client.Do(req)
}
<|endoftext|>


<commit_before>// runoutput

// Copyright 2011 The Go Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Generate test of channel operations and simple selects.
// The output of this program is compiled and run to do the
// actual test.

// Each test does only one real send or receive at a time, but phrased
// in various ways that the compiler may or may not rewrite
// into simpler expressions.

package main

import (
	"bufio"
	"fmt"
	"io"
	"os"
	"text/template"
)

func main() {
	out := bufio.NewWriter(os.Stdout)
	fmt.Fprintln(out, header)
	a := new(arg)

	// Generate each kind of test as a separate function to avoid
	// hitting the 6g optimizer with one enormous function.
	// If we name all the functions init we don't have to
	// maintain a list of which ones to run.
	do := func(t *template.Template) {
		fmt.Fprintln(out, `func init() {`)
		for ; next(); a.reset() {
			run(t, a, out)
		}
		fmt.Fprintln(out, `}`)
	}

	do(recv)
	do(send)
	do(recvOrder)
	do(sendOrder)
	do(nonblock)

	fmt.Fprintln(out, "//", a.nreset, "cases")
	out.Flush()
}

func run(t *template.Template, a interface{}, out io.Writer) {
	if err := t.Execute(out, a); err != nil {
		panic(err)
	}
}

type arg struct {
	def    bool
	nreset int
}

func (a *arg) Maybe() bool {
	return maybe()
}

func (a *arg) MaybeDefault() bool {
	if a.def {
		return false
	}
	a.def = maybe()
	return a.def
}

func (a *arg) MustDefault() bool {
	return !a.def
}

func (a *arg) reset() {
	a.def = false
	a.nreset++
}

const header = `// GENERATED BY select5.go; DO NOT EDIT

package main

// channel is buffered so test is single-goroutine.
// we are not interested in the concurrency aspects
// of select, just testing that the right calls happen.
var c = make(chan int, 1)
var nilch chan int
var n = 1
var x int
var i interface{}
var dummy = make(chan int)
var m = make(map[int]int)
var order = 0

func f(p *int) *int {
	return p
}

// check order of operations by ensuring that
// successive calls to checkorder have increasing o values.
func checkorder(o int) {
	if o <= order {
		println("invalid order", o, "after", order)
		panic("order")
	}
	order = o
}

func fc(c chan int, o int) chan int {
	checkorder(o)
	return c
}

func fp(p *int, o int) *int {
	checkorder(o)
	return p
}

func fn(n, o int) int {
	checkorder(o)
	return n
}

func die(x int) {
	println("have", x, "want", n)
	panic("chan")
}

func main() {
	// everything happens in init funcs
}
`

func parse(name, s string) *template.Template {
	t, err := template.New(name).Parse(s)
	if err != nil {
		panic(fmt.Sprintf("%q: %s", name, err))
	}
	return t
}

var recv = parse("recv", `
	{{/*  Send n, receive it one way or another into x, check that they match. */}}
	c <- n
	{{if .Maybe}}
	x = <-c
	{{else}}
	select {
	{{/*  Blocking or non-blocking, before the receive. */}}
	{{/*  The compiler implements two-case select where one is default with custom code, */}}
	{{/*  so test the default branch both before and after the send. */}}
	{{if .MaybeDefault}}
	default:
		panic("nonblock")
	{{end}}
	{{/*  Receive from c.  Different cases are direct, indirect, :=, interface, and map assignment. */}}
	{{if .Maybe}}
	case x = <-c:
	{{else}}{{if .Maybe}}
	case *f(&x) = <-c:
	{{else}}{{if .Maybe}}
	case y := <-c:
		x = y
	{{else}}{{if .Maybe}}
	case i = <-c:
		x = i.(int)
	{{else}}
	case m[13] = <-c:
		x = m[13]
	{{end}}{{end}}{{end}}{{end}}
	{{/*  Blocking or non-blocking again, after the receive. */}}
	{{if .MaybeDefault}}
	default:
		panic("nonblock")
	{{end}}
	{{/*  Dummy send, receive to keep compiler from optimizing select. */}}
	{{if .Maybe}}
	case dummy <- 1:
		panic("dummy send")
	{{end}}
	{{if .Maybe}}
	case <-dummy:
		panic("dummy receive")
	{{end}}
	{{/*  Nil channel send, receive to keep compiler from optimizing select. */}}
	{{if .Maybe}}
	case nilch <- 1:
		panic("nilch send")
	{{end}}
	{{if .Maybe}}
	case <-nilch:
		panic("nilch recv")
	{{end}}
	}
	{{end}}
	if x != n {
		die(x)
	}
	n++
`)

var recvOrder = parse("recvOrder", `
	{{/*  Send n, receive it one way or another into x, check that they match. */}}
	{{/*  Check order of operations along the way by calling functions that check */}}
	{{/*  that the argument sequence is strictly increasing. */}}
	order = 0
	c <- n
	{{if .Maybe}}
	{{/*  Outside of select, left-to-right rule applies. */}}
	{{/*  (Inside select, assignment waits until case is chosen, */}}
	{{/*  so right hand side happens before anything on left hand side. */}}
	*fp(&x, 1) = <-fc(c, 2)
	{{else}}{{if .Maybe}}
	m[fn(13, 1)] = <-fc(c, 2)
	x = m[13]
	{{else}}
	select {
	{{/*  Blocking or non-blocking, before the receive. */}}
	{{/*  The compiler implements two-case select where one is default with custom code, */}}
	{{/*  so test the default branch both before and after the send. */}}
	{{if .MaybeDefault}}
	default:
		panic("nonblock")
	{{end}}
	{{/*  Receive from c.  Different cases are direct, indirect, :=, interface, and map assignment. */}}
	{{if .Maybe}}
	case *fp(&x, 100) = <-fc(c, 1):
	{{else}}{{if .Maybe}}
	case y := <-fc(c, 1):
		x = y
	{{else}}{{if .Maybe}}
	case i = <-fc(c, 1):
		x = i.(int)
	{{else}}
	case m[fn(13, 100)] = <-fc(c, 1):
		x = m[13]
	{{end}}{{end}}{{end}}
	{{/*  Blocking or non-blocking again, after the receive. */}}
	{{if .MaybeDefault}}
	default:
		panic("nonblock")
	{{end}}
	{{/*  Dummy send, receive to keep compiler from optimizing select. */}}
	{{if .Maybe}}
	case fc(dummy, 2) <- fn(1, 3):
		panic("dummy send")
	{{end}}
	{{if .Maybe}}
	case <-fc(dummy, 4):
		panic("dummy receive")
	{{end}}
	{{/*  Nil channel send, receive to keep compiler from optimizing select. */}}
	{{if .Maybe}}
	case fc(nilch, 5) <- fn(1, 6):
		panic("nilch send")
	{{end}}
	{{if .Maybe}}
	case <-fc(nilch, 7):
		panic("nilch recv")
	{{end}}
	}
	{{end}}{{end}}
	if x != n {
		die(x)
	}
	n++
`)

var send = parse("send", `
	{{/*  Send n one way or another, receive it into x, check that they match. */}}
	{{if .Maybe}}
	c <- n
	{{else}}
	select {
	{{/*  Blocking or non-blocking, before the receive (same reason as in recv). */}}
	{{if .MaybeDefault}}
	default:
		panic("nonblock")
	{{end}}
	{{/*  Send c <- n.  No real special cases here, because no values come back */}}
	{{/*  from the send operation. */}}
	case c <- n:
	{{/*  Blocking or non-blocking. */}}
	{{if .MaybeDefault}}
	default:
		panic("nonblock")
	{{end}}
	{{/*  Dummy send, receive to keep compiler from optimizing select. */}}
	{{if .Maybe}}
	case dummy <- 1:
		panic("dummy send")
	{{end}}
	{{if .Maybe}}
	case <-dummy:
		panic("dummy receive")
	{{end}}
	{{/*  Nil channel send, receive to keep compiler from optimizing select. */}}
	{{if .Maybe}}
	case nilch <- 1:
		panic("nilch send")
	{{end}}
	{{if .Maybe}}
	case <-nilch:
		panic("nilch recv")
	{{end}}
	}
	{{end}}
	x = <-c
	if x != n {
		die(x)
	}
	n++
`)

var sendOrder = parse("sendOrder", `
	{{/*  Send n one way or another, receive it into x, check that they match. */}}
	{{/*  Check order of operations along the way by calling functions that check */}}
	{{/*  that the argument sequence is strictly increasing. */}}
	order = 0
	{{if .Maybe}}
	fc(c, 1) <- fn(n, 2)
	{{else}}
	select {
	{{/*  Blocking or non-blocking, before the receive (same reason as in recv). */}}
	{{if .MaybeDefault}}
	default:
		panic("nonblock")
	{{end}}
	{{/*  Send c <- n.  No real special cases here, because no values come back */}}
	{{/*  from the send operation. */}}
	case fc(c, 1) <- fn(n, 2):
	{{/*  Blocking or non-blocking. */}}
	{{if .MaybeDefault}}
	default:
		panic("nonblock")
	{{end}}
	{{/*  Dummy send, receive to keep compiler from optimizing select. */}}
	{{if .Maybe}}
	case fc(dummy, 3) <- fn(1, 4):
		panic("dummy send")
	{{end}}
	{{if .Maybe}}
	case <-fc(dummy, 5):
		panic("dummy receive")
	{{end}}
	{{/*  Nil channel send, receive to keep compiler from optimizing select. */}}
	{{if .Maybe}}
	case fc(nilch, 6) <- fn(1, 7):
		panic("nilch send")
	{{end}}
	{{if .Maybe}}
	case <-fc(nilch, 8):
		panic("nilch recv")
	{{end}}
	}
	{{end}}
	x = <-c
	if x != n {
		die(x)
	}
	n++
`)

var nonblock = parse("nonblock", `
	x = n
	{{/*  Test various combinations of non-blocking operations. */}}
	{{/*  Receive assignments must not edit or even attempt to compute the address of the lhs. */}}
	select {
	{{if .MaybeDefault}}
	default:
	{{end}}
	{{if .Maybe}}
	case dummy <- 1:
		panic("dummy <- 1")
	{{end}}
	{{if .Maybe}}
	case nilch <- 1:
		panic("nilch <- 1")
	{{end}}
	{{if .Maybe}}
	case <-dummy:
		panic("<-dummy")
	{{end}}
	{{if .Maybe}}
	case x = <-dummy:
		panic("<-dummy x")
	{{end}}
	{{if .Maybe}}
	case **(**int)(nil) = <-dummy:
		panic("<-dummy (and didn't crash saving result!)")
	{{end}}
	{{if .Maybe}}
	case <-nilch:
		panic("<-nilch")
	{{end}}
	{{if .Maybe}}
	case x = <-nilch:
		panic("<-nilch x")
	{{end}}
	{{if .Maybe}}
	case **(**int)(nil) = <-nilch:
		panic("<-nilch (and didn't crash saving result!)")
	{{end}}
	{{if .MustDefault}}
	default:
	{{end}}
	}
	if x != n {
		die(x)
	}
	n++
`)

// Code for enumerating all possible paths through
// some logic.  The logic should call choose(n) when
// it wants to choose between n possibilities.
// On successive runs through the logic, choose(n)
// will return 0, 1, ..., n-1.  The helper maybe() is
// similar but returns true and then false.
//
// Given a function gen that generates an output
// using choose and maybe, code can generate all
// possible outputs using
//
//	for next() {
//		gen()
//	}

type choice struct {
	i, n int
}

var choices []choice
var cp int = -1

func maybe() bool {
	return choose(2) == 0
}

func choose(n int) int {
	if cp >= len(choices) {
		// never asked this before: start with 0.
		choices = append(choices, choice{0, n})
		cp = len(choices)
		return 0
	}
	// otherwise give recorded answer
	if n != choices[cp].n {
		panic("inconsistent choices")
	}
	i := choices[cp].i
	cp++
	return i
}

func next() bool {
	if cp < 0 {
		// start a new round
		cp = 0
		return true
	}

	// increment last choice sequence
	cp = len(choices) - 1
	for cp >= 0 && choices[cp].i == choices[cp].n-1 {
		cp--
	}
	if cp < 0 {
		choices = choices[:0]
		return false
	}
	choices[cp].i++
	choices = choices[:cp+1]
	cp = 0
	return true
}
<commit_msg>test: speed up chan/select5<commit_after>// runoutput

// Copyright 2011 The Go Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Generate test of channel operations and simple selects.
// The output of this program is compiled and run to do the
// actual test.

// Each test does only one real send or receive at a time, but phrased
// in various ways that the compiler may or may not rewrite
// into simpler expressions.

package main

import (
	"bufio"
	"fmt"
	"io"
	"os"
	"text/template"
)

func main() {
	out := bufio.NewWriter(os.Stdout)
	fmt.Fprintln(out, header)
	a := new(arg)

	// Generate each test as a separate function to avoid
	// hitting the 6g optimizer with one enormous function.
	// If we name all the functions init we don't have to
	// maintain a list of which ones to run.
	do := func(t *template.Template) {
		for ; next(); a.reset() {
			fmt.Fprintln(out, `func init() {`)
			run(t, a, out)
			fmt.Fprintln(out, `}`)
		}
	}

	do(recv)
	do(send)
	do(recvOrder)
	do(sendOrder)
	do(nonblock)

	fmt.Fprintln(out, "//", a.nreset, "cases")
	out.Flush()
}

func run(t *template.Template, a interface{}, out io.Writer) {
	if err := t.Execute(out, a); err != nil {
		panic(err)
	}
}

type arg struct {
	def    bool
	nreset int
}

func (a *arg) Maybe() bool {
	return maybe()
}

func (a *arg) MaybeDefault() bool {
	if a.def {
		return false
	}
	a.def = maybe()
	return a.def
}

func (a *arg) MustDefault() bool {
	return !a.def
}

func (a *arg) reset() {
	a.def = false
	a.nreset++
}

const header = `// GENERATED BY select5.go; DO NOT EDIT

package main

// channel is buffered so test is single-goroutine.
// we are not interested in the concurrency aspects
// of select, just testing that the right calls happen.
var c = make(chan int, 1)
var nilch chan int
var n = 1
var x int
var i interface{}
var dummy = make(chan int)
var m = make(map[int]int)
var order = 0

func f(p *int) *int {
	return p
}

// check order of operations by ensuring that
// successive calls to checkorder have increasing o values.
func checkorder(o int) {
	if o <= order {
		println("invalid order", o, "after", order)
		panic("order")
	}
	order = o
}

func fc(c chan int, o int) chan int {
	checkorder(o)
	return c
}

func fp(p *int, o int) *int {
	checkorder(o)
	return p
}

func fn(n, o int) int {
	checkorder(o)
	return n
}

func die(x int) {
	println("have", x, "want", n)
	panic("chan")
}

func main() {
	// everything happens in init funcs
}
`

func parse(name, s string) *template.Template {
	t, err := template.New(name).Parse(s)
	if err != nil {
		panic(fmt.Sprintf("%q: %s", name, err))
	}
	return t
}

var recv = parse("recv", `
	{{/*  Send n, receive it one way or another into x, check that they match. */}}
	c <- n
	{{if .Maybe}}
	x = <-c
	{{else}}
	select {
	{{/*  Blocking or non-blocking, before the receive. */}}
	{{/*  The compiler implements two-case select where one is default with custom code, */}}
	{{/*  so test the default branch both before and after the send. */}}
	{{if .MaybeDefault}}
	default:
		panic("nonblock")
	{{end}}
	{{/*  Receive from c.  Different cases are direct, indirect, :=, interface, and map assignment. */}}
	{{if .Maybe}}
	case x = <-c:
	{{else}}{{if .Maybe}}
	case *f(&x) = <-c:
	{{else}}{{if .Maybe}}
	case y := <-c:
		x = y
	{{else}}{{if .Maybe}}
	case i = <-c:
		x = i.(int)
	{{else}}
	case m[13] = <-c:
		x = m[13]
	{{end}}{{end}}{{end}}{{end}}
	{{/*  Blocking or non-blocking again, after the receive. */}}
	{{if .MaybeDefault}}
	default:
		panic("nonblock")
	{{end}}
	{{/*  Dummy send, receive to keep compiler from optimizing select. */}}
	{{if .Maybe}}
	case dummy <- 1:
		panic("dummy send")
	{{end}}
	{{if .Maybe}}
	case <-dummy:
		panic("dummy receive")
	{{end}}
	{{/*  Nil channel send, receive to keep compiler from optimizing select. */}}
	{{if .Maybe}}
	case nilch <- 1:
		panic("nilch send")
	{{end}}
	{{if .Maybe}}
	case <-nilch:
		panic("nilch recv")
	{{end}}
	}
	{{end}}
	if x != n {
		die(x)
	}
	n++
`)

var recvOrder = parse("recvOrder", `
	{{/*  Send n, receive it one way or another into x, check that they match. */}}
	{{/*  Check order of operations along the way by calling functions that check */}}
	{{/*  that the argument sequence is strictly increasing. */}}
	order = 0
	c <- n
	{{if .Maybe}}
	{{/*  Outside of select, left-to-right rule applies. */}}
	{{/*  (Inside select, assignment waits until case is chosen, */}}
	{{/*  so right hand side happens before anything on left hand side. */}}
	*fp(&x, 1) = <-fc(c, 2)
	{{else}}{{if .Maybe}}
	m[fn(13, 1)] = <-fc(c, 2)
	x = m[13]
	{{else}}
	select {
	{{/*  Blocking or non-blocking, before the receive. */}}
	{{/*  The compiler implements two-case select where one is default with custom code, */}}
	{{/*  so test the default branch both before and after the send. */}}
	{{if .MaybeDefault}}
	default:
		panic("nonblock")
	{{end}}
	{{/*  Receive from c.  Different cases are direct, indirect, :=, interface, and map assignment. */}}
	{{if .Maybe}}
	case *fp(&x, 100) = <-fc(c, 1):
	{{else}}{{if .Maybe}}
	case y := <-fc(c, 1):
		x = y
	{{else}}{{if .Maybe}}
	case i = <-fc(c, 1):
		x = i.(int)
	{{else}}
	case m[fn(13, 100)] = <-fc(c, 1):
		x = m[13]
	{{end}}{{end}}{{end}}
	{{/*  Blocking or non-blocking again, after the receive. */}}
	{{if .MaybeDefault}}
	default:
		panic("nonblock")
	{{end}}
	{{/*  Dummy send, receive to keep compiler from optimizing select. */}}
	{{if .Maybe}}
	case fc(dummy, 2) <- fn(1, 3):
		panic("dummy send")
	{{end}}
	{{if .Maybe}}
	case <-fc(dummy, 4):
		panic("dummy receive")
	{{end}}
	{{/*  Nil channel send, receive to keep compiler from optimizing select. */}}
	{{if .Maybe}}
	case fc(nilch, 5) <- fn(1, 6):
		panic("nilch send")
	{{end}}
	{{if .Maybe}}
	case <-fc(nilch, 7):
		panic("nilch recv")
	{{end}}
	}
	{{end}}{{end}}
	if x != n {
		die(x)
	}
	n++
`)

var send = parse("send", `
	{{/*  Send n one way or another, receive it into x, check that they match. */}}
	{{if .Maybe}}
	c <- n
	{{else}}
	select {
	{{/*  Blocking or non-blocking, before the receive (same reason as in recv). */}}
	{{if .MaybeDefault}}
	default:
		panic("nonblock")
	{{end}}
	{{/*  Send c <- n.  No real special cases here, because no values come back */}}
	{{/*  from the send operation. */}}
	case c <- n:
	{{/*  Blocking or non-blocking. */}}
	{{if .MaybeDefault}}
	default:
		panic("nonblock")
	{{end}}
	{{/*  Dummy send, receive to keep compiler from optimizing select. */}}
	{{if .Maybe}}
	case dummy <- 1:
		panic("dummy send")
	{{end}}
	{{if .Maybe}}
	case <-dummy:
		panic("dummy receive")
	{{end}}
	{{/*  Nil channel send, receive to keep compiler from optimizing select. */}}
	{{if .Maybe}}
	case nilch <- 1:
		panic("nilch send")
	{{end}}
	{{if .Maybe}}
	case <-nilch:
		panic("nilch recv")
	{{end}}
	}
	{{end}}
	x = <-c
	if x != n {
		die(x)
	}
	n++
`)

var sendOrder = parse("sendOrder", `
	{{/*  Send n one way or another, receive it into x, check that they match. */}}
	{{/*  Check order of operations along the way by calling functions that check */}}
	{{/*  that the argument sequence is strictly increasing. */}}
	order = 0
	{{if .Maybe}}
	fc(c, 1) <- fn(n, 2)
	{{else}}
	select {
	{{/*  Blocking or non-blocking, before the receive (same reason as in recv). */}}
	{{if .MaybeDefault}}
	default:
		panic("nonblock")
	{{end}}
	{{/*  Send c <- n.  No real special cases here, because no values come back */}}
	{{/*  from the send operation. */}}
	case fc(c, 1) <- fn(n, 2):
	{{/*  Blocking or non-blocking. */}}
	{{if .MaybeDefault}}
	default:
		panic("nonblock")
	{{end}}
	{{/*  Dummy send, receive to keep compiler from optimizing select. */}}
	{{if .Maybe}}
	case fc(dummy, 3) <- fn(1, 4):
		panic("dummy send")
	{{end}}
	{{if .Maybe}}
	case <-fc(dummy, 5):
		panic("dummy receive")
	{{end}}
	{{/*  Nil channel send, receive to keep compiler from optimizing select. */}}
	{{if .Maybe}}
	case fc(nilch, 6) <- fn(1, 7):
		panic("nilch send")
	{{end}}
	{{if .Maybe}}
	case <-fc(nilch, 8):
		panic("nilch recv")
	{{end}}
	}
	{{end}}
	x = <-c
	if x != n {
		die(x)
	}
	n++
`)

var nonblock = parse("nonblock", `
	x = n
	{{/*  Test various combinations of non-blocking operations. */}}
	{{/*  Receive assignments must not edit or even attempt to compute the address of the lhs. */}}
	select {
	{{if .MaybeDefault}}
	default:
	{{end}}
	{{if .Maybe}}
	case dummy <- 1:
		panic("dummy <- 1")
	{{end}}
	{{if .Maybe}}
	case nilch <- 1:
		panic("nilch <- 1")
	{{end}}
	{{if .Maybe}}
	case <-dummy:
		panic("<-dummy")
	{{end}}
	{{if .Maybe}}
	case x = <-dummy:
		panic("<-dummy x")
	{{end}}
	{{if .Maybe}}
	case **(**int)(nil) = <-dummy:
		panic("<-dummy (and didn't crash saving result!)")
	{{end}}
	{{if .Maybe}}
	case <-nilch:
		panic("<-nilch")
	{{end}}
	{{if .Maybe}}
	case x = <-nilch:
		panic("<-nilch x")
	{{end}}
	{{if .Maybe}}
	case **(**int)(nil) = <-nilch:
		panic("<-nilch (and didn't crash saving result!)")
	{{end}}
	{{if .MustDefault}}
	default:
	{{end}}
	}
	if x != n {
		die(x)
	}
	n++
`)

// Code for enumerating all possible paths through
// some logic.  The logic should call choose(n) when
// it wants to choose between n possibilities.
// On successive runs through the logic, choose(n)
// will return 0, 1, ..., n-1.  The helper maybe() is
// similar but returns true and then false.
//
// Given a function gen that generates an output
// using choose and maybe, code can generate all
// possible outputs using
//
//	for next() {
//		gen()
//	}

type choice struct {
	i, n int
}

var choices []choice
var cp int = -1

func maybe() bool {
	return choose(2) == 0
}

func choose(n int) int {
	if cp >= len(choices) {
		// never asked this before: start with 0.
		choices = append(choices, choice{0, n})
		cp = len(choices)
		return 0
	}
	// otherwise give recorded answer
	if n != choices[cp].n {
		panic("inconsistent choices")
	}
	i := choices[cp].i
	cp++
	return i
}

func next() bool {
	if cp < 0 {
		// start a new round
		cp = 0
		return true
	}

	// increment last choice sequence
	cp = len(choices) - 1
	for cp >= 0 && choices[cp].i == choices[cp].n-1 {
		cp--
	}
	if cp < 0 {
		choices = choices[:0]
		return false
	}
	choices[cp].i++
	choices = choices[:cp+1]
	cp = 0
	return true
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/davyxu/gosproto/meta"
)

const luaCodeTemplate = `-- Generated by github.com/davyxu/gosproto/sprotogen
-- DO NOT EDIT!
{{if .EnumValueGroup}}
ResultToString = function ( result )
	if result == 0 then
		return "OK"
	end

	local str = ResultByID[result]
	if str == nil then
		return string.format("unknown result: %d", result )
	end

	return str
end

ResultByID = {
	{{range $a, $enumObj := .Enums}} {{if .IsResultEnum}} {{range .Fields}} {{if ne .TagNumber 0}}
	[{{.TagNumber}}] = "{{$enumObj.Name}}.{{.Name}}", {{end}} {{end}} {{end}} {{end}}
}

{{end}}

Enum = {
{{range $a, $enumObj := .Enums}}
	{{$enumObj.Name}} = { {{range .Fields}}
		["{{.Name}}"] = {{.TagNumber}}, {{end}}
	},
	{{end}}
}

local sproto = {
	Schema = [[
{{range .Structs}}
.{{.Name}} {	{{range .StFields}}
	{{.Name}} {{.TagNumber}} : {{.CompatibleTypeString}} {{end}}
}
{{end}}
	]],

	NameByID = { {{range .Structs}}
		[{{.MsgID}}] = "{{.Name}}",{{end}}
	},

	IDByName = {},

	ResetByID = { {{range .Structs}}
		[{{.MsgID}}] = function( obj ) -- {{.Name}}
			if obj == nil then return end {{range .StFields}}
			obj.{{.Name}} = {{.LuaDefaultValueString}} {{end}}
		end, {{end}}
	},
}

local t = sproto.IDByName
for k, v in pairs(sproto.NameByID) do
	t[v] = k
end

return sproto

`

func (self *fieldModel) LuaDefaultValueString() string {

	if self.Repeatd {
		return "nil"
	}

	switch self.Type {
	case meta.FieldType_Bool:
		return "false"
	case meta.FieldType_Int32,
		meta.FieldType_Int64,
		meta.FieldType_UInt32,
		meta.FieldType_UInt64,
		meta.FieldType_Integer,
		meta.FieldType_Float32,
		meta.FieldType_Float64,
		meta.FieldType_Enum:
		return "0"
	case meta.FieldType_String:
		return "\"\""
	case meta.FieldType_Struct,
		meta.FieldType_Bytes:
		return "nil"
	}

	return "unknown type" + self.Type.String()
}

func gen_lua(fm *fileModel, filename string) {

	addData(fm, "lua")

	generateCode("sp->lua", luaCodeTemplate, filename, fm, nil)

}
<commit_msg>fixed: lua table style<commit_after>package main

import (
	"github.com/davyxu/gosproto/meta"
)

const luaCodeTemplate = `-- Generated by github.com/davyxu/gosproto/sprotogen
-- DO NOT EDIT!
{{if .EnumValueGroup}}
ResultToString = function ( result )
	if result == 0 then
		return "OK"
	end

	local str = ResultByID[result]
	if str == nil then
		return string.format("unknown result: %d", result )
	end

	return str
end

ResultByID = {
	{{range $a, $enumObj := .Enums}} {{if .IsResultEnum}} {{range .Fields}} {{if ne .TagNumber 0}}
	[{{.TagNumber}}] = "{{$enumObj.Name}}.{{.Name}}", {{end}} {{end}} {{end}} {{end}}
}

{{end}}

Enum = {
{{range $a, $enumObj := .Enums}}
	{{$enumObj.Name}} = { {{range .Fields}}
		{{.Name}} = {{.TagNumber}}, {{end}}
	},
	{{end}}
}

local sproto = {
	Schema = [[
{{range .Structs}}
.{{.Name}} {	{{range .StFields}}
	{{.Name}} {{.TagNumber}} : {{.CompatibleTypeString}} {{end}}
}
{{end}}
	]],

	NameByID = { {{range .Structs}}
		[{{.MsgID}}] = "{{.Name}}",{{end}}
	},

	IDByName = {},

	ResetByID = { {{range .Structs}}
		[{{.MsgID}}] = function( obj ) -- {{.Name}}
			if obj == nil then return end {{range .StFields}}
			obj.{{.Name}} = {{.LuaDefaultValueString}} {{end}}
		end, {{end}}
	},
}

local t = sproto.IDByName
for k, v in pairs(sproto.NameByID) do
	t[v] = k
end

return sproto

`

func (self *fieldModel) LuaDefaultValueString() string {

	if self.Repeatd {
		return "nil"
	}

	switch self.Type {
	case meta.FieldType_Bool:
		return "false"
	case meta.FieldType_Int32,
		meta.FieldType_Int64,
		meta.FieldType_UInt32,
		meta.FieldType_UInt64,
		meta.FieldType_Integer,
		meta.FieldType_Float32,
		meta.FieldType_Float64,
		meta.FieldType_Enum:
		return "0"
	case meta.FieldType_String:
		return "\"\""
	case meta.FieldType_Struct,
		meta.FieldType_Bytes:
		return "nil"
	}

	return "unknown type" + self.Type.String()
}

func gen_lua(fm *fileModel, filename string) {

	addData(fm, "lua")

	generateCode("sp->lua", luaCodeTemplate, filename, fm, nil)

}
<|endoftext|>


<commit_before>// Copyright 2018 by caixw, All rights reserved.
// Use of this source code is governed by a MIT
// license that can be found in the LICENSE file.

package sqlbuilder

import (
	"database/sql"
	"sort"
)

// UpdateStmt 更新语句
type UpdateStmt struct {
	*execStmt

	table  string
	where  *WhereStmt
	values []*updateSet

	occColumn string      // 乐观锁的列名
	occValue  interface{} // 乐观锁的当前值
}

// 表示一条 SET 语句。比如 set key=val
type updateSet struct {
	column string
	value  interface{}
	typ    byte // 类型，可以是 + 自增类型，- 自减类型，或是空值表示正常表达式
}

// Update 声明一条 UPDATE 的 SQL 语句
func Update(e Engine, d Dialect) *UpdateStmt {
	stmt := &UpdateStmt{values: []*updateSet{}}
	stmt.execStmt = newExecStmt(e, d, stmt)
	stmt.where = newWhere(stmt.l, stmt.r)

	return stmt
}

// Table 指定表名
func (stmt *UpdateStmt) Table(table string) *UpdateStmt {
	stmt.table = table
	return stmt
}

// Set 设置值，若 col 相同，则会覆盖
func (stmt *UpdateStmt) Set(col string, val interface{}) *UpdateStmt {
	stmt.values = append(stmt.values, &updateSet{
		column: col,
		value:  val,
		typ:    0,
	})
	return stmt
}

// Increase 给列增加值
func (stmt *UpdateStmt) Increase(col string, val interface{}) *UpdateStmt {
	stmt.values = append(stmt.values, &updateSet{
		column: col,
		value:  val,
		typ:    '+',
	})
	return stmt
}

// Decrease 给钱减少值
func (stmt *UpdateStmt) Decrease(col string, val interface{}) *UpdateStmt {
	stmt.values = append(stmt.values, &updateSet{
		column: col,
		value:  val,
		typ:    '-',
	})
	return stmt
}

// OCC 指定一个用于乐观锁的字段。
//
// val 表示乐观锁原始的值，更新时如果值不等于 val，将更新失败。
func (stmt *UpdateStmt) OCC(col string, val interface{}) *UpdateStmt {
	stmt.occColumn = col
	stmt.occValue = val
	stmt.Increase(col, 1)
	return stmt
}

// WhereStmt 实现 WhereStmter 接口
func (stmt *UpdateStmt) WhereStmt() *WhereStmt {
	return stmt.where
}

// Reset 重置语句
func (stmt *UpdateStmt) Reset() *UpdateStmt {
	stmt.table = ""
	stmt.where.Reset()
	stmt.values = stmt.values[:0]

	stmt.occColumn = ""
	stmt.occValue = nil

	return stmt
}

// SQL 获取 SQL 语句以及对应的参数
func (stmt *UpdateStmt) SQL() (string, []interface{}, error) {
	if err := stmt.checkErrors(); err != nil {
		return "", nil, err
	}

	buf := New("UPDATE ")
	buf.WriteString(stmt.table)
	buf.WriteString(" SET ")

	args := make([]interface{}, 0, len(stmt.values))

	for _, val := range stmt.values {
		buf.Quote(val.column, stmt.l, stmt.r)
		buf.WriteBytes('=')

		if val.typ != 0 {
			buf.Quote(val.column, stmt.l, stmt.r).WriteBytes(val.typ)
		}

		if named, ok := val.value.(sql.NamedArg); ok && named.Name != "" {
			buf.WriteBytes('@')
			buf.WriteString(named.Name)
		} else {
			buf.WriteBytes('?')
		}
		buf.WriteBytes(',')
		args = append(args, val.value)
	}
	buf.TruncateLast(1)

	wq, wa, err := stmt.getWhereSQL()
	if err != nil {
		return "", nil, err
	}

	if wq != "" {
		buf.WriteString(" WHERE ")
		buf.WriteString(wq)
		args = append(args, wa...)
	}

	return buf.String(), args, nil
}

func (stmt *UpdateStmt) getWhereSQL() (string, []interface{}, error) {
	if stmt.occColumn == "" {
		return stmt.where.SQL()
	}

	w := newWhere(stmt.l, stmt.r)
	w.appendGroup(true, stmt.where)

	occ := w.AndGroup()
	if named, ok := stmt.occValue.(sql.NamedArg); ok && named.Name != "" {
		occ.And(stmt.occColumn+"=@"+named.Name, stmt.occValue)
	} else {
		occ.And(stmt.occColumn+"=?", stmt.occValue)
	}

	q, a, err := w.SQL()
	println(q)

	return q, a, err
}

// 检测列名是否存在重复，先排序，再与后一元素比较。
func (stmt *UpdateStmt) checkErrors() error {
	if stmt.table == "" {
		return ErrTableIsEmpty
	}

	if len(stmt.values) == 0 {
		return ErrValueIsEmpty
	}

	if stmt.columnsHasDup() {
		return ErrDupColumn
	}

	return nil
}

// 检测列名是否存在重复，先排序，再与后一元素比较。
func (stmt *UpdateStmt) columnsHasDup() bool {
	sort.SliceStable(stmt.values, func(i, j int) bool {
		return stmt.values[i].column < stmt.values[j].column
	})

	for index, col := range stmt.values {
		if index+1 >= len(stmt.values) {
			return false
		}

		if col.column == stmt.values[index+1].column {
			return true
		}
	}

	return false
}

// Where UpdateStmt.And 的别名
func (stmt *UpdateStmt) Where(cond string, args ...interface{}) *UpdateStmt {
	return stmt.And(cond, args...)
}

// And 添加一条 and 语句
func (stmt *UpdateStmt) And(cond string, args ...interface{}) *UpdateStmt {
	stmt.where.And(cond, args...)
	return stmt
}

// Or 添加一条 OR 语句
func (stmt *UpdateStmt) Or(cond string, args ...interface{}) *UpdateStmt {
	stmt.where.Or(cond, args...)
	return stmt
}

// AndIsNull 指定 WHERE ... AND col IS NULL
func (stmt *UpdateStmt) AndIsNull(col string) *UpdateStmt {
	stmt.where.AndIsNull(col)
	return stmt
}

// OrIsNull 指定 WHERE ... OR col IS NULL
func (stmt *UpdateStmt) OrIsNull(col string) *UpdateStmt {
	stmt.where.OrIsNull(col)
	return stmt
}

// AndIsNotNull 指定 WHERE ... AND col IS NOT NULL
func (stmt *UpdateStmt) AndIsNotNull(col string) *UpdateStmt {
	stmt.where.AndIsNotNull(col)
	return stmt
}

// OrIsNotNull 指定 WHERE ... OR col IS NOT NULL
func (stmt *UpdateStmt) OrIsNotNull(col string) *UpdateStmt {
	stmt.where.OrIsNotNull(col)
	return stmt
}

// AndBetween 指定 WHERE ... AND col BETWEEN v1 AND v2
func (stmt *UpdateStmt) AndBetween(col string, v1, v2 interface{}) *UpdateStmt {
	stmt.where.AndBetween(col, v1, v2)
	return stmt
}

// OrBetween 指定 WHERE ... OR col BETWEEN v1 AND v2
func (stmt *UpdateStmt) OrBetween(col string, v1, v2 interface{}) *UpdateStmt {
	stmt.where.OrBetween(col, v1, v2)
	return stmt
}

// AndNotBetween 指定 WHERE ... AND col NOT BETWEEN v1 AND v2
func (stmt *UpdateStmt) AndNotBetween(col string, v1, v2 interface{}) *UpdateStmt {
	stmt.where.AndNotBetween(col, v1, v2)
	return stmt
}

// OrNotBetween 指定 WHERE ... OR col BETWEEN v1 AND v2
func (stmt *UpdateStmt) OrNotBetween(col string, v1, v2 interface{}) *UpdateStmt {
	stmt.where.OrNotBetween(col, v1, v2)
	return stmt
}

// AndLike 指定 WHERE ... AND col LIKE content
func (stmt *UpdateStmt) AndLike(col string, content interface{}) *UpdateStmt {
	stmt.where.AndLike(col, content)
	return stmt
}

// OrLike 指定 WHERE ... OR col LIKE content
func (stmt *UpdateStmt) OrLike(col string, content interface{}) *UpdateStmt {
	stmt.where.OrLike(col, content)
	return stmt
}

// AndNotLike 指定 WHERE ... AND col NOT LIKE content
func (stmt *UpdateStmt) AndNotLike(col string, content interface{}) *UpdateStmt {
	stmt.where.AndNotLike(col, content)
	return stmt
}

// OrNotLike 指定 WHERE ... OR col NOT LIKE content
func (stmt *UpdateStmt) OrNotLike(col string, content interface{}) *UpdateStmt {
	stmt.where.OrNotLike(col, content)
	return stmt
}

// AndIn 指定 WHERE ... AND col IN(v...)
func (stmt *UpdateStmt) AndIn(col string, v ...interface{}) *UpdateStmt {
	stmt.where.AndIn(col, v...)
	return stmt
}

// OrIn 指定 WHERE ... OR col IN(v...)
func (stmt *UpdateStmt) OrIn(col string, v ...interface{}) *UpdateStmt {
	stmt.where.OrIn(col, v...)
	return stmt
}

// AndNotIn 指定 WHERE ... AND col NOT IN(v...)
func (stmt *UpdateStmt) AndNotIn(col string, v ...interface{}) *UpdateStmt {
	stmt.where.AndNotIn(col, v...)
	return stmt
}

// OrNotIn 指定 WHERE ... OR col IN(v...)
func (stmt *UpdateStmt) OrNotIn(col string, v ...interface{}) *UpdateStmt {
	stmt.where.OrNotIn(col, v...)
	return stmt
}

// AndGroup 开始一个子条件语句
func (stmt *UpdateStmt) AndGroup() *WhereStmt {
	return stmt.where.AndGroup()
}

// OrGroup 开始一个子条件语句
func (stmt *UpdateStmt) OrGroup() *WhereStmt {
	return stmt.where.OrGroup()
}
<commit_msg>[sqlbuilder] 删除无用的代码<commit_after>// Copyright 2018 by caixw, All rights reserved.
// Use of this source code is governed by a MIT
// license that can be found in the LICENSE file.

package sqlbuilder

import (
	"database/sql"
	"sort"
)

// UpdateStmt 更新语句
type UpdateStmt struct {
	*execStmt

	table  string
	where  *WhereStmt
	values []*updateSet

	occColumn string      // 乐观锁的列名
	occValue  interface{} // 乐观锁的当前值
}

// 表示一条 SET 语句。比如 set key=val
type updateSet struct {
	column string
	value  interface{}
	typ    byte // 类型，可以是 + 自增类型，- 自减类型，或是空值表示正常表达式
}

// Update 声明一条 UPDATE 的 SQL 语句
func Update(e Engine, d Dialect) *UpdateStmt {
	stmt := &UpdateStmt{values: []*updateSet{}}
	stmt.execStmt = newExecStmt(e, d, stmt)
	stmt.where = newWhere(stmt.l, stmt.r)

	return stmt
}

// Table 指定表名
func (stmt *UpdateStmt) Table(table string) *UpdateStmt {
	stmt.table = table
	return stmt
}

// Set 设置值，若 col 相同，则会覆盖
func (stmt *UpdateStmt) Set(col string, val interface{}) *UpdateStmt {
	stmt.values = append(stmt.values, &updateSet{
		column: col,
		value:  val,
		typ:    0,
	})
	return stmt
}

// Increase 给列增加值
func (stmt *UpdateStmt) Increase(col string, val interface{}) *UpdateStmt {
	stmt.values = append(stmt.values, &updateSet{
		column: col,
		value:  val,
		typ:    '+',
	})
	return stmt
}

// Decrease 给钱减少值
func (stmt *UpdateStmt) Decrease(col string, val interface{}) *UpdateStmt {
	stmt.values = append(stmt.values, &updateSet{
		column: col,
		value:  val,
		typ:    '-',
	})
	return stmt
}

// OCC 指定一个用于乐观锁的字段。
//
// val 表示乐观锁原始的值，更新时如果值不等于 val，将更新失败。
func (stmt *UpdateStmt) OCC(col string, val interface{}) *UpdateStmt {
	stmt.occColumn = col
	stmt.occValue = val
	stmt.Increase(col, 1)
	return stmt
}

// WhereStmt 实现 WhereStmter 接口
func (stmt *UpdateStmt) WhereStmt() *WhereStmt {
	return stmt.where
}

// Reset 重置语句
func (stmt *UpdateStmt) Reset() *UpdateStmt {
	stmt.table = ""
	stmt.where.Reset()
	stmt.values = stmt.values[:0]

	stmt.occColumn = ""
	stmt.occValue = nil

	return stmt
}

// SQL 获取 SQL 语句以及对应的参数
func (stmt *UpdateStmt) SQL() (string, []interface{}, error) {
	if err := stmt.checkErrors(); err != nil {
		return "", nil, err
	}

	buf := New("UPDATE ")
	buf.WriteString(stmt.table)
	buf.WriteString(" SET ")

	args := make([]interface{}, 0, len(stmt.values))

	for _, val := range stmt.values {
		buf.Quote(val.column, stmt.l, stmt.r)
		buf.WriteBytes('=')

		if val.typ != 0 {
			buf.Quote(val.column, stmt.l, stmt.r).WriteBytes(val.typ)
		}

		if named, ok := val.value.(sql.NamedArg); ok && named.Name != "" {
			buf.WriteBytes('@')
			buf.WriteString(named.Name)
		} else {
			buf.WriteBytes('?')
		}
		buf.WriteBytes(',')
		args = append(args, val.value)
	}
	buf.TruncateLast(1)

	wq, wa, err := stmt.getWhereSQL()
	if err != nil {
		return "", nil, err
	}

	if wq != "" {
		buf.WriteString(" WHERE ")
		buf.WriteString(wq)
		args = append(args, wa...)
	}

	return buf.String(), args, nil
}

func (stmt *UpdateStmt) getWhereSQL() (string, []interface{}, error) {
	if stmt.occColumn == "" {
		return stmt.where.SQL()
	}

	w := newWhere(stmt.l, stmt.r)
	w.appendGroup(true, stmt.where)

	occ := w.AndGroup()
	if named, ok := stmt.occValue.(sql.NamedArg); ok && named.Name != "" {
		occ.And(stmt.occColumn+"=@"+named.Name, stmt.occValue)
	} else {
		occ.And(stmt.occColumn+"=?", stmt.occValue)
	}

	q, a, err := w.SQL()

	return q, a, err
}

// 检测列名是否存在重复，先排序，再与后一元素比较。
func (stmt *UpdateStmt) checkErrors() error {
	if stmt.table == "" {
		return ErrTableIsEmpty
	}

	if len(stmt.values) == 0 {
		return ErrValueIsEmpty
	}

	if stmt.columnsHasDup() {
		return ErrDupColumn
	}

	return nil
}

// 检测列名是否存在重复，先排序，再与后一元素比较。
func (stmt *UpdateStmt) columnsHasDup() bool {
	sort.SliceStable(stmt.values, func(i, j int) bool {
		return stmt.values[i].column < stmt.values[j].column
	})

	for index, col := range stmt.values {
		if index+1 >= len(stmt.values) {
			return false
		}

		if col.column == stmt.values[index+1].column {
			return true
		}
	}

	return false
}

// Where UpdateStmt.And 的别名
func (stmt *UpdateStmt) Where(cond string, args ...interface{}) *UpdateStmt {
	return stmt.And(cond, args...)
}

// And 添加一条 and 语句
func (stmt *UpdateStmt) And(cond string, args ...interface{}) *UpdateStmt {
	stmt.where.And(cond, args...)
	return stmt
}

// Or 添加一条 OR 语句
func (stmt *UpdateStmt) Or(cond string, args ...interface{}) *UpdateStmt {
	stmt.where.Or(cond, args...)
	return stmt
}

// AndIsNull 指定 WHERE ... AND col IS NULL
func (stmt *UpdateStmt) AndIsNull(col string) *UpdateStmt {
	stmt.where.AndIsNull(col)
	return stmt
}

// OrIsNull 指定 WHERE ... OR col IS NULL
func (stmt *UpdateStmt) OrIsNull(col string) *UpdateStmt {
	stmt.where.OrIsNull(col)
	return stmt
}

// AndIsNotNull 指定 WHERE ... AND col IS NOT NULL
func (stmt *UpdateStmt) AndIsNotNull(col string) *UpdateStmt {
	stmt.where.AndIsNotNull(col)
	return stmt
}

// OrIsNotNull 指定 WHERE ... OR col IS NOT NULL
func (stmt *UpdateStmt) OrIsNotNull(col string) *UpdateStmt {
	stmt.where.OrIsNotNull(col)
	return stmt
}

// AndBetween 指定 WHERE ... AND col BETWEEN v1 AND v2
func (stmt *UpdateStmt) AndBetween(col string, v1, v2 interface{}) *UpdateStmt {
	stmt.where.AndBetween(col, v1, v2)
	return stmt
}

// OrBetween 指定 WHERE ... OR col BETWEEN v1 AND v2
func (stmt *UpdateStmt) OrBetween(col string, v1, v2 interface{}) *UpdateStmt {
	stmt.where.OrBetween(col, v1, v2)
	return stmt
}

// AndNotBetween 指定 WHERE ... AND col NOT BETWEEN v1 AND v2
func (stmt *UpdateStmt) AndNotBetween(col string, v1, v2 interface{}) *UpdateStmt {
	stmt.where.AndNotBetween(col, v1, v2)
	return stmt
}

// OrNotBetween 指定 WHERE ... OR col BETWEEN v1 AND v2
func (stmt *UpdateStmt) OrNotBetween(col string, v1, v2 interface{}) *UpdateStmt {
	stmt.where.OrNotBetween(col, v1, v2)
	return stmt
}

// AndLike 指定 WHERE ... AND col LIKE content
func (stmt *UpdateStmt) AndLike(col string, content interface{}) *UpdateStmt {
	stmt.where.AndLike(col, content)
	return stmt
}

// OrLike 指定 WHERE ... OR col LIKE content
func (stmt *UpdateStmt) OrLike(col string, content interface{}) *UpdateStmt {
	stmt.where.OrLike(col, content)
	return stmt
}

// AndNotLike 指定 WHERE ... AND col NOT LIKE content
func (stmt *UpdateStmt) AndNotLike(col string, content interface{}) *UpdateStmt {
	stmt.where.AndNotLike(col, content)
	return stmt
}

// OrNotLike 指定 WHERE ... OR col NOT LIKE content
func (stmt *UpdateStmt) OrNotLike(col string, content interface{}) *UpdateStmt {
	stmt.where.OrNotLike(col, content)
	return stmt
}

// AndIn 指定 WHERE ... AND col IN(v...)
func (stmt *UpdateStmt) AndIn(col string, v ...interface{}) *UpdateStmt {
	stmt.where.AndIn(col, v...)
	return stmt
}

// OrIn 指定 WHERE ... OR col IN(v...)
func (stmt *UpdateStmt) OrIn(col string, v ...interface{}) *UpdateStmt {
	stmt.where.OrIn(col, v...)
	return stmt
}

// AndNotIn 指定 WHERE ... AND col NOT IN(v...)
func (stmt *UpdateStmt) AndNotIn(col string, v ...interface{}) *UpdateStmt {
	stmt.where.AndNotIn(col, v...)
	return stmt
}

// OrNotIn 指定 WHERE ... OR col IN(v...)
func (stmt *UpdateStmt) OrNotIn(col string, v ...interface{}) *UpdateStmt {
	stmt.where.OrNotIn(col, v...)
	return stmt
}

// AndGroup 开始一个子条件语句
func (stmt *UpdateStmt) AndGroup() *WhereStmt {
	return stmt.where.AndGroup()
}

// OrGroup 开始一个子条件语句
func (stmt *UpdateStmt) OrGroup() *WhereStmt {
	return stmt.where.OrGroup()
}
<|endoftext|>