<commit_before>package bsdiff

import (
	"fmt"
	"os"
	"runtime"
	"sync/atomic"
	"time"

	humanize "github.com/dustin/go-humanize"
	"github.com/itchio/wharf/state"
)

type chunk struct {
	addOldStart int
	addNewStart int
	addLength   int
	copyStart   int
	copyEnd     int
	offset      int
	eoc         bool
}

type SendChunkFunc func(c chunk)

type blockWorkerState struct {
	consumed chan bool
	work     chan int
	chunks   chan chunk
}

func (ctx *DiffContext) doPartitioned(obuf []byte, obuflen int, nbuf []byte, nbuflen int, memstats *runtime.MemStats, writeMessage WriteMessageFunc, consumer *state.Consumer) error {
	var err error

	partitions := ctx.Partitions
	if partitions >= len(obuf)-1 {
		partitions = 1
	}

	consumer.ProgressLabel(fmt.Sprintf("Sorting %s...", humanize.IBytes(uint64(obuflen))))
	consumer.Progress(0.0)

	startTime := time.Now()

	pmemstats := &runtime.MemStats{}
	runtime.ReadMemStats(pmemstats)
	oldAlloc := pmemstats.TotalAlloc

	if ctx.I == nil {
		ctx.I = make([]int, len(obuf))
		beforeAlloc := time.Now()
		fmt.Fprintf(os.Stderr, "\nAllocated %d-int I in %s\n", len(obuf), time.Since(beforeAlloc))
	} else {
		for len(ctx.I) < len(obuf) {
			lenBefore := len(ctx.I)
			beforeAlloc := time.Now()
			ctx.I = make([]int, len(obuf))
			fmt.Fprintf(os.Stderr, "\nGrown I from %d to %d in %s\n", lenBefore, len(ctx.I), time.Since(beforeAlloc))
		}
	}

	psa := NewPSA(partitions, obuf, ctx.I)

	runtime.ReadMemStats(pmemstats)
	newAlloc := pmemstats.TotalAlloc
	fmt.Fprintf(os.Stderr, "\nAlloc difference after PSA: %s. Size of I: %s\n", humanize.IBytes(uint64(newAlloc-oldAlloc)), humanize.IBytes(uint64(8*len(psa.I))))

	if ctx.Stats != nil {
		ctx.Stats.TimeSpentSorting += time.Since(startTime)
	}

	if ctx.MeasureMem {
		runtime.ReadMemStats(memstats)
		fmt.Fprintf(os.Stderr, "\nAllocated bytes after qsufsort: %s (%s total)", humanize.IBytes(uint64(memstats.Alloc)), humanize.IBytes(uint64(memstats.TotalAlloc)))
	}

	bsdc := &Control{}

	consumer.ProgressLabel(fmt.Sprintf("Preparing to scan %s...", humanize.IBytes(uint64(nbuflen))))
	consumer.Progress(0.0)

	startTime = time.Now()

	var waitConsumeTime int64
	var waitEnqueueTime int64
	var idleWorkerTime int64
	var sendingWorkerTime int64

	analyzeBlock := func(nbuflen int, nbuf []byte, offset int, chunks chan chunk) {
		var lenf int

		// Compute the differences, writing ctrl as we go
		var scan, pos, length int
		var lastscan, lastpos, lastoffset int

		for scan < nbuflen {
			var oldscore int
			scan += length

			for scsc := scan; scan < nbuflen; scan++ {
				pos, length = psa.search(nbuf[scan:])

				for ; scsc < scan+length; scsc++ {
					if scsc+lastoffset < obuflen &&
						obuf[scsc+lastoffset] == nbuf[scsc] {
						oldscore++
					}
				}

				if (length == oldscore && length != 0) || length > oldscore+8 {
					break
				}

				if scan+lastoffset < obuflen && obuf[scan+lastoffset] == nbuf[scan] {
					oldscore--
				}
			}

			if length != oldscore || scan == nbuflen {
				var s, Sf int
				lenf = 0
				for i := int(0); lastscan+i < scan && lastpos+i < obuflen; {
					if obuf[lastpos+i] == nbuf[lastscan+i] {
						s++
					}
					i++
					if s*2-i > Sf*2-lenf {
						Sf = s
						lenf = i
					}
				}

				lenb := 0
				if scan < nbuflen {
					var s, Sb int
					for i := int(1); (scan >= lastscan+i) && (pos >= i); i++ {
						if obuf[pos-i] == nbuf[scan-i] {
							s++
						}
						if s*2-i > Sb*2-lenb {
							Sb = s
							lenb = i
						}
					}
				}

				if lastscan+lenf > scan-lenb {
					overlap := (lastscan + lenf) - (scan - lenb)
					s := int(0)
					Ss := int(0)
					lens := int(0)
					for i := int(0); i < overlap; i++ {
						if nbuf[lastscan+lenf-overlap+i] == obuf[lastpos+lenf-overlap+i] {
							s++
						}
						if nbuf[scan-lenb+i] == obuf[pos-lenb+i] {
							s--
						}
						if s > Ss {
							Ss = s
							lens = i + 1
						}
					}

					lenf += lens - overlap
					lenb -= lens
				}

				c := chunk{
					addOldStart: lastpos,
					addNewStart: lastscan,
					addLength:   lenf,
					copyStart:   lastscan + lenf,
					copyEnd:     scan - lenb,
					offset:      offset,
				}

				if c.addLength > 0 || (c.copyEnd != c.copyStart) {
					// if not a no-op, send
					beforeSend := time.Now()
					chunks <- c
					atomic.AddInt64(&sendingWorkerTime, int64(time.Since(beforeSend)))
				}

				lastscan = scan - lenb
				lastpos = pos - lenb
				lastoffset = pos - scan
			}
		}

		beforeSend := time.Now()
		chunks <- chunk{eoc: true}
		atomic.AddInt64(&sendingWorkerTime, int64(time.Since(beforeSend)))
	}

	blockSize := 128 * 1024
	numBlocks := (nbuflen + blockSize - 1) / blockSize

	if numBlocks < partitions {
		blockSize = nbuflen / partitions
		numBlocks = (nbuflen + blockSize - 1) / blockSize
	}

	numWorkers := partitions * 16
	if numWorkers > numBlocks {
		numWorkers = numBlocks
	}

	// fmt.Fprintf(os.Stderr, "Divvying %s in %d block(s) of %s (with %d workers)\n",
	// 	humanize.IBytes(uint64(nbuflen)),
	// 	numBlocks,
	// 	humanize.IBytes(uint64(blockSize)),
	// 	partitions,
	// )

	blockWorkersState := make([]blockWorkerState, numWorkers)

	// initialize all channels
	for i := 0; i < numWorkers; i++ {
		blockWorkersState[i].work = make(chan int, 1)
		blockWorkersState[i].chunks = make(chan chunk, 256)
		blockWorkersState[i].consumed = make(chan bool, 1)
		blockWorkersState[i].consumed <- true
	}

	for i := 0; i < numWorkers; i++ {
		go func(workerState blockWorkerState, workerIndex int) {
			lastWorkTime := time.Now()
			for blockIndex := range workerState.work {
				atomic.AddInt64(&idleWorkerTime, int64(time.Since(lastWorkTime)))

				// fmt.Fprintf(os.Stderr, "\nWorker %d should analyze block %d", workerIndex, blockIndex)
				boundary := blockSize * blockIndex
				realBlockSize := blockSize
				if blockIndex == numBlocks-1 {
					realBlockSize = nbuflen - boundary
				}
				// fmt.Fprintf(os.Stderr, "Analyzing %s block at %d\n", humanize.IBytes(uint64(realBlockSize)), i)

				analyzeBlock(realBlockSize, nbuf[boundary:boundary+realBlockSize], boundary, workerState.chunks)
				// fmt.Fprintf(os.Stderr, "\nWorker %d done analyzing block %d", workerIndex, blockIndex)

				lastWorkTime = time.Now()
			}
		}(blockWorkersState[i], i)
	}

	go func() {
		workerIndex := 0

		for i := 0; i < numBlocks; i++ {
			beforeConsume := time.Now()
			<-blockWorkersState[workerIndex].consumed
			atomic.AddInt64(&waitConsumeTime, int64(time.Since(beforeConsume)))

			beforeEnqueue := time.Now()
			blockWorkersState[workerIndex].work <- i
			atomic.AddInt64(&waitEnqueueTime, int64(time.Since(beforeEnqueue)))

			workerIndex = (workerIndex + 1) % numWorkers
		}

		for workerIndex := 0; workerIndex < numWorkers; workerIndex++ {
			close(blockWorkersState[workerIndex].work)
		}
		// fmt.Fprintf(os.Stderr, "Sent all blockworks\n")
	}()

	if ctx.MeasureMem {
		runtime.ReadMemStats(memstats)
		fmt.Fprintf(os.Stderr, "\nAllocated bytes after scan-prepare: %s (%s total)", humanize.IBytes(uint64(memstats.Alloc)), humanize.IBytes(uint64(memstats.TotalAlloc)))
	}

	var prevChunk chunk
	first := true

	consumer.ProgressLabel(fmt.Sprintf("Scanning %s (%d blocks of %s)...", humanize.IBytes(uint64(nbuflen)), numBlocks, humanize.IBytes(uint64(blockSize))))

	allChunks := make(chan chunk, 1024)

	go func() {
		workerIndex := 0
		for blockIndex := 0; blockIndex < numBlocks; blockIndex++ {
			consumer.Progress(float64(blockIndex) / float64(numBlocks))

			// fmt.Fprintf(os.Stderr, "\nWaiting on worker %d for block %d", workerIndex, blockIndex)
			consumer.Progress(float64(blockIndex) / float64(numBlocks))
			state := blockWorkersState[workerIndex]

			for chunk := range state.chunks {
				// fmt.Fprintf(os.Stderr, "\nFor block %d, received chunk %#v", blockIndex, chunk)
				if chunk.eoc {
					break
				}

				allChunks <- chunk
			}

			state.consumed <- true
			workerIndex = (workerIndex + 1) % numWorkers
		}

		close(allChunks)
	}()

	for chunk := range allChunks {
		// fmt.Fprintf(os.Stderr, "\nWaiting on worker %d for block %d", workerIndex, blockIndex)
		// fmt.Fprintf(os.Stderr, "\nFor block %d, received chunk %#v", blockIndex, chunk)
		if chunk.eoc {
			break
		}

		if first {
			first = false
		} else {
			bsdc.Seek = int64(chunk.addOldStart - (prevChunk.addOldStart + prevChunk.addLength))

			// fmt.Fprintf(os.Stderr, "%d bytes add, %d bytes copy\n", len(bsdc.Add), len(bsdc.Copy))

			err := writeMessage(bsdc)
			if err != nil {
				return err
			}
		}

		ctx.db.Reset()
		ctx.db.Grow(chunk.addLength)

		addNewStart := chunk.addNewStart + chunk.offset

		for i := 0; i < chunk.addLength; i++ {
			ctx.db.WriteByte(nbuf[addNewStart+i] - obuf[chunk.addOldStart+i])
		}

		bsdc.Add = ctx.db.Bytes()
		bsdc.Copy = nbuf[chunk.offset+chunk.copyStart : chunk.offset+chunk.copyEnd]

		if ctx.Stats != nil && ctx.Stats.BiggestAdd < int64(len(bsdc.Add)) {
			ctx.Stats.BiggestAdd = int64(len(bsdc.Add))
		}

		prevChunk = chunk
	}

	// fmt.Fprintf(os.Stderr, "%d bytes add, %d bytes copy\n", len(bsdc.Add), len(bsdc.Copy))

	bsdc.Seek = 0
	err = writeMessage(bsdc)
	if err != nil {
		return err
	}

	if ctx.Stats != nil {
		ctx.Stats.TimeSpentScanning += time.Since(startTime)
	}

	if ctx.MeasureMem {
		runtime.ReadMemStats(memstats)
		consumer.Debugf("\nAllocated bytes after scan: %s (%s total)", humanize.IBytes(uint64(memstats.Alloc)), humanize.IBytes(uint64(memstats.TotalAlloc)))
		fmt.Fprintf(os.Stderr, "\nAllocated bytes after scan: %s (%s total)", humanize.IBytes(uint64(memstats.Alloc)), humanize.IBytes(uint64(memstats.TotalAlloc)))
	}

	fmt.Fprintf(os.Stderr, "\nStats: waitConsume %s, waitEnqueue %s, idleWorker %s, sendingWorker %s\n",
		time.Duration(atomic.LoadInt64(&waitConsumeTime)),
		time.Duration(atomic.LoadInt64(&waitEnqueueTime)),
		time.Duration(atomic.LoadInt64(&idleWorkerTime)),
		time.Duration(atomic.LoadInt64(&sendingWorkerTime)),
	)

	bsdc.Reset()
	bsdc.Eof = true
	err = writeMessage(bsdc)
	if err != nil {
		return err
	}

	return nil
}
<commit_msg>Clean up a little<commit_after>package bsdiff

import (
	"fmt"
	"os"
	"runtime"
	"time"

	humanize "github.com/dustin/go-humanize"
	"github.com/itchio/wharf/state"
)

type chunk struct {
	addOldStart int
	addNewStart int
	addLength   int
	copyStart   int
	copyEnd     int
	offset      int
	eoc         bool
}

type blockWorkerState struct {
	consumed chan bool
	work     chan int
	chunks   chan chunk
}

func (ctx *DiffContext) doPartitioned(obuf []byte, obuflen int, nbuf []byte, nbuflen int, memstats *runtime.MemStats, writeMessage WriteMessageFunc, consumer *state.Consumer) error {
	var err error

	partitions := ctx.Partitions
	if partitions >= len(obuf)-1 {
		partitions = 1
	}

	consumer.ProgressLabel(fmt.Sprintf("Sorting %s...", humanize.IBytes(uint64(obuflen))))
	consumer.Progress(0.0)

	startTime := time.Now()

	pmemstats := &runtime.MemStats{}
	runtime.ReadMemStats(pmemstats)
	oldAlloc := pmemstats.TotalAlloc

	if ctx.I == nil {
		ctx.I = make([]int, len(obuf))
		beforeAlloc := time.Now()
		fmt.Fprintf(os.Stderr, "\nAllocated %d-int I in %s\n", len(obuf), time.Since(beforeAlloc))
	} else {
		for len(ctx.I) < len(obuf) {
			lenBefore := len(ctx.I)
			beforeAlloc := time.Now()
			ctx.I = make([]int, len(obuf))
			fmt.Fprintf(os.Stderr, "\nGrown I from %d to %d in %s\n", lenBefore, len(ctx.I), time.Since(beforeAlloc))
		}
	}

	psa := NewPSA(partitions, obuf, ctx.I)

	runtime.ReadMemStats(pmemstats)
	newAlloc := pmemstats.TotalAlloc
	fmt.Fprintf(os.Stderr, "\nAlloc difference after PSA: %s. Size of I: %s\n", humanize.IBytes(uint64(newAlloc-oldAlloc)), humanize.IBytes(uint64(8*len(psa.I))))

	if ctx.Stats != nil {
		ctx.Stats.TimeSpentSorting += time.Since(startTime)
	}

	if ctx.MeasureMem {
		runtime.ReadMemStats(memstats)
		fmt.Fprintf(os.Stderr, "\nAllocated bytes after qsufsort: %s (%s total)", humanize.IBytes(uint64(memstats.Alloc)), humanize.IBytes(uint64(memstats.TotalAlloc)))
	}

	bsdc := &Control{}

	consumer.ProgressLabel(fmt.Sprintf("Preparing to scan %s...", humanize.IBytes(uint64(nbuflen))))
	consumer.Progress(0.0)

	startTime = time.Now()

	analyzeBlock := func(nbuflen int, nbuf []byte, offset int, chunks chan chunk) {
		var lenf int

		// Compute the differences, writing ctrl as we go
		var scan, pos, length int
		var lastscan, lastpos, lastoffset int

		for scan < nbuflen {
			var oldscore int
			scan += length

			for scsc := scan; scan < nbuflen; scan++ {
				pos, length = psa.search(nbuf[scan:])

				for ; scsc < scan+length; scsc++ {
					if scsc+lastoffset < obuflen &&
						obuf[scsc+lastoffset] == nbuf[scsc] {
						oldscore++
					}
				}

				if (length == oldscore && length != 0) || length > oldscore+8 {
					break
				}

				if scan+lastoffset < obuflen && obuf[scan+lastoffset] == nbuf[scan] {
					oldscore--
				}
			}

			if length != oldscore || scan == nbuflen {
				var s, Sf int
				lenf = 0
				for i := int(0); lastscan+i < scan && lastpos+i < obuflen; {
					if obuf[lastpos+i] == nbuf[lastscan+i] {
						s++
					}
					i++
					if s*2-i > Sf*2-lenf {
						Sf = s
						lenf = i
					}
				}

				lenb := 0
				if scan < nbuflen {
					var s, Sb int
					for i := int(1); (scan >= lastscan+i) && (pos >= i); i++ {
						if obuf[pos-i] == nbuf[scan-i] {
							s++
						}
						if s*2-i > Sb*2-lenb {
							Sb = s
							lenb = i
						}
					}
				}

				if lastscan+lenf > scan-lenb {
					overlap := (lastscan + lenf) - (scan - lenb)
					s := int(0)
					Ss := int(0)
					lens := int(0)
					for i := int(0); i < overlap; i++ {
						if nbuf[lastscan+lenf-overlap+i] == obuf[lastpos+lenf-overlap+i] {
							s++
						}
						if nbuf[scan-lenb+i] == obuf[pos-lenb+i] {
							s--
						}
						if s > Ss {
							Ss = s
							lens = i + 1
						}
					}

					lenf += lens - overlap
					lenb -= lens
				}

				c := chunk{
					addOldStart: lastpos,
					addNewStart: lastscan,
					addLength:   lenf,
					copyStart:   lastscan + lenf,
					copyEnd:     scan - lenb,
					offset:      offset,
				}

				if c.addLength > 0 || (c.copyEnd != c.copyStart) {
					// if not a no-op, send
					chunks <- c
				}

				lastscan = scan - lenb
				lastpos = pos - lenb
				lastoffset = pos - scan
			}
		}

		chunks <- chunk{eoc: true}
	}

	blockSize := 128 * 1024
	numBlocks := (nbuflen + blockSize - 1) / blockSize

	if numBlocks < partitions {
		blockSize = nbuflen / partitions
		numBlocks = (nbuflen + blockSize - 1) / blockSize
	}

	numWorkers := partitions * 16
	if numWorkers > numBlocks {
		numWorkers = numBlocks
	}

	// fmt.Fprintf(os.Stderr, "Divvying %s in %d block(s) of %s (with %d workers)\n",
	// 	humanize.IBytes(uint64(nbuflen)),
	// 	numBlocks,
	// 	humanize.IBytes(uint64(blockSize)),
	// 	partitions,
	// )

	blockWorkersState := make([]blockWorkerState, numWorkers)

	// initialize all channels
	for i := 0; i < numWorkers; i++ {
		blockWorkersState[i].work = make(chan int, 1)
		blockWorkersState[i].chunks = make(chan chunk, 256)
		blockWorkersState[i].consumed = make(chan bool, 1)
		blockWorkersState[i].consumed <- true
	}

	for i := 0; i < numWorkers; i++ {
		go func(workerState blockWorkerState, workerIndex int) {
			for blockIndex := range workerState.work {
				// fmt.Fprintf(os.Stderr, "\nWorker %d should analyze block %d", workerIndex, blockIndex)
				boundary := blockSize * blockIndex
				realBlockSize := blockSize
				if blockIndex == numBlocks-1 {
					realBlockSize = nbuflen - boundary
				}
				// fmt.Fprintf(os.Stderr, "Analyzing %s block at %d\n", humanize.IBytes(uint64(realBlockSize)), i)

				analyzeBlock(realBlockSize, nbuf[boundary:boundary+realBlockSize], boundary, workerState.chunks)
				// fmt.Fprintf(os.Stderr, "\nWorker %d done analyzing block %d", workerIndex, blockIndex)
			}
		}(blockWorkersState[i], i)
	}

	go func() {
		workerIndex := 0

		for i := 0; i < numBlocks; i++ {
			<-blockWorkersState[workerIndex].consumed
			blockWorkersState[workerIndex].work <- i

			workerIndex = (workerIndex + 1) % numWorkers
		}

		for workerIndex := 0; workerIndex < numWorkers; workerIndex++ {
			close(blockWorkersState[workerIndex].work)
		}
		// fmt.Fprintf(os.Stderr, "Sent all blockworks\n")
	}()

	if ctx.MeasureMem {
		runtime.ReadMemStats(memstats)
		fmt.Fprintf(os.Stderr, "\nAllocated bytes after scan-prepare: %s (%s total)", humanize.IBytes(uint64(memstats.Alloc)), humanize.IBytes(uint64(memstats.TotalAlloc)))
	}

	var prevChunk chunk
	first := true

	consumer.ProgressLabel(fmt.Sprintf("Scanning %s (%d blocks of %s)...", humanize.IBytes(uint64(nbuflen)), numBlocks, humanize.IBytes(uint64(blockSize))))

	allChunks := make(chan chunk, 1024)

	go func() {
		workerIndex := 0
		for blockIndex := 0; blockIndex < numBlocks; blockIndex++ {
			consumer.Progress(float64(blockIndex) / float64(numBlocks))

			// fmt.Fprintf(os.Stderr, "\nWaiting on worker %d for block %d", workerIndex, blockIndex)
			consumer.Progress(float64(blockIndex) / float64(numBlocks))
			state := blockWorkersState[workerIndex]

			for chunk := range state.chunks {
				// fmt.Fprintf(os.Stderr, "\nFor block %d, received chunk %#v", blockIndex, chunk)
				if chunk.eoc {
					break
				}

				allChunks <- chunk
			}

			state.consumed <- true
			workerIndex = (workerIndex + 1) % numWorkers
		}

		close(allChunks)
	}()

	for chunk := range allChunks {
		// fmt.Fprintf(os.Stderr, "\nWaiting on worker %d for block %d", workerIndex, blockIndex)
		// fmt.Fprintf(os.Stderr, "\nFor block %d, received chunk %#v", blockIndex, chunk)
		if chunk.eoc {
			break
		}

		if first {
			first = false
		} else {
			bsdc.Seek = int64(chunk.addOldStart - (prevChunk.addOldStart + prevChunk.addLength))

			// fmt.Fprintf(os.Stderr, "%d bytes add, %d bytes copy\n", len(bsdc.Add), len(bsdc.Copy))

			err := writeMessage(bsdc)
			if err != nil {
				return err
			}
		}

		ctx.db.Reset()
		ctx.db.Grow(chunk.addLength)

		addNewStart := chunk.addNewStart + chunk.offset

		for i := 0; i < chunk.addLength; i++ {
			ctx.db.WriteByte(nbuf[addNewStart+i] - obuf[chunk.addOldStart+i])
		}

		bsdc.Add = ctx.db.Bytes()
		bsdc.Copy = nbuf[chunk.offset+chunk.copyStart : chunk.offset+chunk.copyEnd]

		if ctx.Stats != nil && ctx.Stats.BiggestAdd < int64(len(bsdc.Add)) {
			ctx.Stats.BiggestAdd = int64(len(bsdc.Add))
		}

		prevChunk = chunk
	}

	// fmt.Fprintf(os.Stderr, "%d bytes add, %d bytes copy\n", len(bsdc.Add), len(bsdc.Copy))

	bsdc.Seek = 0
	err = writeMessage(bsdc)
	if err != nil {
		return err
	}

	if ctx.Stats != nil {
		ctx.Stats.TimeSpentScanning += time.Since(startTime)
	}

	if ctx.MeasureMem {
		runtime.ReadMemStats(memstats)
		consumer.Debugf("\nAllocated bytes after scan: %s (%s total)", humanize.IBytes(uint64(memstats.Alloc)), humanize.IBytes(uint64(memstats.TotalAlloc)))
		fmt.Fprintf(os.Stderr, "\nAllocated bytes after scan: %s (%s total)", humanize.IBytes(uint64(memstats.Alloc)), humanize.IBytes(uint64(memstats.TotalAlloc)))
	}

	bsdc.Reset()
	bsdc.Eof = true
	err = writeMessage(bsdc)
	if err != nil {
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>/*
Copyright 2019 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 upgrade

import (
	"context"
	"testing"

	"k8s.io/api/networking/v1beta1"
	"k8s.io/apimachinery/pkg/util/intstr"
	"k8s.io/ingress-gce/pkg/e2e"
	"k8s.io/ingress-gce/pkg/fuzz"
	"k8s.io/ingress-gce/pkg/utils/common"
)

var (
	port80  = intstr.FromInt(80)
	ingName = "ing1"
)

// Finalizer implements e2e.UpgradeTest interface.
type BasicHTTP struct {
	t         *testing.T
	s         *e2e.Sandbox
	framework *e2e.Framework
	crud      e2e.IngressCRUD
	ing       *v1beta1.Ingress
}

// NewBasicHTTPUpgradeTest returns an upgrade test that tests the basic behavior
// of an ingress with http load-balancer.
func NewBasicHTTPUpgradeTest() e2e.UpgradeTest {
	return &BasicHTTP{}
}

// Name implements e2e.UpgradeTest.Init.
func (bh *BasicHTTP) Name() string {
	return "BasicHTTPUpgrade"
}

// Init implements e2e.UpgradeTest.Init.
func (bh *BasicHTTP) Init(t *testing.T, s *e2e.Sandbox, framework *e2e.Framework) error {
	bh.t = t
	bh.s = s
	bh.framework = framework
	return nil
}

// PreUpgrade implements e2e.UpgradeTest.PreUpgrade.
func (bh *BasicHTTP) PreUpgrade() error {
	_, err := e2e.CreateEchoService(bh.s, svcName, nil)
	if err != nil {
		bh.t.Fatalf("error creating echo service: %v", err)
	}
	bh.t.Logf("Echo service created (%s/%s)", bh.s.Namespace, svcName)

	ing := fuzz.NewIngressBuilder(bh.s.Namespace, ingName, "").
		AddPath("foo.com", "/", svcName, port80).
		Build()
	ingKey := common.NamespacedName(ing)
	bh.crud = e2e.IngressCRUD{C: bh.framework.Clientset}
	if _, err := bh.crud.Create(ing); err != nil {
		bh.t.Fatalf("error creating Ingress %s: %v", ingKey, err)
	}
	bh.t.Logf("Ingress created (%s)", ingKey)

	if bh.ing, err = e2e.UpgradeTestWaitForIngress(bh.s, ing, &e2e.WaitForIngressOptions{ExpectUnreachable: true}); err != nil {
		bh.t.Fatalf("error waiting for Ingress %s to stabilize: %v", ingKey, err)
	}
	bh.t.Logf("GCLB resources created (%s)", ingKey)

	if _, err := e2e.WhiteboxTest(bh.ing, bh.s, bh.framework.Cloud, ""); err != nil {
		bh.t.Fatalf("e2e.WhiteboxTest(%s, ...) = %v, want nil", ingKey, err)
	}
	return nil
}

// DuringUpgrade implements e2e.UpgradeTest.DuringUpgrade.
func (bh *BasicHTTP) DuringUpgrade() error {
	return nil
}

// PostUpgrade implements e2e.UpgradeTest.PostUpgrade
func (bh *BasicHTTP) PostUpgrade() error {
	// force ingress update. only add path once
	newIng := fuzz.NewIngressBuilderFromExisting(bh.ing).
		AddPath("bar.com", "/", "service-1", port80).
		Build()
	ingKey := common.NamespacedName(newIng)
	// TODO: does the path need to be different for each upgrade
	if _, err := bh.crud.Update(newIng); err != nil {
		bh.t.Fatalf("error updating Ingress %s: %v", ingKey, err)
	} else {
		// If Ingress upgrade succeeds, we update the status on this Ingress
		// to Unstable. It is set back to Stable after WaitForIngress below
		// finishes successfully.
		bh.s.PutStatus(e2e.Unstable)
	}

	// Verify the Ingress has stabilized after the master upgrade and we
	// trigger an Ingress update
	ing, err := e2e.WaitForIngress(bh.s, bh.ing, &e2e.WaitForIngressOptions{ExpectUnreachable: true})
	if err != nil {
		bh.t.Fatalf("error waiting for Ingress %s to stabilize: %v", ingKey, err)
	}
	bh.s.PutStatus(e2e.Stable)
	bh.t.Logf("GCLB is stable (%s)", ingKey)
	gclb, err := e2e.WhiteboxTest(ing, bh.s, bh.framework.Cloud, "")
	if err != nil {
		bh.t.Fatalf("e2e.WhiteboxTest(%s, ...) = %v, want nil", ingKey, err)
	}

	// If the Master has upgraded and the Ingress is stable,
	// we delete the Ingress and exit out of the loop to indicate that
	// the test is done.
	deleteOptions := &fuzz.GCLBDeleteOptions{
		SkipDefaultBackend: true,
	}

	if err := e2e.WaitForIngressDeletion(context.Background(), gclb, bh.s, ing, deleteOptions); err != nil { // Sometimes times out waiting
		bh.t.Errorf("e2e.WaitForIngressDeletion(..., %q, nil) = %v, want nil", ingKey, err)
	}
	return nil
}
<commit_msg>Bugfix: Fix service name in upgrade test<commit_after>/*
Copyright 2019 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 upgrade

import (
	"context"
	"testing"

	"k8s.io/api/networking/v1beta1"
	"k8s.io/apimachinery/pkg/util/intstr"
	"k8s.io/ingress-gce/pkg/e2e"
	"k8s.io/ingress-gce/pkg/fuzz"
	"k8s.io/ingress-gce/pkg/utils/common"
)

var (
	port80  = intstr.FromInt(80)
	ingName = "ing1"
)

// Finalizer implements e2e.UpgradeTest interface.
type BasicHTTP struct {
	t         *testing.T
	s         *e2e.Sandbox
	framework *e2e.Framework
	crud      e2e.IngressCRUD
	ing       *v1beta1.Ingress
}

// NewBasicHTTPUpgradeTest returns an upgrade test that tests the basic behavior
// of an ingress with http load-balancer.
func NewBasicHTTPUpgradeTest() e2e.UpgradeTest {
	return &BasicHTTP{}
}

// Name implements e2e.UpgradeTest.Init.
func (bh *BasicHTTP) Name() string {
	return "BasicHTTPUpgrade"
}

// Init implements e2e.UpgradeTest.Init.
func (bh *BasicHTTP) Init(t *testing.T, s *e2e.Sandbox, framework *e2e.Framework) error {
	bh.t = t
	bh.s = s
	bh.framework = framework
	return nil
}

// PreUpgrade implements e2e.UpgradeTest.PreUpgrade.
func (bh *BasicHTTP) PreUpgrade() error {
	_, err := e2e.CreateEchoService(bh.s, svcName, nil)
	if err != nil {
		bh.t.Fatalf("error creating echo service: %v", err)
	}
	bh.t.Logf("Echo service created (%s/%s)", bh.s.Namespace, svcName)

	ing := fuzz.NewIngressBuilder(bh.s.Namespace, ingName, "").
		AddPath("foo.com", "/", svcName, port80).
		Build()
	ingKey := common.NamespacedName(ing)
	bh.crud = e2e.IngressCRUD{C: bh.framework.Clientset}
	if _, err := bh.crud.Create(ing); err != nil {
		bh.t.Fatalf("error creating Ingress %s: %v", ingKey, err)
	}
	bh.t.Logf("Ingress created (%s)", ingKey)

	if bh.ing, err = e2e.UpgradeTestWaitForIngress(bh.s, ing, &e2e.WaitForIngressOptions{ExpectUnreachable: true}); err != nil {
		bh.t.Fatalf("error waiting for Ingress %s to stabilize: %v", ingKey, err)
	}
	bh.t.Logf("GCLB resources created (%s)", ingKey)

	if _, err := e2e.WhiteboxTest(bh.ing, bh.s, bh.framework.Cloud, ""); err != nil {
		bh.t.Fatalf("e2e.WhiteboxTest(%s, ...) = %v, want nil", ingKey, err)
	}
	return nil
}

// DuringUpgrade implements e2e.UpgradeTest.DuringUpgrade.
func (bh *BasicHTTP) DuringUpgrade() error {
	return nil
}

// PostUpgrade implements e2e.UpgradeTest.PostUpgrade
func (bh *BasicHTTP) PostUpgrade() error {
	// force ingress update. only add path once
	newIng := fuzz.NewIngressBuilderFromExisting(bh.ing).
		AddPath("bar.com", "/", svcName, port80).
		Build()
	ingKey := common.NamespacedName(newIng)
	// TODO: does the path need to be different for each upgrade
	if _, err := bh.crud.Update(newIng); err != nil {
		bh.t.Fatalf("error updating Ingress %s: %v", ingKey, err)
	} else {
		// If Ingress upgrade succeeds, we update the status on this Ingress
		// to Unstable. It is set back to Stable after WaitForIngress below
		// finishes successfully.
		bh.s.PutStatus(e2e.Unstable)
	}

	// Verify the Ingress has stabilized after the master upgrade and we
	// trigger an Ingress update
	ing, err := e2e.WaitForIngress(bh.s, bh.ing, &e2e.WaitForIngressOptions{ExpectUnreachable: true})
	if err != nil {
		bh.t.Fatalf("error waiting for Ingress %s to stabilize: %v", ingKey, err)
	}
	bh.s.PutStatus(e2e.Stable)
	bh.t.Logf("GCLB is stable (%s)", ingKey)
	gclb, err := e2e.WhiteboxTest(ing, bh.s, bh.framework.Cloud, "")
	if err != nil {
		bh.t.Fatalf("e2e.WhiteboxTest(%s, ...) = %v, want nil", ingKey, err)
	}

	// If the Master has upgraded and the Ingress is stable,
	// we delete the Ingress and exit out of the loop to indicate that
	// the test is done.
	deleteOptions := &fuzz.GCLBDeleteOptions{
		SkipDefaultBackend: true,
	}

	if err := e2e.WaitForIngressDeletion(context.Background(), gclb, bh.s, ing, deleteOptions); err != nil { // Sometimes times out waiting
		bh.t.Errorf("e2e.WaitForIngressDeletion(..., %q, nil) = %v, want nil", ingKey, err)
	}
	return nil
}
<|endoftext|>


<commit_before>package highlight

import (
	"regexp"
	"strings"

	"github.com/dlclark/regexp2"
)

func combineLineMatch(src, dst LineMatch) LineMatch {
	for k, v := range src {
		if g, ok := dst[k]; ok {
			if g == 0 {
				dst[k] = v
			}
		} else {
			dst[k] = v
		}
	}
	return dst
}

// A State represents the region at the end of a line
type State *Region

// LineStates is an interface for a buffer-like object which can also store the states and matches for every line
type LineStates interface {
	Line(n int) string
	LinesNum() int
	State(lineN int) State
	SetState(lineN int, s State)
	SetMatch(lineN int, m LineMatch)
}

// A Highlighter contains the information needed to highlight a string
type Highlighter struct {
	lastRegion *Region
	def        *Def
}

// NewHighlighter returns a new highlighter from the given syntax definition
func NewHighlighter(def *Def) *Highlighter {
	h := new(Highlighter)
	h.def = def
	return h
}

// LineMatch represents the syntax highlighting matches for one line. Each index where the coloring is changed is marked with that
// color's group (represented as one byte)
type LineMatch map[int]uint8

func findIndex(regex *regexp2.Regexp, str []byte, canMatchStart, canMatchEnd bool) []int {
	regexStr := regex.String()
	if strings.Contains(regexStr, "^") {
		if !canMatchStart {
			return nil
		}
	}
	if strings.Contains(regexStr, "$") {
		if !canMatchEnd {
			return nil
		}
	}
	match, _ := regex.FindStringMatch(string(str))
	if match == nil {
		return nil
	}
	return []int{match.Index, match.Index + match.Length}
}

func findAllIndex(regex *regexp.Regexp, str []byte, canMatchStart, canMatchEnd bool) [][]int {
	regexStr := regex.String()
	if strings.Contains(regexStr, "^") {
		if !canMatchStart {
			return nil
		}
	}
	if strings.Contains(regexStr, "$") {
		if !canMatchEnd {
			return nil
		}
	}
	return regex.FindAllIndex(str, -1)
}

func (h *Highlighter) highlightRegion(highlights LineMatch, start int, canMatchEnd bool, lineNum int, line []byte, region *Region, statesOnly bool) LineMatch {
	// highlights := make(LineMatch)

	if start == 0 {
		if !statesOnly {
			highlights[0] = region.group
		}
	}

	loc := findIndex(region.end, line, start == 0, canMatchEnd)
	if loc != nil {
		if !statesOnly {
			highlights[start+loc[1]-1] = region.group
		}
		if region.parent == nil {
			if !statesOnly {
				highlights[start+loc[1]] = 0
				h.highlightRegion(highlights, start, false, lineNum, line[:loc[0]], region, statesOnly)
			}
			h.highlightEmptyRegion(highlights, start+loc[1], canMatchEnd, lineNum, line[loc[1]:], statesOnly)
			return highlights
		}
		if !statesOnly {
			highlights[start+loc[1]] = region.parent.group
			h.highlightRegion(highlights, start, false, lineNum, line[:loc[0]], region, statesOnly)
		}
		h.highlightRegion(highlights, start+loc[1], canMatchEnd, lineNum, line[loc[1]:], region.parent, statesOnly)
		return highlights
	}

	if len(line) == 0 || statesOnly {
		if canMatchEnd {
			h.lastRegion = region
		}

		return highlights
	}

	firstLoc := []int{len(line), 0}
	var firstRegion *Region
	for _, r := range region.rules.regions {
		loc := findIndex(r.start, line, start == 0, canMatchEnd)
		if loc != nil {
			if loc[0] < firstLoc[0] {
				firstLoc = loc
				firstRegion = r
			}
		}
	}
	if firstLoc[0] != len(line) {
		highlights[start+firstLoc[0]] = firstRegion.group
		h.highlightRegion(highlights, start, false, lineNum, line[:firstLoc[0]], region, statesOnly)
		h.highlightRegion(highlights, start+firstLoc[1], canMatchEnd, lineNum, line[firstLoc[1]:], firstRegion, statesOnly)
		return highlights
	}

	fullHighlights := make([]uint8, len([]rune(string(line))))
	for i := 0; i < len(fullHighlights); i++ {
		fullHighlights[i] = region.group
	}

	for _, p := range region.rules.patterns {
		matches := findAllIndex(p.regex, line, start == 0, canMatchEnd)
		for _, m := range matches {
			for i := m[0]; i < m[1]; i++ {
				fullHighlights[i] = p.group
			}
		}
	}
	for i, h := range fullHighlights {
		if i == 0 || h != fullHighlights[i-1] {
			if _, ok := highlights[start+i]; !ok {
				highlights[start+i] = h
			}
		}
	}

	if canMatchEnd {
		h.lastRegion = region
	}

	return highlights
}

func (h *Highlighter) highlightEmptyRegion(highlights LineMatch, start int, canMatchEnd bool, lineNum int, line []byte, statesOnly bool) LineMatch {
	if len(line) == 0 {
		if canMatchEnd {
			h.lastRegion = nil
		}
		return highlights
	}

	firstLoc := []int{len(line), 0}
	var firstRegion *Region
	for _, r := range h.def.rules.regions {
		loc := findIndex(r.start, line, start == 0, canMatchEnd)
		if loc != nil {
			if loc[0] < firstLoc[0] {
				firstLoc = loc
				firstRegion = r
			}
		}
	}
	if firstLoc[0] != len(line) {
		if !statesOnly {
			highlights[start+firstLoc[0]] = firstRegion.group
		}
		h.highlightEmptyRegion(highlights, start, false, lineNum, line[:firstLoc[0]], statesOnly)
		h.highlightRegion(highlights, start+firstLoc[1], canMatchEnd, lineNum, line[firstLoc[1]:], firstRegion, statesOnly)
		return highlights
	}

	if statesOnly {
		if canMatchEnd {
			h.lastRegion = nil
		}

		return highlights
	}

	fullHighlights := make([]uint8, len(line))
	for _, p := range h.def.rules.patterns {
		matches := findAllIndex(p.regex, line, start == 0, canMatchEnd)
		for _, m := range matches {
			for i := m[0]; i < m[1]; i++ {
				fullHighlights[i] = p.group
			}
		}
	}
	for i, h := range fullHighlights {
		if i == 0 || h != fullHighlights[i-1] {
			if _, ok := highlights[start+i]; !ok {
				highlights[start+i] = h
			}
		}
	}

	if canMatchEnd {
		h.lastRegion = nil
	}

	return highlights
}

// HighlightString syntax highlights a string
// Use this function for simple syntax highlighting and use the other functions for
// more advanced syntax highlighting. They are optimized for quick rehighlighting of the same
// text with minor changes made
func (h *Highlighter) HighlightString(input string) []LineMatch {
	lines := strings.Split(input, "\n")
	var lineMatches []LineMatch

	for i := 0; i < len(lines); i++ {
		line := []byte(lines[i])
		highlights := make(LineMatch)

		if i == 0 || h.lastRegion == nil {
			lineMatches = append(lineMatches, h.highlightEmptyRegion(highlights, 0, true, i, line, false))
		} else {
			lineMatches = append(lineMatches, h.highlightRegion(highlights, 0, true, i, line, h.lastRegion, false))
		}
	}

	return lineMatches
}

// HighlightStates correctly sets all states for the buffer
func (h *Highlighter) HighlightStates(input LineStates) {
	for i := 0; i < input.LinesNum(); i++ {
		line := []byte(input.Line(i))
		// highlights := make(LineMatch)

		if i == 0 || h.lastRegion == nil {
			h.highlightEmptyRegion(nil, 0, true, i, line, true)
		} else {
			h.highlightRegion(nil, 0, true, i, line, h.lastRegion, true)
		}

		curState := h.lastRegion

		input.SetState(i, curState)
	}
}

// HighlightMatches sets the matches for each line in between startline and endline
// It sets all other matches in the buffer to nil to conserve memory
// This assumes that all the states are set correctly
func (h *Highlighter) HighlightMatches(input LineStates, startline, endline int) {
	for i := startline; i < endline; i++ {
		if i >= input.LinesNum() {
			break
		}

		line := []byte(input.Line(i))
		highlights := make(LineMatch)

		var match LineMatch
		if i == 0 || input.State(i-1) == nil {
			match = h.highlightEmptyRegion(highlights, 0, true, i, line, false)
		} else {
			match = h.highlightRegion(highlights, 0, true, i, line, input.State(i-1), false)
		}

		input.SetMatch(i, match)
	}
}

// ReHighlightStates will scan down from `startline` and set the appropriate end of line state
// for each line until it comes across the same state in two consecutive lines
func (h *Highlighter) ReHighlightStates(input LineStates, startline int) {
	// lines := input.LineData()

	h.lastRegion = nil
	if startline > 0 {
		h.lastRegion = input.State(startline - 1)
	}
	for i := startline; i < input.LinesNum(); i++ {
		line := []byte(input.Line(i))
		// highlights := make(LineMatch)

		// var match LineMatch
		if i == 0 || h.lastRegion == nil {
			h.highlightEmptyRegion(nil, 0, true, i, line, true)
		} else {
			h.highlightRegion(nil, 0, true, i, line, h.lastRegion, true)
		}
		curState := h.lastRegion
		lastState := input.State(i)

		input.SetState(i, curState)

		if curState == lastState {
			break
		}
	}
}

// ReHighlightLine will rehighlight the state and match for a single line
func (h *Highlighter) ReHighlightLine(input LineStates, lineN int) {
	line := []byte(input.Line(lineN))
	highlights := make(LineMatch)

	h.lastRegion = nil
	if lineN > 0 {
		h.lastRegion = input.State(lineN - 1)
	}

	var match LineMatch
	if lineN == 0 || h.lastRegion == nil {
		match = h.highlightEmptyRegion(highlights, 0, true, lineN, line, false)
	} else {
		match = h.highlightRegion(highlights, 0, true, lineN, line, h.lastRegion, false)
	}
	curState := h.lastRegion

	input.SetMatch(lineN, match)
	input.SetState(lineN, curState)
}
<commit_msg>Proper unicode support<commit_after>package highlight

import (
	"regexp"
	"strings"

	"github.com/dlclark/regexp2"
)

func combineLineMatch(src, dst LineMatch) LineMatch {
	for k, v := range src {
		if g, ok := dst[k]; ok {
			if g == 0 {
				dst[k] = v
			}
		} else {
			dst[k] = v
		}
	}
	return dst
}

// A State represents the region at the end of a line
type State *Region

// LineStates is an interface for a buffer-like object which can also store the states and matches for every line
type LineStates interface {
	Line(n int) string
	LinesNum() int
	State(lineN int) State
	SetState(lineN int, s State)
	SetMatch(lineN int, m LineMatch)
}

// A Highlighter contains the information needed to highlight a string
type Highlighter struct {
	lastRegion *Region
	def        *Def
}

// NewHighlighter returns a new highlighter from the given syntax definition
func NewHighlighter(def *Def) *Highlighter {
	h := new(Highlighter)
	h.def = def
	return h
}

// LineMatch represents the syntax highlighting matches for one line. Each index where the coloring is changed is marked with that
// color's group (represented as one byte)
type LineMatch map[int]uint8

func findIndex(regex *regexp2.Regexp, str []rune, canMatchStart, canMatchEnd bool) []int {
	regexStr := regex.String()
	if strings.Contains(regexStr, "^") {
		if !canMatchStart {
			return nil
		}
	}
	if strings.Contains(regexStr, "$") {
		if !canMatchEnd {
			return nil
		}
	}
	match, _ := regex.FindStringMatch(string(str))
	if match == nil {
		return nil
	}
	return []int{match.Index, match.Index + match.Length}
}

func findAllIndex(regex *regexp.Regexp, str []rune, canMatchStart, canMatchEnd bool) [][]int {
	regexStr := regex.String()
	if strings.Contains(regexStr, "^") {
		if !canMatchStart {
			return nil
		}
	}
	if strings.Contains(regexStr, "$") {
		if !canMatchEnd {
			return nil
		}
	}
	return regex.FindAllIndex([]byte(string(str)), -1)
}

func (h *Highlighter) highlightRegion(highlights LineMatch, start int, canMatchEnd bool, lineNum int, line []rune, region *Region, statesOnly bool) LineMatch {
	// highlights := make(LineMatch)

	if start == 0 {
		if !statesOnly {
			highlights[0] = region.group
		}
	}

	loc := findIndex(region.end, line, start == 0, canMatchEnd)
	if loc != nil {
		if !statesOnly {
			highlights[start+loc[1]-1] = region.group
		}
		if region.parent == nil {
			if !statesOnly {
				highlights[start+loc[1]] = 0
				h.highlightRegion(highlights, start, false, lineNum, line[:loc[0]], region, statesOnly)
			}
			h.highlightEmptyRegion(highlights, start+loc[1], canMatchEnd, lineNum, line[loc[1]:], statesOnly)
			return highlights
		}
		if !statesOnly {
			highlights[start+loc[1]] = region.parent.group
			h.highlightRegion(highlights, start, false, lineNum, line[:loc[0]], region, statesOnly)
		}
		h.highlightRegion(highlights, start+loc[1], canMatchEnd, lineNum, line[loc[1]:], region.parent, statesOnly)
		return highlights
	}

	if len(line) == 0 || statesOnly {
		if canMatchEnd {
			h.lastRegion = region
		}

		return highlights
	}

	firstLoc := []int{len(line), 0}
	var firstRegion *Region
	for _, r := range region.rules.regions {
		loc := findIndex(r.start, line, start == 0, canMatchEnd)
		if loc != nil {
			if loc[0] < firstLoc[0] {
				firstLoc = loc
				firstRegion = r
			}
		}
	}
	if firstLoc[0] != len(line) {
		highlights[start+firstLoc[0]] = firstRegion.group
		h.highlightRegion(highlights, start, false, lineNum, line[:firstLoc[0]], region, statesOnly)
		h.highlightRegion(highlights, start+firstLoc[1], canMatchEnd, lineNum, line[firstLoc[1]:], firstRegion, statesOnly)
		return highlights
	}

	fullHighlights := make([]uint8, len([]rune(string(line))))
	for i := 0; i < len(fullHighlights); i++ {
		fullHighlights[i] = region.group
	}

	for _, p := range region.rules.patterns {
		matches := findAllIndex(p.regex, line, start == 0, canMatchEnd)
		for _, m := range matches {
			for i := m[0]; i < m[1]; i++ {
				fullHighlights[i] = p.group
			}
		}
	}
	for i, h := range fullHighlights {
		if i == 0 || h != fullHighlights[i-1] {
			if _, ok := highlights[start+i]; !ok {
				highlights[start+i] = h
			}
		}
	}

	if canMatchEnd {
		h.lastRegion = region
	}

	return highlights
}

func (h *Highlighter) highlightEmptyRegion(highlights LineMatch, start int, canMatchEnd bool, lineNum int, line []rune, statesOnly bool) LineMatch {
	if len(line) == 0 {
		if canMatchEnd {
			h.lastRegion = nil
		}
		return highlights
	}

	firstLoc := []int{len(line), 0}
	var firstRegion *Region
	for _, r := range h.def.rules.regions {
		loc := findIndex(r.start, line, start == 0, canMatchEnd)
		if loc != nil {
			if loc[0] < firstLoc[0] {
				firstLoc = loc
				firstRegion = r
			}
		}
	}
	if firstLoc[0] != len(line) {
		if !statesOnly {
			highlights[start+firstLoc[0]] = firstRegion.group
		}
		h.highlightEmptyRegion(highlights, start, false, lineNum, line[:firstLoc[0]], statesOnly)
		h.highlightRegion(highlights, start+firstLoc[1], canMatchEnd, lineNum, line[firstLoc[1]:], firstRegion, statesOnly)
		return highlights
	}

	if statesOnly {
		if canMatchEnd {
			h.lastRegion = nil
		}

		return highlights
	}

	fullHighlights := make([]uint8, len(line))
	for _, p := range h.def.rules.patterns {
		matches := findAllIndex(p.regex, line, start == 0, canMatchEnd)
		for _, m := range matches {
			for i := m[0]; i < m[1]; i++ {
				fullHighlights[i] = p.group
			}
		}
	}
	for i, h := range fullHighlights {
		if i == 0 || h != fullHighlights[i-1] {
			if _, ok := highlights[start+i]; !ok {
				highlights[start+i] = h
			}
		}
	}

	if canMatchEnd {
		h.lastRegion = nil
	}

	return highlights
}

// HighlightString syntax highlights a string
// Use this function for simple syntax highlighting and use the other functions for
// more advanced syntax highlighting. They are optimized for quick rehighlighting of the same
// text with minor changes made
func (h *Highlighter) HighlightString(input string) []LineMatch {
	lines := strings.Split(input, "\n")
	var lineMatches []LineMatch

	for i := 0; i < len(lines); i++ {
		line := []rune(lines[i])
		highlights := make(LineMatch)

		if i == 0 || h.lastRegion == nil {
			lineMatches = append(lineMatches, h.highlightEmptyRegion(highlights, 0, true, i, line, false))
		} else {
			lineMatches = append(lineMatches, h.highlightRegion(highlights, 0, true, i, line, h.lastRegion, false))
		}
	}

	return lineMatches
}

// HighlightStates correctly sets all states for the buffer
func (h *Highlighter) HighlightStates(input LineStates) {
	for i := 0; i < input.LinesNum(); i++ {
		line := []rune(input.Line(i))
		// highlights := make(LineMatch)

		if i == 0 || h.lastRegion == nil {
			h.highlightEmptyRegion(nil, 0, true, i, line, true)
		} else {
			h.highlightRegion(nil, 0, true, i, line, h.lastRegion, true)
		}

		curState := h.lastRegion

		input.SetState(i, curState)
	}
}

// HighlightMatches sets the matches for each line in between startline and endline
// It sets all other matches in the buffer to nil to conserve memory
// This assumes that all the states are set correctly
func (h *Highlighter) HighlightMatches(input LineStates, startline, endline int) {
	for i := startline; i < endline; i++ {
		if i >= input.LinesNum() {
			break
		}

		line := []rune(input.Line(i))
		highlights := make(LineMatch)

		var match LineMatch
		if i == 0 || input.State(i-1) == nil {
			match = h.highlightEmptyRegion(highlights, 0, true, i, line, false)
		} else {
			match = h.highlightRegion(highlights, 0, true, i, line, input.State(i-1), false)
		}

		input.SetMatch(i, match)
	}
}

// ReHighlightStates will scan down from `startline` and set the appropriate end of line state
// for each line until it comes across the same state in two consecutive lines
func (h *Highlighter) ReHighlightStates(input LineStates, startline int) {
	// lines := input.LineData()

	h.lastRegion = nil
	if startline > 0 {
		h.lastRegion = input.State(startline - 1)
	}
	for i := startline; i < input.LinesNum(); i++ {
		line := []rune(input.Line(i))
		// highlights := make(LineMatch)

		// var match LineMatch
		if i == 0 || h.lastRegion == nil {
			h.highlightEmptyRegion(nil, 0, true, i, line, true)
		} else {
			h.highlightRegion(nil, 0, true, i, line, h.lastRegion, true)
		}
		curState := h.lastRegion
		lastState := input.State(i)

		input.SetState(i, curState)

		if curState == lastState {
			break
		}
	}
}

// ReHighlightLine will rehighlight the state and match for a single line
func (h *Highlighter) ReHighlightLine(input LineStates, lineN int) {
	line := []rune(input.Line(lineN))
	highlights := make(LineMatch)

	h.lastRegion = nil
	if lineN > 0 {
		h.lastRegion = input.State(lineN - 1)
	}

	var match LineMatch
	if lineN == 0 || h.lastRegion == nil {
		match = h.highlightEmptyRegion(highlights, 0, true, lineN, line, false)
	} else {
		match = h.highlightRegion(highlights, 0, true, lineN, line, h.lastRegion, false)
	}
	curState := h.lastRegion

	input.SetMatch(lineN, match)
	input.SetState(lineN, curState)
}
<|endoftext|>


<commit_before>package highlight

import (
	"regexp"
	"strings"
	"unicode/utf8"
)

// RunePos returns the rune index of a given byte index
// This could cause problems if the byte index is between code points
func runePos(p int, str string) int {
	if p < 0 {
		return 0
	}
	if p >= len(str) {
		return utf8.RuneCountInString(str)
	}
	return utf8.RuneCountInString(str[:p])
}

func combineLineMatch(src, dst LineMatch) LineMatch {
	for k, v := range src {
		if g, ok := dst[k]; ok {
			if g == 0 {
				dst[k] = v
			}
		} else {
			dst[k] = v
		}
	}
	return dst
}

// A State represents the region at the end of a line
type State *region

// LineStates is an interface for a buffer-like object which can also store the states and matches for every line
type LineStates interface {
	Line(n int) string
	LinesNum() int
	State(lineN int) State
	SetState(lineN int, s State)
	SetMatch(lineN int, m LineMatch)
}

// A Highlighter contains the information needed to highlight a string
type Highlighter struct {
	lastRegion *region
	Def        *Def
}

// NewHighlighter returns a new highlighter from the given syntax definition
func NewHighlighter(def *Def) *Highlighter {
	h := new(Highlighter)
	h.Def = def
	return h
}

// LineMatch represents the syntax highlighting matches for one line. Each index where the coloring is changed is marked with that
// color's group (represented as one byte)
type LineMatch map[int]Group

func findIndex(regex *regexp.Regexp, skip *regexp.Regexp, str []rune, canMatchStart, canMatchEnd bool) []int {
	regexStr := regex.String()
	if strings.Contains(regexStr, "^") {
		if !canMatchStart {
			return nil
		}
	}
	if strings.Contains(regexStr, "$") {
		if !canMatchEnd {
			return nil
		}
	}

	var strbytes []byte
	if skip != nil {
		strbytes = skip.ReplaceAllFunc(strbytes, func(match []byte) []byte {
			res := make([]byte, utf8.RuneCount(match))
			return res
		})
	} else {
		strbytes = []byte(string(str))
	}

	match := regex.FindIndex(strbytes)
	if match == nil {
		return nil
	}
	// return []int{match.Index, match.Index + match.Length}
	return []int{runePos(match[0], string(str)), runePos(match[1], string(str))}
}

func findAllIndex(regex *regexp.Regexp, str []rune, canMatchStart, canMatchEnd bool) [][]int {
	regexStr := regex.String()
	if strings.Contains(regexStr, "^") {
		if !canMatchStart {
			return nil
		}
	}
	if strings.Contains(regexStr, "$") {
		if !canMatchEnd {
			return nil
		}
	}
	matches := regex.FindAllIndex([]byte(string(str)), -1)
	for i, m := range matches {
		matches[i][0] = runePos(m[0], string(str))
		matches[i][1] = runePos(m[1], string(str))
	}
	return matches
}

func (h *Highlighter) highlightRegion(highlights LineMatch, start int, canMatchEnd bool, lineNum int, line []rune, curRegion *region, statesOnly bool) LineMatch {
	// highlights := make(LineMatch)

	if start == 0 {
		if !statesOnly {
			highlights[0] = curRegion.group
		}
	}

	loc := findIndex(curRegion.end, curRegion.skip, line, start == 0, canMatchEnd)
	if loc != nil {
		if !statesOnly {
			highlights[start+loc[1]-1] = curRegion.group
		}
		if curRegion.parent == nil {
			if !statesOnly {
				highlights[start+loc[1]] = 0
				h.highlightRegion(highlights, start, false, lineNum, line[:loc[0]], curRegion, statesOnly)
			}
			h.highlightEmptyRegion(highlights, start+loc[1], canMatchEnd, lineNum, line[loc[1]:], statesOnly)
			return highlights
		}
		if !statesOnly {
			highlights[start+loc[1]] = curRegion.parent.group
			h.highlightRegion(highlights, start, false, lineNum, line[:loc[0]], curRegion, statesOnly)
		}
		h.highlightRegion(highlights, start+loc[1], canMatchEnd, lineNum, line[loc[1]:], curRegion.parent, statesOnly)
		return highlights
	}

	if len(line) == 0 || statesOnly {
		if canMatchEnd {
			h.lastRegion = curRegion
		}

		return highlights
	}

	firstLoc := []int{len(line), 0}

	var firstRegion *region
	for _, r := range curRegion.rules.regions {
		loc := findIndex(r.start, nil, line, start == 0, canMatchEnd)
		if loc != nil {
			if loc[0] < firstLoc[0] {
				firstLoc = loc
				firstRegion = r
			}
		}
	}
	if firstLoc[0] != len(line) {
		highlights[start+firstLoc[0]] = firstRegion.group
		h.highlightRegion(highlights, start, false, lineNum, line[:firstLoc[0]], curRegion, statesOnly)
		h.highlightRegion(highlights, start+firstLoc[1], canMatchEnd, lineNum, line[firstLoc[1]:], firstRegion, statesOnly)
		return highlights
	}

	fullHighlights := make([]Group, len([]rune(string(line))))
	for i := 0; i < len(fullHighlights); i++ {
		fullHighlights[i] = curRegion.group
	}

	for _, p := range curRegion.rules.patterns {
		matches := findAllIndex(p.regex, line, start == 0, canMatchEnd)
		for _, m := range matches {
			for i := m[0]; i < m[1]; i++ {
				fullHighlights[i] = p.group
			}
		}
	}
	for i, h := range fullHighlights {
		if i == 0 || h != fullHighlights[i-1] {
			if _, ok := highlights[start+i]; !ok {
				highlights[start+i] = h
			}
		}
	}

	if canMatchEnd {
		h.lastRegion = curRegion
	}

	return highlights
}

func (h *Highlighter) highlightEmptyRegion(highlights LineMatch, start int, canMatchEnd bool, lineNum int, line []rune, statesOnly bool) LineMatch {
	if len(line) == 0 {
		if canMatchEnd {
			h.lastRegion = nil
		}
		return highlights
	}

	firstLoc := []int{len(line), 0}
	var firstRegion *region
	for _, r := range h.Def.rules.regions {
		loc := findIndex(r.start, nil, line, start == 0, canMatchEnd)
		if loc != nil {
			if loc[0] < firstLoc[0] {
				firstLoc = loc
				firstRegion = r
			}
		}
	}
	if firstLoc[0] != len(line) {
		if !statesOnly {
			highlights[start+firstLoc[0]] = firstRegion.group
		}
		h.highlightEmptyRegion(highlights, start, false, lineNum, line[:firstLoc[0]], statesOnly)
		h.highlightRegion(highlights, start+firstLoc[1], canMatchEnd, lineNum, line[firstLoc[1]:], firstRegion, statesOnly)
		return highlights
	}

	if statesOnly {
		if canMatchEnd {
			h.lastRegion = nil
		}

		return highlights
	}

	fullHighlights := make([]Group, len(line))
	for _, p := range h.Def.rules.patterns {
		matches := findAllIndex(p.regex, line, start == 0, canMatchEnd)
		for _, m := range matches {
			for i := m[0]; i < m[1]; i++ {
				fullHighlights[i] = p.group
			}
		}
	}
	for i, h := range fullHighlights {
		if i == 0 || h != fullHighlights[i-1] {
			if _, ok := highlights[start+i]; !ok {
				highlights[start+i] = h
			}
		}
	}

	if canMatchEnd {
		h.lastRegion = nil
	}

	return highlights
}

// HighlightString syntax highlights a string
// Use this function for simple syntax highlighting and use the other functions for
// more advanced syntax highlighting. They are optimized for quick rehighlighting of the same
// text with minor changes made
func (h *Highlighter) HighlightString(input string) []LineMatch {
	lines := strings.Split(input, "\n")
	var lineMatches []LineMatch

	for i := 0; i < len(lines); i++ {
		line := []rune(lines[i])
		highlights := make(LineMatch)

		if i == 0 || h.lastRegion == nil {
			lineMatches = append(lineMatches, h.highlightEmptyRegion(highlights, 0, true, i, line, false))
		} else {
			lineMatches = append(lineMatches, h.highlightRegion(highlights, 0, true, i, line, h.lastRegion, false))
		}
	}

	return lineMatches
}

// HighlightStates correctly sets all states for the buffer
func (h *Highlighter) HighlightStates(input LineStates) {
	for i := 0; i < input.LinesNum(); i++ {
		line := []rune(input.Line(i))
		// highlights := make(LineMatch)

		if i == 0 || h.lastRegion == nil {
			h.highlightEmptyRegion(nil, 0, true, i, line, true)
		} else {
			h.highlightRegion(nil, 0, true, i, line, h.lastRegion, true)
		}

		curState := h.lastRegion

		input.SetState(i, curState)
	}
}

// HighlightMatches sets the matches for each line in between startline and endline
// It sets all other matches in the buffer to nil to conserve memory
// This assumes that all the states are set correctly
func (h *Highlighter) HighlightMatches(input LineStates, startline, endline int) {
	for i := startline; i < endline; i++ {
		if i >= input.LinesNum() {
			break
		}

		line := []rune(input.Line(i))
		highlights := make(LineMatch)

		var match LineMatch
		if i == 0 || input.State(i-1) == nil {
			match = h.highlightEmptyRegion(highlights, 0, true, i, line, false)
		} else {
			match = h.highlightRegion(highlights, 0, true, i, line, input.State(i-1), false)
		}

		input.SetMatch(i, match)
	}
}

// ReHighlightStates will scan down from `startline` and set the appropriate end of line state
// for each line until it comes across the same state in two consecutive lines
func (h *Highlighter) ReHighlightStates(input LineStates, startline int) {
	// lines := input.LineData()

	h.lastRegion = nil
	if startline > 0 {
		h.lastRegion = input.State(startline - 1)
	}
	for i := startline; i < input.LinesNum(); i++ {
		line := []rune(input.Line(i))
		// highlights := make(LineMatch)

		// var match LineMatch
		if i == 0 || h.lastRegion == nil {
			h.highlightEmptyRegion(nil, 0, true, i, line, true)
		} else {
			h.highlightRegion(nil, 0, true, i, line, h.lastRegion, true)
		}
		curState := h.lastRegion
		lastState := input.State(i)

		input.SetState(i, curState)

		if curState == lastState {
			break
		}
	}
}

// ReHighlightLine will rehighlight the state and match for a single line
func (h *Highlighter) ReHighlightLine(input LineStates, lineN int) {
	line := []rune(input.Line(lineN))
	highlights := make(LineMatch)

	h.lastRegion = nil
	if lineN > 0 {
		h.lastRegion = input.State(lineN - 1)
	}

	var match LineMatch
	if lineN == 0 || h.lastRegion == nil {
		match = h.highlightEmptyRegion(highlights, 0, true, lineN, line, false)
	} else {
		match = h.highlightRegion(highlights, 0, true, lineN, line, h.lastRegion, false)
	}
	curState := h.lastRegion

	input.SetMatch(lineN, match)
	input.SetState(lineN, curState)
}
<commit_msg>Slight improvements to included region highlighting<commit_after>package highlight

import (
	"regexp"
	"strings"
	"unicode/utf8"
)

// RunePos returns the rune index of a given byte index
// This could cause problems if the byte index is between code points
func runePos(p int, str string) int {
	if p < 0 {
		return 0
	}
	if p >= len(str) {
		return utf8.RuneCountInString(str)
	}
	return utf8.RuneCountInString(str[:p])
}

func combineLineMatch(src, dst LineMatch) LineMatch {
	for k, v := range src {
		if g, ok := dst[k]; ok {
			if g == 0 {
				dst[k] = v
			}
		} else {
			dst[k] = v
		}
	}
	return dst
}

// A State represents the region at the end of a line
type State *region

// LineStates is an interface for a buffer-like object which can also store the states and matches for every line
type LineStates interface {
	Line(n int) string
	LinesNum() int
	State(lineN int) State
	SetState(lineN int, s State)
	SetMatch(lineN int, m LineMatch)
}

// A Highlighter contains the information needed to highlight a string
type Highlighter struct {
	lastRegion *region
	Def        *Def
}

// NewHighlighter returns a new highlighter from the given syntax definition
func NewHighlighter(def *Def) *Highlighter {
	h := new(Highlighter)
	h.Def = def
	return h
}

// LineMatch represents the syntax highlighting matches for one line. Each index where the coloring is changed is marked with that
// color's group (represented as one byte)
type LineMatch map[int]Group

func findIndex(regex *regexp.Regexp, skip *regexp.Regexp, str []rune, canMatchStart, canMatchEnd bool) []int {
	regexStr := regex.String()
	if strings.Contains(regexStr, "^") {
		if !canMatchStart {
			return nil
		}
	}
	if strings.Contains(regexStr, "$") {
		if !canMatchEnd {
			return nil
		}
	}

	var strbytes []byte
	if skip != nil {
		strbytes = skip.ReplaceAllFunc(strbytes, func(match []byte) []byte {
			res := make([]byte, utf8.RuneCount(match))
			return res
		})
	} else {
		strbytes = []byte(string(str))
	}

	match := regex.FindIndex(strbytes)
	if match == nil {
		return nil
	}
	// return []int{match.Index, match.Index + match.Length}
	return []int{runePos(match[0], string(str)), runePos(match[1], string(str))}
}

func findAllIndex(regex *regexp.Regexp, str []rune, canMatchStart, canMatchEnd bool) [][]int {
	regexStr := regex.String()
	if strings.Contains(regexStr, "^") {
		if !canMatchStart {
			return nil
		}
	}
	if strings.Contains(regexStr, "$") {
		if !canMatchEnd {
			return nil
		}
	}
	matches := regex.FindAllIndex([]byte(string(str)), -1)
	for i, m := range matches {
		matches[i][0] = runePos(m[0], string(str))
		matches[i][1] = runePos(m[1], string(str))
	}
	return matches
}

func (h *Highlighter) highlightRegion(highlights LineMatch, start int, canMatchEnd bool, lineNum int, line []rune, curRegion *region, statesOnly bool) LineMatch {
	// highlights := make(LineMatch)

	if start == 0 {
		if !statesOnly {
			highlights[0] = curRegion.group
		}
	}

	loc := findIndex(curRegion.end, curRegion.skip, line, start == 0, canMatchEnd)
	if loc != nil {
		if !statesOnly {
			highlights[start+loc[1]-1] = curRegion.group
		}
		if curRegion.parent == nil {
			if !statesOnly {
				highlights[start+loc[1]] = 0
				h.highlightRegion(highlights, start, false, lineNum, line[:loc[0]], curRegion, statesOnly)
			}
			h.highlightEmptyRegion(highlights, start+loc[1], canMatchEnd, lineNum, line[loc[1]:], statesOnly)
			return highlights
		}
		if !statesOnly {
			highlights[start+loc[1]] = curRegion.parent.group
			h.highlightRegion(highlights, start, false, lineNum, line[:loc[0]], curRegion, statesOnly)
		}
		h.highlightRegion(highlights, start+loc[1], canMatchEnd, lineNum, line[loc[1]:], curRegion.parent, statesOnly)
		return highlights
	}

	if len(line) == 0 || statesOnly {
		if canMatchEnd {
			h.lastRegion = curRegion
		}

		return highlights
	}

	firstLoc := []int{len(line), 0}

	var firstRegion *region
	for _, r := range curRegion.rules.regions {
		loc := findIndex(r.start, nil, line, start == 0, canMatchEnd)
		if loc != nil {
			if loc[0] < firstLoc[0] {
				firstLoc = loc
				firstRegion = r
			}
		}
	}
	if firstLoc[0] != len(line) {
		highlights[start+firstLoc[0]] = firstRegion.group
		h.highlightRegion(highlights, start, false, lineNum, line[:firstLoc[0]], curRegion, statesOnly)
		h.highlightRegion(highlights, start+firstLoc[1], canMatchEnd, lineNum, line[firstLoc[1]:], firstRegion, statesOnly)
		return highlights
	}

	fullHighlights := make([]Group, len([]rune(string(line))))
	for i := 0; i < len(fullHighlights); i++ {
		fullHighlights[i] = curRegion.group
	}

	for _, p := range curRegion.rules.patterns {
		matches := findAllIndex(p.regex, line, start == 0, canMatchEnd)
		for _, m := range matches {
			for i := m[0]; i < m[1]; i++ {
				fullHighlights[i] = p.group
			}
		}
	}
	for i, h := range fullHighlights {
		if i == 0 || h != fullHighlights[i-1] {
			// if _, ok := highlights[start+i]; !ok {
			highlights[start+i] = h
			// }
		}
	}

	if canMatchEnd {
		h.lastRegion = curRegion
	}

	return highlights
}

func (h *Highlighter) highlightEmptyRegion(highlights LineMatch, start int, canMatchEnd bool, lineNum int, line []rune, statesOnly bool) LineMatch {
	if len(line) == 0 {
		if canMatchEnd {
			h.lastRegion = nil
		}
		return highlights
	}

	firstLoc := []int{len(line), 0}
	var firstRegion *region
	for _, r := range h.Def.rules.regions {
		loc := findIndex(r.start, nil, line, start == 0, canMatchEnd)
		if loc != nil {
			if loc[0] < firstLoc[0] {
				firstLoc = loc
				firstRegion = r
			}
		}
	}
	if firstLoc[0] != len(line) {
		if !statesOnly {
			highlights[start+firstLoc[0]] = firstRegion.group
		}
		h.highlightEmptyRegion(highlights, start, false, lineNum, line[:firstLoc[0]], statesOnly)
		h.highlightRegion(highlights, start+firstLoc[1], canMatchEnd, lineNum, line[firstLoc[1]:], firstRegion, statesOnly)
		return highlights
	}

	if statesOnly {
		if canMatchEnd {
			h.lastRegion = nil
		}

		return highlights
	}

	fullHighlights := make([]Group, len(line))
	for _, p := range h.Def.rules.patterns {
		matches := findAllIndex(p.regex, line, start == 0, canMatchEnd)
		for _, m := range matches {
			for i := m[0]; i < m[1]; i++ {
				fullHighlights[i] = p.group
			}
		}
	}
	for i, h := range fullHighlights {
		if i == 0 || h != fullHighlights[i-1] {
			// if _, ok := highlights[start+i]; !ok {
			highlights[start+i] = h
			// }
		}
	}

	if canMatchEnd {
		h.lastRegion = nil
	}

	return highlights
}

// HighlightString syntax highlights a string
// Use this function for simple syntax highlighting and use the other functions for
// more advanced syntax highlighting. They are optimized for quick rehighlighting of the same
// text with minor changes made
func (h *Highlighter) HighlightString(input string) []LineMatch {
	lines := strings.Split(input, "\n")
	var lineMatches []LineMatch

	for i := 0; i < len(lines); i++ {
		line := []rune(lines[i])
		highlights := make(LineMatch)

		if i == 0 || h.lastRegion == nil {
			lineMatches = append(lineMatches, h.highlightEmptyRegion(highlights, 0, true, i, line, false))
		} else {
			lineMatches = append(lineMatches, h.highlightRegion(highlights, 0, true, i, line, h.lastRegion, false))
		}
	}

	return lineMatches
}

// HighlightStates correctly sets all states for the buffer
func (h *Highlighter) HighlightStates(input LineStates) {
	for i := 0; i < input.LinesNum(); i++ {
		line := []rune(input.Line(i))
		// highlights := make(LineMatch)

		if i == 0 || h.lastRegion == nil {
			h.highlightEmptyRegion(nil, 0, true, i, line, true)
		} else {
			h.highlightRegion(nil, 0, true, i, line, h.lastRegion, true)
		}

		curState := h.lastRegion

		input.SetState(i, curState)
	}
}

// HighlightMatches sets the matches for each line in between startline and endline
// It sets all other matches in the buffer to nil to conserve memory
// This assumes that all the states are set correctly
func (h *Highlighter) HighlightMatches(input LineStates, startline, endline int) {
	for i := startline; i < endline; i++ {
		if i >= input.LinesNum() {
			break
		}

		line := []rune(input.Line(i))
		highlights := make(LineMatch)

		var match LineMatch
		if i == 0 || input.State(i-1) == nil {
			match = h.highlightEmptyRegion(highlights, 0, true, i, line, false)
		} else {
			match = h.highlightRegion(highlights, 0, true, i, line, input.State(i-1), false)
		}

		input.SetMatch(i, match)
	}
}

// ReHighlightStates will scan down from `startline` and set the appropriate end of line state
// for each line until it comes across the same state in two consecutive lines
func (h *Highlighter) ReHighlightStates(input LineStates, startline int) {
	// lines := input.LineData()

	h.lastRegion = nil
	if startline > 0 {
		h.lastRegion = input.State(startline - 1)
	}
	for i := startline; i < input.LinesNum(); i++ {
		line := []rune(input.Line(i))
		// highlights := make(LineMatch)

		// var match LineMatch
		if i == 0 || h.lastRegion == nil {
			h.highlightEmptyRegion(nil, 0, true, i, line, true)
		} else {
			h.highlightRegion(nil, 0, true, i, line, h.lastRegion, true)
		}
		curState := h.lastRegion
		lastState := input.State(i)

		input.SetState(i, curState)

		if curState == lastState {
			break
		}
	}
}

// ReHighlightLine will rehighlight the state and match for a single line
func (h *Highlighter) ReHighlightLine(input LineStates, lineN int) {
	line := []rune(input.Line(lineN))
	highlights := make(LineMatch)

	h.lastRegion = nil
	if lineN > 0 {
		h.lastRegion = input.State(lineN - 1)
	}

	var match LineMatch
	if lineN == 0 || h.lastRegion == nil {
		match = h.highlightEmptyRegion(highlights, 0, true, lineN, line, false)
	} else {
		match = h.highlightRegion(highlights, 0, true, lineN, line, h.lastRegion, false)
	}
	curState := h.lastRegion

	input.SetMatch(lineN, match)
	input.SetState(lineN, curState)
}
<|endoftext|>


<commit_before>package cmd

import (
	"fmt"
	"time"

	"github.com/grafana/metrictank/cmd/mt-fakemetrics/metricbuilder"
	"github.com/grafana/metrictank/cmd/mt-fakemetrics/out"
	"github.com/grafana/metrictank/cmd/mt-fakemetrics/policy"
	"github.com/grafana/metrictank/schema"
	"github.com/raintank/worldping-api/pkg/log"
)

// examples (everything perOrg)
// num metrics - flush (s) - period (s) - speedup -> ratePerSPerOrg      -> ratePerFlushPerOrg
// mpo         -                                     mpo*speedup/period  -> mpo*speedup*flush/period
// mpo         - 0.1         1            1          mpo                    mpo/10           1/10 -> flush 1 of 10 fractions runDivided <-- default
// mpo         - 0.1         1            2          mpo*2                  mpo/5            1/5 -> flush 1 of 5 fractions   runDivided
// mpo         - 0.1         1            100        mpo*100                mpo*10           1/5 -> flush 1x                 runMultiplied
// mpo         - 1           1            2          mpo*2                  mpo*2            2x -> flush 2x                  runMultiplied
// mpo         - 2           1            1          mpo                    mpo*2            2x -> flush 2x                  runMultiplied
// mpo         - 2           1            2          mpo*2                  mpo*2            4x -> flush 4x                  runMultiplied

// dataFeed supports both realtime, as backfill, with speedup
// important:
// period in seconds
// flush  in ms
// offset in seconds
func dataFeed(out out.Out, orgs, mpo, period, flush, offset, speedup int, stopAtNow bool, builder metricbuilder.Builder, vp policy.ValuePolicy) {
	flushDur := time.Duration(flush) * time.Millisecond

	if mpo*speedup%period != 0 {
		panic("not a good fit. mpo*speedup must divide by period, to compute clean rate/s/org")
	}
	ratePerSPerOrg := mpo * speedup / period

	if mpo*speedup*flush%(1000*period) != 0 {
		panic("not a good fit. mpo*speedup*flush must divide by period, to compute clean rate/flush/org")
	}
	ratePerFlushPerOrg := ratePerSPerOrg * flush / 1000

	if addTags && len(customTags) > 0 {
		panic("cannot use regular-tags and custom-tags at the same time")
	}

	if numUniqueTags > 10 || numUniqueTags < 0 {
		panic(fmt.Sprintf("num-unique-tags must be a value between 0 and 10, you entered %d", numUniqueTags))
	}

	if numUniqueCustomTags > len(customTags) || numUniqueCustomTags < 0 {
		panic(fmt.Sprintf("num-unique-custom-tags must be a value between 0 and %d, you entered %d", len(customTags), numUniqueCustomTags))
	}

	ratePerS := ratePerSPerOrg * orgs
	ratePerFlush := ratePerFlushPerOrg * orgs

	tmpl := `params: %s, orgs=%d, mpo=%d, period=%d, flush=%d, offset=%d, speedup=%d, stopAtNow=%t
per org:         each %s, flushing %d metrics so rate of %d Hz. (%d total unique series)
times %4d orgs: each %s, flushing %d metrics so rate of %d Hz. (%d total unique series)
`
	fmt.Printf(tmpl, builder.Info(), orgs, mpo, period, flush, offset, speedup, stopAtNow,
		flushDur, ratePerFlush, ratePerS, orgs*mpo,
		orgs, flushDur, ratePerFlush, ratePerS, orgs*mpo)

	tick := time.NewTicker(flushDur)

	metrics := builder.Build(orgs, mpo, period)

	// set initial conditions
	mp := int64(period)
	// set start to now-offset because we add mp back every time we start a cycle going through metrics[o]
	ts := time.Now().Unix() - int64(offset) - mp
	startFrom := 0

	// huh what if we increment ts beyond the now ts?
	// this can only happen if we repeatedly loop, and bump ts each time
	// let's say we loop 5 times, so:
	// ratePerFlushPerOrg == 5 * mpo
	// then last ts = ts+4*period
	// (loops-1)*period < flush
	// (ceil(ratePerFlushPerOrg/mpo)-1)*period < flush
	// (ceil(mpo * speedup * flush /period /mpo)-1)*period < flush
	// (ceil(speedup * flush /period)-1)*period < flush
	// (ceil(speedup * flush - period ) < flush

	for nowT := range tick.C {
		now := nowT.Unix()
		var data []*schema.MetricData

		for o := 0; o < len(metrics); o++ {
			// for each org, we need to flush ratePerFlushPerOrg,
			// starting at wherever a previous flush (if any) left off.
			var m int
			for num := 0; num < ratePerFlushPerOrg; num++ {
				// note that ratePerFlushPerOrg may be any of >, =, < mpo
				// it all depends on what the user requested
				// we mainly need to ensure both cases properly bump the timestamp
				m = (startFrom + num) % mpo
				metricData := metrics[o][m]
				// every time we cycle through metrics[o], we bump timestamp
				if m == 0 {
					ts += mp
				}
				metricData.Time = ts
				metricData.Value = vp.Value(ts)

				data = append(data, &metricData)
			}
			// next metrics iteration should start where we left off... but..
			// it looks like this will affect the next org during the current iteration :?
			startFrom = (m + 1) % mpo
		}

		preFlush := time.Now()
		err := out.Flush(data)
		if err != nil {
			log.Error(0, err.Error())
		}
		flushDuration.Value(time.Since(preFlush))

		if ts >= now && stopAtNow {
			return
		}
	}
}
<commit_msg>fix datafeed loop starting to lag behind if ticks are missed<commit_after>package cmd

import (
	"fmt"
	"time"

	"github.com/grafana/metrictank/clock"
	"github.com/grafana/metrictank/cmd/mt-fakemetrics/metricbuilder"
	"github.com/grafana/metrictank/cmd/mt-fakemetrics/out"
	"github.com/grafana/metrictank/cmd/mt-fakemetrics/policy"
	"github.com/grafana/metrictank/schema"
	"github.com/raintank/worldping-api/pkg/log"
)

// examples (everything perOrg)
// num metrics - flush (s) - period (s) - speedup -> ratePerSPerOrg      -> ratePerFlushPerOrg
// mpo         -                                     mpo*speedup/period  -> mpo*speedup*flush/period
// mpo         - 0.1         1            1          mpo                    mpo/10           1/10 -> flush 1 of 10 fractions runDivided <-- default
// mpo         - 0.1         1            2          mpo*2                  mpo/5            1/5 -> flush 1 of 5 fractions   runDivided
// mpo         - 0.1         1            100        mpo*100                mpo*10           1/5 -> flush 1x                 runMultiplied
// mpo         - 1           1            2          mpo*2                  mpo*2            2x -> flush 2x                  runMultiplied
// mpo         - 2           1            1          mpo                    mpo*2            2x -> flush 2x                  runMultiplied
// mpo         - 2           1            2          mpo*2                  mpo*2            4x -> flush 4x                  runMultiplied

// dataFeed supports both realtime, as backfill, with speedup
// important:
// period in seconds
// flush  in ms
// offset in seconds
func dataFeed(out out.Out, orgs, mpo, period, flush, offset, speedup int, stopAtNow bool, builder metricbuilder.Builder, vp policy.ValuePolicy) {
	flushDur := time.Duration(flush) * time.Millisecond

	if mpo*speedup%period != 0 {
		panic("not a good fit. mpo*speedup must divide by period, to compute clean rate/s/org")
	}
	ratePerSPerOrg := mpo * speedup / period

	if mpo*speedup*flush%(1000*period) != 0 {
		panic("not a good fit. mpo*speedup*flush must divide by period, to compute clean rate/flush/org")
	}
	ratePerFlushPerOrg := ratePerSPerOrg * flush / 1000

	if addTags && len(customTags) > 0 {
		panic("cannot use regular-tags and custom-tags at the same time")
	}

	if numUniqueTags > 10 || numUniqueTags < 0 {
		panic(fmt.Sprintf("num-unique-tags must be a value between 0 and 10, you entered %d", numUniqueTags))
	}

	if numUniqueCustomTags > len(customTags) || numUniqueCustomTags < 0 {
		panic(fmt.Sprintf("num-unique-custom-tags must be a value between 0 and %d, you entered %d", len(customTags), numUniqueCustomTags))
	}

	ratePerS := ratePerSPerOrg * orgs
	ratePerFlush := ratePerFlushPerOrg * orgs

	tmpl := `params: %s, orgs=%d, mpo=%d, period=%d, flush=%d, offset=%d, speedup=%d, stopAtNow=%t
per org:         each %s, flushing %d metrics so rate of %d Hz. (%d total unique series)
times %4d orgs: each %s, flushing %d metrics so rate of %d Hz. (%d total unique series)
`
	fmt.Printf(tmpl, builder.Info(), orgs, mpo, period, flush, offset, speedup, stopAtNow,
		flushDur, ratePerFlush, ratePerS, orgs*mpo,
		orgs, flushDur, ratePerFlush, ratePerS, orgs*mpo)

	metrics := builder.Build(orgs, mpo, period)

	// set initial conditions
	mp := int64(period)
	// set start to now-offset because we add mp back every time we start a cycle going through metrics[o]
	ts := time.Now().Unix() - int64(offset) - mp
	startFrom := 0

	// huh what if we increment ts beyond the now ts?
	// this can only happen if we repeatedly loop, and bump ts each time
	// let's say we loop 5 times, so:
	// ratePerFlushPerOrg == 5 * mpo
	// then last ts = ts+4*period
	// (loops-1)*period < flush
	// (ceil(ratePerFlushPerOrg/mpo)-1)*period < flush
	// (ceil(mpo * speedup * flush /period /mpo)-1)*period < flush
	// (ceil(speedup * flush /period)-1)*period < flush
	// (ceil(speedup * flush - period ) < flush

	for nowT := range clock.AlignedTickLossless(flushDur) {
		now := nowT.Unix()
		var data []*schema.MetricData

		for o := 0; o < len(metrics); o++ {
			// for each org, we need to flush ratePerFlushPerOrg,
			// starting at wherever a previous flush (if any) left off.
			var m int
			for num := 0; num < ratePerFlushPerOrg; num++ {
				// note that ratePerFlushPerOrg may be any of >, =, < mpo
				// it all depends on what the user requested
				// we mainly need to ensure both cases properly bump the timestamp
				m = (startFrom + num) % mpo
				metricData := metrics[o][m]
				// every time we cycle through metrics[o], we bump timestamp
				if m == 0 {
					ts += mp
				}
				metricData.Time = ts
				metricData.Value = vp.Value(ts)

				data = append(data, &metricData)
			}
			// next metrics iteration should start where we left off... but..
			// it looks like this will affect the next org during the current iteration :?
			startFrom = (m + 1) % mpo
		}

		preFlush := time.Now()
		err := out.Flush(data)
		if err != nil {
			log.Error(0, err.Error())
		}
		flushDuration.Value(time.Since(preFlush))

		if ts >= now && stopAtNow {
			return
		}
	}
}
<|endoftext|>


<commit_before>package cli

import (
	"fmt"
	"github.com/qiniu/api/auth/digest"
	fio "github.com/qiniu/api/io"
	rio "github.com/qiniu/api/resumable/io"
	"github.com/qiniu/api/rs"
	"github.com/qiniu/log"
	"os"
	"sort"
)

func FormPut(cmd string, params ...string) {
	if len(params) == 3 || len(params) == 4 {
		bucket := params[0]
		key := params[1]
		localFile := params[2]
		mimeType := ""
		if len(params) == 4 {
			mimeType = params[3]
		}
		accountS.Get()
		mac := digest.Mac{accountS.AccessKey, []byte(accountS.SecretKey)}
		policy := rs.PutPolicy{}
		policy.Scope = bucket
		putExtra := fio.PutExtra{}
		if mimeType != "" {
			putExtra.MimeType = mimeType
		}
		uptoken := policy.Token(&mac)
		putRet := fio.PutRet{}
		err := fio.PutFile(nil, &putRet, uptoken, key, localFile, &putExtra)
		if err != nil {
			log.Error("Put file error", err)
		} else {
			fmt.Println("Put file", localFile, "=>", bucket, ":", putRet.Key, "(", putRet.Hash, ")", "success!")
		}
	} else {
		Help(cmd)
	}
}

func ResumablePut(cmd string, params ...string) {
	if len(params) == 3 || len(params) == 4 {
		bucket := params[0]
		key := params[1]
		localFile := params[2]
		mimeType := ""
		if len(params) == 4 {
			mimeType = params[3]
		}
		accountS.Get()
		mac := digest.Mac{accountS.AccessKey, []byte(accountS.SecretKey)}
		policy := rs.PutPolicy{}
		policy.Scope = bucket
		putExtra := rio.PutExtra{}
		if mimeType != "" {
			putExtra.MimeType = mimeType
		}
		progressHandler := ProgressHandler{
			BlockIndices:    make([]int, 0),
			BlockProgresses: make(map[int]float32),
		}
		putExtra.Notify = progressHandler.Notify
		putExtra.NotifyErr = progressHandler.NotifyErr
		uptoken := policy.Token(&mac)
		putRet := rio.PutRet{}
		err := rio.PutFile(nil, &putRet, uptoken, key, localFile, &putExtra)
		if err != nil {
			log.Error("Put file error", err)
		} else {
			fmt.Println("\r\nPut file", localFile, "=>", bucket, ":", putRet.Key, "(", putRet.Hash, ")", "success!")
		}
	} else {
		Help(cmd)
	}
}

type ProgressHandler struct {
	BlockIndices    []int
	BlockProgresses map[int]float32
}

func (this *ProgressHandler) Notify(blkIdx int, blkSize int, ret *rio.BlkputRet) {
	offset := ret.Offset
	perent := float32(offset) * 100 / float32(blkSize)
	if _, ok := this.BlockProgresses[blkIdx]; !ok {
		this.BlockIndices = append(this.BlockIndices, blkIdx)
		sort.Ints(this.BlockIndices)
	}
	this.BlockProgresses[blkIdx] = perent
	output := fmt.Sprintf("\r")
	for _, blockIndex := range this.BlockIndices {
		blockProgress := this.BlockProgresses[blockIndex]
		output += fmt.Sprintf("Block %d=>%.2f%%\t", blockIndex+1, blockProgress)
	}
	fmt.Print(output)
	os.Stdout.Sync()
}
func (this *ProgressHandler) NotifyErr(blkIdx int, blkSize int, err error) {

}
<commit_msg>A little fix.<commit_after>package cli

import (
	"fmt"
	"github.com/qiniu/api/auth/digest"
	fio "github.com/qiniu/api/io"
	rio "github.com/qiniu/api/resumable/io"
	"github.com/qiniu/api/rs"
	"github.com/qiniu/log"
	"os"
	"sort"
)

func FormPut(cmd string, params ...string) {
	if len(params) == 3 || len(params) == 4 {
		bucket := params[0]
		key := params[1]
		localFile := params[2]
		mimeType := ""
		if len(params) == 4 {
			mimeType = params[3]
		}
		accountS.Get()
		mac := digest.Mac{accountS.AccessKey, []byte(accountS.SecretKey)}
		policy := rs.PutPolicy{}
		policy.Scope = bucket
		putExtra := fio.PutExtra{}
		if mimeType != "" {
			putExtra.MimeType = mimeType
		}
		uptoken := policy.Token(&mac)
		putRet := fio.PutRet{}
		err := fio.PutFile(nil, &putRet, uptoken, key, localFile, &putExtra)
		if err != nil {
			log.Error("Put file error", err)
		} else {
			fmt.Println("Put file", localFile, "=>", bucket, ":", putRet.Key, "(", putRet.Hash, ")", "success!")
		}
	} else {
		Help(cmd)
	}
}

func ResumablePut(cmd string, params ...string) {
	if len(params) == 3 || len(params) == 4 {
		bucket := params[0]
		key := params[1]
		localFile := params[2]
		mimeType := ""
		if len(params) == 4 {
			mimeType = params[3]
		}
		accountS.Get()
		mac := digest.Mac{accountS.AccessKey, []byte(accountS.SecretKey)}
		policy := rs.PutPolicy{}
		policy.Scope = bucket
		putExtra := rio.PutExtra{}
		if mimeType != "" {
			putExtra.MimeType = mimeType
		}
		progressHandler := ProgressHandler{
			BlockIndices:    make([]int, 0),
			BlockProgresses: make(map[int]float32),
		}
		putExtra.Notify = progressHandler.Notify
		putExtra.NotifyErr = progressHandler.NotifyErr
		uptoken := policy.Token(&mac)
		putRet := rio.PutRet{}
		err := rio.PutFile(nil, &putRet, uptoken, key, localFile, &putExtra)
		if err != nil {
			log.Error("Put file error", err)
		} else {
			fmt.Println("\r\nPut file", localFile, "=>", bucket, ":", putRet.Key, "(", putRet.Hash, ")", "success!")
		}
	} else {
		Help(cmd)
	}
}

type ProgressHandler struct {
	BlockIndices    []int
	BlockProgresses map[int]float32
}

func (this *ProgressHandler) Notify(blkIdx int, blkSize int, ret *rio.BlkputRet) {
	offset := ret.Offset
	perent := float32(offset) * 100 / float32(blkSize)
	if _, ok := this.BlockProgresses[blkIdx]; !ok {
		this.BlockIndices = append(this.BlockIndices, blkIdx)
		sort.Ints(this.BlockIndices)
	}
	this.BlockProgresses[blkIdx] = perent
	output := fmt.Sprintf("\r")
	for _, blockIndex := range this.BlockIndices {
		blockProgress := this.BlockProgresses[blockIndex]
		output += fmt.Sprintf("[Block %d=>%.2f%%], ", blockIndex+1, blockProgress)
	}
	fmt.Print(output)
	os.Stdout.Sync()
}
func (this *ProgressHandler) NotifyErr(blkIdx int, blkSize int, err error) {

}
<|endoftext|>


<commit_before>// Package jantar is a lightweight mvc web framework with emphasis on security written in golang.
//
// It has been largely inspired by Martini(https://github.com/codegangsta/martini) but prefers performance over
// syntactic sugar and aims to provide crucial security settings and features right out of the box.
package jantar

import (
	"crypto/tls"
	"fmt"
	"github.com/tsurai/jantar/context"
	"net"
	"net/http"
	"os"
	"os/signal"
	"path/filepath"
	"strings"
	"sync"
	"time"
)

// Log is a package global Log instance using the prefix "[jantar] " on outputs
var (
	Log *JLogger
)

// Jantar is the top level application type
type Jantar struct {
	closing    bool
	wg         sync.WaitGroup
	listener   net.Listener
	config     *Config
	middleware []IMiddleware
	tm         *TemplateManager
	router     *router
}

// TLSConfig can be given to Jantar to enable tls support
type TLSConfig struct {
	CertFile string
	KeyFile  string
	CertPem  []byte
	KeyPem   []byte
	cert     tls.Certificate
}

// Config is the main configuration struct for jantar
type Config struct {
	Hostname string
	Port     int
	TLS      *TLSConfig
}

// New creates a new Jantar instance ready to listen on a given hostname and port.
// Choosing a port small than 1 will cause Jantar to use the standard ports.
func New(config *Config) *Jantar {
	// create Log
	Log = NewJLogger(os.Stdout, "", LogLevelInfo)

	if config == nil {
		Log.Fatal("no config given")
	}

	j := &Jantar{
		config:     config,
		tm:         newTemplateManager("views"),
		router:     newRouter(),
		middleware: nil,
		closing:    false,
	}

	if j.config.Port < 1 {
		if j.config.TLS == nil {
			j.config.Port = 80
		} else {
			j.config.Port = 443
		}
	}

	// load default middleware
	j.AddMiddleware(&csrf{})

	// load ssl certificate
	if config.TLS != nil {
		if err := loadTLSCertificate(config.TLS); err != nil {
			Log.Fatald(JLData{"error": err}, "failed to load x509 certificate")
		}
	}

	setModule(ModuleTemplateManager, j.tm)
	setModule(ModuleRouter, j.router)

	return j
}

// AddMiddleware adds a given middleware to the current middleware list. Middlewares are executed
// once for every request before the actual route handler is called
func (j *Jantar) AddMiddleware(mware IMiddleware) {
	if len(j.middleware) > 0 {
		j.middleware[len(j.middleware)-1].setNext(&mware)
	}
	j.middleware = append(j.middleware, mware)
}

func (j *Jantar) initMiddleware() {
	for _, mw := range j.middleware {
		mw.Initialize()
	}
}

func (j *Jantar) cleanupMiddleware() {
	for _, mw := range j.middleware {
		mw.Cleanup()
	}
}

func (j *Jantar) callMiddleware(respw http.ResponseWriter, req *http.Request) bool {
	for _, mw := range j.middleware {
		if !mw.Call(respw, req) {
			return false
		}

		if mw.doesYield() {
			break
		}
	}
	return true
}

// AddRoute adds a route with given method, pattern and handler to the Router
func (j *Jantar) AddRoute(method string, pattern string, handler interface{}) *route {
	return j.router.addRoute(method, pattern, handler)
}

func (j *Jantar) listenForSignals() {
	sigChan := make(chan os.Signal, 1)

	signal.Notify(sigChan, os.Interrupt, os.Kill)

	s := <-sigChan
	if s == os.Kill {
		Log.Fatal("Got SIGKILL")
	}

	j.Stop()
}

func serveStatic(respw http.ResponseWriter, req *http.Request) bool {
	if file, stat := getFile("/", "views/.static", req.URL.Path); file != nil {
		http.ServeContent(respw, req, req.URL.Path, stat.ModTime(), file)
		file.Close()

		return true
	}

	return false
}

func servePublic(respw http.ResponseWriter, req *http.Request) bool {
	if strings.HasPrefix(req.URL.Path, "/public/") {
		if file, stat := getFile("/public/", "public", req.URL.Path); file != nil {
			http.ServeContent(respw, req, req.URL.Path, stat.ModTime(), file)
			file.Close()

			return true
		}

		Log.Errord(JLData{"file": req.URL.Path}, "failed to serve public file")
		http.NotFound(respw, req)

		return true
	}
	return false
}

func getFile(prefix string, root string, request string) (http.File, os.FileInfo) {
	var file http.File
	var stat os.FileInfo
	var publicpath, publicfilepath string
	var err error

	fname := request[len(prefix):]

	if !strings.HasPrefix(fname, ".") {
		if publicpath, err = filepath.Abs(root); err == nil {
			if publicfilepath, err = filepath.Abs(root + "/" + fname); err == nil {
				if strings.HasPrefix(publicfilepath, publicpath) {
					if file, err = http.Dir(root).Open(fname); err == nil {
						if stat, err = file.Stat(); err == nil {
							if !stat.IsDir() {
								return file, stat
							}
						}
					}
				}
			}
		}
	}

	return nil, nil
}

func (j *Jantar) listenAndServe(addr string, handler http.Handler) error {
	var err error

	if addr == "" {
		addr = ":http"
	}

	if j.config.TLS != nil {
		// configure tls with secure settings
		tlsConfig.Certificates = []tls.Certificate{j.config.TLS.cert}
		j.listener, err = tls.Listen("tcp", addr, tlsConfig)

		// listen redirect port 80 to 443 if using the standard port
		if j.config.Port == 443 {
			go http.ListenAndServe(fmt.Sprintf("%s:%d", j.config.Hostname, 80), http.HandlerFunc(
				func(respw http.ResponseWriter, req *http.Request) {
					http.Redirect(respw, req, "https://"+j.config.Hostname+req.RequestURI, 301)
				}))
		}
	} else {
		j.listener, err = net.Listen("tcp", addr)
	}

	if err != nil {
		return err
	}

	server := &http.Server{Addr: addr, Handler: handler}
	if err = server.Serve(j.listener); !j.closing {
		return err
	}

	return nil
}

// ServeTemplate is a simple http.HandlerFunc wrapper serving a template. It is intended for cases in which a Controller would be too much.
func (j *Jantar) ServeTemplate(name string) func(http.ResponseWriter, *http.Request) {
	return func(respw http.ResponseWriter, req *http.Request) {
		j.tm.RenderTemplate(respw, req, name, make(map[string]interface{}))
	}
}

// ServeHTTP implements the http.Handler interface
func (j *Jantar) ServeHTTP(respw http.ResponseWriter, req *http.Request) {
	j.wg.Add(1)

	t0 := time.Now()

	if method := req.FormValue("_method"); method != "" {
		req.Method = method
	}

	Log.Infof("%s %s", req.Method, req.URL.Path)

	// set security header
	respw.Header().Set("Strict-Transport-Security", "max-age=31536000;includeSubDomains")
	respw.Header().Set("X-Frame-Options", "sameorigin")
	respw.Header().Set("X-XSS-Protection", "1;mode=block")
	respw.Header().Set("X-Content-Type-Options", "nosniff")

	if !servePublic(respw, req) {
		if route := j.router.searchRoute(req); route != nil {
			context.Set(req, "renderArgs", make(map[string]interface{}), true)
			if j.callMiddleware(respw, req) {
				route.handler(respw, req)
			}
		} else if !servePublic(respw, req) {
			Log.Info("404 page not found")
			http.NotFound(respw, req)
		}
	}

	context.ClearData(req)
	Log.Infof("completed in %v", time.Since(t0))

	j.wg.Done()
}

// Stop closes the listener and stops the server when all pending requests have been finished
func (j *Jantar) Stop() {
	j.closing = true

	// stop listening for new connections
	j.listener.Close()

	// wait until all pending requests have been finished
	j.wg.Wait()

	j.cleanupMiddleware()
}

// Run starts the http server and listens on the hostname and port given to New
func (j *Jantar) Run() {
	j.initMiddleware()

	if err := j.tm.loadTemplates(); err != nil {
		Log.Fataldf(JLData{"error": err}, "failed to load templates")
	}

	go j.listenForSignals()

	Log.Infod(JLData{"hostname": j.config.Hostname, "port": j.config.Port, "TLS": j.config.TLS != nil}, "starting server & listening")

	if err := j.listenAndServe(fmt.Sprintf("%s:%d", j.config.Hostname, j.config.Port), j); err != nil {
		Log.Fatal(err)
	}

	Log.Info("stopping server")
}
<commit_msg>fixed typo<commit_after>// Package jantar is a lightweight mvc web framework with emphasis on security written in golang.
//
// It has been largely inspired by Martini(https://github.com/codegangsta/martini) but prefers performance over
// syntactic sugar and aims to provide crucial security settings and features right out of the box.
package jantar

import (
	"crypto/tls"
	"fmt"
	"github.com/tsurai/jantar/context"
	"net"
	"net/http"
	"os"
	"os/signal"
	"path/filepath"
	"strings"
	"sync"
	"time"
)

// Log is a package global Log instance using the prefix "[jantar] " on outputs
var (
	Log *JLogger
)

// Jantar is the top level application type
type Jantar struct {
	closing    bool
	wg         sync.WaitGroup
	listener   net.Listener
	config     *Config
	middleware []IMiddleware
	tm         *TemplateManager
	router     *router
}

// TLSConfig can be given to Jantar to enable tls support
type TLSConfig struct {
	CertFile string
	KeyFile  string
	CertPem  []byte
	KeyPem   []byte
	cert     tls.Certificate
}

// Config is the main configuration struct for jantar
type Config struct {
	Hostname string
	Port     int
	TLS      *TLSConfig
}

// New creates a new Jantar instance ready to listen on a given hostname and port.
// Choosing a port small than 1 will cause Jantar to use the standard ports.
func New(config *Config) *Jantar {
	// create Log
	Log = NewJLogger(os.Stdout, "", LogLevelInfo)

	if config == nil {
		Log.Fatal("no config given")
	}

	j := &Jantar{
		config:     config,
		tm:         newTemplateManager("views"),
		router:     newRouter(),
		middleware: nil,
		closing:    false,
	}

	if j.config.Port < 1 {
		if j.config.TLS == nil {
			j.config.Port = 80
		} else {
			j.config.Port = 443
		}
	}

	// load default middleware
	j.AddMiddleware(&csrf{})

	// load ssl certificate
	if config.TLS != nil {
		if err := loadTLSCertificate(config.TLS); err != nil {
			Log.Fatald(JLData{"error": err}, "failed to load x509 certificate")
		}
	}

	setModule(ModuleTemplateManager, j.tm)
	setModule(ModuleRouter, j.router)

	return j
}

// AddMiddleware adds a given middleware to the current middleware list. Middlewares are executed
// once for every request before the actual route handler is called
func (j *Jantar) AddMiddleware(mware IMiddleware) {
	if len(j.middleware) > 0 {
		j.middleware[len(j.middleware)-1].setNext(&mware)
	}
	j.middleware = append(j.middleware, mware)
}

func (j *Jantar) initMiddleware() {
	for _, mw := range j.middleware {
		mw.Initialize()
	}
}

func (j *Jantar) cleanupMiddleware() {
	for _, mw := range j.middleware {
		mw.Cleanup()
	}
}

func (j *Jantar) callMiddleware(respw http.ResponseWriter, req *http.Request) bool {
	for _, mw := range j.middleware {
		if !mw.Call(respw, req) {
			return false
		}

		if mw.doesYield() {
			break
		}
	}
	return true
}

// AddRoute adds a route with given method, pattern and handler to the Router
func (j *Jantar) AddRoute(method string, pattern string, handler interface{}) *route {
	return j.router.addRoute(method, pattern, handler)
}

func (j *Jantar) listenForSignals() {
	sigChan := make(chan os.Signal, 1)

	signal.Notify(sigChan, os.Interrupt, os.Kill)

	s := <-sigChan
	if s == os.Kill {
		Log.Fatal("Got SIGKILL")
	}

	j.Stop()
}

func serveStatic(respw http.ResponseWriter, req *http.Request) bool {
	if file, stat := getFile("/", "views/.static", req.URL.Path); file != nil {
		http.ServeContent(respw, req, req.URL.Path, stat.ModTime(), file)
		file.Close()

		return true
	}

	return false
}

func servePublic(respw http.ResponseWriter, req *http.Request) bool {
	if strings.HasPrefix(req.URL.Path, "/public/") {
		if file, stat := getFile("/public/", "public", req.URL.Path); file != nil {
			http.ServeContent(respw, req, req.URL.Path, stat.ModTime(), file)
			file.Close()

			return true
		}

		Log.Errord(JLData{"file": req.URL.Path}, "failed to serve public file")
		http.NotFound(respw, req)

		return true
	}
	return false
}

func getFile(prefix string, root string, request string) (http.File, os.FileInfo) {
	var file http.File
	var stat os.FileInfo
	var publicpath, publicfilepath string
	var err error

	fname := request[len(prefix):]

	if !strings.HasPrefix(fname, ".") {
		if publicpath, err = filepath.Abs(root); err == nil {
			if publicfilepath, err = filepath.Abs(root + "/" + fname); err == nil {
				if strings.HasPrefix(publicfilepath, publicpath) {
					if file, err = http.Dir(root).Open(fname); err == nil {
						if stat, err = file.Stat(); err == nil {
							if !stat.IsDir() {
								return file, stat
							}
						}
					}
				}
			}
		}
	}

	return nil, nil
}

func (j *Jantar) listenAndServe(addr string, handler http.Handler) error {
	var err error

	if addr == "" {
		addr = ":http"
	}

	if j.config.TLS != nil {
		// configure tls with secure settings
		tlsConfig.Certificates = []tls.Certificate{j.config.TLS.cert}
		j.listener, err = tls.Listen("tcp", addr, tlsConfig)

		// listen redirect port 80 to 443 if using the standard port
		if j.config.Port == 443 {
			go http.ListenAndServe(fmt.Sprintf("%s:%d", j.config.Hostname, 80), http.HandlerFunc(
				func(respw http.ResponseWriter, req *http.Request) {
					http.Redirect(respw, req, "https://"+j.config.Hostname+req.RequestURI, 301)
				}))
		}
	} else {
		j.listener, err = net.Listen("tcp", addr)
	}

	if err != nil {
		return err
	}

	server := &http.Server{Addr: addr, Handler: handler}
	if err = server.Serve(j.listener); !j.closing {
		return err
	}

	return nil
}

// ServeTemplate is a simple http.HandlerFunc wrapper serving a template. It is intended for cases in which a Controller would be too much.
func (j *Jantar) ServeTemplate(name string) func(http.ResponseWriter, *http.Request) {
	return func(respw http.ResponseWriter, req *http.Request) {
		j.tm.RenderTemplate(respw, req, name, make(map[string]interface{}))
	}
}

// ServeHTTP implements the http.Handler interface
func (j *Jantar) ServeHTTP(respw http.ResponseWriter, req *http.Request) {
	j.wg.Add(1)

	t0 := time.Now()

	if method := req.FormValue("_method"); method != "" {
		req.Method = method
	}

	Log.Infof("%s %s", req.Method, req.URL.Path)

	// set security header
	respw.Header().Set("Strict-Transport-Security", "max-age=31536000;includeSubDomains")
	respw.Header().Set("X-Frame-Options", "sameorigin")
	respw.Header().Set("X-XSS-Protection", "1;mode=block")
	respw.Header().Set("X-Content-Type-Options", "nosniff")

	if !servePublic(respw, req) {
		if route := j.router.searchRoute(req); route != nil {
			context.Set(req, "renderArgs", make(map[string]interface{}), true)
			if j.callMiddleware(respw, req) {
				route.handler(respw, req)
			}
		} else if !serveStatic(respw, req) {
			Log.Info("404 page not found")
			http.NotFound(respw, req)
		}
	}

	context.ClearData(req)
	Log.Infof("completed in %v", time.Since(t0))

	j.wg.Done()
}

// Stop closes the listener and stops the server when all pending requests have been finished
func (j *Jantar) Stop() {
	j.closing = true

	// stop listening for new connections
	j.listener.Close()

	// wait until all pending requests have been finished
	j.wg.Wait()

	j.cleanupMiddleware()
}

// Run starts the http server and listens on the hostname and port given to New
func (j *Jantar) Run() {
	j.initMiddleware()

	if err := j.tm.loadTemplates(); err != nil {
		Log.Fataldf(JLData{"error": err}, "failed to load templates")
	}

	go j.listenForSignals()

	Log.Infod(JLData{"hostname": j.config.Hostname, "port": j.config.Port, "TLS": j.config.TLS != nil}, "starting server & listening")

	if err := j.listenAndServe(fmt.Sprintf("%s:%d", j.config.Hostname, j.config.Port), j); err != nil {
		Log.Fatal(err)
	}

	Log.Info("stopping server")
}
<|endoftext|>