<commit_before>// Copyright © 2017 The Kubicorn 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 resources

import (
	"encoding/base64"
	"fmt"
	"strconv"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/service/autoscaling"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/kris-nova/kubicorn/apis/cluster"
	"github.com/kris-nova/kubicorn/cloud"
	"github.com/kris-nova/kubicorn/cutil/compare"
	"github.com/kris-nova/kubicorn/cutil/defaults"
	"github.com/kris-nova/kubicorn/cutil/logger"
	"github.com/kris-nova/kubicorn/cutil/script"
)

var _ cloud.Resource = &Lc{}

type Lc struct {
	Shared
	InstanceType     string
	Image            string
	SpotPrice        string
	InstanceProfile  string
	ServerPool       *cluster.ServerPool
	BootstrapScripts []string
}

const (
	MasterIPAttempts               = 40
	MasterIPSleepSecondsPerAttempt = 3
)

func (r *Lc) Actual(immutable *cluster.Cluster) (*cluster.Cluster, cloud.Resource, error) {
	logger.Debug("lc.Actual")
	newResource := &Lc{
		Shared: Shared{
			Name: r.Name,
			Tags: make(map[string]string),
		},
	}
	if r.ServerPool.Identifier != "" {
		lcInput := &autoscaling.DescribeLaunchConfigurationsInput{
			LaunchConfigurationNames: []*string{&r.ServerPool.Identifier},
		}
		lcOutput, err := Sdk.ASG.DescribeLaunchConfigurations(lcInput)
		if err != nil {
			return nil, nil, err
		}
		llc := len(lcOutput.LaunchConfigurations)
		if llc != 1 {
			return nil, nil, fmt.Errorf("Found [%d] Launch Configurations for ID [%s]", llc, r.ServerPool.Identifier)
		}
		lc := lcOutput.LaunchConfigurations[0]
		newResource.Image = *lc.ImageId
		if lc.SpotPrice != nil {
			newResource.SpotPrice = *lc.SpotPrice
		}
		newResource.Identifier = *lc.LaunchConfigurationName
		if lc.IamInstanceProfile != nil {
			newResource.InstanceProfile = *lc.IamInstanceProfile
		}
		newResource.Tags = map[string]string{
			"Name":              r.Name,
			"KubernetesCluster": immutable.Name,
		}
	} else {
		newResource.Image = r.ServerPool.Image
		newResource.InstanceType = r.ServerPool.Size
		if r.ServerPool.Type == cluster.ServerPoolTypeNode {
			newResource.SpotPrice = r.ServerPool.AwsConfiguration.SpotPrice
		}
	}
	newResource.BootstrapScripts = r.ServerPool.BootstrapScripts

	newCluster := r.immutableRender(newResource, immutable)
	return newCluster, newResource, nil
}

func (r *Lc) Expected(immutable *cluster.Cluster) (*cluster.Cluster, cloud.Resource, error) {
	logger.Debug("asg.Expected")
	newResource := &Lc{
		Shared: Shared{
			Tags: map[string]string{
				"Name":              r.Name,
				"KubernetesCluster": immutable.Name,
			},
			Identifier: r.ServerPool.Identifier,
			Name:       r.Name,
		},
		InstanceType:     r.ServerPool.Size,
		Image:            r.ServerPool.Image,
		BootstrapScripts: r.ServerPool.BootstrapScripts,
	}
	if r.ServerPool.InstanceProfile != nil {
		newResource.InstanceProfile = r.ServerPool.InstanceProfile.Name
	}
	if r.ServerPool.Type == cluster.ServerPoolTypeNode {
		newResource.SpotPrice = r.ServerPool.AwsConfiguration.SpotPrice
	}
	newCluster := r.immutableRender(newResource, immutable)
	return newCluster, newResource, nil
}

func (r *Lc) Apply(actual, expected cloud.Resource, immutable *cluster.Cluster) (*cluster.Cluster, cloud.Resource, error) {
	logger.Debug("lc.Apply")
	applyResource := expected.(*Lc)
	isEqual, err := compare.IsEqual(actual.(*Lc), expected.(*Lc))
	if err != nil {
		return nil, nil, err
	}
	if isEqual {
		return immutable, applyResource, nil
	}
	logger.Debug("Actual: %#v", actual)
	logger.Debug("Expectd: %#v", expected)
	var sgs []*string
	found := false
	for _, serverPool := range immutable.ServerPools {
		if serverPool.Name == expected.(*Lc).Name {
			for _, firewall := range serverPool.Firewalls {
				sgs = append(sgs, &firewall.Identifier)
			}
			found = true
		}
	}
	if !found {
		return nil, nil, fmt.Errorf("Unable to lookup serverpool for Launch Configuration %s", r.Name)
	}

	// --- Hack in here for master IP
	privip := ""
	pubip := ""
	if strings.Contains(r.ServerPool.Name, "node") {
		found := false
		logger.Debug("Tag query: [%s] %s", "Name", fmt.Sprintf("%s.master", immutable.Name))
		logger.Debug("Tag query: [%s] %s", "KubernetesCluster", immutable.Name)
		for i := 0; i < MasterIPAttempts; i++ {
			logger.Debug("Attempting to lookup master IP for node registration..")
			input := &ec2.DescribeInstancesInput{
				Filters: []*ec2.Filter{
					{
						Name:   S("tag:Name"),
						Values: []*string{S(fmt.Sprintf("%s.master", immutable.Name))},
					},
					{
						Name:   S("tag:KubernetesCluster"),
						Values: []*string{S(immutable.Name)},
					},
				},
			}
			output, err := Sdk.Ec2.DescribeInstances(input)
			if err != nil {
				return nil, nil, err
			}
			lr := len(output.Reservations)
			if lr == 0 {
				logger.Debug("Found %d Reservations, hanging ", lr)
				time.Sleep(time.Duration(MasterIPSleepSecondsPerAttempt) * time.Second)
				continue
			}
			for _, reservation := range output.Reservations {
				for _, instance := range reservation.Instances {
					if instance.PublicIpAddress != nil {
						privip = *instance.PrivateIpAddress
						pubip = *instance.PublicIpAddress
						immutable.Values.ItemMap["INJECTEDMASTER"] = fmt.Sprintf("%s:%s", privip, immutable.KubernetesAPI.Port)
						immutable.KubernetesAPI.Endpoint = pubip
						logger.Info("Found public IP for master: [%s]", pubip)
						found = true
					}
				}
			}
			if found == true {
				break
			}
			time.Sleep(time.Duration(MasterIPSleepSecondsPerAttempt) * time.Second)
		}
		if !found {
			return nil, nil, fmt.Errorf("Unable to find Master IP")
		}
	}

	immutable.Values.ItemMap["INJECTEDPORT"] = immutable.KubernetesAPI.Port

	newResource := &Lc{}
	userData, err := script.BuildBootstrapScript(r.ServerPool.BootstrapScripts, immutable)
	if err != nil {
		return nil, nil, err
	}

	b64data := base64.StdEncoding.EncodeToString(userData)
	lcInput := &autoscaling.CreateLaunchConfigurationInput{
		AssociatePublicIpAddress: B(true),
		LaunchConfigurationName:  &expected.(*Lc).Name,
		ImageId:                  &expected.(*Lc).Image,
		InstanceType:             &expected.(*Lc).InstanceType,
		KeyName:                  &immutable.SSH.Identifier,
		SecurityGroups:           sgs,
		UserData:                 &b64data,
	}
	if expected.(*Lc).InstanceProfile != "" {
		lcInput.IamInstanceProfile = &expected.(*Lc).InstanceProfile
	}
	spotPrice, err := strconv.ParseFloat(*&expected.(*Lc).SpotPrice, 64)
	if *&expected.(*Lc).InstanceType != cluster.ServerPoolTypeMaster && err == nil && spotPrice > 0 {
		lcInput.SpotPrice = &expected.(*Lc).SpotPrice
	}
	//Make it repeatable due to InstanceProfile
	for i := 0; i < 10; i++ {
		_, err = Sdk.ASG.CreateLaunchConfiguration(lcInput)
		if err != nil {
			if awserr, ok := err.(awserr.Error); ok {
				switch awserr.Code() {
				case autoscaling.ErrCodeAlreadyExistsFault:
					logger.Debug(autoscaling.ErrCodeAlreadyExistsFault, awserr.Error())
				case autoscaling.ErrCodeLimitExceededFault:
					logger.Debug(autoscaling.ErrCodeLimitExceededFault, awserr.Error())
				case autoscaling.ErrCodeResourceContentionFault:
					logger.Debug(autoscaling.ErrCodeResourceContentionFault, awserr.Error())
				default:
					logger.Debug(awserr.Error())
				}
			} else {
				logger.Debug(err.Error())
			}
			if strings.Contains(err.Error(), "Invalid IamInstanceProfile") {
				logger.Debug("InstanceProfile missing waiting...")
				time.Sleep(time.Duration(i) * time.Second)
				continue
			}

		} else {
			break
		}
		return nil, nil, err
	}
	logger.Info("Created Launch Configuration [%s]", r.Name)
	newResource.Image = expected.(*Lc).Image
	newResource.InstanceType = expected.(*Lc).InstanceType
	newResource.Name = expected.(*Lc).Name
	newResource.Identifier = expected.(*Lc).Name
	newResource.BootstrapScripts = r.ServerPool.BootstrapScripts
	newResource.InstanceProfile = expected.(*Lc).InstanceProfile

	newCluster := r.immutableRender(newResource, immutable)
	return newCluster, newResource, nil
}

func (r *Lc) Delete(actual cloud.Resource, immutable *cluster.Cluster) (*cluster.Cluster, cloud.Resource, error) {
	logger.Debug("lc.Delete")
	deleteResource := actual.(*Lc)
	if deleteResource.Name == "" {
		return nil, nil, fmt.Errorf("Unable to delete Launch Configuration resource without Name [%s]", deleteResource.Name)
	}
	input := &autoscaling.DeleteLaunchConfigurationInput{
		LaunchConfigurationName: &actual.(*Lc).Name,
	}
	_, err := Sdk.ASG.DeleteLaunchConfiguration(input)
	if err != nil {
		return nil, nil, err
	}
	logger.Info("Deleted Launch Configuration [%s]", actual.(*Lc).Name)

	// Kubernetes API
	// Todo (@kris-nova) this obviously isn't immutable
	immutable.KubernetesAPI.Endpoint = ""

	newResource := &Lc{}
	newResource.Name = actual.(*Lc).Name
	newResource.Tags = actual.(*Lc).Tags
	newResource.Image = actual.(*Lc).Image
	newResource.InstanceType = actual.(*Lc).InstanceType
	newResource.BootstrapScripts = actual.(*Lc).BootstrapScripts
	newResource.InstanceProfile = actual.(*Lc).InstanceProfile

	newCluster := r.immutableRender(newResource, immutable)
	return newCluster, newResource, nil
}

func (r *Lc) immutableRender(newResource cloud.Resource, inaccurateCluster *cluster.Cluster) *cluster.Cluster {
	logger.Debug("lc.Render")
	newCluster := defaults.NewClusterDefaults(inaccurateCluster)
	serverPool := &cluster.ServerPool{}
	serverPool.Image = newResource.(*Lc).Image
	serverPool.Size = newResource.(*Lc).InstanceType
	serverPool.BootstrapScripts = newResource.(*Lc).BootstrapScripts
	found := false
	for i := 0; i < len(newCluster.ServerPools); i++ {
		if newCluster.ServerPools[i].Name == newResource.(*Lc).Name {
			newCluster.ServerPools[i].Image = newResource.(*Lc).Image
			newCluster.ServerPools[i].Size = newResource.(*Lc).InstanceType
			newCluster.ServerPools[i].BootstrapScripts = newResource.(*Lc).BootstrapScripts
			found = true
		}
	}
	if !found {
		newCluster.ServerPools = append(newCluster.ServerPools, serverPool)
	}

	return newCluster
}
<commit_msg>AWS spot price check accidentally removed by IAM feature.<commit_after>// Copyright © 2017 The Kubicorn 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 resources

import (
	"encoding/base64"
	"fmt"
	"strconv"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/service/autoscaling"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/kris-nova/kubicorn/apis/cluster"
	"github.com/kris-nova/kubicorn/cloud"
	"github.com/kris-nova/kubicorn/cutil/compare"
	"github.com/kris-nova/kubicorn/cutil/defaults"
	"github.com/kris-nova/kubicorn/cutil/logger"
	"github.com/kris-nova/kubicorn/cutil/script"
)

var _ cloud.Resource = &Lc{}

type Lc struct {
	Shared
	InstanceType     string
	Image            string
	SpotPrice        string
	ServerPool       *cluster.ServerPool
	BootstrapScripts []string
}

const (
	MasterIPAttempts               = 40
	MasterIPSleepSecondsPerAttempt = 3
)

func (r *Lc) Actual(immutable *cluster.Cluster) (*cluster.Cluster, cloud.Resource, error) {
	logger.Debug("lc.Actual")
	newResource := &Lc{
		Shared: Shared{
			Name: r.Name,
			Tags: make(map[string]string),
		},
	}
	if r.ServerPool.Identifier != "" {
		lcInput := &autoscaling.DescribeLaunchConfigurationsInput{
			LaunchConfigurationNames: []*string{&r.ServerPool.Identifier},
		}
		lcOutput, err := Sdk.ASG.DescribeLaunchConfigurations(lcInput)
		if err != nil {
			return nil, nil, err
		}
		llc := len(lcOutput.LaunchConfigurations)
		if llc != 1 {
			return nil, nil, fmt.Errorf("Found [%d] Launch Configurations for ID [%s]", llc, r.ServerPool.Identifier)
		}
		lc := lcOutput.LaunchConfigurations[0]
		newResource.Image = *lc.ImageId
		if lc.SpotPrice != nil {
			newResource.SpotPrice = *lc.SpotPrice
		}
		newResource.Identifier = *lc.LaunchConfigurationName
		newResource.Tags = map[string]string{
			"Name":              r.Name,
			"KubernetesCluster": immutable.Name,
		}
	} else {
		newResource.Image = r.ServerPool.Image
		newResource.InstanceType = r.ServerPool.Size
		if r.ServerPool.Type == cluster.ServerPoolTypeNode && r.ServerPool.AwsConfiguration != nil {
			newResource.SpotPrice = r.ServerPool.AwsConfiguration.SpotPrice
		}
	}
	newResource.BootstrapScripts = r.ServerPool.BootstrapScripts

	newCluster := r.immutableRender(newResource, immutable)
	return newCluster, newResource, nil
}

func (r *Lc) Expected(immutable *cluster.Cluster) (*cluster.Cluster, cloud.Resource, error) {
	logger.Debug("asg.Expected")
	newResource := &Lc{
		Shared: Shared{
			Tags: map[string]string{
				"Name":              r.Name,
				"KubernetesCluster": immutable.Name,
			},
			Identifier: r.ServerPool.Identifier,
			Name:       r.Name,
		},
		InstanceType:     r.ServerPool.Size,
		Image:            r.ServerPool.Image,
		BootstrapScripts: r.ServerPool.BootstrapScripts,
	}
	if r.ServerPool.Type == cluster.ServerPoolTypeNode && r.ServerPool.AwsConfiguration != nil {
		newResource.SpotPrice = r.ServerPool.AwsConfiguration.SpotPrice
	}
	newCluster := r.immutableRender(newResource, immutable)
	return newCluster, newResource, nil
}

func (r *Lc) Apply(actual, expected cloud.Resource, immutable *cluster.Cluster) (*cluster.Cluster, cloud.Resource, error) {
	logger.Debug("lc.Apply")
	applyResource := expected.(*Lc)
	isEqual, err := compare.IsEqual(actual.(*Lc), expected.(*Lc))
	if err != nil {
		return nil, nil, err
	}
	if isEqual {
		return immutable, applyResource, nil
	}
	var sgs []*string
	found := false
	for _, serverPool := range immutable.ServerPools {
		if serverPool.Name == expected.(*Lc).Name {
			for _, firewall := range serverPool.Firewalls {
				sgs = append(sgs, &firewall.Identifier)
			}
			found = true
		}
	}
	if !found {
		return nil, nil, fmt.Errorf("Unable to lookup serverpool for Launch Configuration %s", r.Name)
	}

	// --- Hack in here for master IP
	privip := ""
	pubip := ""
	if strings.Contains(r.ServerPool.Name, "node") {
		found := false
		logger.Debug("Tag query: [%s] %s", "Name", fmt.Sprintf("%s.master", immutable.Name))
		logger.Debug("Tag query: [%s] %s", "KubernetesCluster", immutable.Name)
		for i := 0; i < MasterIPAttempts; i++ {
			logger.Debug("Attempting to lookup master IP for node registration..")
			input := &ec2.DescribeInstancesInput{
				Filters: []*ec2.Filter{
					{
						Name:   S("tag:Name"),
						Values: []*string{S(fmt.Sprintf("%s.master", immutable.Name))},
					},
					{
						Name:   S("tag:KubernetesCluster"),
						Values: []*string{S(immutable.Name)},
					},
				},
			}
			output, err := Sdk.Ec2.DescribeInstances(input)
			if err != nil {
				return nil, nil, err
			}
			lr := len(output.Reservations)
			if lr == 0 {
				logger.Debug("Found %d Reservations, hanging ", lr)
				time.Sleep(time.Duration(MasterIPSleepSecondsPerAttempt) * time.Second)
				continue
			}
			for _, reservation := range output.Reservations {
				for _, instance := range reservation.Instances {
					if instance.PublicIpAddress != nil {
						privip = *instance.PrivateIpAddress
						pubip = *instance.PublicIpAddress
						immutable.Values.ItemMap["INJECTEDMASTER"] = fmt.Sprintf("%s:%s", privip, immutable.KubernetesAPI.Port)
						immutable.KubernetesAPI.Endpoint = pubip
						logger.Info("Found public IP for master: [%s]", pubip)
						found = true
					}
				}
			}
			if found == true {
				break
			}
			time.Sleep(time.Duration(MasterIPSleepSecondsPerAttempt) * time.Second)
		}
		if !found {
			return nil, nil, fmt.Errorf("Unable to find Master IP")
		}
	}

	immutable.Values.ItemMap["INJECTEDPORT"] = immutable.KubernetesAPI.Port

	newResource := &Lc{}
	userData, err := script.BuildBootstrapScript(r.ServerPool.BootstrapScripts, immutable)
	if err != nil {
		return nil, nil, err
	}

	b64data := base64.StdEncoding.EncodeToString(userData)
	lcInput := &autoscaling.CreateLaunchConfigurationInput{
		AssociatePublicIpAddress: B(true),
		LaunchConfigurationName:  &expected.(*Lc).Name,
		ImageId:                  &expected.(*Lc).Image,
		InstanceType:             &expected.(*Lc).InstanceType,
		KeyName:                  &immutable.SSH.Identifier,
		SecurityGroups:           sgs,
		UserData:                 &b64data,
	}

	spotPrice, err := strconv.ParseFloat(*&expected.(*Lc).SpotPrice, 64)
	if *&expected.(*Lc).InstanceType != cluster.ServerPoolTypeMaster && err == nil && spotPrice > 0 {
		lcInput.SpotPrice = &expected.(*Lc).SpotPrice
	}
	_, err = Sdk.ASG.CreateLaunchConfiguration(lcInput)
	if err != nil {
		if awserr, ok := err.(awserr.Error); ok {
			switch awserr.Code() {
			case autoscaling.ErrCodeAlreadyExistsFault:
				logger.Debug(autoscaling.ErrCodeAlreadyExistsFault, awserr.Error())
			case autoscaling.ErrCodeLimitExceededFault:
				logger.Debug(autoscaling.ErrCodeLimitExceededFault, awserr.Error())
			case autoscaling.ErrCodeResourceContentionFault:
				logger.Debug(autoscaling.ErrCodeResourceContentionFault, awserr.Error())
			default:
				logger.Debug(awserr.Error())
			}
		} else {
			logger.Debug(err.Error())
		}
		return nil, nil, err
	}
	logger.Info("Created Launch Configuration [%s]", r.Name)
	newResource.Image = expected.(*Lc).Image
	newResource.InstanceType = expected.(*Lc).InstanceType
	newResource.Name = expected.(*Lc).Name
	newResource.Identifier = expected.(*Lc).Name
	newResource.BootstrapScripts = r.ServerPool.BootstrapScripts

	newCluster := r.immutableRender(newResource, immutable)
	return newCluster, newResource, nil
}

func (r *Lc) Delete(actual cloud.Resource, immutable *cluster.Cluster) (*cluster.Cluster, cloud.Resource, error) {
	logger.Debug("lc.Delete")
	deleteResource := actual.(*Lc)
	if deleteResource.Name == "" {
		return nil, nil, fmt.Errorf("Unable to delete Launch Configuration resource without Name [%s]", deleteResource.Name)
	}
	input := &autoscaling.DeleteLaunchConfigurationInput{
		LaunchConfigurationName: &actual.(*Lc).Name,
	}
	_, err := Sdk.ASG.DeleteLaunchConfiguration(input)
	if err != nil {
		return nil, nil, err
	}
	logger.Info("Deleted Launch Configuration [%s]", actual.(*Lc).Name)

	// Kubernetes API
	// Todo (@kris-nova) this obviously isn't immutable
	immutable.KubernetesAPI.Endpoint = ""

	newResource := &Lc{}
	newResource.Name = actual.(*Lc).Name
	newResource.Tags = actual.(*Lc).Tags
	newResource.Image = actual.(*Lc).Image
	newResource.InstanceType = actual.(*Lc).InstanceType
	newResource.BootstrapScripts = actual.(*Lc).BootstrapScripts

	newCluster := r.immutableRender(newResource, immutable)
	return newCluster, newResource, nil
}

func (r *Lc) immutableRender(newResource cloud.Resource, inaccurateCluster *cluster.Cluster) *cluster.Cluster {
	logger.Debug("lc.Render")
	newCluster := defaults.NewClusterDefaults(inaccurateCluster)
	serverPool := &cluster.ServerPool{}
	serverPool.Image = newResource.(*Lc).Image
	serverPool.Size = newResource.(*Lc).InstanceType
	serverPool.BootstrapScripts = newResource.(*Lc).BootstrapScripts
	found := false
	for i := 0; i < len(newCluster.ServerPools); i++ {
		if newCluster.ServerPools[i].Name == newResource.(*Lc).Name {
			newCluster.ServerPools[i].Image = newResource.(*Lc).Image
			newCluster.ServerPools[i].Size = newResource.(*Lc).InstanceType
			newCluster.ServerPools[i].BootstrapScripts = newResource.(*Lc).BootstrapScripts
			found = true
		}
	}
	if !found {
		newCluster.ServerPools = append(newCluster.ServerPools, serverPool)
	}

	return newCluster
}
<|endoftext|>


<commit_before>package edit

import (
	"bytes"
	"fmt"
	"os"
	"strings"
	"unicode"
	"unicode/utf8"

	"github.com/xiaq/elvish/edit/tty"
	"github.com/xiaq/elvish/util"
)

// cell is an indivisible unit on the screen. It is not necessarily 1 column
// wide.
type cell struct {
	rune
	width byte
	attr  string
}

// pos is the position within a buffer.
type pos struct {
	line, col int
}

// buffer reflects a continuous range of lines on the terminal. The Unix
// terminal API provides only awkward ways of querying the terminal buffer, so
// we keep an internal reflection and do one-way synchronizations (buffer ->
// terminal, and not the other way around). This requires us to exactly match
// the terminal's idea of the width of characters (wcwidth) and where to
// insert soft carriage returns, so there could be bugs.
type buffer struct {
	width, col, indent int
	newlineWhenFull    bool
	cells              [][]cell // cells reflect len(cells) lines on the terminal.
	dot                pos      // dot is what the user perceives as the cursor.
}

func newBuffer(width int) *buffer {
	return &buffer{width: width, cells: [][]cell{make([]cell, 0, width)}}
}

func (b *buffer) appendCell(c cell) {
	n := len(b.cells)
	b.cells[n-1] = append(b.cells[n-1], c)
	b.col += int(c.width)
}

func (b *buffer) appendLine() {
	b.cells = append(b.cells, make([]cell, 0, b.width))
	b.col = 0
}

func (b *buffer) newline() {
	b.appendLine()

	if b.indent > 0 {
		for i := 0; i < b.indent; i++ {
			b.appendCell(cell{rune: ' ', width: 1})
		}
	}
}

func (b *buffer) extend(b2 *buffer) {
	if b2 != nil && b2.cells != nil {
		b.cells = append(b.cells, b2.cells...)
		b.col = b2.col
	}
}

// write appends a single rune to a buffer.
func (b *buffer) write(r rune, attr string) {
	if r == '\n' {
		b.newline()
		return
	} else if !unicode.IsPrint(r) {
		// BUG(xiaq): buffer.write drops unprintable runes silently
		return
	}
	wd := wcwidth(r)
	c := cell{r, byte(wd), attr}

	if b.col+wd > b.width {
		b.newline()
		b.appendCell(c)
	} else {
		b.appendCell(c)
		if b.col == b.width && b.newlineWhenFull {
			b.newline()
		}
	}
}

func (b *buffer) writes(s string, attr string) {
	for _, r := range s {
		b.write(r, attr)
	}
}

func (b *buffer) writePadding(w int, attr string) {
	b.writes(strings.Repeat(" ", w), attr)
}

func (b *buffer) line() int {
	return len(b.cells) - 1
}

func (b *buffer) cursor() pos {
	return pos{len(b.cells) - 1, b.col}
}

func (b *buffer) trimToLines(low, high int) {
	for i := 0; i < low; i++ {
		b.cells[i] = nil
	}
	for i := high; i < len(b.cells); i++ {
		b.cells[i] = nil
	}
	b.cells = b.cells[low:high]
	b.dot.line -= low
}

// writer is the part of an Editor responsible for keeping the status of and
// updating the screen.
type writer struct {
	file   *os.File
	oldBuf *buffer
}

func newWriter(f *os.File) *writer {
	writer := &writer{file: f, oldBuf: newBuffer(0)}
	return writer
}

// deltaPos calculates the escape sequence needed to move the cursor from one
// position to another.
func deltaPos(from, to pos) []byte {
	buf := new(bytes.Buffer)
	if from.line < to.line {
		// move down
		fmt.Fprintf(buf, "\033[%dB", to.line-from.line)
	} else if from.line > to.line {
		// move up
		fmt.Fprintf(buf, "\033[%dA", from.line-to.line)
	}
	fmt.Fprintf(buf, "\033[%dG", to.col+1)
	return buf.Bytes()
}

// commitBuffer updates the terminal display to reflect current buffer.
// TODO Instead of erasing w.oldBuf entirely and then draw buf, compute a
// delta between w.oldBuf and buf
func (w *writer) commitBuffer(buf *buffer) error {
	bytesBuf := new(bytes.Buffer)

	pLine := w.oldBuf.dot.line
	if pLine > 0 {
		fmt.Fprintf(bytesBuf, "\033[%dA", pLine)
	}
	bytesBuf.WriteString("\r\033[J")

	attr := ""
	for i, line := range buf.cells {
		if i > 0 {
			bytesBuf.WriteString("\n")
		}
		for _, c := range line {
			if c.width > 0 && c.attr != attr {
				fmt.Fprintf(bytesBuf, "\033[m\033[%sm", c.attr)
				attr = c.attr
			}
			bytesBuf.WriteString(string(c.rune))
		}
	}
	if attr != "" {
		bytesBuf.WriteString("\033[m")
	}
	cursor := buf.cursor()
	if cursor.col == buf.width {
		cursor.col--
	}
	bytesBuf.Write(deltaPos(cursor, buf.dot))

	_, err := w.file.Write(bytesBuf.Bytes())
	if err != nil {
		return err
	}

	w.oldBuf = buf
	return nil
}

func lines(bufs ...*buffer) (l int) {
	for _, buf := range bufs {
		if buf != nil {
			l += len(buf.cells)
		}
	}
	return
}

// findWindow finds a window of lines around the selected line in a total
// number of height lines, that is at most max lines.
func findWindow(height, selected, max int) (low, high int) {
	if height > max {
		low = selected - max/2
		high = low + max
		switch {
		case low < 0:
			// Near top of the list, move the window down
			low = 0
			high = low + max
		case high > height:
			// Near bottom of the list, move the window down
			high = height
			low = high - max
		}
		return
	} else {
		return 0, height
	}
}

func trimToWindow(s []string, selected, max int) ([]string, int) {
	low, high := findWindow(len(s), selected, max)
	return s[low:high], low
}

// refresh redraws the line editor. The dot is passed as an index into text;
// the corresponding position will be calculated.
func (w *writer) refresh(bs *editorState) error {
	winsize := tty.GetWinsize(int(w.file.Fd()))
	width, height := int(winsize.Col), int(winsize.Row)

	var bufLine, bufMode, bufTips, bufListing, buf *buffer
	// bufLine
	b := newBuffer(width)
	bufLine = b

	b.newlineWhenFull = true

	b.writes(bs.prompt, attrForPrompt)

	if b.line() == 0 && b.col*2 < b.width {
		b.indent = b.col
	}

	// i keeps track of number of bytes written.
	i := 0
	if bs.dot == 0 {
		b.dot = b.cursor()
	}

	comp := bs.completion
	var suppress = false

tokens:
	for _, token := range bs.tokens {
		for _, r := range token.Val {
			if suppress && i < comp.end {
				// Silence the part that is being completed
			} else {
				b.write(r, attrForType[token.Typ])
			}
			i += utf8.RuneLen(r)
			if comp != nil && comp.current != -1 && i == comp.start {
				// Put the current candidate and instruct text up to comp.end
				// to be suppressed. The cursor should be placed correctly
				// (i.e. right after the candidate)
				for _, part := range comp.candidates[comp.current].parts {
					attr := attrForType[comp.typ]
					if part.completed {
						attr += attrForCompleted
					}
					b.writes(part.text, attr)
				}
				suppress = true
			}
			if bs.mode == modeHistory && i == len(bs.history.prefix) {
				break tokens
			}
			if bs.dot == i {
				b.dot = b.cursor()
			}
		}
	}

	if bs.mode == modeHistory {
		// Put the rest of current history, position the cursor at the
		// end of the line, and finish writing
		h := bs.history
		b.writes(h.items[h.current][len(h.prefix):], attrForCompletedHistory)
		b.dot = b.cursor()
	}

	// Write rprompt
	padding := b.width - b.col - wcwidths(bs.rprompt)
	if padding >= 1 {
		b.newlineWhenFull = false
		b.writePadding(padding, "")
		b.writes(bs.rprompt, attrForRprompt)
	}

	// bufMode
	if bs.mode != modeInsert {
		b := newBuffer(width)
		bufMode = b
		text := ""
		switch bs.mode {
		case modeCommand:
			text = "Command"
			b.writes(trimWcwidth("Command", width), attrForMode)
		case modeCompletion:
			text = fmt.Sprintf("Completing %s", bs.line[comp.start:comp.end])
		case modeNavigation:
			text = "Navigating"
		case modeHistory:
			text = fmt.Sprintf("History #%d", bs.history.current)
		}
		b.writes(trimWcwidth(text, width), attrForMode)
	}

	// bufTips
	// TODO tips is assumed to contain no newlines.
	if len(bs.tips) > 0 {
		b := newBuffer(width)
		bufTips = b
		b.writes(trimWcwidth(strings.Join(bs.tips, ", "), width), attrForTip)
	}

	listingHeight := 0
	// Trim lines and determine the maximum height for bufListing
	switch {
	case height >= lines(bufLine, bufMode, bufTips):
		listingHeight = height - lines(bufLine, bufMode, bufTips)
	case height >= lines(bufLine, bufTips):
		bufMode, bufListing = nil, nil
	case height >= lines(bufLine):
		bufTips, bufMode, bufListing = nil, nil, nil
	case height >= 1:
		bufTips, bufMode, bufListing = nil, nil, nil
		dotLine := bufLine.dot.line
		bufLine.trimToLines(dotLine+1-height, dotLine+1)
	default:
		bufLine, bufTips, bufMode, bufListing = nil, nil, nil, nil
	}

	// Render bufListing under the maximum height constraint
	nav := bs.navigation
	if listingHeight > 0 && comp != nil || nav != nil {
		b := newBuffer(width)
		bufListing = b
		// Completion listing
		if comp != nil {
			// Layout candidates in multiple columns
			cands := comp.candidates

			// First decide the shape (# of rows and columns)
			colWidth := 0
			colMargin := 2
			for _, cand := range cands {
				width := wcwidths(cand.text)
				if colWidth < width {
					colWidth = width
				}
			}

			cols := (b.width + colMargin) / (colWidth + colMargin)
			if cols == 0 {
				cols = 1
			}
			lines := util.CeilDiv(len(cands), cols)
			bs.completionLines = lines

			// Determine the window to show.
			low, high := findWindow(lines, comp.current%lines, listingHeight)
			for i := low; i < high; i++ {
				if i > low {
					b.newline()
				}
				for j := 0; j < cols; j++ {
					k := j*lines + i
					if k >= len(cands) {
						continue
					}
					attr := ""
					if k == comp.current {
						attr = attrForCurrentCompletion
					}
					text := cands[k].text
					b.writes(text, attr)
					b.writePadding(colWidth-wcwidths(text), attr)
					b.writePadding(colMargin, "")
				}
			}
		}

		// Navigation listing
		if nav != nil {
			b := newBuffer(width)
			bufListing = b

			filenames, low := trimToWindow(nav.current.names, nav.current.selected, listingHeight)
			parentFilenames, parentLow := trimToWindow(nav.parent.names, nav.parent.selected, listingHeight)

			// TODO(xiaq): When laying out the navigation listing, determine
			// the width of two columns more intelligently instead of
			// allocating half of screen for each. Maybe the algorithm used by
			// ranger could be pirated.
			colMargin := 1
			parentWidth := (width + colMargin) / 2
			currentWidth := width - colMargin - parentWidth
			for i := 0; i < len(filenames) || i < len(parentFilenames); i++ {
				if i > 0 {
					b.newline()
				}
				text, attr := "", ""
				if i < len(parentFilenames) {
					text = parentFilenames[i]
				}
				if i+parentLow == nav.parent.selected {
					attr = attrForSelectedFile
				}
				b.writes(trimWcwidth(text, parentWidth), attr)
				b.writePadding(parentWidth-wcwidths(text), attr)
				b.writePadding(colMargin, "")

				if i < len(filenames) {
					attr := ""
					if i+low == nav.current.selected {
						attr = attrForSelectedFile
					}
					text := filenames[i]
					b.writes(trimWcwidth(text, currentWidth), attr)
					b.writePadding(currentWidth-wcwidths(text), attr)
				}
			}
		}
		// Trim bufListing.
		// XXX This algorithm only works for completion listing.

	}

	// Combine buffers (reusing bufLine)
	buf = bufLine
	buf.extend(bufMode)
	buf.extend(bufTips)
	buf.extend(bufListing)

	return w.commitBuffer(buf)
}
<commit_msg>Much cleaner navigation listing rendering<commit_after>package edit

import (
	"bytes"
	"fmt"
	"os"
	"strings"
	"unicode"
	"unicode/utf8"

	"github.com/xiaq/elvish/edit/tty"
	"github.com/xiaq/elvish/util"
)

// cell is an indivisible unit on the screen. It is not necessarily 1 column
// wide.
type cell struct {
	rune
	width byte
	attr  string
}

// pos is the position within a buffer.
type pos struct {
	line, col int
}

// buffer reflects a continuous range of lines on the terminal. The Unix
// terminal API provides only awkward ways of querying the terminal buffer, so
// we keep an internal reflection and do one-way synchronizations (buffer ->
// terminal, and not the other way around). This requires us to exactly match
// the terminal's idea of the width of characters (wcwidth) and where to
// insert soft carriage returns, so there could be bugs.
type buffer struct {
	width, col, indent int
	newlineWhenFull    bool
	cells              [][]cell // cells reflect len(cells) lines on the terminal.
	dot                pos      // dot is what the user perceives as the cursor.
}

func newBuffer(width int) *buffer {
	return &buffer{width: width, cells: [][]cell{make([]cell, 0, width)}}
}

func (b *buffer) appendCell(c cell) {
	n := len(b.cells)
	b.cells[n-1] = append(b.cells[n-1], c)
	b.col += int(c.width)
}

func (b *buffer) appendLine() {
	b.cells = append(b.cells, make([]cell, 0, b.width))
	b.col = 0
}

func (b *buffer) newline() {
	b.appendLine()

	if b.indent > 0 {
		for i := 0; i < b.indent; i++ {
			b.appendCell(cell{rune: ' ', width: 1})
		}
	}
}

func (b *buffer) extend(b2 *buffer) {
	if b2 != nil && b2.cells != nil {
		b.cells = append(b.cells, b2.cells...)
		b.col = b2.col
	}
}

func makePadding(n int) []cell {
	s := make([]cell, n)
	for i := 0; i < n; i++ {
		s[i].rune = ' '
		s[i].width = 1
	}
	return s
}

// extendHorizontal extends b horizontally, appending each line in b2 to b,
// preceeded by m margin. If b2 has more lines than b, the last len(b2) -
// len(b) lines are first filled with paddings of width w.
func (b *buffer) extendHorizontal(b2 *buffer, w, m int) {
	i := 0
	margin := makePadding(m)
	padding := makePadding(w + m)
	for ; i < len(b.cells) && i < len(b2.cells); i++ {
		b.cells[i] = append(append(b.cells[i], margin...), b2.cells[i]...)
	}
	for ; i < len(b2.cells); i++ {
		row := make([]cell, 0, w+m+len(b2.cells[i]))
		row = append(append(row, padding...), b2.cells[i]...)
		b.cells = append(b.cells, row)
	}
}

// write appends a single rune to a buffer.
func (b *buffer) write(r rune, attr string) {
	if r == '\n' {
		b.newline()
		return
	} else if !unicode.IsPrint(r) {
		// BUG(xiaq): buffer.write drops unprintable runes silently
		return
	}
	wd := wcwidth(r)
	c := cell{r, byte(wd), attr}

	if b.col+wd > b.width {
		b.newline()
		b.appendCell(c)
	} else {
		b.appendCell(c)
		if b.col == b.width && b.newlineWhenFull {
			b.newline()
		}
	}
}

func (b *buffer) writes(s string, attr string) {
	for _, r := range s {
		b.write(r, attr)
	}
}

func (b *buffer) writePadding(w int, attr string) {
	b.writes(strings.Repeat(" ", w), attr)
}

func (b *buffer) line() int {
	return len(b.cells) - 1
}

func (b *buffer) cursor() pos {
	return pos{len(b.cells) - 1, b.col}
}

func (b *buffer) trimToLines(low, high int) {
	for i := 0; i < low; i++ {
		b.cells[i] = nil
	}
	for i := high; i < len(b.cells); i++ {
		b.cells[i] = nil
	}
	b.cells = b.cells[low:high]
	b.dot.line -= low
}

// writer is the part of an Editor responsible for keeping the status of and
// updating the screen.
type writer struct {
	file   *os.File
	oldBuf *buffer
}

func newWriter(f *os.File) *writer {
	writer := &writer{file: f, oldBuf: newBuffer(0)}
	return writer
}

// deltaPos calculates the escape sequence needed to move the cursor from one
// position to another.
func deltaPos(from, to pos) []byte {
	buf := new(bytes.Buffer)
	if from.line < to.line {
		// move down
		fmt.Fprintf(buf, "\033[%dB", to.line-from.line)
	} else if from.line > to.line {
		// move up
		fmt.Fprintf(buf, "\033[%dA", from.line-to.line)
	}
	fmt.Fprintf(buf, "\033[%dG", to.col+1)
	return buf.Bytes()
}

// commitBuffer updates the terminal display to reflect current buffer.
// TODO Instead of erasing w.oldBuf entirely and then draw buf, compute a
// delta between w.oldBuf and buf
func (w *writer) commitBuffer(buf *buffer) error {
	bytesBuf := new(bytes.Buffer)

	pLine := w.oldBuf.dot.line
	if pLine > 0 {
		fmt.Fprintf(bytesBuf, "\033[%dA", pLine)
	}
	bytesBuf.WriteString("\r\033[J")

	attr := ""
	for i, line := range buf.cells {
		if i > 0 {
			bytesBuf.WriteString("\n")
		}
		for _, c := range line {
			if c.width > 0 && c.attr != attr {
				fmt.Fprintf(bytesBuf, "\033[m\033[%sm", c.attr)
				attr = c.attr
			}
			bytesBuf.WriteString(string(c.rune))
		}
	}
	if attr != "" {
		bytesBuf.WriteString("\033[m")
	}
	cursor := buf.cursor()
	if cursor.col == buf.width {
		cursor.col--
	}
	bytesBuf.Write(deltaPos(cursor, buf.dot))

	_, err := w.file.Write(bytesBuf.Bytes())
	if err != nil {
		return err
	}

	w.oldBuf = buf
	return nil
}

func lines(bufs ...*buffer) (l int) {
	for _, buf := range bufs {
		if buf != nil {
			l += len(buf.cells)
		}
	}
	return
}

// findWindow finds a window of lines around the selected line in a total
// number of height lines, that is at most max lines.
func findWindow(height, selected, max int) (low, high int) {
	if height > max {
		low = selected - max/2
		high = low + max
		switch {
		case low < 0:
			// Near top of the list, move the window down
			low = 0
			high = low + max
		case high > height:
			// Near bottom of the list, move the window down
			high = height
			low = high - max
		}
		return
	} else {
		return 0, height
	}
}

func trimToWindow(s []string, selected, max int) ([]string, int) {
	low, high := findWindow(len(s), selected, max)
	return s[low:high], low
}

func renderNavColumn(nc navColumn, w, h int) *buffer {
	b := newBuffer(w)
	low, high := findWindow(len(nc.names), nc.selected, h)
	for i := low; i < high; i++ {
		if i > low {
			b.newline()
		}
		text := nc.names[i]
		attr := ""
		if i == nc.selected {
			attr = attrForSelectedFile
		}
		b.writes(trimWcwidth(text, w), attr)
		b.writePadding(w-wcwidths(text), attr)
	}
	return b
}

// refresh redraws the line editor. The dot is passed as an index into text;
// the corresponding position will be calculated.
func (w *writer) refresh(bs *editorState) error {
	winsize := tty.GetWinsize(int(w.file.Fd()))
	width, height := int(winsize.Col), int(winsize.Row)

	var bufLine, bufMode, bufTips, bufListing, buf *buffer
	// bufLine
	b := newBuffer(width)
	bufLine = b

	b.newlineWhenFull = true

	b.writes(bs.prompt, attrForPrompt)

	if b.line() == 0 && b.col*2 < b.width {
		b.indent = b.col
	}

	// i keeps track of number of bytes written.
	i := 0
	if bs.dot == 0 {
		b.dot = b.cursor()
	}

	comp := bs.completion
	var suppress = false

tokens:
	for _, token := range bs.tokens {
		for _, r := range token.Val {
			if suppress && i < comp.end {
				// Silence the part that is being completed
			} else {
				b.write(r, attrForType[token.Typ])
			}
			i += utf8.RuneLen(r)
			if comp != nil && comp.current != -1 && i == comp.start {
				// Put the current candidate and instruct text up to comp.end
				// to be suppressed. The cursor should be placed correctly
				// (i.e. right after the candidate)
				for _, part := range comp.candidates[comp.current].parts {
					attr := attrForType[comp.typ]
					if part.completed {
						attr += attrForCompleted
					}
					b.writes(part.text, attr)
				}
				suppress = true
			}
			if bs.mode == modeHistory && i == len(bs.history.prefix) {
				break tokens
			}
			if bs.dot == i {
				b.dot = b.cursor()
			}
		}
	}

	if bs.mode == modeHistory {
		// Put the rest of current history, position the cursor at the
		// end of the line, and finish writing
		h := bs.history
		b.writes(h.items[h.current][len(h.prefix):], attrForCompletedHistory)
		b.dot = b.cursor()
	}

	// Write rprompt
	padding := b.width - b.col - wcwidths(bs.rprompt)
	if padding >= 1 {
		b.newlineWhenFull = false
		b.writePadding(padding, "")
		b.writes(bs.rprompt, attrForRprompt)
	}

	// bufMode
	if bs.mode != modeInsert {
		b := newBuffer(width)
		bufMode = b
		text := ""
		switch bs.mode {
		case modeCommand:
			text = "Command"
			b.writes(trimWcwidth("Command", width), attrForMode)
		case modeCompletion:
			text = fmt.Sprintf("Completing %s", bs.line[comp.start:comp.end])
		case modeNavigation:
			text = "Navigating"
		case modeHistory:
			text = fmt.Sprintf("History #%d", bs.history.current)
		}
		b.writes(trimWcwidth(text, width), attrForMode)
	}

	// bufTips
	// TODO tips is assumed to contain no newlines.
	if len(bs.tips) > 0 {
		b := newBuffer(width)
		bufTips = b
		b.writes(trimWcwidth(strings.Join(bs.tips, ", "), width), attrForTip)
	}

	listingHeight := 0
	// Trim lines and determine the maximum height for bufListing
	switch {
	case height >= lines(bufLine, bufMode, bufTips):
		listingHeight = height - lines(bufLine, bufMode, bufTips)
	case height >= lines(bufLine, bufTips):
		bufMode, bufListing = nil, nil
	case height >= lines(bufLine):
		bufTips, bufMode, bufListing = nil, nil, nil
	case height >= 1:
		bufTips, bufMode, bufListing = nil, nil, nil
		dotLine := bufLine.dot.line
		bufLine.trimToLines(dotLine+1-height, dotLine+1)
	default:
		bufLine, bufTips, bufMode, bufListing = nil, nil, nil, nil
	}

	// Render bufListing under the maximum height constraint
	nav := bs.navigation
	if listingHeight > 0 && comp != nil || nav != nil {
		b := newBuffer(width)
		bufListing = b
		// Completion listing
		if comp != nil {
			// Layout candidates in multiple columns
			cands := comp.candidates

			// First decide the shape (# of rows and columns)
			colWidth := 0
			colMargin := 2
			for _, cand := range cands {
				width := wcwidths(cand.text)
				if colWidth < width {
					colWidth = width
				}
			}

			cols := (b.width + colMargin) / (colWidth + colMargin)
			if cols == 0 {
				cols = 1
			}
			lines := util.CeilDiv(len(cands), cols)
			bs.completionLines = lines

			// Determine the window to show.
			low, high := findWindow(lines, comp.current%lines, listingHeight)
			for i := low; i < high; i++ {
				if i > low {
					b.newline()
				}
				for j := 0; j < cols; j++ {
					k := j*lines + i
					if k >= len(cands) {
						continue
					}
					attr := ""
					if k == comp.current {
						attr = attrForCurrentCompletion
					}
					text := cands[k].text
					b.writes(text, attr)
					b.writePadding(colWidth-wcwidths(text), attr)
					b.writePadding(colMargin, "")
				}
			}
		}

		// Navigation listing
		if nav != nil {
			// TODO(xiaq): When laying out the navigation listing, determine
			// the width of two columns more intelligently instead of
			// allocating half of screen for each. Maybe the algorithm used by
			// ranger could be pirated.
			colMargin := 1
			wParent := (width + colMargin) / 2
			wCurrent := width - colMargin - wParent

			b := renderNavColumn(nav.parent, wParent, listingHeight)
			bufListing = b

			bCurrent := renderNavColumn(nav.current, wCurrent, listingHeight)
			b.extendHorizontal(bCurrent, wParent, colMargin)
		}
	}

	// Combine buffers (reusing bufLine)
	buf = bufLine
	buf.extend(bufMode)
	buf.extend(bufTips)
	buf.extend(bufListing)

	return w.commitBuffer(buf)
}
<|endoftext|>


<commit_before>package editor

import (
	"github.com/atotto/clipboard"
	"github.com/gdamore/tcell"
	"github.com/jantb/olive/goPlugin"
	"github.com/rivo/tview"
	"log"
)

// View implements the main editor view.
type View struct {
	dataView map[string]*Dataview
	Lines    [][]Block
	*tview.Box
	*Editor
	offy, offx int
}

type Block struct {
	Rune  rune
	Style tcell.Style
}

// NewView returns a new main view primitive.
func NewView() *View {
	view := View{
		Box:   tview.NewBox().SetBorder(false),
		Lines: [][]Block{},
	}
	return &view
}

// Draw draws this primitive onto the screen.
func (m *View) Draw(screen tcell.Screen) {
	_, bg, _ := defaultStyle.Decompose()
	m.Box.SetBackgroundColor(bg).Draw(screen)
	dataview := m.dataView[m.curViewID]
	if dataview == nil {
		return
	}
	lines := dataview.Lines()
	m.Lines = [][]Block{}
	if len(lines) < m.offy {
		return
	}
	for y, line := range lines[m.offy:] {
		var blocks []Block
		m.Lines = append(m.Lines, blocks)
		for x, r := range line.Text {
			var style = defaultStyle
			if line.StyleIds[x] != nil {
				for _, value := range line.StyleIds[x] {
					s := styles[value]
					if value == 0 {
						s = s.Background(tcell.NewRGBColor(m.Editor.theme.Selection.ToRGB()))
					}

					fg, bg, _ := s.Decompose()

					if fg != tcell.ColorDefault {
						style = style.Foreground(fg)
					}
					if bg != tcell.ColorDefault {
						style = style.Background(bg)
					}
				}
			}
			m.Lines[y] = append(m.Lines[y], Block{Rune: r, Style: style})
		}
	}

	for y, line := range m.Lines {
		offX := 0
		for x, block := range line {
			if block.Rune == '\t' {
				m.draw(screen, x, y, block)
				m.draw(screen, x+1, y, block)
				m.draw(screen, x+2, y, block)
				m.draw(screen, x+3, y, block)
				offX += 3
				continue
			}
			m.draw(screen, x+offX, y, block)
		}
	}

	// Draw cursors
	for y, line := range lines[m.offy:] {
		for _, cursor := range line.Cursors {
			x := GetCursorVisualX(cursor, line.Text)
			content := m.getContent(screen, x, y)
			content.Style = content.Style.Reverse(true)
			m.draw(screen, x, y, content)
		}
	}
}

func (m *View) draw(screen tcell.Screen, x int, y int, b Block) {

	xMin, yMin, width, height := m.Box.GetInnerRect()
	x = xMin + x - m.offx
	y = yMin + y

	if x < xMin || y < yMin || x >= width+xMin || y >= height+yMin {
		return
	}
	screen.SetContent(x, y, b.Rune, nil, b.Style)
}

func (m *View) getContent(screen tcell.Screen, x int, y int) Block {

	xMin, yMin, width, height := m.Box.GetInnerRect()
	x = xMin + x - m.offx
	y = yMin + y

	if x < xMin || y < yMin || x >= width+xMin || y >= height+yMin {
		return Block{}
	}
	mainc, _, style, _ := screen.GetContent(x, y)
	return Block{Rune: mainc, Style: style}
}

func (m *View) MakeVisible(x, y int) {
	lines := m.dataView[m.curViewID].Lines()
	if len(lines) <= y {
		return
	}
	x = GetCursorVisualX(x, lines[y].Text)
	_, _, width, height := m.Box.GetInnerRect()

	if y >= m.offy+height {
		m.offy = y - (height - 1)
	}

	if y >= 0 && y < m.offy {
		m.offy = y
	}

	if x >= m.offx+width {
		m.offx = x - (width - 1)
	}
	if x >= 0 && x < m.offx {
		m.offx = x
	}
}

// InputHandler returns the handler for this primitive.
func (m *View) InputHandler() func(event *tcell.EventKey, setFocus func(p tview.Primitive)) {
	return m.WrapInputHandler(func(event *tcell.EventKey, setFocus func(p tview.Primitive)) {
		dataview := m.dataView[m.curViewID]
		ctrl := event.Modifiers()&tcell.ModCtrl != 0
		alt := event.Modifiers()&tcell.ModAlt != 0
		shift := event.Modifiers()&tcell.ModShift != 0
		if !ctrl && !alt && !shift {
			switch event.Key() {
			case tcell.KeyEsc:
				dataview.CancelOperation()
			case tcell.KeyUp:
				dataview.MoveUp()
			case tcell.KeyDown:
				dataview.MoveDown()
			case tcell.KeyLeft:
				dataview.MoveLeft()
			case tcell.KeyEnter:
				dataview.Newline()
			case tcell.KeyRight:
				dataview.MoveRight()
			case tcell.KeyRune:
				dataview.Insert(string(event.Rune()))
			case tcell.KeyTab:
				dataview.Tab()
			case tcell.KeyBS:
				dataview.DeleteForward()
			case tcell.KeyDEL:
				dataview.DeleteBackward()
			case tcell.KeyPgUp:
				dataview.ScrollPageUp()
			case tcell.KeyPgDn:
				dataview.ScrollPageDown()
			default:
				log.Println(event.Name())
			}
		}
		if !ctrl && !alt && shift {
			switch event.Key() {
			case tcell.KeyRune:
				dataview.Insert(string(event.Rune()))
			case tcell.KeyPgUp:
				dataview.ScrollPageUpAndModifySelection()
			case tcell.KeyPgDn:
				dataview.ScrollPageDownAndModifySelection()
			}
			switch event.Name() {
			case "Shift+Right":
				dataview.MoveRightAndModifySelection()
			case "Shift+Left":
				dataview.MoveLeftAndModifySelection()
			case "Shift+Up":
				dataview.MoveUpAndModifySelection()
			case "Shift+Down":
				dataview.MoveDownAndModifySelection()
			default:
				log.Println(event.Name())
			}
		}
		if ctrl && alt {
			switch event.Name() {
			case "Alt+Ctrl+Z":
				dataview.Redo()
			case "Alt+Ctrl+L":
				if m.footer.language == "Go" {
					dataview.CancelOperation()
					dataview.SelectAll()
					src := dataview.Copy()
					dataview.Insert(goPlugin.Format(src))
					dataview.CancelOperation()
				} else if m.footer.language == "XML" {
					dataview.CancelOperation()
					dataview.SelectAll()
					src := dataview.Copy()
					dataview.Insert(goPlugin.FormatXml(src))
					dataview.CancelOperation()
				} else if m.footer.language == "JSON" {
					dataview.CancelOperation()
					dataview.SelectAll()
					src := dataview.Copy()
					dataview.Insert(goPlugin.FormatJson(src))
					dataview.CancelOperation()
				}

			}
		}
		if ctrl && !alt && !shift {
			switch event.Key() {
			case tcell.KeyCtrlS:
				dataview.Save()
			case tcell.KeyCtrlA:
				dataview.SelectAll()
			case tcell.KeyCtrlZ:
				dataview.Undo()
			case tcell.KeyCtrlQ:
				dataview.Save()
				dataview.Close()
				m.curViewID = ""
				m.focusFileselector()
			case tcell.KeyCtrlD:
				dataview.DuplicateLine()
			case tcell.KeyCtrlV:
				s, e := clipboard.ReadAll()
				if e != nil {
					log.Println(e)
					return
				}
				dataview.Insert(s)
			default:
				log.Println(event.Name())
			}
		}
		if !ctrl && alt {
			switch event.Name() {
			case "Alt+Up":
				dataview.AddSelectionAbove()
			case "Alt+Down":
				dataview.AddSelectionBelow()
			case "Alt+Rune[0]":
				m.focusFileselector()
			case "Alt+Rune[c]":
				clipboard.WriteAll(dataview.Copy())
			case "Alt+Rune[x]":
				clipboard.WriteAll(dataview.Cut())
			case "Alt+Rune[v]":
				s, e := clipboard.ReadAll()
				if e != nil {
					log.Println(e)
					return
				}
				dataview.Insert(s)
			default:
				log.Println(event.Name())
			}
		}
		if ctrl && alt {
			switch event.Key() {

			default:
				log.Println(event.Name())
			}
		}
		log.Println(event.Name())
	})
}
<commit_msg>Delete key<commit_after>package editor

import (
	"github.com/atotto/clipboard"
	"github.com/gdamore/tcell"
	"github.com/jantb/olive/goPlugin"
	"github.com/rivo/tview"
	"log"
)

// View implements the main editor view.
type View struct {
	dataView map[string]*Dataview
	Lines    [][]Block
	*tview.Box
	*Editor
	offy, offx int
}

type Block struct {
	Rune  rune
	Style tcell.Style
}

// NewView returns a new main view primitive.
func NewView() *View {
	view := View{
		Box:   tview.NewBox().SetBorder(false),
		Lines: [][]Block{},
	}
	return &view
}

// Draw draws this primitive onto the screen.
func (m *View) Draw(screen tcell.Screen) {
	_, bg, _ := defaultStyle.Decompose()
	m.Box.SetBackgroundColor(bg).Draw(screen)
	dataview := m.dataView[m.curViewID]
	if dataview == nil {
		return
	}
	lines := dataview.Lines()
	m.Lines = [][]Block{}
	if len(lines) < m.offy {
		return
	}
	for y, line := range lines[m.offy:] {
		var blocks []Block
		m.Lines = append(m.Lines, blocks)
		for x, r := range line.Text {
			var style = defaultStyle
			if line.StyleIds[x] != nil {
				for _, value := range line.StyleIds[x] {
					s := styles[value]
					if value == 0 {
						s = s.Background(tcell.NewRGBColor(m.Editor.theme.Selection.ToRGB()))
					}

					fg, bg, _ := s.Decompose()

					if fg != tcell.ColorDefault {
						style = style.Foreground(fg)
					}
					if bg != tcell.ColorDefault {
						style = style.Background(bg)
					}
				}
			}
			m.Lines[y] = append(m.Lines[y], Block{Rune: r, Style: style})
		}
	}

	for y, line := range m.Lines {
		offX := 0
		for x, block := range line {
			if block.Rune == '\t' {
				m.draw(screen, x, y, block)
				m.draw(screen, x+1, y, block)
				m.draw(screen, x+2, y, block)
				m.draw(screen, x+3, y, block)
				offX += 3
				continue
			}
			m.draw(screen, x+offX, y, block)
		}
	}

	// Draw cursors
	for y, line := range lines[m.offy:] {
		for _, cursor := range line.Cursors {
			x := GetCursorVisualX(cursor, line.Text)
			content := m.getContent(screen, x, y)
			content.Style = content.Style.Reverse(true)
			m.draw(screen, x, y, content)
		}
	}
}

func (m *View) draw(screen tcell.Screen, x int, y int, b Block) {

	xMin, yMin, width, height := m.Box.GetInnerRect()
	x = xMin + x - m.offx
	y = yMin + y

	if x < xMin || y < yMin || x >= width+xMin || y >= height+yMin {
		return
	}
	screen.SetContent(x, y, b.Rune, nil, b.Style)
}

func (m *View) getContent(screen tcell.Screen, x int, y int) Block {

	xMin, yMin, width, height := m.Box.GetInnerRect()
	x = xMin + x - m.offx
	y = yMin + y

	if x < xMin || y < yMin || x >= width+xMin || y >= height+yMin {
		return Block{}
	}
	mainc, _, style, _ := screen.GetContent(x, y)
	return Block{Rune: mainc, Style: style}
}

func (m *View) MakeVisible(x, y int) {
	lines := m.dataView[m.curViewID].Lines()
	if len(lines) <= y {
		return
	}
	x = GetCursorVisualX(x, lines[y].Text)
	_, _, width, height := m.Box.GetInnerRect()

	if y >= m.offy+height {
		m.offy = y - (height - 1)
	}

	if y >= 0 && y < m.offy {
		m.offy = y
	}

	if x >= m.offx+width {
		m.offx = x - (width - 1)
	}
	if x >= 0 && x < m.offx {
		m.offx = x
	}
}

// InputHandler returns the handler for this primitive.
func (m *View) InputHandler() func(event *tcell.EventKey, setFocus func(p tview.Primitive)) {
	return m.WrapInputHandler(func(event *tcell.EventKey, setFocus func(p tview.Primitive)) {
		dataview := m.dataView[m.curViewID]
		ctrl := event.Modifiers()&tcell.ModCtrl != 0
		alt := event.Modifiers()&tcell.ModAlt != 0
		shift := event.Modifiers()&tcell.ModShift != 0
		if !ctrl && !alt && !shift {
			switch event.Key() {
			case tcell.KeyEsc:
				dataview.CancelOperation()
			case tcell.KeyUp:
				dataview.MoveUp()
			case tcell.KeyDown:
				dataview.MoveDown()
			case tcell.KeyLeft:
				dataview.MoveLeft()
			case tcell.KeyEnter:
				dataview.Newline()
			case tcell.KeyRight:
				dataview.MoveRight()
			case tcell.KeyRune:
				dataview.Insert(string(event.Rune()))
			case tcell.KeyTab:
				dataview.Tab()
			case tcell.KeyBS:
				dataview.DeleteForward()
			case tcell.KeyDelete:
				dataview.DeleteBackward()
			case tcell.KeyPgUp:
				dataview.ScrollPageUp()
			case tcell.KeyPgDn:
				dataview.ScrollPageDown()
			default:
				log.Println(event.Name())
			}
		}
		if !ctrl && !alt && shift {
			switch event.Key() {
			case tcell.KeyRune:
				dataview.Insert(string(event.Rune()))
			case tcell.KeyPgUp:
				dataview.ScrollPageUpAndModifySelection()
			case tcell.KeyPgDn:
				dataview.ScrollPageDownAndModifySelection()
			}
			switch event.Name() {
			case "Shift+Right":
				dataview.MoveRightAndModifySelection()
			case "Shift+Left":
				dataview.MoveLeftAndModifySelection()
			case "Shift+Up":
				dataview.MoveUpAndModifySelection()
			case "Shift+Down":
				dataview.MoveDownAndModifySelection()
			default:
				log.Println(event.Name())
			}
		}
		if ctrl && alt {
			switch event.Name() {
			case "Alt+Ctrl+Z":
				dataview.Redo()
			case "Alt+Ctrl+L":
				if m.footer.language == "Go" {
					dataview.CancelOperation()
					dataview.SelectAll()
					src := dataview.Copy()
					dataview.Insert(goPlugin.Format(src))
					dataview.CancelOperation()
				} else if m.footer.language == "XML" {
					dataview.CancelOperation()
					dataview.SelectAll()
					src := dataview.Copy()
					dataview.Insert(goPlugin.FormatXml(src))
					dataview.CancelOperation()
				} else if m.footer.language == "JSON" {
					dataview.CancelOperation()
					dataview.SelectAll()
					src := dataview.Copy()
					dataview.Insert(goPlugin.FormatJson(src))
					dataview.CancelOperation()
				}

			}
		}
		if ctrl && !alt && !shift {
			switch event.Key() {
			case tcell.KeyCtrlS:
				dataview.Save()
			case tcell.KeyCtrlA:
				dataview.SelectAll()
			case tcell.KeyCtrlZ:
				dataview.Undo()
			case tcell.KeyCtrlQ:
				dataview.Save()
				dataview.Close()
				m.curViewID = ""
				m.focusFileselector()
			case tcell.KeyCtrlD:
				dataview.DuplicateLine()
			case tcell.KeyCtrlV:
				s, e := clipboard.ReadAll()
				if e != nil {
					log.Println(e)
					return
				}
				dataview.Insert(s)
			default:
				log.Println(event.Name())
			}
		}
		if !ctrl && alt {
			switch event.Name() {
			case "Alt+Up":
				dataview.AddSelectionAbove()
			case "Alt+Down":
				dataview.AddSelectionBelow()
			case "Alt+Rune[0]":
				m.focusFileselector()
			case "Alt+Rune[c]":
				clipboard.WriteAll(dataview.Copy())
			case "Alt+Rune[x]":
				clipboard.WriteAll(dataview.Cut())
			case "Alt+Rune[v]":
				s, e := clipboard.ReadAll()
				if e != nil {
					log.Println(e)
					return
				}
				dataview.Insert(s)
			default:
				log.Println(event.Name())
			}
		}
		if ctrl && alt {
			switch event.Key() {

			default:
				log.Println(event.Name())
			}
		}
		log.Println(event.Name())
	})
}
<|endoftext|>


<commit_before>package timer

import (
	"testing"
	"time"
)

func TestStartStop(t *testing.T) {
	var now = time.Now()
	var timer = AfterFunc(time.Second, func() {
		println(time.Now().Sub(now))
		if time.Now().Sub(now) < time.Second*2 {
			t.Fatal("Early func")
		}
	})
	timer.Start()
	time.Sleep(time.Microsecond * 500)
	timer.Pause()
	time.Sleep(time.Second)
	timer.Start()
	time.Sleep(time.Second)
}
<commit_msg>More tests<commit_after>package timer

import (
	"testing"
	"time"
)

func TestStart(t *testing.T) {
	var now = time.Now()
	var timer = AfterFunc(time.Second, func() {
		if time.Now().Sub(now) < time.Second {
			t.Fatal("Early func")
		}
	})
	timer.Start()
	time.Sleep(time.Second * 2)
}

func TestStop(t *testing.T) {
	var timer = AfterFunc(time.Second, func() {
		t.Fatal("Stopping is not working")
	})
	timer.Start()
	timer.Stop()
	time.Sleep(time.Second * 2)
}

func TestPause(t *testing.T) {
	var now = time.Now()
	var timer = AfterFunc(time.Second, func() {
		if time.Now().Sub(now) < time.Second*2 {
			t.Fatal("Early func")
		}
	})
	timer.Start()
	time.Sleep(time.Microsecond * 500)
	timer.Pause()
	time.Sleep(time.Second)
	timer.Start()
	time.Sleep(time.Second)
}

func TestChanStart(t *testing.T) {
	var now = time.Now()
	var timer = NewTimer(time.Second)
	go func() {
		<-timer.C
		if time.Now().Sub(now) < time.Second {
			t.Fatal("Early func")
		}
	}()
	timer.Start()
	time.Sleep(time.Second * 2)
}

func TestChanStop(t *testing.T) {
	var timer = NewTimer(time.Second)
	go func() {
		<-timer.C
		t.Fatal("Stopping is not working")
	}()
	timer.Start()
	timer.Stop()
	time.Sleep(time.Second * 2)
}

func TestChanPause(t *testing.T) {
	var now = time.Now()
	var timer = NewTimer(time.Second)
	go func() {
		<-timer.C
		if time.Now().Sub(now) < time.Second*2 {
			t.Fatal("Early func")
		}
	}()
	timer.Start()
	time.Sleep(time.Microsecond * 500)
	timer.Pause()
	time.Sleep(time.Second)
	timer.Start()
	time.Sleep(time.Second)
}
<|endoftext|>


<commit_before>// Copyright 2015 Google Inc. All Rights Reserved.
//
// 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 forgetfs

import (
	"fmt"
	"os"

	"github.com/jacobsa/fuse"
	"github.com/jacobsa/fuse/fuseops"
	"github.com/jacobsa/fuse/fuseutil"
	"github.com/jacobsa/gcloud/syncutil"
)

// Create a file system whose sole contents are a file named "foo" and a
// directory named "bar".
//
// The file "foo" may be opened for reading and/or writing, but reads and
// writes aren't supported. Additionally, any non-existent file or directory
// name may be created within any directory, but the resulting inode will
// appear to have been unlinked immediately.
//
// The file system maintains reference counts for the inodes involved. It will
// panic if a reference count becomes negative or if an inode ID is re-used
// after we expect it to be dead. Its Check method may be used to check that
// there are no inodes with unexpected reference counts remaining, after
// unmounting.
func NewFileSystem() (fs *ForgetFS) {
	// Set up the actual file system.
	impl := &fsImpl{
		inodes: map[fuseops.InodeID]*inode{
			cannedID_Root: &inode{
				lookupCount: 1,
				attributes: fuseops.InodeAttributes{
					Nlink: 1,
					Mode:  0777 | os.ModeDir,
				},
			},
			cannedID_Foo: &inode{
				attributes: fuseops.InodeAttributes{
					Nlink: 1,
					Mode:  0777,
				},
			},
			cannedID_Bar: &inode{
				attributes: fuseops.InodeAttributes{
					Nlink: 1,
					Mode:  0777 | os.ModeDir,
				},
			},
		},
		nextInodeID: cannedID_Next,
	}

	impl.mu = syncutil.NewInvariantMutex(impl.checkInvariants)

	// Set up a wrapper that exposes only certain methods.
	fs = &ForgetFS{
		impl:   impl,
		server: fuseutil.NewFileSystemServer(impl),
	}

	return
}

////////////////////////////////////////////////////////////////////////
// ForgetFS
////////////////////////////////////////////////////////////////////////

type ForgetFS struct {
	impl   *fsImpl
	server fuse.Server
}

func (fs *ForgetFS) ServeOps(c *fuse.Connection) {
	fs.server.ServeOps(c)
}

// Panic if there are any inodes that have a non-zero reference count. For use
// after unmounting.
func (fs *ForgetFS) Check() {
	fs.impl.Check()
}

////////////////////////////////////////////////////////////////////////
// Actual implementation
////////////////////////////////////////////////////////////////////////

const (
	cannedID_Root = fuseops.RootInodeID + iota
	cannedID_Foo
	cannedID_Bar
	cannedID_Next
)

type fsImpl struct {
	fuseutil.NotImplementedFileSystem

	/////////////////////////
	// Mutable state
	/////////////////////////

	mu syncutil.InvariantMutex

	// An index of inode by ID, for all IDs we have issued.
	//
	// INVARIANT: For each v in inodes, v.lookupCount >= 0
	//
	// GUARDED_BY(mu)
	inodes map[fuseops.InodeID]*inode

	// The next ID to issue.
	//
	// INVARIANT: For each k in inodes, k < nextInodeID
	//
	// GUARDED_BY(mu)
	nextInodeID fuseops.InodeID
}

////////////////////////////////////////////////////////////////////////
// inode
////////////////////////////////////////////////////////////////////////

type inode struct {
	attributes fuseops.InodeAttributes

	// The current lookup count.
	lookupCount int

	// true if lookupCount has ever been positive.
	lookedUp bool
}

func (in *inode) Forgotten() bool {
	return in.lookedUp && in.lookupCount <= 0
}

func (in *inode) IncrementLookupCount() {
	in.lookupCount++
	in.lookedUp = true
}

////////////////////////////////////////////////////////////////////////
// Helpers
////////////////////////////////////////////////////////////////////////

// LOCKS_REQUIRED(fs.mu)
func (fs *fsImpl) checkInvariants() {
	// INVARIANT: For each v in inodes, v.lookupCount >= 0
	for _, v := range fs.inodes {
		if !(v.lookupCount >= 0) {
			panic("Negative lookup count")
		}
	}

	// INVARIANT: For each k in inodes, k < nextInodeID
	for k, _ := range fs.inodes {
		if !(k < fs.nextInodeID) {
			panic("Unexpectedly large inode ID")
		}
	}
}

// LOCKS_EXCLUDED(fs.mu)
func (fs *fsImpl) Check() {
	fs.mu.Lock()
	defer fs.mu.Unlock()

	for k, v := range fs.inodes {
		// Special case: the root inode should have an implicit lookup count of 1.
		if k == fuseops.RootInodeID {
			if v.lookupCount != 1 {
				panic(fmt.Sprintf("Root has lookup count %v", v.lookupCount))
			}

			continue
		}

		// Check other inodes.
		if v.lookupCount != 0 {
			panic(fmt.Sprintf("Inode %v has lookup count %v", k, v.lookupCount))
		}
	}
}

// Look up the inode and verify it hasn't been forgotten.
//
// LOCKS_REQUIRED(fs.mu)
func (fs *fsImpl) findInodeByID(id fuseops.InodeID) (in *inode) {
	in = fs.inodes[id]
	if in == nil {
		panic(fmt.Sprintf("Unknown inode: %v", id))
	}

	if in.Forgotten() {
		panic(fmt.Sprintf("Forgotten inode: %v", id))
	}

	return
}

////////////////////////////////////////////////////////////////////////
// FileSystem methods
////////////////////////////////////////////////////////////////////////

func (fs *fsImpl) Init(
	op *fuseops.InitOp) {
	var err error
	defer fuseutil.RespondToOp(op, &err)

	return
}

func (fs *fsImpl) LookUpInode(
	op *fuseops.LookUpInodeOp) {
	var err error
	defer fuseutil.RespondToOp(op, &err)

	fs.mu.Lock()
	defer fs.mu.Unlock()

	// Make sure the parent exists and has not been forgotten.
	_ = fs.findInodeByID(op.Parent)

	// Handle the names we support.
	var childID fuseops.InodeID
	switch {
	case op.Parent == cannedID_Root && op.Name == "foo":
		childID = cannedID_Foo

	case op.Parent == cannedID_Root && op.Name == "bar":
		childID = cannedID_Bar

	default:
		err = fuse.ENOENT
		return
	}

	// Look up the child.
	child := fs.findInodeByID(childID)
	child.IncrementLookupCount()

	// Return an appropriate entry.
	op.Entry = fuseops.ChildInodeEntry{
		Child:      childID,
		Attributes: child.attributes,
	}

	return
}

func (fs *fsImpl) GetInodeAttributes(
	op *fuseops.GetInodeAttributesOp) {
	var err error
	defer fuseutil.RespondToOp(op, &err)

	fs.mu.Lock()
	defer fs.mu.Unlock()

	// Find the inode, verifying that it has not been forgotten.
	in := fs.findInodeByID(op.Inode)

	// Return appropriate attributes.
	op.Attributes = in.attributes

	return
}

func (fs *fsImpl) OpenFile(
	op *fuseops.OpenFileOp) {
	var err error
	defer fuseutil.RespondToOp(op, &err)

	fs.mu.Lock()
	defer fs.mu.Unlock()

	// Verify that the inode has not been forgotten.
	_ = fs.findInodeByID(op.Inode)

	return
}
<commit_msg>fsImpl.ForgetInode<commit_after>// Copyright 2015 Google Inc. All Rights Reserved.
//
// 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 forgetfs

import (
	"fmt"
	"os"

	"github.com/jacobsa/fuse"
	"github.com/jacobsa/fuse/fuseops"
	"github.com/jacobsa/fuse/fuseutil"
	"github.com/jacobsa/gcloud/syncutil"
)

// Create a file system whose sole contents are a file named "foo" and a
// directory named "bar".
//
// The file "foo" may be opened for reading and/or writing, but reads and
// writes aren't supported. Additionally, any non-existent file or directory
// name may be created within any directory, but the resulting inode will
// appear to have been unlinked immediately.
//
// The file system maintains reference counts for the inodes involved. It will
// panic if a reference count becomes negative or if an inode ID is re-used
// after we expect it to be dead. Its Check method may be used to check that
// there are no inodes with unexpected reference counts remaining, after
// unmounting.
func NewFileSystem() (fs *ForgetFS) {
	// Set up the actual file system.
	impl := &fsImpl{
		inodes: map[fuseops.InodeID]*inode{
			cannedID_Root: &inode{
				lookupCount: 1,
				attributes: fuseops.InodeAttributes{
					Nlink: 1,
					Mode:  0777 | os.ModeDir,
				},
			},
			cannedID_Foo: &inode{
				attributes: fuseops.InodeAttributes{
					Nlink: 1,
					Mode:  0777,
				},
			},
			cannedID_Bar: &inode{
				attributes: fuseops.InodeAttributes{
					Nlink: 1,
					Mode:  0777 | os.ModeDir,
				},
			},
		},
		nextInodeID: cannedID_Next,
	}

	impl.mu = syncutil.NewInvariantMutex(impl.checkInvariants)

	// Set up a wrapper that exposes only certain methods.
	fs = &ForgetFS{
		impl:   impl,
		server: fuseutil.NewFileSystemServer(impl),
	}

	return
}

////////////////////////////////////////////////////////////////////////
// ForgetFS
////////////////////////////////////////////////////////////////////////

type ForgetFS struct {
	impl   *fsImpl
	server fuse.Server
}

func (fs *ForgetFS) ServeOps(c *fuse.Connection) {
	fs.server.ServeOps(c)
}

// Panic if there are any inodes that have a non-zero reference count. For use
// after unmounting.
func (fs *ForgetFS) Check() {
	fs.impl.Check()
}

////////////////////////////////////////////////////////////////////////
// Actual implementation
////////////////////////////////////////////////////////////////////////

const (
	cannedID_Root = fuseops.RootInodeID + iota
	cannedID_Foo
	cannedID_Bar
	cannedID_Next
)

type fsImpl struct {
	fuseutil.NotImplementedFileSystem

	/////////////////////////
	// Mutable state
	/////////////////////////

	mu syncutil.InvariantMutex

	// An index of inode by ID, for all IDs we have issued.
	//
	// GUARDED_BY(mu)
	inodes map[fuseops.InodeID]*inode

	// The next ID to issue.
	//
	// INVARIANT: For each k in inodes, k < nextInodeID
	//
	// GUARDED_BY(mu)
	nextInodeID fuseops.InodeID
}

////////////////////////////////////////////////////////////////////////
// inode
////////////////////////////////////////////////////////////////////////

type inode struct {
	attributes fuseops.InodeAttributes

	// The current lookup count.
	lookupCount uint64

	// true if lookupCount has ever been positive.
	lookedUp bool
}

func (in *inode) Forgotten() bool {
	return in.lookedUp && in.lookupCount == 0
}

func (in *inode) IncrementLookupCount() {
	in.lookupCount++
	in.lookedUp = true
}

func (in *inode) DecrementLookupCount(n uint64) {
	if in.lookupCount < n {
		panic(fmt.Sprintf(
			"Overly large decrement: %v, %v",
			in.lookupCount,
			n))
	}

	in.lookupCount -= n
}

////////////////////////////////////////////////////////////////////////
// Helpers
////////////////////////////////////////////////////////////////////////

// LOCKS_REQUIRED(fs.mu)
func (fs *fsImpl) checkInvariants() {
	// INVARIANT: For each k in inodes, k < nextInodeID
	for k, _ := range fs.inodes {
		if !(k < fs.nextInodeID) {
			panic("Unexpectedly large inode ID")
		}
	}
}

// LOCKS_EXCLUDED(fs.mu)
func (fs *fsImpl) Check() {
	fs.mu.Lock()
	defer fs.mu.Unlock()

	for k, v := range fs.inodes {
		// Special case: the root inode should have an implicit lookup count of 1.
		if k == fuseops.RootInodeID {
			if v.lookupCount != 1 {
				panic(fmt.Sprintf("Root has lookup count %v", v.lookupCount))
			}

			continue
		}

		// Check other inodes.
		if v.lookupCount != 0 {
			panic(fmt.Sprintf("Inode %v has lookup count %v", k, v.lookupCount))
		}
	}
}

// Look up the inode and verify it hasn't been forgotten.
//
// LOCKS_REQUIRED(fs.mu)
func (fs *fsImpl) findInodeByID(id fuseops.InodeID) (in *inode) {
	in = fs.inodes[id]
	if in == nil {
		panic(fmt.Sprintf("Unknown inode: %v", id))
	}

	if in.Forgotten() {
		panic(fmt.Sprintf("Forgotten inode: %v", id))
	}

	return
}

////////////////////////////////////////////////////////////////////////
// FileSystem methods
////////////////////////////////////////////////////////////////////////

func (fs *fsImpl) Init(
	op *fuseops.InitOp) {
	var err error
	defer fuseutil.RespondToOp(op, &err)

	return
}

func (fs *fsImpl) LookUpInode(
	op *fuseops.LookUpInodeOp) {
	var err error
	defer fuseutil.RespondToOp(op, &err)

	fs.mu.Lock()
	defer fs.mu.Unlock()

	// Make sure the parent exists and has not been forgotten.
	_ = fs.findInodeByID(op.Parent)

	// Handle the names we support.
	var childID fuseops.InodeID
	switch {
	case op.Parent == cannedID_Root && op.Name == "foo":
		childID = cannedID_Foo

	case op.Parent == cannedID_Root && op.Name == "bar":
		childID = cannedID_Bar

	default:
		err = fuse.ENOENT
		return
	}

	// Look up the child.
	child := fs.findInodeByID(childID)
	child.IncrementLookupCount()

	// Return an appropriate entry.
	op.Entry = fuseops.ChildInodeEntry{
		Child:      childID,
		Attributes: child.attributes,
	}

	return
}

func (fs *fsImpl) GetInodeAttributes(
	op *fuseops.GetInodeAttributesOp) {
	var err error
	defer fuseutil.RespondToOp(op, &err)

	fs.mu.Lock()
	defer fs.mu.Unlock()

	// Find the inode, verifying that it has not been forgotten.
	in := fs.findInodeByID(op.Inode)

	// Return appropriate attributes.
	op.Attributes = in.attributes

	return
}

func (fs *fsImpl) ForgetInode(
	op *fuseops.ForgetInodeOp) {
	var err error
	defer fuseutil.RespondToOp(op, &err)

	fs.mu.Lock()
	defer fs.mu.Unlock()

	// Find the inode and decrement its count.
	in := fs.findInodeByID(op.Inode)
	in.DecrementLookupCount(op.N)

	return
}

func (fs *fsImpl) OpenFile(
	op *fuseops.OpenFileOp) {
	var err error
	defer fuseutil.RespondToOp(op, &err)

	fs.mu.Lock()
	defer fs.mu.Unlock()

	// Verify that the inode has not been forgotten.
	_ = fs.findInodeByID(op.Inode)

	return
}
<|endoftext|>


<commit_before>// Copyright 2015 Google Inc. All Rights Reserved.
//
// 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 forgetfs

import (
	"fmt"

	"github.com/jacobsa/fuse"
	"github.com/jacobsa/fuse/fuseops"
	"github.com/jacobsa/fuse/fuseutil"
	"github.com/jacobsa/gcloud/syncutil"
)

// Create a file system whose sole contents are a file named "foo" and a
// directory named "bar".
//
// The file "foo" may be opened for reading and/or writing, but reads and
// writes aren't supported. Additionally, any non-existent file or directory
// name may be created within any directory, but the resulting inode will
// appear to have been unlinked immediately.
//
// The file system maintains reference counts for the inodes involved. It will
// panic if a reference count becomes negative or if an inode ID is re-used
// after we expect it to be dead. Its Check method may be used to check that
// there are no inodes with non-zero reference counts remaining, after
// unmounting.
func NewFileSystem() (fs *ForgetFS) {
	// Set up the actual file system.
	impl := &fsImpl{
		inodes: map[fuseops.InodeID]*inode{
			cannedID_Root: &inode{lookupCount: 1},
			cannedID_Foo:  &inode{},
			cannedID_Bar:  &inode{},
		},
		nextInodeID: cannedID_Next,
	}

	impl.mu = syncutil.NewInvariantMutex(impl.checkInvariants)

	// Set up a wrapper that exposes only certain methods.
	fs = &ForgetFS{
		impl:   impl,
		server: fuseutil.NewFileSystemServer(impl),
	}

	return
}

////////////////////////////////////////////////////////////////////////
// ForgetFS
////////////////////////////////////////////////////////////////////////

type ForgetFS struct {
	impl   *fsImpl
	server fuse.Server
}

func (fs *ForgetFS) ServeOps(c *fuse.Connection) {
	fs.server.ServeOps(c)
}

// Panic if there are any inodes that have a non-zero reference count. For use
// after unmounting.
func (fs *ForgetFS) Check() {
	fs.impl.Check()
}

////////////////////////////////////////////////////////////////////////
// Actual implementation
////////////////////////////////////////////////////////////////////////

const (
	cannedID_Root = fuseops.RootInodeID + iota
	cannedID_Foo
	cannedID_Bar
	cannedID_Next
)

type fsImpl struct {
	fuseutil.NotImplementedFileSystem

	/////////////////////////
	// Mutable state
	/////////////////////////

	mu syncutil.InvariantMutex

	// An index of inode by ID, for all IDs we have issued.
	//
	// INVARIANT: For each v in inodes, v.lookupCount >= 0
	//
	// GUARDED_BY(mu)
	inodes map[fuseops.InodeID]*inode

	// The next ID to issue.
	//
	// INVARIANT: For each k in inodes, k < nextInodeID
	//
	// GUARDED_BY(mu)
	nextInodeID fuseops.InodeID
}

type inode struct {
	// The current lookup count.
	lookupCount int
}

// LOCKS_REQUIRED(fs.mu)
func (fs *fsImpl) checkInvariants() {
	// INVARIANT: For each v in inodes, v.lookupCount >= 0
	for _, v := range fs.inodes {
		if !(v.lookupCount >= 0) {
			panic("Negative lookup count")
		}
	}

	// INVARIANT: For each k in inodes, k < nextInodeID
	for k, _ := range fs.inodes {
		if !(k < fs.nextInodeID) {
			panic("Unexpectedly large inode ID")
		}
	}
}

// LOCKS_EXCLUDED(fs.mu)
func (fs *fsImpl) Check() {
	fs.mu.Lock()
	defer fs.mu.Unlock()

	for k, v := range fs.inodes {
		if v.lookupCount != 0 {
			panic(fmt.Sprintf("Inode %v has lookup count %v", k, v.lookupCount))
		}
	}
}

func (fs *fsImpl) Init(
	op *fuseops.InitOp) {
	var err error
	defer fuseutil.RespondToOp(op, &err)

	return
}
<commit_msg>fsImpl.GetInodeAttributes<commit_after>// Copyright 2015 Google Inc. All Rights Reserved.
//
// 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 forgetfs

import (
	"fmt"
	"os"

	"github.com/jacobsa/fuse"
	"github.com/jacobsa/fuse/fuseops"
	"github.com/jacobsa/fuse/fuseutil"
	"github.com/jacobsa/gcloud/syncutil"
)

// Create a file system whose sole contents are a file named "foo" and a
// directory named "bar".
//
// The file "foo" may be opened for reading and/or writing, but reads and
// writes aren't supported. Additionally, any non-existent file or directory
// name may be created within any directory, but the resulting inode will
// appear to have been unlinked immediately.
//
// The file system maintains reference counts for the inodes involved. It will
// panic if a reference count becomes negative or if an inode ID is re-used
// after we expect it to be dead. Its Check method may be used to check that
// there are no inodes with non-zero reference counts remaining, after
// unmounting.
func NewFileSystem() (fs *ForgetFS) {
	// Set up the actual file system.
	impl := &fsImpl{
		inodes: map[fuseops.InodeID]*inode{
			cannedID_Root: &inode{
				lookupCount: 1,
				attributes: fuseops.InodeAttributes{
					Nlink: 1,
					Mode:  0777 | os.ModeDir,
				},
			},
			cannedID_Foo: &inode{
				attributes: fuseops.InodeAttributes{
					Nlink: 1,
					Mode:  0777,
				},
			},
			cannedID_Bar: &inode{
				attributes: fuseops.InodeAttributes{
					Nlink: 1,
					Mode:  0777 | os.ModeDir,
				},
			},
		},
		nextInodeID: cannedID_Next,
	}

	impl.mu = syncutil.NewInvariantMutex(impl.checkInvariants)

	// Set up a wrapper that exposes only certain methods.
	fs = &ForgetFS{
		impl:   impl,
		server: fuseutil.NewFileSystemServer(impl),
	}

	return
}

////////////////////////////////////////////////////////////////////////
// ForgetFS
////////////////////////////////////////////////////////////////////////

type ForgetFS struct {
	impl   *fsImpl
	server fuse.Server
}

func (fs *ForgetFS) ServeOps(c *fuse.Connection) {
	fs.server.ServeOps(c)
}

// Panic if there are any inodes that have a non-zero reference count. For use
// after unmounting.
func (fs *ForgetFS) Check() {
	fs.impl.Check()
}

////////////////////////////////////////////////////////////////////////
// Actual implementation
////////////////////////////////////////////////////////////////////////

const (
	cannedID_Root = fuseops.RootInodeID + iota
	cannedID_Foo
	cannedID_Bar
	cannedID_Next
)

type fsImpl struct {
	fuseutil.NotImplementedFileSystem

	/////////////////////////
	// Mutable state
	/////////////////////////

	mu syncutil.InvariantMutex

	// An index of inode by ID, for all IDs we have issued.
	//
	// INVARIANT: For each v in inodes, v.lookupCount >= 0
	//
	// GUARDED_BY(mu)
	inodes map[fuseops.InodeID]*inode

	// The next ID to issue.
	//
	// INVARIANT: For each k in inodes, k < nextInodeID
	//
	// GUARDED_BY(mu)
	nextInodeID fuseops.InodeID
}

type inode struct {
	attributes fuseops.InodeAttributes

	// The current lookup count.
	lookupCount int
}

// LOCKS_REQUIRED(fs.mu)
func (fs *fsImpl) checkInvariants() {
	// INVARIANT: For each v in inodes, v.lookupCount >= 0
	for _, v := range fs.inodes {
		if !(v.lookupCount >= 0) {
			panic("Negative lookup count")
		}
	}

	// INVARIANT: For each k in inodes, k < nextInodeID
	for k, _ := range fs.inodes {
		if !(k < fs.nextInodeID) {
			panic("Unexpectedly large inode ID")
		}
	}
}

// LOCKS_EXCLUDED(fs.mu)
func (fs *fsImpl) Check() {
	fs.mu.Lock()
	defer fs.mu.Unlock()

	for k, v := range fs.inodes {
		if v.lookupCount != 0 {
			panic(fmt.Sprintf("Inode %v has lookup count %v", k, v.lookupCount))
		}
	}
}

func (fs *fsImpl) Init(
	op *fuseops.InitOp) {
	var err error
	defer fuseutil.RespondToOp(op, &err)

	return
}

func (fs *fsImpl) GetInodeAttributes(
	op *fuseops.GetInodeAttributesOp) {
	var err error
	defer fuseutil.RespondToOp(op, &err)

	fs.mu.Lock()
	defer fs.mu.Unlock()

	// Find the inode, verifying that it has not been forgotten.
	in := fs.inodes[op.Inode]
	if in == nil {
		panic(fmt.Sprintf("Unknown inode: %v", op.Inode))
	}

	if in.lookupCount <= 0 {
		panic(fmt.Sprintf("Forgotten inode: %v", op.Inode))
	}

	// Return appropriate attributes.
	op.Attributes = in.attributes

	return
}
<|endoftext|>


<commit_before>// Copyright (c) 2016 The btcsuite developers
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package txsizes

import (
	"github.com/btcsuite/btcd/blockchain"
	"github.com/btcsuite/btcd/txscript"
	"github.com/btcsuite/btcd/wire"
)

// Worst case script and input/output size estimates.
const (
	// RedeemP2PKHSigScriptSize is the worst case (largest) serialize size
	// of a transaction input script that redeems a compressed P2PKH output.
	// It is calculated as:
	//
	//   - OP_DATA_73
	//   - 72 bytes DER signature + 1 byte sighash
	//   - OP_DATA_33
	//   - 33 bytes serialized compressed pubkey
	RedeemP2PKHSigScriptSize = 1 + 73 + 1 + 33

	// P2PKHPkScriptSize is the size of a transaction output script that
	// pays to a compressed pubkey hash.  It is calculated as:
	//
	//   - OP_DUP
	//   - OP_HASH160
	//   - OP_DATA_20
	//   - 20 bytes pubkey hash
	//   - OP_EQUALVERIFY
	//   - OP_CHECKSIG
	P2PKHPkScriptSize = 1 + 1 + 1 + 20 + 1 + 1

	// RedeemP2PKHInputSize is the worst case (largest) serialize size of a
	// transaction input redeeming a compressed P2PKH output.  It is
	// calculated as:
	//
	//   - 32 bytes previous tx
	//   - 4 bytes output index
	//   - 1 byte compact int encoding value 107
	//   - 107 bytes signature script
	//   - 4 bytes sequence
	RedeemP2PKHInputSize = 32 + 4 + 1 + RedeemP2PKHSigScriptSize + 4

	// P2PKHOutputSize is the serialize size of a transaction output with a
	// P2PKH output script.  It is calculated as:
	//
	//   - 8 bytes output value
	//   - 1 byte compact int encoding value 25
	//   - 25 bytes P2PKH output script
	P2PKHOutputSize = 8 + 1 + P2PKHPkScriptSize

	// P2WPKHPkScriptSize is the size of a transaction output script that
	// pays to a witness pubkey hash. It is calculated as:
	//
	//   - OP_0
	//   - OP_DATA_20
	//   - 20 bytes pubkey hash
	P2WPKHPkScriptSize = 1 + 1 + 20

	// P2WPKHOutputSize is the serialize size of a transaction output with a
	// P2WPKH output script. It is calculated as:
	//
	//   - 8 bytes output value
	//   - 1 byte compact int encoding value 22
	//   - 22 bytes P2PKH output script
	P2WPKHOutputSize = 8 + 1 + P2WPKHPkScriptSize

	// RedeemP2WPKHScriptSize is the size of a transaction input script
	// that spends a pay-to-witness-public-key hash (P2WPKH). The redeem
	// script for P2WPKH spends MUST be empty.
	RedeemP2WPKHScriptSize = 0

	// RedeemP2WPKHInputSize is the worst case size of a transaction
	// input redeeming a P2WPKH output. It is calculated as:
	//
	//   - 32 bytes previous tx
	//   - 4 bytes output index
	//   - 1 byte encoding empty redeem script
	//   - 0 bytes redeem script
	//   - 4 bytes sequence
	RedeemP2WPKHInputSize = 32 + 4 + 1 + RedeemP2WPKHScriptSize + 4

	// NestedP2WPKHPkScriptSize is the size of a transaction output script
	// that pays to a pay-to-witness-key hash nested in P2SH (P2SH-P2WPKH).
	// It is calculated as:
	//
	//   - OP_HASH160
	//   - OP_DATA_20
	//   - 20 bytes script hash
	//   - OP_EQUAL
	NestedP2WPKHPkScriptSize = 1 + 1 + 20 + 1

	// RedeemNestedP2WPKHScriptSize is the worst case size of a transaction
	// input script that redeems a pay-to-witness-key hash nested in P2SH
	// (P2SH-P2WPKH). It is calculated as:
	//
	//   - 1 byte compact int encoding value 22
	//   - OP_0
	//   - 1 byte compact int encoding value 20
	//   - 20 byte key hash
	RedeemNestedP2WPKHScriptSize = 1 + 1 + 1 + 20

	// RedeemNestedP2WPKHInputSize is the worst case size of a
	// transaction input redeeming a P2SH-P2WPKH output. It is
	// calculated as:
	//
	//   - 32 bytes previous tx
	//   - 4 bytes output index
	//   - 1 byte compact int encoding value 23
	//   - 23 bytes redeem script (scriptSig)
	//   - 4 bytes sequence
	RedeemNestedP2WPKHInputSize = 32 + 4 + 1 +
		RedeemNestedP2WPKHScriptSize + 4

	// RedeemP2WPKHInputWitnessWeight is the worst case weight of
	// a witness for spending P2WPKH and nested P2WPKH outputs. It
	// is calculated as:
	//
	//   - 1 wu compact int encoding value 2 (number of items)
	//   - 1 wu compact int encoding value 73
	//   - 72 wu DER signature + 1 wu sighash
	//   - 1 wu compact int encoding value 33
	//   - 33 wu serialized compressed pubkey
	RedeemP2WPKHInputWitnessWeight = 1 + 1 + 73 + 1 + 33
)

// SumOutputSerializeSizes sums up the serialized size of the supplied outputs.
func SumOutputSerializeSizes(outputs []*wire.TxOut) (serializeSize int) {
	for _, txOut := range outputs {
		serializeSize += txOut.SerializeSize()
	}
	return serializeSize
}

// EstimateSerializeSize returns a worst case serialize size estimate for a
// signed transaction that spends inputCount number of compressed P2PKH outputs
// and contains each transaction output from txOuts.  The estimated size is
// incremented for an additional P2PKH change output if addChangeOutput is true.
func EstimateSerializeSize(inputCount int, txOuts []*wire.TxOut, addChangeOutput bool) int {
	changeSize := 0
	outputCount := len(txOuts)
	if addChangeOutput {
		changeSize = P2PKHOutputSize
		outputCount++
	}

	// 8 additional bytes are for version and locktime
	return 8 + wire.VarIntSerializeSize(uint64(inputCount)) +
		wire.VarIntSerializeSize(uint64(outputCount)) +
		inputCount*RedeemP2PKHInputSize +
		SumOutputSerializeSizes(txOuts) +
		changeSize
}

// EstimateVirtualSize returns a worst case virtual size estimate for a
// signed transaction that spends the given number of P2PKH, P2WPKH and
// (nested) P2SH-P2WPKH outputs, and contains each transaction output
// from txOuts. The estimate is incremented for an additional P2PKH
// change output if addChangeOutput is true.
func EstimateVirtualSize(numP2PKHIns, numP2WPKHIns, numNestedP2WPKHIns int,
	txOuts []*wire.TxOut, changeScriptSize int) int {
	outputCount := len(txOuts)

	changeOutputSize := 0
	if changeScriptSize > 0 {
		changeOutputSize = 8 +
			wire.VarIntSerializeSize(uint64(changeScriptSize)) +
			changeScriptSize
		outputCount++
	}

	// Version 4 bytes + LockTime 4 bytes + Serialized var int size for the
	// number of transaction inputs and outputs + size of redeem scripts +
	// the size out the serialized outputs and change.
	baseSize := 8 +
		wire.VarIntSerializeSize(
			uint64(numP2PKHIns+numP2WPKHIns+numNestedP2WPKHIns)) +
		wire.VarIntSerializeSize(uint64(len(txOuts))) +
		numP2PKHIns*RedeemP2PKHInputSize +
		numP2WPKHIns*RedeemP2WPKHInputSize +
		numNestedP2WPKHIns*RedeemNestedP2WPKHInputSize +
		SumOutputSerializeSizes(txOuts) +
		changeOutputSize

	// If this transaction has any witness inputs, we must count the
	// witness data.
	witnessWeight := 0
	if numP2WPKHIns+numNestedP2WPKHIns > 0 {
		// Additional 2 weight units for segwit marker + flag.
		witnessWeight = 2 +
			wire.VarIntSerializeSize(
				uint64(numP2WPKHIns+numNestedP2WPKHIns)) +
			numP2WPKHIns*RedeemP2WPKHInputWitnessWeight +
			numNestedP2WPKHIns*RedeemP2WPKHInputWitnessWeight
	}

	// We add 3 to the witness weight to make sure the result is
	// always rounded up.
	return baseSize + (witnessWeight+3)/blockchain.WitnessScaleFactor
}

// GetMinInputVirtualSize returns the minimum number of vbytes that this input
// adds to a transaction.
func GetMinInputVirtualSize(pkScript []byte) int {
	var baseSize, witnessWeight int
	switch {
	// If this is a p2sh output, we assume this is a
	// nested P2WKH.
	case txscript.IsPayToScriptHash(pkScript):
		baseSize = RedeemNestedP2WPKHInputSize
		witnessWeight = RedeemP2WPKHInputWitnessWeight

	case txscript.IsPayToWitnessPubKeyHash(pkScript):
		baseSize = RedeemP2WPKHInputSize
		witnessWeight = RedeemP2WPKHInputWitnessWeight

	default:
		baseSize = RedeemP2PKHInputSize
	}

	return baseSize +
		(witnessWeight+blockchain.WitnessScaleFactor-1)/
			blockchain.WitnessScaleFactor
}
<commit_msg>Fix size estimates for taproot inputs<commit_after>// Copyright (c) 2016 The btcsuite developers
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package txsizes

import (
	"github.com/btcsuite/btcd/blockchain"
	"github.com/btcsuite/btcd/txscript"
	"github.com/btcsuite/btcd/wire"
)

// Worst case script and input/output size estimates.
const (
	// RedeemP2PKHSigScriptSize is the worst case (largest) serialize size
	// of a transaction input script that redeems a compressed P2PKH output.
	// It is calculated as:
	//
	//   - OP_DATA_73
	//   - 72 bytes DER signature + 1 byte sighash
	//   - OP_DATA_33
	//   - 33 bytes serialized compressed pubkey
	RedeemP2PKHSigScriptSize = 1 + 73 + 1 + 33

	// P2PKHPkScriptSize is the size of a transaction output script that
	// pays to a compressed pubkey hash.  It is calculated as:
	//
	//   - OP_DUP
	//   - OP_HASH160
	//   - OP_DATA_20
	//   - 20 bytes pubkey hash
	//   - OP_EQUALVERIFY
	//   - OP_CHECKSIG
	P2PKHPkScriptSize = 1 + 1 + 1 + 20 + 1 + 1

	// RedeemP2PKHInputSize is the worst case (largest) serialize size of a
	// transaction input redeeming a compressed P2PKH output.  It is
	// calculated as:
	//
	//   - 32 bytes previous tx
	//   - 4 bytes output index
	//   - 1 byte compact int encoding value 107
	//   - 107 bytes signature script
	//   - 4 bytes sequence
	RedeemP2PKHInputSize = 32 + 4 + 1 + RedeemP2PKHSigScriptSize + 4

	// P2PKHOutputSize is the serialize size of a transaction output with a
	// P2PKH output script.  It is calculated as:
	//
	//   - 8 bytes output value
	//   - 1 byte compact int encoding value 25
	//   - 25 bytes P2PKH output script
	P2PKHOutputSize = 8 + 1 + P2PKHPkScriptSize

	// P2WPKHPkScriptSize is the size of a transaction output script that
	// pays to a witness pubkey hash. It is calculated as:
	//
	//   - OP_0
	//   - OP_DATA_20
	//   - 20 bytes pubkey hash
	P2WPKHPkScriptSize = 1 + 1 + 20

	// P2WPKHOutputSize is the serialize size of a transaction output with a
	// P2WPKH output script. It is calculated as:
	//
	//   - 8 bytes output value
	//   - 1 byte compact int encoding value 22
	//   - 22 bytes P2PKH output script
	P2WPKHOutputSize = 8 + 1 + P2WPKHPkScriptSize

	// RedeemP2WPKHScriptSize is the size of a transaction input script
	// that spends a pay-to-witness-public-key hash (P2WPKH). The redeem
	// script for P2WPKH spends MUST be empty.
	RedeemP2WPKHScriptSize = 0

	// RedeemP2WPKHInputSize is the worst case size of a transaction
	// input redeeming a P2WPKH output. It is calculated as:
	//
	//   - 32 bytes previous tx
	//   - 4 bytes output index
	//   - 1 byte encoding empty redeem script
	//   - 0 bytes redeem script
	//   - 4 bytes sequence
	RedeemP2WPKHInputSize = 32 + 4 + 1 + RedeemP2WPKHScriptSize + 4

	// P2TRPkScriptSize is the size of a transaction output script that
	// pays to a taproot pubkey. It is calculated as:
	//
	//   - OP_1
	//   - OP_DATA_32
	//   - 32 bytes pubkey
	P2TRPkScriptSize = 1 + 1 + 32

	// P2TROutputSize is the serialize size of a transaction output with a
	// P2TR output script. It is calculated as:
	//
	//   - 8 bytes output value
	//   - 1 byte compact int encoding value 34
	//   - 34 bytes P2TR output script
	P2TROutputSize = 8 + 1 + P2TRPkScriptSize

	// RedeemP2TRScriptSize is the size of a transaction input script
	// that spends a pay-to-taproot hash (P2TR). The redeem
	// script for P2TR spends MUST be empty.
	RedeemP2TRScriptSize = 0

	// RedeemP2TRInputSize is the worst case size of a transaction
	// input redeeming a P2TR output. It is calculated as:
	//
	//   - 32 bytes previous tx
	//   - 4 bytes output index
	//   - 1 byte encoding empty redeem script
	//   - 0 bytes redeem script
	//   - 4 bytes sequence
	RedeemP2TRInputSize = 32 + 4 + 1 + RedeemP2TRScriptSize + 4

	// NestedP2WPKHPkScriptSize is the size of a transaction output script
	// that pays to a pay-to-witness-key hash nested in P2SH (P2SH-P2WPKH).
	// It is calculated as:
	//
	//   - OP_HASH160
	//   - OP_DATA_20
	//   - 20 bytes script hash
	//   - OP_EQUAL
	NestedP2WPKHPkScriptSize = 1 + 1 + 20 + 1

	// RedeemNestedP2WPKHScriptSize is the worst case size of a transaction
	// input script that redeems a pay-to-witness-key hash nested in P2SH
	// (P2SH-P2WPKH). It is calculated as:
	//
	//   - 1 byte compact int encoding value 22
	//   - OP_0
	//   - 1 byte compact int encoding value 20
	//   - 20 byte key hash
	RedeemNestedP2WPKHScriptSize = 1 + 1 + 1 + 20

	// RedeemNestedP2WPKHInputSize is the worst case size of a
	// transaction input redeeming a P2SH-P2WPKH output. It is
	// calculated as:
	//
	//   - 32 bytes previous tx
	//   - 4 bytes output index
	//   - 1 byte compact int encoding value 23
	//   - 23 bytes redeem script (scriptSig)
	//   - 4 bytes sequence
	RedeemNestedP2WPKHInputSize = 32 + 4 + 1 +
		RedeemNestedP2WPKHScriptSize + 4

	// RedeemP2WPKHInputWitnessWeight is the worst case weight of
	// a witness for spending P2WPKH and nested P2WPKH outputs. It
	// is calculated as:
	//
	//   - 1 wu compact int encoding value 2 (number of items)
	//   - 1 wu compact int encoding value 73
	//   - 72 wu DER signature + 1 wu sighash
	//   - 1 wu compact int encoding value 33
	//   - 33 wu serialized compressed pubkey
	RedeemP2WPKHInputWitnessWeight = 1 + 1 + 73 + 1 + 33

	// RedeemP2TRInputWitnessWeight is the worst case weight of
	// a witness for spending P2TR outputs. It
	// is calculated as:
	//
	//   - 1 wu compact int encoding value 65
	//   - 64 wu BIP-340 schnorr signature + 1 wu sighash
	RedeemP2TRInputWitnessWeight = 1 + 65
)

// SumOutputSerializeSizes sums up the serialized size of the supplied outputs.
func SumOutputSerializeSizes(outputs []*wire.TxOut) (serializeSize int) {
	for _, txOut := range outputs {
		serializeSize += txOut.SerializeSize()
	}
	return serializeSize
}

// EstimateSerializeSize returns a worst case serialize size estimate for a
// signed transaction that spends inputCount number of compressed P2PKH outputs
// and contains each transaction output from txOuts.  The estimated size is
// incremented for an additional P2PKH change output if addChangeOutput is true.
func EstimateSerializeSize(inputCount int, txOuts []*wire.TxOut, addChangeOutput bool) int {
	changeSize := 0
	outputCount := len(txOuts)
	if addChangeOutput {
		changeSize = P2PKHOutputSize
		outputCount++
	}

	// 8 additional bytes are for version and locktime
	return 8 + wire.VarIntSerializeSize(uint64(inputCount)) +
		wire.VarIntSerializeSize(uint64(outputCount)) +
		inputCount*RedeemP2PKHInputSize +
		SumOutputSerializeSizes(txOuts) +
		changeSize
}

// EstimateVirtualSize returns a worst case virtual size estimate for a
// signed transaction that spends the given number of P2PKH, P2TR, P2WPKH and
// (nested) P2SH-P2WPKH outputs, and contains each transaction output
// from txOuts. The estimate is incremented for an additional P2PKH
// change output if addChangeOutput is true.
func EstimateVirtualSize(numP2PKHIns, numP2TRIns, numP2WPKHIns, numNestedP2WPKHIns int,
	txOuts []*wire.TxOut, changeScriptSize int) int {
	outputCount := len(txOuts)

	changeOutputSize := 0
	if changeScriptSize > 0 {
		changeOutputSize = 8 +
			wire.VarIntSerializeSize(uint64(changeScriptSize)) +
			changeScriptSize
		outputCount++
	}

	// Version 4 bytes + LockTime 4 bytes + Serialized var int size for the
	// number of transaction inputs and outputs + size of redeem scripts +
	// the size out the serialized outputs and change.
	baseSize := 8 +
		wire.VarIntSerializeSize(
			uint64(numP2PKHIns+numP2TRIns+numP2WPKHIns+numNestedP2WPKHIns)) +
		wire.VarIntSerializeSize(uint64(len(txOuts))) +
		numP2PKHIns*RedeemP2PKHInputSize +
		numP2WPKHIns*RedeemP2WPKHInputSize +
		numP2TRIns*RedeemP2TRInputSize +
		numNestedP2WPKHIns*RedeemNestedP2WPKHInputSize +
		SumOutputSerializeSizes(txOuts) +
		changeOutputSize

	// If this transaction has any witness inputs, we must count the
	// witness data.
	witnessWeight := 0
	if numP2WPKHIns+numNestedP2WPKHIns+numP2TRIns > 0 {
		// Additional 2 weight units for segwit marker + flag.
		witnessWeight = 2 +
			wire.VarIntSerializeSize(
				uint64(numP2WPKHIns+numNestedP2WPKHIns+numP2TRIns)) +
			numP2WPKHIns*RedeemP2WPKHInputWitnessWeight +
			numP2TRIns*RedeemP2TRInputWitnessWeight +
			numNestedP2WPKHIns*RedeemP2WPKHInputWitnessWeight
	}

	// We add 3 to the witness weight to make sure the result is
	// always rounded up.
	return baseSize + (witnessWeight+3)/blockchain.WitnessScaleFactor
}

// GetMinInputVirtualSize returns the minimum number of vbytes that this input
// adds to a transaction.
func GetMinInputVirtualSize(pkScript []byte) int {
	var baseSize, witnessWeight int
	switch {
	// If this is a p2sh output, we assume this is a
	// nested P2WKH.
	case txscript.IsPayToScriptHash(pkScript):
		baseSize = RedeemNestedP2WPKHInputSize
		witnessWeight = RedeemP2WPKHInputWitnessWeight

	case txscript.IsPayToWitnessPubKeyHash(pkScript):
		baseSize = RedeemP2WPKHInputSize
		witnessWeight = RedeemP2WPKHInputWitnessWeight

	case txscript.IsPayToTaproot(pkScript):
		baseSize = RedeemP2TRInputSize
		witnessWeight = RedeemP2TRInputWitnessWeight

	default:
		baseSize = RedeemP2PKHInputSize
	}

	return baseSize +
		(witnessWeight+blockchain.WitnessScaleFactor-1)/
			blockchain.WitnessScaleFactor
}
<|endoftext|>