<commit_before>package core

import (
	"errors"

	"github.com/MG-RAST/AWE/lib/logger"
	"sort"
	//"sync"
)

type WorkQueue struct {
	//sync.RWMutex
	//workMap  map[string]*Workunit //all parsed workunits
	all      WorkunitMap
	Queue    WorkunitMap // WORK_STAT_QUEUED - waiting workunits
	Checkout WorkunitMap // WORK_STAT_CHECKOUT - workunits being checked out
	Suspend  WorkunitMap // WORK_STAT_SUSPEND - suspended workunits
}

func NewWorkQueue() *WorkQueue {
	wq := &WorkQueue{
		all:      *NewWorkunitMap(),
		Queue:    *NewWorkunitMap(),
		Checkout: *NewWorkunitMap(),
		Suspend:  *NewWorkunitMap(),
	}

	wq.all.Init("WorkQueue/workMap")
	wq.Queue.Init("WorkQueue/Queue")
	wq.Checkout.Init("WorkQueue/Checkout")
	wq.Suspend.Init("WorkQueue/Suspend")

	return wq
}

//--------accessor methods-------

func (wq *WorkQueue) Len() (int, error) {
	return wq.all.Len()
}

func (wq *WorkQueue) Add(workunit *Workunit) (err error) {
	if workunit.Id == "" {
		return errors.New("try to push a workunit with an empty id")
	}
	logger.Debug(3, "(WorkQueue/Add)")
	//id := workunit.Id

	err = wq.all.Set(workunit)
	if err != nil {
		return
	}

	err = wq.StatusChange("", workunit, WORK_STAT_QUEUED)
	if err != nil {
		return
	}

	return
}

func (wq *WorkQueue) Get(id string) (w *Workunit, ok bool, err error) {

	w, ok, err = wq.all.Get(id)

	return

}

func (wq *WorkQueue) GetForJob(jobid string) (worklist []*Workunit, err error) {

	workunits, err := wq.all.GetWorkunits()
	if err != nil {
		return
	}
	for _, work := range workunits {

		if jobid == getParentJobId(work.Id) {
			worklist = append(worklist, work)
		}
	}
	return
}

func (wq *WorkQueue) GetAll() (worklist []*Workunit, err error) {

	return wq.all.GetWorkunits()
}

func (wq *WorkQueue) Clean() (workids []string) {
	//wq.Lock()
	//defer wq.Unlock()
	workunt_list, err := wq.all.GetWorkunits()
	if err != nil {
		return
	}
	for _, work := range workunt_list {
		id := work.Id
		if work == nil || work.Info == nil {
			workids = append(workids, id)
			wq.Queue.Delete(id)
			wq.Checkout.Delete(id)
			wq.Suspend.Delete(id)
			wq.all.Delete(id)
			logger.Error("error: in WorkQueue workunit %s is nil, deleted from queue", id)
		}
	}
	return
}

func (wq *WorkQueue) Delete(id string) (err error) {

	err = wq.Queue.Delete(id)
	if err != nil {
		return
	}

	err = wq.Checkout.Delete(id)

	if err != nil {
		return
	}

	err = wq.Suspend.Delete(id)
	if err != nil {
		return
	}

	err = wq.all.Delete(id)
	if err != nil {
		return
	}

	return

}

func (wq *WorkQueue) Has(id string) (has bool, err error) {

	_, has, err = wq.all.Get(id)

	return
}

//--------end of accessors-------

func (wq *WorkQueue) StatusChange(id string, workunit *Workunit, new_status string) (err error) {
	//move workunit id between maps. no need to care about the old status because
	//delete function will do nothing if the operated map has no such key.

	if workunit == nil {
		var ok bool
		workunit, ok, err = wq.all.Get(id)
		if err != nil {
			return
		}
		if !ok {
			return errors.New("WQueue.statusChange: invalid workunit id:" + id)
		}
	}

	if workunit.State == new_status {
		return
	}
	if workunit.State != WORK_STAT_CHECKOUT && workunit.State != WORK_STAT_RESERVED {
		workunit.Client = ""
	}

	switch new_status {
	case WORK_STAT_CHECKOUT:
		wq.Queue.Delete(id)
		wq.Suspend.Delete(id)
		workunit.SetState(new_status)
		wq.Checkout.Set(workunit)
	case WORK_STAT_QUEUED:
		wq.Checkout.Delete(id)
		wq.Suspend.Delete(id)
		workunit.SetState(new_status)
		wq.Queue.Set(workunit)

	case WORK_STAT_SUSPEND:
		wq.Checkout.Delete(id)
		wq.Queue.Delete(id)
		workunit.SetState(new_status)
		wq.Suspend.Set(workunit)

	default:
		wq.Checkout.Delete(id)
		wq.Queue.Delete(id)
		wq.Suspend.Delete(id)
		workunit.SetState(new_status)

	}

	return
}

//select workunits, return a slice of ids based on given queuing policy and requested count
//if available is a positive value, filter by workunit input size
func (wq *WorkQueue) selectWorkunits(workunits WorkList, policy string, available int64, count int) (selected []*Workunit, err error) {
	logger.Debug(3, "starting selectWorkunits")

	if policy == "FCFS" {
		sort.Sort(byFCFS{workunits})
	}
	added := 0
	for _, work := range workunits {
		if added == count {
			break
		}

		inputSize := int64(0)
		for _, input := range work.Inputs {
			inputSize = inputSize + input.Size
		}
		// skip work that is too large for client
		if (available < 0) || (available > inputSize) {
			selected = append(selected, work)
			added = added + 1
		}

	}

	if len(selected) == 0 {
		err = errors.New(e.NoEligibleWorkunitFound)
		return
	}

	logger.Debug(3, "done with selectWorkunits")
	return
}

//queuing/prioritizing related functions
type WorkList []*Workunit

func (wl WorkList) Len() int      { return len(wl) }
func (wl WorkList) Swap(i, j int) { wl[i], wl[j] = wl[j], wl[i] }

type byFCFS struct{ WorkList }

//compare priority first, then FCFS (if priorities are the same)
func (s byFCFS) Less(i, j int) (ret bool) {
	p_i := s.WorkList[i].Info.Priority
	p_j := s.WorkList[j].Info.Priority
	switch {
	case p_i > p_j:
		return true
	case p_i < p_j:
		return false
	case p_i == p_j:
		return s.WorkList[i].Info.SubmitTime.Before(s.WorkList[j].Info.SubmitTime)
	}
	return
}
<commit_msg>added missing import<commit_after>package core

import (
	"errors"
	e "github.com/MG-RAST/AWE/lib/errors"
	"github.com/MG-RAST/AWE/lib/logger"
	"sort"
	//"sync"
)

type WorkQueue struct {
	//sync.RWMutex
	//workMap  map[string]*Workunit //all parsed workunits
	all      WorkunitMap
	Queue    WorkunitMap // WORK_STAT_QUEUED - waiting workunits
	Checkout WorkunitMap // WORK_STAT_CHECKOUT - workunits being checked out
	Suspend  WorkunitMap // WORK_STAT_SUSPEND - suspended workunits
}

func NewWorkQueue() *WorkQueue {
	wq := &WorkQueue{
		all:      *NewWorkunitMap(),
		Queue:    *NewWorkunitMap(),
		Checkout: *NewWorkunitMap(),
		Suspend:  *NewWorkunitMap(),
	}

	wq.all.Init("WorkQueue/workMap")
	wq.Queue.Init("WorkQueue/Queue")
	wq.Checkout.Init("WorkQueue/Checkout")
	wq.Suspend.Init("WorkQueue/Suspend")

	return wq
}

//--------accessor methods-------

func (wq *WorkQueue) Len() (int, error) {
	return wq.all.Len()
}

func (wq *WorkQueue) Add(workunit *Workunit) (err error) {
	if workunit.Id == "" {
		return errors.New("try to push a workunit with an empty id")
	}
	logger.Debug(3, "(WorkQueue/Add)")
	//id := workunit.Id

	err = wq.all.Set(workunit)
	if err != nil {
		return
	}

	err = wq.StatusChange("", workunit, WORK_STAT_QUEUED)
	if err != nil {
		return
	}

	return
}

func (wq *WorkQueue) Get(id string) (w *Workunit, ok bool, err error) {

	w, ok, err = wq.all.Get(id)

	return

}

func (wq *WorkQueue) GetForJob(jobid string) (worklist []*Workunit, err error) {

	workunits, err := wq.all.GetWorkunits()
	if err != nil {
		return
	}
	for _, work := range workunits {

		if jobid == getParentJobId(work.Id) {
			worklist = append(worklist, work)
		}
	}
	return
}

func (wq *WorkQueue) GetAll() (worklist []*Workunit, err error) {

	return wq.all.GetWorkunits()
}

func (wq *WorkQueue) Clean() (workids []string) {
	//wq.Lock()
	//defer wq.Unlock()
	workunt_list, err := wq.all.GetWorkunits()
	if err != nil {
		return
	}
	for _, work := range workunt_list {
		id := work.Id
		if work == nil || work.Info == nil {
			workids = append(workids, id)
			wq.Queue.Delete(id)
			wq.Checkout.Delete(id)
			wq.Suspend.Delete(id)
			wq.all.Delete(id)
			logger.Error("error: in WorkQueue workunit %s is nil, deleted from queue", id)
		}
	}
	return
}

func (wq *WorkQueue) Delete(id string) (err error) {

	err = wq.Queue.Delete(id)
	if err != nil {
		return
	}

	err = wq.Checkout.Delete(id)

	if err != nil {
		return
	}

	err = wq.Suspend.Delete(id)
	if err != nil {
		return
	}

	err = wq.all.Delete(id)
	if err != nil {
		return
	}

	return

}

func (wq *WorkQueue) Has(id string) (has bool, err error) {

	_, has, err = wq.all.Get(id)

	return
}

//--------end of accessors-------

func (wq *WorkQueue) StatusChange(id string, workunit *Workunit, new_status string) (err error) {
	//move workunit id between maps. no need to care about the old status because
	//delete function will do nothing if the operated map has no such key.

	if workunit == nil {
		var ok bool
		workunit, ok, err = wq.all.Get(id)
		if err != nil {
			return
		}
		if !ok {
			return errors.New("WQueue.statusChange: invalid workunit id:" + id)
		}
	}

	if workunit.State == new_status {
		return
	}
	if workunit.State != WORK_STAT_CHECKOUT && workunit.State != WORK_STAT_RESERVED {
		workunit.Client = ""
	}

	switch new_status {
	case WORK_STAT_CHECKOUT:
		wq.Queue.Delete(id)
		wq.Suspend.Delete(id)
		workunit.SetState(new_status)
		wq.Checkout.Set(workunit)
	case WORK_STAT_QUEUED:
		wq.Checkout.Delete(id)
		wq.Suspend.Delete(id)
		workunit.SetState(new_status)
		wq.Queue.Set(workunit)

	case WORK_STAT_SUSPEND:
		wq.Checkout.Delete(id)
		wq.Queue.Delete(id)
		workunit.SetState(new_status)
		wq.Suspend.Set(workunit)

	default:
		wq.Checkout.Delete(id)
		wq.Queue.Delete(id)
		wq.Suspend.Delete(id)
		workunit.SetState(new_status)

	}

	return
}

//select workunits, return a slice of ids based on given queuing policy and requested count
//if available is a positive value, filter by workunit input size
func (wq *WorkQueue) selectWorkunits(workunits WorkList, policy string, available int64, count int) (selected []*Workunit, err error) {
	logger.Debug(3, "starting selectWorkunits")

	if policy == "FCFS" {
		sort.Sort(byFCFS{workunits})
	}
	added := 0
	for _, work := range workunits {
		if added == count {
			break
		}

		inputSize := int64(0)
		for _, input := range work.Inputs {
			inputSize = inputSize + input.Size
		}
		// skip work that is too large for client
		if (available < 0) || (available > inputSize) {
			selected = append(selected, work)
			added = added + 1
		}

	}

	if len(selected) == 0 {
		err = errors.New(e.NoEligibleWorkunitFound)
		return
	}

	logger.Debug(3, "done with selectWorkunits")
	return
}

//queuing/prioritizing related functions
type WorkList []*Workunit

func (wl WorkList) Len() int      { return len(wl) }
func (wl WorkList) Swap(i, j int) { wl[i], wl[j] = wl[j], wl[i] }

type byFCFS struct{ WorkList }

//compare priority first, then FCFS (if priorities are the same)
func (s byFCFS) Less(i, j int) (ret bool) {
	p_i := s.WorkList[i].Info.Priority
	p_j := s.WorkList[j].Info.Priority
	switch {
	case p_i > p_j:
		return true
	case p_i < p_j:
		return false
	case p_i == p_j:
		return s.WorkList[i].Info.SubmitTime.Before(s.WorkList[j].Info.SubmitTime)
	}
	return
}
<|endoftext|>


<commit_before>// Copyright 2019, OpenTelemetry 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 ddsketch // import "go.opentelemetry.io/otel/sdk/metric/aggregator/ddsketch"

import (
	"context"
	"math"
	"sync"

	sdk "github.com/DataDog/sketches-go/ddsketch"

	"go.opentelemetry.io/otel/api/core"

	export "go.opentelemetry.io/otel/sdk/export/metric"
	"go.opentelemetry.io/otel/sdk/export/metric/aggregator"
)

// Config is an alias for the underlying DDSketch config object.
type Config = sdk.Config

// Aggregator aggregates measure events.
type Aggregator struct {
	lock       sync.Mutex
	cfg        *Config
	kind       core.NumberKind
	current    *sdk.DDSketch
	checkpoint *sdk.DDSketch
}

var _ export.Aggregator = &Aggregator{}
var _ aggregator.MinMaxSumCount = &Aggregator{}
var _ aggregator.Distribution = &Aggregator{}

// New returns a new DDSketch aggregator.
func New(cfg *Config, desc *export.Descriptor) *Aggregator {
	return &Aggregator{
		cfg:     cfg,
		kind:    desc.NumberKind(),
		current: sdk.NewDDSketch(cfg),
	}
}

// NewDefaultConfig returns a new, default DDSketch config.
//
// TODO: Should the Config constructor set minValue to -Inf to
// when the descriptor has absolute=false?  This requires providing
// values for alpha and maxNumBins, apparently.
func NewDefaultConfig() *Config {
	return sdk.NewDefaultConfig()
}

// Sum returns the sum of values in the checkpoint.
func (c *Aggregator) Sum() (core.Number, error) {
	return c.toNumber(c.checkpoint.Sum()), nil
}

// Count returns the number of values in the checkpoint.
func (c *Aggregator) Count() (int64, error) {
	return c.checkpoint.Count(), nil
}

// Max returns the maximum value in the checkpoint.
func (c *Aggregator) Max() (core.Number, error) {
	return c.Quantile(1)
}

// Min returns the minimum value in the checkpoint.
func (c *Aggregator) Min() (core.Number, error) {
	return c.Quantile(0)
}

// Quantile returns the estimated quantile of data in the checkpoint.
// It is an error if `q` is less than 0 or greated than 1.
func (c *Aggregator) Quantile(q float64) (core.Number, error) {
	if c.checkpoint.Count() == 0 {
		return core.Number(0), aggregator.ErrEmptyDataSet
	}
	f := c.checkpoint.Quantile(q)
	if math.IsNaN(f) {
		return core.Number(0), aggregator.ErrInvalidQuantile
	}
	return c.toNumber(f), nil
}

func (c *Aggregator) toNumber(f float64) core.Number {
	if c.kind == core.Float64NumberKind {
		return core.NewFloat64Number(f)
	}
	return core.NewInt64Number(int64(f))
}

// Checkpoint saves the current state and resets the current state to
// the empty set, taking a lock to prevent concurrent Update() calls.
func (c *Aggregator) Checkpoint(ctx context.Context, _ *export.Descriptor) {
	replace := sdk.NewDDSketch(c.cfg)

	c.lock.Lock()
	c.checkpoint = c.current
	c.current = replace
	c.lock.Unlock()
}

// Update adds the recorded measurement to the current data set.
// Update takes a lock to prevent concurrent Update() and Checkpoint()
// calls.
func (c *Aggregator) Update(_ context.Context, number core.Number, desc *export.Descriptor) error {
	c.lock.Lock()
	defer c.lock.Unlock()
	c.current.Add(number.CoerceToFloat64(desc.NumberKind()))
	return nil
}

// Merge combines two sketches into one.
func (c *Aggregator) Merge(oa export.Aggregator, d *export.Descriptor) error {
	o, _ := oa.(*Aggregator)
	if o == nil {
		return aggregator.NewInconsistentMergeError(c, oa)
	}

	c.checkpoint.Merge(o.checkpoint)
	return nil
}
<commit_msg>initialize checkpoint when creating ddsketch aggregator to prevent panic when merging into a empty one (#402) (#403)<commit_after>// Copyright 2019, OpenTelemetry 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 ddsketch // import "go.opentelemetry.io/otel/sdk/metric/aggregator/ddsketch"

import (
	"context"
	"math"
	"sync"

	sdk "github.com/DataDog/sketches-go/ddsketch"

	"go.opentelemetry.io/otel/api/core"

	export "go.opentelemetry.io/otel/sdk/export/metric"
	"go.opentelemetry.io/otel/sdk/export/metric/aggregator"
)

// Config is an alias for the underlying DDSketch config object.
type Config = sdk.Config

// Aggregator aggregates measure events.
type Aggregator struct {
	lock       sync.Mutex
	cfg        *Config
	kind       core.NumberKind
	current    *sdk.DDSketch
	checkpoint *sdk.DDSketch
}

var _ export.Aggregator = &Aggregator{}
var _ aggregator.MinMaxSumCount = &Aggregator{}
var _ aggregator.Distribution = &Aggregator{}

// New returns a new DDSketch aggregator.
func New(cfg *Config, desc *export.Descriptor) *Aggregator {
	return &Aggregator{
		cfg:        cfg,
		kind:       desc.NumberKind(),
		current:    sdk.NewDDSketch(cfg),
		checkpoint: sdk.NewDDSketch(cfg),
	}
}

// NewDefaultConfig returns a new, default DDSketch config.
//
// TODO: Should the Config constructor set minValue to -Inf to
// when the descriptor has absolute=false?  This requires providing
// values for alpha and maxNumBins, apparently.
func NewDefaultConfig() *Config {
	return sdk.NewDefaultConfig()
}

// Sum returns the sum of values in the checkpoint.
func (c *Aggregator) Sum() (core.Number, error) {
	return c.toNumber(c.checkpoint.Sum()), nil
}

// Count returns the number of values in the checkpoint.
func (c *Aggregator) Count() (int64, error) {
	return c.checkpoint.Count(), nil
}

// Max returns the maximum value in the checkpoint.
func (c *Aggregator) Max() (core.Number, error) {
	return c.Quantile(1)
}

// Min returns the minimum value in the checkpoint.
func (c *Aggregator) Min() (core.Number, error) {
	return c.Quantile(0)
}

// Quantile returns the estimated quantile of data in the checkpoint.
// It is an error if `q` is less than 0 or greated than 1.
func (c *Aggregator) Quantile(q float64) (core.Number, error) {
	if c.checkpoint.Count() == 0 {
		return core.Number(0), aggregator.ErrEmptyDataSet
	}
	f := c.checkpoint.Quantile(q)
	if math.IsNaN(f) {
		return core.Number(0), aggregator.ErrInvalidQuantile
	}
	return c.toNumber(f), nil
}

func (c *Aggregator) toNumber(f float64) core.Number {
	if c.kind == core.Float64NumberKind {
		return core.NewFloat64Number(f)
	}
	return core.NewInt64Number(int64(f))
}

// Checkpoint saves the current state and resets the current state to
// the empty set, taking a lock to prevent concurrent Update() calls.
func (c *Aggregator) Checkpoint(ctx context.Context, _ *export.Descriptor) {
	replace := sdk.NewDDSketch(c.cfg)

	c.lock.Lock()
	c.checkpoint = c.current
	c.current = replace
	c.lock.Unlock()
}

// Update adds the recorded measurement to the current data set.
// Update takes a lock to prevent concurrent Update() and Checkpoint()
// calls.
func (c *Aggregator) Update(_ context.Context, number core.Number, desc *export.Descriptor) error {
	c.lock.Lock()
	defer c.lock.Unlock()
	c.current.Add(number.CoerceToFloat64(desc.NumberKind()))
	return nil
}

// Merge combines two sketches into one.
func (c *Aggregator) Merge(oa export.Aggregator, d *export.Descriptor) error {
	o, _ := oa.(*Aggregator)
	if o == nil {
		return aggregator.NewInconsistentMergeError(c, oa)
	}

	c.checkpoint.Merge(o.checkpoint)
	return nil
}
<|endoftext|>


<commit_before>// (c) 2021, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package proposervm

import (
	"errors"
	"time"

	"github.com/ava-labs/avalanchego/ids"
	"github.com/ava-labs/avalanchego/snow/choices"
	"github.com/ava-labs/avalanchego/snow/consensus/snowman"
	"github.com/ava-labs/avalanchego/vms/proposervm/block"
)

var (
	errExpectedNoProposer = errors.New("expected no proposer to be named")
	errPChainHeightTooLow = errors.New("block P-chain height is too low")

	_ Block = &preForkBlock{}
)

// preForkBlock implements proposervm.Block
type preForkBlock struct {
	snowman.Block
	vm *VM
}

func (b *preForkBlock) Parent() ids.ID {
	return b.Block.Parent()
}

func (b *preForkBlock) Verify() error {
	b.vm.ctx.Log.Debug("Snowman++ calling verify on %s", b.ID())

	parent, err := b.vm.getBlock(b.Block.Parent())
	if err != nil {
		return err
	}
	return parent.verifyPreForkChild(b)
}

func (b *preForkBlock) Options() ([2]snowman.Block, error) {
	oracleBlk, ok := b.Block.(snowman.OracleBlock)
	if !ok {
		return [2]snowman.Block{}, snowman.ErrNotOracle
	}

	options, err := oracleBlk.Options()
	if err != nil {
		return [2]snowman.Block{}, err
	}
	// A pre-fork block's child options are always pre-fork blocks
	return [2]snowman.Block{
		&preForkBlock{
			Block: options[0],
			vm:    b.vm,
		},
		&preForkBlock{
			Block: options[1],
			vm:    b.vm,
		},
	}, nil
}

func (b *preForkBlock) verifyPreForkChild(child *preForkBlock) error {
	parentTimestamp := b.Timestamp()
	if !parentTimestamp.Before(b.vm.activationTime) {
		if err := verifyIsOracleBlock(b.Block); err != nil {
			b.vm.ctx.Log.Debug("a pre-fork block is only allowed after the fork time if the parent is an oracle block")
			return err
		}
		b.vm.ctx.Log.Debug("allowing pre-fork block after the fork time because the parent is an oracle block")
	}

	return child.Block.Verify()
}

func (b *preForkBlock) getInnerBlk() snowman.Block {
	return b.Block
}

// This method only returns nil once (during the transition)
func (b *preForkBlock) verifyPostForkChild(child *postForkBlock) error {
	if err := verifyIsNotOracleBlock(b.Block); err != nil {
		b.vm.ctx.Log.Debug("option block shouldn't have a signed child")
		return err
	}

	childID := child.ID()
	childPChainHeight := child.PChainHeight()
	currentPChainHeight, err := b.vm.PChainHeight()
	if err != nil {
		b.vm.ctx.Log.Error("couldn't retrieve current P-Chain height while verifying %s: %s", childID, err)
		return err
	}
	if childPChainHeight > currentPChainHeight {
		b.vm.ctx.Log.Warn("Snowman++ verify - dropped post-fork block; expected chid's P-Chain height to be <=%d but got %d",
			currentPChainHeight, childPChainHeight)
		return errPChainHeightNotReached
	}
	if childPChainHeight < b.vm.minimumPChainHeight {
		b.vm.ctx.Log.Warn("Snowman++ verify - dropped post-fork block; expected chid's P-Chain height to be >=%d but got %d",
			b.vm.minimumPChainHeight, childPChainHeight)
		return errPChainHeightTooLow
	}

	// Make sure [b] is the parent of [child]'s inner block
	expectedInnerParentID := b.ID()
	innerParentID := child.innerBlk.Parent()
	if innerParentID != expectedInnerParentID {
		b.vm.ctx.Log.Warn("Snowman++ verify - dropped post-fork block; expected inner parent %s but got %s",
			expectedInnerParentID, innerParentID)
		return errInnerParentMismatch
	}

	// A *preForkBlock can only have a *postForkBlock child
	// if the *preForkBlock is the last *preForkBlock before activation takes effect
	// (its timestamp is at or after the activation time)
	parentTimestamp := b.Timestamp()
	if parentTimestamp.Before(b.vm.activationTime) {
		return errProposersNotActivated
	}

	// Child's timestamp must be at or after its parent's timestamp
	childTimestamp := child.Timestamp()
	if childTimestamp.Before(parentTimestamp) {
		b.vm.ctx.Log.Warn("Snowman++ verify - dropped post-fork block; expected child's timestamp (%s) to be at or after parent's timestamp (%s)",
			childTimestamp, parentTimestamp)
		return errTimeNotMonotonic
	}

	// Child timestamp can't be too far in the future
	maxTimestamp := b.vm.Time().Add(maxSkew)
	if childTimestamp.After(maxTimestamp) {
		b.vm.ctx.Log.Warn("Snowman++ verify - dropped post-fork block; block's timestamp (%s) is after the synchrony bound (%s)",
			childTimestamp, maxTimestamp)
		return errTimeTooAdvanced
	}

	proposer := child.SignedBlock.Proposer()
	if proposer != ids.ShortEmpty {
		return errExpectedNoProposer
	}

	// Verify the lack of signature on the node
	if err := child.SignedBlock.Verify(b.vm.ctx.ChainID); err != nil {
		return err
	}

	// If inner block's Verify returned true, don't call it again.
	// Note that if [child.innerBlk.Verify] returns nil,
	// this method returns nil. This must always remain the case to
	// maintain the inner block's invariant that if it's Verify()
	// returns nil, it is eventually accepted/rejected.
	if !b.vm.Tree.Contains(child.innerBlk) {
		if err := child.innerBlk.Verify(); err != nil {
			return err
		}
		b.vm.Tree.Add(child.innerBlk)
	}

	b.vm.verifiedBlocks[childID] = child
	return nil
}

func (b *preForkBlock) verifyPostForkOption(child *postForkOption) error {
	b.vm.ctx.Log.Debug("post-fork option has pre-fork block as parent")
	return errUnexpectedBlockType
}

func (b *preForkBlock) buildChild(innerBlock snowman.Block) (Block, error) {
	parentTimestamp := b.Timestamp()
	if parentTimestamp.Before(b.vm.activationTime) {
		// The chain hasn't forked yet
		res := &preForkBlock{
			Block: innerBlock,
			vm:    b.vm,
		}
		b.vm.ctx.Log.Debug("Snowman++ build pre-fork block %s - timestamp parent block %v",
			res.ID(), b.Timestamp())

		return res, nil
	}

	// The chain is currently forking

	parentID := b.ID()
	newTimestamp := b.vm.Time().Truncate(time.Second)
	if newTimestamp.Before(parentTimestamp) {
		newTimestamp = parentTimestamp
	}

	pChainHeight, err := b.vm.PChainHeight()
	if err != nil {
		return nil, err
	}

	statelessBlock, err := block.BuildUnsigned(
		parentID,
		newTimestamp,
		pChainHeight,
		innerBlock.Bytes(),
	)
	if err != nil {
		return nil, err
	}

	blk := &postForkBlock{
		SignedBlock: statelessBlock,
		postForkCommonComponents: postForkCommonComponents{
			vm:       b.vm,
			innerBlk: innerBlock,
			status:   choices.Processing,
		},
	}

	b.vm.ctx.Log.Debug("Snowman++ build post-fork block %s - parent timestamp %v, expected delay NA, block timestamp %v.",
		blk.ID(), parentTimestamp, newTimestamp)
	return blk, b.vm.storePostForkBlock(blk.SignedBlock, blk.status)
}

func (b *preForkBlock) pChainHeight() (uint64, error) {
	return 0, nil
}
<commit_msg>standardize method order<commit_after>// (c) 2021, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package proposervm

import (
	"errors"
	"time"

	"github.com/ava-labs/avalanchego/ids"
	"github.com/ava-labs/avalanchego/snow/choices"
	"github.com/ava-labs/avalanchego/snow/consensus/snowman"
	"github.com/ava-labs/avalanchego/vms/proposervm/block"
)

var (
	errExpectedNoProposer = errors.New("expected no proposer to be named")
	errPChainHeightTooLow = errors.New("block P-chain height is too low")

	_ Block = &preForkBlock{}
)

// preForkBlock implements proposervm.Block
type preForkBlock struct {
	snowman.Block
	vm *VM
}

func (b *preForkBlock) Parent() ids.ID {
	return b.Block.Parent()
}

func (b *preForkBlock) Verify() error {
	b.vm.ctx.Log.Debug("Snowman++ calling verify on %s", b.ID())

	parent, err := b.vm.getBlock(b.Block.Parent())
	if err != nil {
		return err
	}
	return parent.verifyPreForkChild(b)
}

func (b *preForkBlock) Options() ([2]snowman.Block, error) {
	oracleBlk, ok := b.Block.(snowman.OracleBlock)
	if !ok {
		return [2]snowman.Block{}, snowman.ErrNotOracle
	}

	options, err := oracleBlk.Options()
	if err != nil {
		return [2]snowman.Block{}, err
	}
	// A pre-fork block's child options are always pre-fork blocks
	return [2]snowman.Block{
		&preForkBlock{
			Block: options[0],
			vm:    b.vm,
		},
		&preForkBlock{
			Block: options[1],
			vm:    b.vm,
		},
	}, nil
}

func (b *preForkBlock) getInnerBlk() snowman.Block {
	return b.Block
}

func (b *preForkBlock) verifyPreForkChild(child *preForkBlock) error {
	parentTimestamp := b.Timestamp()
	if !parentTimestamp.Before(b.vm.activationTime) {
		if err := verifyIsOracleBlock(b.Block); err != nil {
			b.vm.ctx.Log.Debug("a pre-fork block is only allowed after the fork time if the parent is an oracle block")
			return err
		}
		b.vm.ctx.Log.Debug("allowing pre-fork block after the fork time because the parent is an oracle block")
	}

	return child.Block.Verify()
}

// This method only returns nil once (during the transition)
func (b *preForkBlock) verifyPostForkChild(child *postForkBlock) error {
	if err := verifyIsNotOracleBlock(b.Block); err != nil {
		b.vm.ctx.Log.Debug("option block shouldn't have a signed child")
		return err
	}

	childID := child.ID()
	childPChainHeight := child.PChainHeight()
	currentPChainHeight, err := b.vm.PChainHeight()
	if err != nil {
		b.vm.ctx.Log.Error("couldn't retrieve current P-Chain height while verifying %s: %s", childID, err)
		return err
	}
	if childPChainHeight > currentPChainHeight {
		b.vm.ctx.Log.Warn("Snowman++ verify - dropped post-fork block; expected chid's P-Chain height to be <=%d but got %d",
			currentPChainHeight, childPChainHeight)
		return errPChainHeightNotReached
	}
	if childPChainHeight < b.vm.minimumPChainHeight {
		b.vm.ctx.Log.Warn("Snowman++ verify - dropped post-fork block; expected chid's P-Chain height to be >=%d but got %d",
			b.vm.minimumPChainHeight, childPChainHeight)
		return errPChainHeightTooLow
	}

	// Make sure [b] is the parent of [child]'s inner block
	expectedInnerParentID := b.ID()
	innerParentID := child.innerBlk.Parent()
	if innerParentID != expectedInnerParentID {
		b.vm.ctx.Log.Warn("Snowman++ verify - dropped post-fork block; expected inner parent %s but got %s",
			expectedInnerParentID, innerParentID)
		return errInnerParentMismatch
	}

	// A *preForkBlock can only have a *postForkBlock child
	// if the *preForkBlock is the last *preForkBlock before activation takes effect
	// (its timestamp is at or after the activation time)
	parentTimestamp := b.Timestamp()
	if parentTimestamp.Before(b.vm.activationTime) {
		return errProposersNotActivated
	}

	// Child's timestamp must be at or after its parent's timestamp
	childTimestamp := child.Timestamp()
	if childTimestamp.Before(parentTimestamp) {
		b.vm.ctx.Log.Warn("Snowman++ verify - dropped post-fork block; expected child's timestamp (%s) to be at or after parent's timestamp (%s)",
			childTimestamp, parentTimestamp)
		return errTimeNotMonotonic
	}

	// Child timestamp can't be too far in the future
	maxTimestamp := b.vm.Time().Add(maxSkew)
	if childTimestamp.After(maxTimestamp) {
		b.vm.ctx.Log.Warn("Snowman++ verify - dropped post-fork block; block's timestamp (%s) is after the synchrony bound (%s)",
			childTimestamp, maxTimestamp)
		return errTimeTooAdvanced
	}

	proposer := child.SignedBlock.Proposer()
	if proposer != ids.ShortEmpty {
		return errExpectedNoProposer
	}

	// Verify the lack of signature on the node
	if err := child.SignedBlock.Verify(b.vm.ctx.ChainID); err != nil {
		return err
	}

	// If inner block's Verify returned true, don't call it again.
	// Note that if [child.innerBlk.Verify] returns nil,
	// this method returns nil. This must always remain the case to
	// maintain the inner block's invariant that if it's Verify()
	// returns nil, it is eventually accepted/rejected.
	if !b.vm.Tree.Contains(child.innerBlk) {
		if err := child.innerBlk.Verify(); err != nil {
			return err
		}
		b.vm.Tree.Add(child.innerBlk)
	}

	b.vm.verifiedBlocks[childID] = child
	return nil
}

func (b *preForkBlock) verifyPostForkOption(child *postForkOption) error {
	b.vm.ctx.Log.Debug("post-fork option has pre-fork block as parent")
	return errUnexpectedBlockType
}

func (b *preForkBlock) buildChild(innerBlock snowman.Block) (Block, error) {
	parentTimestamp := b.Timestamp()
	if parentTimestamp.Before(b.vm.activationTime) {
		// The chain hasn't forked yet
		res := &preForkBlock{
			Block: innerBlock,
			vm:    b.vm,
		}
		b.vm.ctx.Log.Debug("Snowman++ build pre-fork block %s - timestamp parent block %v",
			res.ID(), b.Timestamp())

		return res, nil
	}

	// The chain is currently forking

	parentID := b.ID()
	newTimestamp := b.vm.Time().Truncate(time.Second)
	if newTimestamp.Before(parentTimestamp) {
		newTimestamp = parentTimestamp
	}

	pChainHeight, err := b.vm.PChainHeight()
	if err != nil {
		return nil, err
	}

	statelessBlock, err := block.BuildUnsigned(
		parentID,
		newTimestamp,
		pChainHeight,
		innerBlock.Bytes(),
	)
	if err != nil {
		return nil, err
	}

	blk := &postForkBlock{
		SignedBlock: statelessBlock,
		postForkCommonComponents: postForkCommonComponents{
			vm:       b.vm,
			innerBlk: innerBlock,
			status:   choices.Processing,
		},
	}

	b.vm.ctx.Log.Debug("Snowman++ build post-fork block %s - parent timestamp %v, expected delay NA, block timestamp %v.",
		blk.ID(), parentTimestamp, newTimestamp)
	return blk, b.vm.storePostForkBlock(blk.SignedBlock, blk.status)
}

func (b *preForkBlock) pChainHeight() (uint64, error) {
	return 0, nil
}
<|endoftext|>


<commit_before>package deb

import (
	"errors"
	"github.com/smira/aptly/database"
	. "launchpad.net/gocheck"
)

type PackageRefListSuite struct {
	// Simple list with "real" packages from stanzas
	list                   *PackageList
	p1, p2, p3, p4, p5, p6 *Package
}

var _ = Suite(&PackageRefListSuite{})

func (s *PackageRefListSuite) SetUpTest(c *C) {
	s.list = NewPackageList()

	s.p1 = NewPackageFromControlFile(packageStanza.Copy())
	s.p2 = NewPackageFromControlFile(packageStanza.Copy())
	stanza := packageStanza.Copy()
	stanza["Package"] = "mars-invaders"
	s.p3 = NewPackageFromControlFile(stanza)
	stanza = packageStanza.Copy()
	stanza["Size"] = "42"
	s.p4 = NewPackageFromControlFile(stanza)
	stanza = packageStanza.Copy()
	stanza["Package"] = "lonely-strangers"
	s.p5 = NewPackageFromControlFile(stanza)
	stanza = packageStanza.Copy()
	stanza["Version"] = "99.1"
	s.p6 = NewPackageFromControlFile(stanza)
}

func (s *PackageRefListSuite) TestNewPackageListFromRefList(c *C) {
	db, _ := database.OpenDB(c.MkDir())
	coll := NewPackageCollection(db)
	coll.Update(s.p1)
	coll.Update(s.p3)

	s.list.Add(s.p1)
	s.list.Add(s.p3)
	s.list.Add(s.p5)
	s.list.Add(s.p6)

	reflist := NewPackageRefListFromPackageList(s.list)

	_, err := NewPackageListFromRefList(reflist, coll, nil)
	c.Assert(err, ErrorMatches, "unable to load package with key.*")

	coll.Update(s.p5)
	coll.Update(s.p6)

	list, err := NewPackageListFromRefList(reflist, coll, nil)
	c.Assert(err, IsNil)
	c.Check(list.Len(), Equals, 4)
	c.Check(list.Add(s.p4), ErrorMatches, "conflict in package.*")

	list, err = NewPackageListFromRefList(nil, coll, nil)
	c.Assert(err, IsNil)
	c.Check(list.Len(), Equals, 0)
}

func (s *PackageRefListSuite) TestNewPackageRefList(c *C) {
	s.list.Add(s.p1)
	s.list.Add(s.p3)
	s.list.Add(s.p5)
	s.list.Add(s.p6)

	reflist := NewPackageRefListFromPackageList(s.list)
	c.Assert(reflist.Len(), Equals, 4)
	c.Check(reflist.Refs[0], DeepEquals, []byte(s.p1.Key("")))
	c.Check(reflist.Refs[1], DeepEquals, []byte(s.p6.Key("")))
	c.Check(reflist.Refs[2], DeepEquals, []byte(s.p5.Key("")))
	c.Check(reflist.Refs[3], DeepEquals, []byte(s.p3.Key("")))

	reflist = NewPackageRefList()
	c.Check(reflist.Len(), Equals, 0)
}

func (s *PackageRefListSuite) TestPackageRefListEncodeDecode(c *C) {
	s.list.Add(s.p1)
	s.list.Add(s.p3)
	s.list.Add(s.p5)
	s.list.Add(s.p6)

	reflist := NewPackageRefListFromPackageList(s.list)

	reflist2 := &PackageRefList{}
	err := reflist2.Decode(reflist.Encode())
	c.Assert(err, IsNil)
	c.Check(reflist2.Len(), Equals, reflist.Len())
	c.Check(reflist2.Refs, DeepEquals, reflist.Refs)
}

func (s *PackageRefListSuite) TestPackageRefListForeach(c *C) {
	s.list.Add(s.p1)
	s.list.Add(s.p3)
	s.list.Add(s.p5)
	s.list.Add(s.p6)

	reflist := NewPackageRefListFromPackageList(s.list)

	Len := 0
	err := reflist.ForEach(func([]byte) error {
		Len++
		return nil
	})

	c.Check(Len, Equals, 4)
	c.Check(err, IsNil)

	e := errors.New("b")

	err = reflist.ForEach(func([]byte) error {
		return e
	})

	c.Check(err, Equals, e)
}

func (s *PackageRefListSuite) TestSubstract(c *C) {
	r1 := []byte("r1")
	r2 := []byte("r2")
	r3 := []byte("r3")
	r4 := []byte("r4")
	r5 := []byte("r5")

	empty := &PackageRefList{Refs: [][]byte{}}
	l1 := &PackageRefList{Refs: [][]byte{r1, r2, r3, r4}}
	l2 := &PackageRefList{Refs: [][]byte{r1, r3}}
	l3 := &PackageRefList{Refs: [][]byte{r2, r4}}
	l4 := &PackageRefList{Refs: [][]byte{r4, r5}}
	l5 := &PackageRefList{Refs: [][]byte{r1, r2, r3}}

	c.Check(l1.Substract(empty), DeepEquals, l1)
	c.Check(l1.Substract(l2), DeepEquals, l3)
	c.Check(l1.Substract(l3), DeepEquals, l2)
	c.Check(l1.Substract(l4), DeepEquals, l5)
	c.Check(empty.Substract(l1), DeepEquals, empty)
	c.Check(l2.Substract(l3), DeepEquals, l2)
}

func (s *PackageRefListSuite) TestDiff(c *C) {
	db, _ := database.OpenDB(c.MkDir())
	coll := NewPackageCollection(db)

	packages := []*Package{
		&Package{Name: "lib", Version: "1.0", Architecture: "i386"},      //0
		&Package{Name: "dpkg", Version: "1.7", Architecture: "i386"},     //1
		&Package{Name: "data", Version: "1.1~bp1", Architecture: "all"},  //2
		&Package{Name: "app", Version: "1.1~bp1", Architecture: "i386"},  //3
		&Package{Name: "app", Version: "1.1~bp2", Architecture: "i386"},  //4
		&Package{Name: "app", Version: "1.1~bp2", Architecture: "amd64"}, //5
		&Package{Name: "xyz", Version: "3.0", Architecture: "sparc"},     //6
	}

	for _, p := range packages {
		coll.Update(p)
	}

	listA := NewPackageList()
	listA.Add(packages[0])
	listA.Add(packages[1])
	listA.Add(packages[2])
	listA.Add(packages[3])
	listA.Add(packages[6])

	listB := NewPackageList()
	listB.Add(packages[0])
	listB.Add(packages[2])
	listB.Add(packages[4])
	listB.Add(packages[5])

	reflistA := NewPackageRefListFromPackageList(listA)
	reflistB := NewPackageRefListFromPackageList(listB)

	diffAA, err := reflistA.Diff(reflistA, coll)
	c.Check(err, IsNil)
	c.Check(diffAA, HasLen, 0)

	diffAB, err := reflistA.Diff(reflistB, coll)
	c.Check(err, IsNil)
	c.Check(diffAB, HasLen, 4)

	c.Check(diffAB[0].Left, IsNil)
	c.Check(diffAB[0].Right.String(), Equals, "app_1.1~bp2_amd64")

	c.Check(diffAB[1].Left.String(), Equals, "app_1.1~bp1_i386")
	c.Check(diffAB[1].Right.String(), Equals, "app_1.1~bp2_i386")

	c.Check(diffAB[2].Left.String(), Equals, "dpkg_1.7_i386")
	c.Check(diffAB[2].Right, IsNil)

	c.Check(diffAB[3].Left.String(), Equals, "xyz_3.0_sparc")
	c.Check(diffAB[3].Right, IsNil)

	diffBA, err := reflistB.Diff(reflistA, coll)
	c.Check(err, IsNil)
	c.Check(diffBA, HasLen, 4)

	c.Check(diffBA[0].Right, IsNil)
	c.Check(diffBA[0].Left.String(), Equals, "app_1.1~bp2_amd64")

	c.Check(diffBA[1].Right.String(), Equals, "app_1.1~bp1_i386")
	c.Check(diffBA[1].Left.String(), Equals, "app_1.1~bp2_i386")

	c.Check(diffBA[2].Right.String(), Equals, "dpkg_1.7_i386")
	c.Check(diffBA[2].Left, IsNil)

	c.Check(diffBA[3].Right.String(), Equals, "xyz_3.0_sparc")
	c.Check(diffBA[3].Left, IsNil)

}

func (s *PackageRefListSuite) TestMerge(c *C) {
	db, _ := database.OpenDB(c.MkDir())
	coll := NewPackageCollection(db)

	packages := []*Package{
		&Package{Name: "lib", Version: "1.0", Architecture: "i386"},      //0
		&Package{Name: "dpkg", Version: "1.7", Architecture: "i386"},     //1
		&Package{Name: "data", Version: "1.1~bp1", Architecture: "all"},  //2
		&Package{Name: "app", Version: "1.1~bp1", Architecture: "i386"},  //3
		&Package{Name: "app", Version: "1.1~bp2", Architecture: "i386"},  //4
		&Package{Name: "app", Version: "1.1~bp2", Architecture: "amd64"}, //5
		&Package{Name: "dpkg", Version: "1.0", Architecture: "i386"},     //6
		&Package{Name: "xyz", Version: "1.0", Architecture: "sparc"},     //7
	}

	for _, p := range packages {
		coll.Update(p)
	}

	listA := NewPackageList()
	listA.Add(packages[0])
	listA.Add(packages[1])
	listA.Add(packages[2])
	listA.Add(packages[3])
	listA.Add(packages[7])

	listB := NewPackageList()
	listB.Add(packages[0])
	listB.Add(packages[2])
	listB.Add(packages[4])
	listB.Add(packages[5])
	listB.Add(packages[6])

	reflistA := NewPackageRefListFromPackageList(listA)
	reflistB := NewPackageRefListFromPackageList(listB)

	toStrSlice := func(reflist *PackageRefList) (result []string) {
		result = make([]string, reflist.Len())
		for i, r := range reflist.Refs {
			result[i] = string(r)
		}
		return
	}

	mergeAB := reflistA.Merge(reflistB, true)
	mergeBA := reflistB.Merge(reflistA, true)

	c.Check(toStrSlice(mergeAB), DeepEquals,
		[]string{"Pall data 1.1~bp1", "Pamd64 app 1.1~bp2", "Pi386 app 1.1~bp2", "Pi386 dpkg 1.0", "Pi386 lib 1.0", "Psparc xyz 1.0"})
	c.Check(toStrSlice(mergeBA), DeepEquals,
		[]string{"Pall data 1.1~bp1", "Pamd64 app 1.1~bp2", "Pi386 app 1.1~bp1", "Pi386 dpkg 1.7", "Pi386 lib 1.0", "Psparc xyz 1.0"})

	mergeABall := reflistA.Merge(reflistB, false)
	mergeBAall := reflistB.Merge(reflistA, false)

	c.Check(mergeABall, DeepEquals, mergeBAall)
	c.Check(toStrSlice(mergeBAall), DeepEquals,
		[]string{"Pall data 1.1~bp1", "Pamd64 app 1.1~bp2", "Pi386 app 1.1~bp1", "Pi386 app 1.1~bp2", "Pi386 dpkg 1.0", "Pi386 dpkg 1.7", "Pi386 lib 1.0", "Psparc xyz 1.0"})
}
<commit_msg>Initial pass at testing merged snapshotting with -latest flag<commit_after>package deb

import (
	"errors"
	"github.com/smira/aptly/database"
	. "launchpad.net/gocheck"
)

type PackageRefListSuite struct {
	// Simple list with "real" packages from stanzas
	list                   *PackageList
	p1, p2, p3, p4, p5, p6 *Package
}

var _ = Suite(&PackageRefListSuite{})

func (s *PackageRefListSuite) SetUpTest(c *C) {
	s.list = NewPackageList()

	s.p1 = NewPackageFromControlFile(packageStanza.Copy())
	s.p2 = NewPackageFromControlFile(packageStanza.Copy())
	stanza := packageStanza.Copy()
	stanza["Package"] = "mars-invaders"
	s.p3 = NewPackageFromControlFile(stanza)
	stanza = packageStanza.Copy()
	stanza["Size"] = "42"
	s.p4 = NewPackageFromControlFile(stanza)
	stanza = packageStanza.Copy()
	stanza["Package"] = "lonely-strangers"
	s.p5 = NewPackageFromControlFile(stanza)
	stanza = packageStanza.Copy()
	stanza["Version"] = "99.1"
	s.p6 = NewPackageFromControlFile(stanza)
}

func (s *PackageRefListSuite) TestNewPackageListFromRefList(c *C) {
	db, _ := database.OpenDB(c.MkDir())
	coll := NewPackageCollection(db)
	coll.Update(s.p1)
	coll.Update(s.p3)

	s.list.Add(s.p1)
	s.list.Add(s.p3)
	s.list.Add(s.p5)
	s.list.Add(s.p6)

	reflist := NewPackageRefListFromPackageList(s.list)

	_, err := NewPackageListFromRefList(reflist, coll, nil)
	c.Assert(err, ErrorMatches, "unable to load package with key.*")

	coll.Update(s.p5)
	coll.Update(s.p6)

	list, err := NewPackageListFromRefList(reflist, coll, nil)
	c.Assert(err, IsNil)
	c.Check(list.Len(), Equals, 4)
	c.Check(list.Add(s.p4), ErrorMatches, "conflict in package.*")

	list, err = NewPackageListFromRefList(nil, coll, nil)
	c.Assert(err, IsNil)
	c.Check(list.Len(), Equals, 0)
}

func (s *PackageRefListSuite) TestNewPackageRefList(c *C) {
	s.list.Add(s.p1)
	s.list.Add(s.p3)
	s.list.Add(s.p5)
	s.list.Add(s.p6)

	reflist := NewPackageRefListFromPackageList(s.list)
	c.Assert(reflist.Len(), Equals, 4)
	c.Check(reflist.Refs[0], DeepEquals, []byte(s.p1.Key("")))
	c.Check(reflist.Refs[1], DeepEquals, []byte(s.p6.Key("")))
	c.Check(reflist.Refs[2], DeepEquals, []byte(s.p5.Key("")))
	c.Check(reflist.Refs[3], DeepEquals, []byte(s.p3.Key("")))

	reflist = NewPackageRefList()
	c.Check(reflist.Len(), Equals, 0)
}

func (s *PackageRefListSuite) TestPackageRefListEncodeDecode(c *C) {
	s.list.Add(s.p1)
	s.list.Add(s.p3)
	s.list.Add(s.p5)
	s.list.Add(s.p6)

	reflist := NewPackageRefListFromPackageList(s.list)

	reflist2 := &PackageRefList{}
	err := reflist2.Decode(reflist.Encode())
	c.Assert(err, IsNil)
	c.Check(reflist2.Len(), Equals, reflist.Len())
	c.Check(reflist2.Refs, DeepEquals, reflist.Refs)
}

func (s *PackageRefListSuite) TestPackageRefListForeach(c *C) {
	s.list.Add(s.p1)
	s.list.Add(s.p3)
	s.list.Add(s.p5)
	s.list.Add(s.p6)

	reflist := NewPackageRefListFromPackageList(s.list)

	Len := 0
	err := reflist.ForEach(func([]byte) error {
		Len++
		return nil
	})

	c.Check(Len, Equals, 4)
	c.Check(err, IsNil)

	e := errors.New("b")

	err = reflist.ForEach(func([]byte) error {
		return e
	})

	c.Check(err, Equals, e)
}

func (s *PackageRefListSuite) TestSubstract(c *C) {
	r1 := []byte("r1")
	r2 := []byte("r2")
	r3 := []byte("r3")
	r4 := []byte("r4")
	r5 := []byte("r5")

	empty := &PackageRefList{Refs: [][]byte{}}
	l1 := &PackageRefList{Refs: [][]byte{r1, r2, r3, r4}}
	l2 := &PackageRefList{Refs: [][]byte{r1, r3}}
	l3 := &PackageRefList{Refs: [][]byte{r2, r4}}
	l4 := &PackageRefList{Refs: [][]byte{r4, r5}}
	l5 := &PackageRefList{Refs: [][]byte{r1, r2, r3}}

	c.Check(l1.Substract(empty), DeepEquals, l1)
	c.Check(l1.Substract(l2), DeepEquals, l3)
	c.Check(l1.Substract(l3), DeepEquals, l2)
	c.Check(l1.Substract(l4), DeepEquals, l5)
	c.Check(empty.Substract(l1), DeepEquals, empty)
	c.Check(l2.Substract(l3), DeepEquals, l2)
}

func (s *PackageRefListSuite) TestDiff(c *C) {
	db, _ := database.OpenDB(c.MkDir())
	coll := NewPackageCollection(db)

	packages := []*Package{
		&Package{Name: "lib", Version: "1.0", Architecture: "i386"},      //0
		&Package{Name: "dpkg", Version: "1.7", Architecture: "i386"},     //1
		&Package{Name: "data", Version: "1.1~bp1", Architecture: "all"},  //2
		&Package{Name: "app", Version: "1.1~bp1", Architecture: "i386"},  //3
		&Package{Name: "app", Version: "1.1~bp2", Architecture: "i386"},  //4
		&Package{Name: "app", Version: "1.1~bp2", Architecture: "amd64"}, //5
		&Package{Name: "xyz", Version: "3.0", Architecture: "sparc"},     //6
	}

	for _, p := range packages {
		coll.Update(p)
	}

	listA := NewPackageList()
	listA.Add(packages[0])
	listA.Add(packages[1])
	listA.Add(packages[2])
	listA.Add(packages[3])
	listA.Add(packages[6])

	listB := NewPackageList()
	listB.Add(packages[0])
	listB.Add(packages[2])
	listB.Add(packages[4])
	listB.Add(packages[5])

	reflistA := NewPackageRefListFromPackageList(listA)
	reflistB := NewPackageRefListFromPackageList(listB)

	diffAA, err := reflistA.Diff(reflistA, coll)
	c.Check(err, IsNil)
	c.Check(diffAA, HasLen, 0)

	diffAB, err := reflistA.Diff(reflistB, coll)
	c.Check(err, IsNil)
	c.Check(diffAB, HasLen, 4)

	c.Check(diffAB[0].Left, IsNil)
	c.Check(diffAB[0].Right.String(), Equals, "app_1.1~bp2_amd64")

	c.Check(diffAB[1].Left.String(), Equals, "app_1.1~bp1_i386")
	c.Check(diffAB[1].Right.String(), Equals, "app_1.1~bp2_i386")

	c.Check(diffAB[2].Left.String(), Equals, "dpkg_1.7_i386")
	c.Check(diffAB[2].Right, IsNil)

	c.Check(diffAB[3].Left.String(), Equals, "xyz_3.0_sparc")
	c.Check(diffAB[3].Right, IsNil)

	diffBA, err := reflistB.Diff(reflistA, coll)
	c.Check(err, IsNil)
	c.Check(diffBA, HasLen, 4)

	c.Check(diffBA[0].Right, IsNil)
	c.Check(diffBA[0].Left.String(), Equals, "app_1.1~bp2_amd64")

	c.Check(diffBA[1].Right.String(), Equals, "app_1.1~bp1_i386")
	c.Check(diffBA[1].Left.String(), Equals, "app_1.1~bp2_i386")

	c.Check(diffBA[2].Right.String(), Equals, "dpkg_1.7_i386")
	c.Check(diffBA[2].Left, IsNil)

	c.Check(diffBA[3].Right.String(), Equals, "xyz_3.0_sparc")
	c.Check(diffBA[3].Left, IsNil)

}

func (s *PackageRefListSuite) TestMerge(c *C) {
	db, _ := database.OpenDB(c.MkDir())
	coll := NewPackageCollection(db)

	packages := []*Package{
		&Package{Name: "lib", Version: "1.0", Architecture: "i386"},      //0
		&Package{Name: "dpkg", Version: "1.7", Architecture: "i386"},     //1
		&Package{Name: "data", Version: "1.1~bp1", Architecture: "all"},  //2
		&Package{Name: "app", Version: "1.1~bp1", Architecture: "i386"},  //3
		&Package{Name: "app", Version: "1.1~bp2", Architecture: "i386"},  //4
		&Package{Name: "app", Version: "1.1~bp2", Architecture: "amd64"}, //5
		&Package{Name: "dpkg", Version: "1.0", Architecture: "i386"},     //6
		&Package{Name: "xyz", Version: "1.0", Architecture: "sparc"},     //7
	}

	for _, p := range packages {
		coll.Update(p)
	}

	listA := NewPackageList()
	listA.Add(packages[0])
	listA.Add(packages[1])
	listA.Add(packages[2])
	listA.Add(packages[3])
	listA.Add(packages[7])

	listB := NewPackageList()
	listB.Add(packages[0])
	listB.Add(packages[2])
	listB.Add(packages[4])
	listB.Add(packages[5])
	listB.Add(packages[6])

	listC := NewPackageList()
	listC.Add(packages[3])
	listC.Add(packages[3])
	listC.Add(packages[4])
	listC.Add(packages[4])
	listC.Add(packages[7])
	listC.Add(packages[7])

	reflistA := NewPackageRefListFromPackageList(listA)
	reflistB := NewPackageRefListFromPackageList(listB)
	reflistC := NewPackageRefListFromPackageList(listC)

	toStrSlice := func(reflist *PackageRefList) (result []string) {
		result = make([]string, reflist.Len())
		for i, r := range reflist.Refs {
			result[i] = string(r)
		}
		return
	}

	mergeAB := reflistA.Merge(reflistB, true, false)
	mergeBA := reflistB.Merge(reflistA, true, false)

	c.Check(toStrSlice(mergeAB), DeepEquals,
		[]string{"Pall data 1.1~bp1", "Pamd64 app 1.1~bp2", "Pi386 app 1.1~bp2", "Pi386 dpkg 1.0", "Pi386 lib 1.0", "Psparc xyz 1.0"})
	c.Check(toStrSlice(mergeBA), DeepEquals,
		[]string{"Pall data 1.1~bp1", "Pamd64 app 1.1~bp2", "Pi386 app 1.1~bp1", "Pi386 dpkg 1.7", "Pi386 lib 1.0", "Psparc xyz 1.0"})

	mergeABall := reflistA.Merge(reflistB, false, false)
	mergeBAall := reflistB.Merge(reflistA, false, false)

	c.Check(mergeABall, DeepEquals, mergeBAall)
	c.Check(toStrSlice(mergeBAall), DeepEquals,
		[]string{"Pall data 1.1~bp1", "Pamd64 app 1.1~bp2", "Pi386 app 1.1~bp1", "Pi386 app 1.1~bp2", "Pi386 dpkg 1.0", "Pi386 dpkg 1.7", "Pi386 lib 1.0", "Psparc xyz 1.0"})

	mergeABClatest := reflistA.Merge(reflistB, false, true)
	mergeABClatest = mergeABClatest.Merge(reflistC, false, true)

	mergeCBAlatest := reflistC.Merge(reflistB, false, true)
	mergeCBAlatest = mergeCBAlatest.Merge(reflistA, false, true)

	c.Check(toStrSlice(mergeABClatest), DeepEquals,
		[]string{"Pall data 1.1~bp1", "Pamd64 app 1.1~bp2", "Pi386 app 1.1~bp2", "Pi386 dpkg 1.7", "Pi386 lib 1.0", "Psparc xyz 1.0"})
	c.Check(toStrSlice(mergeCBAlatest), DeepEquals,
		[]string{"Pall data 1.1~bp1", "Pamd64 app 1.1~bp2", "Pi386 app 1.1~bp2", "Pi386 dpkg 1.7", "Pi386 lib 1.0", "Psparc xyz 1.0"})
}
<|endoftext|>


<commit_before>// Copyright 2017 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 auth contains integration tests for the firebase.google.com/go/auth package.
package auth

import (
	"fmt"
	"reflect"
	"testing"

	"google.golang.org/api/iterator"

	"firebase.google.com/go/auth"

	"golang.org/x/net/context"
)

var testFixtures = struct {
	uidList            []string
	sampleUserBlank    *auth.UserRecord
	sampleUserNil      *auth.UserRecord
	sampleUserWithData *auth.UserRecord
}{}

func TestUserManagement(t *testing.T) {
	t.Run("add some users", createTestUsers)
	t.Run("created users", testCreatedAllUsers)
	t.Run("get user", testGetUser)
	t.Run("user iterator test", testUserIterator)
	t.Run("paging iterator test", testPager)
	t.Run("disable", testDisableUser)

	t.Run("update user", testUpdateUser)
	t.Run("disable", testDisableUser)

	t.Run("remove PhoneNumber, PhotoURL and DisplaName", testRemovePhonePhotoName)

	t.Run("Remove custom claims", testRemoveCustomClaims)
	t.Run("add custom claims", testAddCustomClaims)

	t.Run("delete all users", cleanupUsers)
}

func createTestUsers(t *testing.T) {
	for i := 0; i < 2; i++ {
		u, err := client.CreateUser(context.Background(),
			(&auth.UserToCreate{}).
				UID(fmt.Sprintf("tempTestUserID-%d", i)))
		if err != nil {
			t.Fatal("failed to create user", i, err)
		}
		testFixtures.uidList = append(testFixtures.uidList, u.UID)
	}
	u, err := client.CreateUser(context.Background(), nil)
	if err != nil {
		t.Fatal(err)
	}
	testFixtures.sampleUserNil = u
	testFixtures.uidList = append(testFixtures.uidList, u.UID)

	u, err = client.CreateUser(context.Background(), (&auth.UserToCreate{}))
	if err != nil {
		t.Fatal(err)
	}
	testFixtures.sampleUserBlank = u
	testFixtures.uidList = append(testFixtures.uidList, u.UID)
	uid := "tempUserId1234"

	u, err = client.CreateUser(context.Background(), (&auth.UserToCreate{}).
		UID(uid).
		Email(uid+"email@test.com").
		DisplayName("display_name").
		Password("passawd"))
	if err != nil {
		t.Fatal(err)
	}
	testFixtures.sampleUserWithData = u
	testFixtures.uidList = append(testFixtures.uidList, u.UID)
}

func testCreatedAllUsers(t *testing.T) {
	for _, id := range testFixtures.uidList {
		_, err := client.GetUser(context.Background(), id)
		if err != nil {
			t.Errorf("GetUser(%q) user not found. %s", id, err)
		}
	}
}

func testGetUser(t *testing.T) {
	u, err := client.GetUser(context.Background(), testFixtures.sampleUserWithData.UID)
	if err != nil {
		t.Fatalf("error getting user %s", err)
	}
	if u.UID != testFixtures.sampleUserWithData.UID || u.Email != testFixtures.sampleUserWithData.Email {
		t.Errorf("GetUser() = %#v; want: %#v", testFixtures.sampleUserWithData, u.UserInfo)
	}
	if !reflect.DeepEqual(u, testFixtures.sampleUserWithData) {
		t.Errorf("GetUser(UID) = %#v; want: %#v", u, testFixtures.sampleUserWithData)
	}
}

func testGetUserByPhoneNumber(t *testing.T) {
	u, err := client.GetUserByPhoneNumber(context.Background(), testFixtures.sampleUserWithData.PhoneNumber)
	if err != nil {
		t.Fatalf("error getting user %s", err)
	}
	if u.UID != testFixtures.sampleUserWithData.UID || u.PhoneNumber != testFixtures.sampleUserWithData.PhoneNumber {
		t.Errorf("GetUserByPhoneNumber(%q) = %#v; want: %#v",
			testFixtures.sampleUserWithData.PhoneNumber, u, testFixtures.sampleUserWithData)
	}
	if !reflect.DeepEqual(u, testFixtures.sampleUserWithData) {
		t.Errorf("GetUserByPhoneNumber(%q) = %#v; want: %#v",
			testFixtures.sampleUserWithData.PhoneNumber, u, testFixtures.sampleUserWithData)
	}
}

func testGetUserByEmail(t *testing.T) {
	u, err := client.GetUserByEmail(context.Background(), testFixtures.sampleUserWithData.Email)
	if err != nil {
		t.Fatalf("error getting user %s", err)
	}
	if u.UID != testFixtures.sampleUserWithData.UID || u.Email != testFixtures.sampleUserWithData.Email {
		t.Errorf("GetUserByEmail(%q) = %#v; want: %#v",
			testFixtures.sampleUserWithData.Email, u, testFixtures.sampleUserWithData)
	}
	if !reflect.DeepEqual(u, testFixtures.sampleUserWithData) {
		t.Errorf("GetUserByEmail(%q) = %#v; want: %#v",
			testFixtures.sampleUserWithData.Email, u, testFixtures.sampleUserWithData)
	}
}

func testUserIterator(t *testing.T) {
	iter := client.Users(context.Background(), "")
	uids := map[string]bool{}
	gotCount := 0

	for {
		u, err := iter.Next()
		if err == iterator.Done {
			break
		}
		if err != nil {
			t.Fatal(err)
		}

		gotCount++
		uids[u.UID] = true
	}
	if gotCount < 5 {
		t.Errorf("Users() got %d users; want: at least 5 users", gotCount)
	}
	// verify that all the expected users are present
	for _, uid := range testFixtures.uidList {
		if _, ok := uids[uid]; !ok {
			t.Errorf("Users() missing UID %s; want: UID %s", uid, uid)
		}
	}
}

func testPager(t *testing.T) {
	iter := client.Users(context.Background(), "")
	pager := iterator.NewPager(iter, 2, "")
	userCount := 0
	pageCount := 0

	for {
		pageCount++
		var users []*auth.ExportedUserRecord
		nextPageToken, err := pager.NextPage(&users)
		if err != nil {
			t.Fatalf("paging error %v", err)
		}
		userCount += len(users)
		for _, u := range users {
			// this iterates over users in a page
			if u.UID == "something" {
				// do something
			}
		}
		if nextPageToken == "" {
			break
		}
	}
	if userCount < 5 || pageCount < 3 {
		t.Errorf("Users(), NewPager() got %d pages with %d users;  want at least %d pages with at least %d users,", pageCount, userCount, 3, 5)
	}
}

func testDisableUser(t *testing.T) {
	u, err := client.GetUser(context.Background(), testFixtures.sampleUserBlank.UID)
	if err != nil {
		t.Fatal(err)
	}
	if u.Disabled {
		t.Errorf("GetUser() user is disabled; want: user not disabled")
	}

	_, err = client.UpdateUser(context.Background(), u.UID,
		(&auth.UserToUpdate{}).Disabled(true))
	if err != nil {
		t.Fatal(err)
	}
	u, err = client.GetUser(context.Background(), testFixtures.sampleUserBlank.UID)
	if !u.Disabled {
		t.Errorf("Disabled() user is not disabled; want: user disabled")
	}
	_, err = client.UpdateUser(context.Background(), u.UID,
		(&auth.UserToUpdate{}).Disabled(false))
	if err != nil {
		t.Fatal(err)
	}
	u, err = client.GetUser(context.Background(), testFixtures.sampleUserBlank.UID)
	if err != nil {
		t.Fatal(err)
	}
	if u.Disabled {
		t.Errorf("Disabled() user is disabled; want: user not disabled")
	}
}

func testUpdateUser(t *testing.T) {
	u, err := client.GetUser(context.Background(), testFixtures.sampleUserBlank.UID)

	if err != nil || u == nil {
		t.Fatalf("error getting user %s", err)
	}
	refU := &auth.UserRecord{
		UserInfo: &auth.UserInfo{UID: testFixtures.sampleUserBlank.UID},
		UserMetadata: &auth.UserMetadata{
			CreationTimestamp: testFixtures.sampleUserBlank.UserMetadata.CreationTimestamp,
		},
	}
	if !reflect.DeepEqual(u, refU) {
		t.Errorf("GetUser() got = %s; \nwant: %s", toString(refU), toString(u))
	}
	utup := (&auth.UserToUpdate{}).
		Disabled(false).
		DisplayName("name").
		PhoneNumber("+12345678901").
		PhotoURL("http://photo.png").
		Email("abc@ab.ab").
		EmailVerified(true).
		Password("wordpass").
		CustomClaims(map[string]interface{}{"custom": "claims"})

	u, err = client.UpdateUser(context.Background(), u.UID, utup)
	if err != nil {
		t.Fatal(err)
	}

	refU = &auth.UserRecord{
		UserInfo: &auth.UserInfo{
			UID:         testFixtures.sampleUserBlank.UID,
			DisplayName: "name",
			PhoneNumber: "+12345678901",
			PhotoURL:    "http://photo.png",
			Email:       "abc@ab.ab",
		},
		UserMetadata: &auth.UserMetadata{
			CreationTimestamp: testFixtures.sampleUserBlank.UserMetadata.CreationTimestamp,
		},
		Disabled:      false,
		EmailVerified: true,
		CustomClaims:  map[string]interface{}{"custom": "claims"},
	}
	// compare provider info seperatley since the order of the providers isn't guaranteed.
	testProviderInfo(u.ProviderUserInfo,
		&auth.UserInfo{
			DisplayName: "name",
			Email:       "abc@ab.ab",
			PhotoURL:    "http://photo.png",
			ProviderID:  "password"},
		&auth.UserInfo{
			PhoneNumber: "+12345678901",
			ProviderID:  "phone"},
		t)
	// now compare the rest of the record, without the ProviderInfo
	u.ProviderUserInfo = nil
	if !reflect.DeepEqual(u, refU) {
		t.Errorf("UpdateUser() got = %s\nexpecting: %s", toString(u), toString(refU))
	}
}

func testProviderInfo(pi []*auth.UserInfo, passwordUI, phoneUI *auth.UserInfo, t *testing.T) {
	var compareWith *auth.UserInfo
	for _, ui := range pi {
		switch ui.ProviderID {
		case "password":
			compareWith = passwordUI
		case "phone":
			compareWith = phoneUI
		}
		if !reflect.DeepEqual(ui, compareWith) {
			t.Errorf("UpdateUser()got: %#v; \nwant: %#v", ui, compareWith)
		}
	}
}

func testRemovePhonePhotoName(t *testing.T) {
	u, err := client.GetUser(context.Background(), testFixtures.sampleUserBlank.UID)
	if err != nil {
		t.Fatal(err)
	}
	if u.PhoneNumber == "" {
		t.Errorf("GetUser() expecting non empty PhoneNumber")
	}
	if len(u.ProviderUserInfo) != 2 {
		t.Errorf("GetUser() expecting 2 providers")
	}
	if u.PhotoURL == "" {
		t.Errorf("GetUser() expecting non empty PhotoURL")
	}
	if u.DisplayName == "" {
		t.Errorf("GetUser() expecting non empty DisplayName")
	}

	_, err = client.UpdateUser(context.Background(), u.UID,
		(&auth.UserToUpdate{}).PhoneNumber("").PhotoURL("").DisplayName(""))
	if err != nil {
		t.Fatal(err)
	}

	u, err = client.GetUser(context.Background(), testFixtures.sampleUserBlank.UID)
	if u.PhoneNumber != "" {
		t.Errorf("UpdateUser() [remove PhoneNumber] PhoneNumber = %q; want: PhoneNumber=%q", u.PhoneNumber, "")
	}
	if len(u.ProviderUserInfo) != 1 {
		t.Errorf("UpdateUser() [remove PhoneNumber] got %d ProviderUserInfo records, want:want 1 provider", len(u.ProviderUserInfo))
	}
	if u.DisplayName != "" {
		t.Errorf("UpdateUser() [remove DisplayName] DisplayName = %q; want: DisplayName=%q", u.DisplayName, "")
	}
	if u.PhotoURL != "" {
		t.Errorf("UpdateUser() [remove PhotoURL] PhotoURL = %q; want: PhotoURL=%q", u.PhotoURL, "")
	}
}

func testRemoveCustomClaims(t *testing.T) {
	u, err := client.GetUser(context.Background(), testFixtures.sampleUserBlank.UID)
	if err != nil {
		t.Fatal(err)
	}
	if !reflect.DeepEqual(u.CustomClaims,
		map[string]interface{}{"custom": "claims"}) {
		t.Errorf("expecting CustomClaims")
	}

	err = client.SetCustomUserClaims(context.Background(), u.UID, nil)
	if err != nil {
		t.Fatal(err)
	}
	u, err = client.GetUser(context.Background(), testFixtures.sampleUserBlank.UID)
	if u.CustomClaims != nil {
		t.Errorf("CustomClaims() = %#v; want {}", u.CustomClaims)

	}
}

func testAddCustomClaims(t *testing.T) {
	u, err := client.GetUser(context.Background(), testFixtures.sampleUserBlank.UID)
	if err != nil {
		t.Fatal(err)
	}
	if u.CustomClaims != nil {
		t.Errorf("GetUser(), CustomClaims = %q; want <nil>", u.CustomClaims)
	}

	_, err = client.UpdateUser(context.Background(), u.UID,
		(&auth.UserToUpdate{}).CustomClaims(map[string]interface{}{"2custom": "2claims"}))
	if err != nil {
		t.Fatal(err)
	}
	u, err = client.GetUser(context.Background(), testFixtures.sampleUserBlank.UID)
	if u.CustomClaims == nil {
		t.Errorf("CustomClaims = <nil>; want: {\"2custom\": \"2claims\"}")
	}
}

func provString(e *auth.UserRecord) string {
	providerStr := ""
	if e.ProviderUserInfo != nil {
		for _, info := range e.ProviderUserInfo {
			providerStr += fmt.Sprintf("\n            %#v", info)
		}
	}
	return providerStr
}

func cleanupUsers(t *testing.T) {
	for _, id := range testFixtures.uidList {
		err := client.DeleteUser(context.Background(), id)
		if err != nil {
			t.Errorf("DeleteUser(%s) error deleting user, %s", id, err)
		}
	}
}

func toString(e *auth.UserRecord) string {
	return fmt.Sprintf(
		"    UserRecord: %#v\n"+
			"        UserInfo: %#v\n"+
			"        MetaData: %#v\n"+
			"        CustomClaims: %#v\n"+
			"        ProviderData: %#v %s",
		e,
		e.UserInfo,
		e.UserMetadata,
		e.CustomClaims,
		e.ProviderUserInfo,
		provString(e))
}
<commit_msg>Refactored integration tests (#46)<commit_after>// Copyright 2017 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 auth contains integration tests for the firebase.google.com/go/auth package.
package auth

import (
	"fmt"
	"reflect"
	"testing"

	"google.golang.org/api/iterator"

	"firebase.google.com/go/auth"

	"golang.org/x/net/context"
)

var testFixtures = struct {
	uidList            []string
	sampleUserBlank    *auth.UserRecord
	sampleUserWithData *auth.UserRecord
}{}

func TestUserManagement(t *testing.T) {
	t.Run("Create test users", testCreateUsers)
	t.Run("Get user", testGetUser)
	t.Run("Iterate users", testUserIterator)
	t.Run("Paged iteration", testPager)
	t.Run("Disable user account", testDisableUser)
	t.Run("Update user", testUpdateUser)
	t.Run("Remove user attributes", testRemovePhonePhotoName)
	t.Run("Remove custom claims", testRemoveCustomClaims)
	t.Run("Add custom claims", testAddCustomClaims)
	t.Run("Delete test users", testDeleteUsers)
}

func testCreateUsers(t *testing.T) {
	// Create users with uid
	for i := 0; i < 2; i++ {
		params := (&auth.UserToCreate{}).UID(fmt.Sprintf("tempTestUserID-%d", i))
		u, err := client.CreateUser(context.Background(), params)
		if err != nil {
			t.Fatal("failed to create user", i, err)
		}
		testFixtures.uidList = append(testFixtures.uidList, u.UID)
	}

	// Create user with no parameters (zero-value)
	u, err := client.CreateUser(context.Background(), (&auth.UserToCreate{}))
	if err != nil {
		t.Fatal(err)
	}
	testFixtures.sampleUserBlank = u
	testFixtures.uidList = append(testFixtures.uidList, u.UID)

	// Create user with parameters
	uid := "tempUserId1234"
	params := (&auth.UserToCreate{}).
		UID(uid).
		Email(uid + "email@test.com").
		DisplayName("display_name").
		Password("password")
	u, err = client.CreateUser(context.Background(), params)
	if err != nil {
		t.Fatal(err)
	}
	testFixtures.sampleUserWithData = u
	testFixtures.uidList = append(testFixtures.uidList, u.UID)
}

func testGetUser(t *testing.T) {
	want := testFixtures.sampleUserWithData
	u, err := client.GetUser(context.Background(), want.UID)
	if err != nil {
		t.Fatalf("error getting user %s", err)
	}
	if !reflect.DeepEqual(u, want) {
		t.Errorf("GetUser(UID) = %#v; want = %#v", u, want)
	}
}

func testGetUserByPhoneNumber(t *testing.T) {
	want := testFixtures.sampleUserWithData
	u, err := client.GetUserByPhoneNumber(context.Background(), want.PhoneNumber)
	if err != nil {
		t.Fatalf("error getting user %s", err)
	}
	if !reflect.DeepEqual(u, want) {
		t.Errorf("GetUserByPhoneNumber(%q) = %#v; want = %#v", want.PhoneNumber, u, want)
	}
}

func testGetUserByEmail(t *testing.T) {
	want := testFixtures.sampleUserWithData
	u, err := client.GetUserByEmail(context.Background(), want.Email)
	if err != nil {
		t.Fatalf("error getting user %s", err)
	}
	if !reflect.DeepEqual(u, want) {
		t.Errorf("GetUserByEmail(%q) = %#v; want = %#v", want.Email, u, want)
	}
}

func testUserIterator(t *testing.T) {
	iter := client.Users(context.Background(), "")
	uids := map[string]bool{}
	count := 0

	for {
		u, err := iter.Next()
		if err == iterator.Done {
			break
		}
		if err != nil {
			t.Fatal(err)
		}

		count++
		uids[u.UID] = true
	}
	if count < 5 {
		t.Errorf("Users() count = %d; want >= 5", count)
	}
	// verify that all the expected users are present
	for _, uid := range testFixtures.uidList {
		if _, ok := uids[uid]; !ok {
			t.Errorf("Users() missing uid: %s", uid)
		}
	}
}

func testPager(t *testing.T) {
	iter := client.Users(context.Background(), "")
	pager := iterator.NewPager(iter, 2, "")
	userCount := 0
	pageCount := 0

	for {
		pageCount++
		var users []*auth.ExportedUserRecord
		nextPageToken, err := pager.NextPage(&users)
		if err != nil {
			t.Fatalf("paging error %v", err)
		}
		userCount += len(users)
		for _, u := range users {
			// this iterates over users in a page
			if u.UID == "something" {
				// do something
			}
		}
		if nextPageToken == "" {
			break
		}
	}
	if userCount < 5 {
		t.Errorf("Users() count = %d;  want >= 5", userCount)
	}
	if pageCount < 3 {
		t.Errorf("NewPager() count = %d;  want >= 3", pageCount)
	}
}

func testDisableUser(t *testing.T) {
	want := testFixtures.sampleUserBlank
	u, err := client.GetUser(context.Background(), want.UID)
	if err != nil {
		t.Fatal(err)
	}
	if u.Disabled {
		t.Errorf("GetUser().Disabled = true; want = false")
	}

	params := (&auth.UserToUpdate{}).Disabled(true)
	u, err = client.UpdateUser(context.Background(), u.UID, params)
	if err != nil {
		t.Fatal(err)
	}
	if !u.Disabled {
		t.Errorf("UpdateUser(disable).Disabled = false; want = true")
	}

	params = (&auth.UserToUpdate{}).Disabled(false)
	u, err = client.UpdateUser(context.Background(), u.UID, params)
	if err != nil {
		t.Fatal(err)
	}
	if u.Disabled {
		t.Errorf("UpdateUser(disable).Disabled = true; want = false")
	}
}

func testUpdateUser(t *testing.T) {
	u, err := client.GetUser(context.Background(), testFixtures.sampleUserBlank.UID)
	if u == nil || err != nil {
		t.Fatalf("GetUser() = (%v, %v); want = (user, nil)", u, err)
	}

	want := &auth.UserRecord{
		UserInfo: &auth.UserInfo{UID: testFixtures.sampleUserBlank.UID},
		UserMetadata: &auth.UserMetadata{
			CreationTimestamp: testFixtures.sampleUserBlank.UserMetadata.CreationTimestamp,
		},
	}
	if !reflect.DeepEqual(u, want) {
		t.Errorf("GetUser() = %v; want = %v", u, want)
	}

	params := (&auth.UserToUpdate{}).
		Disabled(false).
		DisplayName("name").
		PhoneNumber("+12345678901").
		PhotoURL("http://photo.png").
		Email("abc@ab.ab").
		EmailVerified(true).
		Password("wordpass").
		CustomClaims(map[string]interface{}{"custom": "claims"})
	u, err = client.UpdateUser(context.Background(), u.UID, params)
	if err != nil {
		t.Fatal(err)
	}

	want = &auth.UserRecord{
		UserInfo: &auth.UserInfo{
			UID:         testFixtures.sampleUserBlank.UID,
			DisplayName: "name",
			PhoneNumber: "+12345678901",
			PhotoURL:    "http://photo.png",
			Email:       "abc@ab.ab",
		},
		UserMetadata: &auth.UserMetadata{
			CreationTimestamp: testFixtures.sampleUserBlank.UserMetadata.CreationTimestamp,
		},
		Disabled:      false,
		EmailVerified: true,
		CustomClaims:  map[string]interface{}{"custom": "claims"},
	}

	testProviderInfo := func(pi []*auth.UserInfo, t *testing.T) {
		passwordUI := &auth.UserInfo{
			DisplayName: "name",
			Email:       "abc@ab.ab",
			PhotoURL:    "http://photo.png",
			ProviderID:  "password",
		}
		phoneUI := &auth.UserInfo{
			PhoneNumber: "+12345678901",
			ProviderID:  "phone",
		}

		var compareWith *auth.UserInfo
		for _, ui := range pi {
			switch ui.ProviderID {
			case "password":
				compareWith = passwordUI
			case "phone":
				compareWith = phoneUI
			}
			if !reflect.DeepEqual(ui, compareWith) {
				t.Errorf("UpdateUser()got: %#v; \nwant: %#v", ui, compareWith)
			}
		}
	}

	// compare provider info seperatley since the order of the providers isn't guaranteed.
	testProviderInfo(u.ProviderUserInfo, t)

	// now compare the rest of the record, without the ProviderInfo
	u.ProviderUserInfo = nil
	if !reflect.DeepEqual(u, want) {
		t.Errorf("UpdateUser() = %v; want = %v", u, want)
	}
}

func testRemovePhonePhotoName(t *testing.T) {
	u, err := client.GetUser(context.Background(), testFixtures.sampleUserBlank.UID)
	if err != nil {
		t.Fatal(err)
	}
	if u.PhoneNumber == "" {
		t.Errorf("GetUser().PhoneNumber = empty; want = non-empty")
	}
	if len(u.ProviderUserInfo) != 2 {
		t.Errorf("GetUser().ProviderUserInfo = %d; want = 2", len(u.ProviderUserInfo))
	}
	if u.PhotoURL == "" {
		t.Errorf("GetUser().PhotoURL = empty; want = non-empty")
	}
	if u.DisplayName == "" {
		t.Errorf("GetUser().DisplayName = empty; want = non-empty")
	}

	params := (&auth.UserToUpdate{}).PhoneNumber("").PhotoURL("").DisplayName("")
	u, err = client.UpdateUser(context.Background(), u.UID, params)
	if err != nil {
		t.Fatal(err)
	}
	if u.PhoneNumber != "" {
		t.Errorf("UpdateUser().PhoneNumber = %q; want: %q", u.PhoneNumber, "")
	}
	if len(u.ProviderUserInfo) != 1 {
		t.Errorf("UpdateUser().ProviderUserInfo = %d, want = 1", len(u.ProviderUserInfo))
	}
	if u.DisplayName != "" {
		t.Errorf("UpdateUser().DisplayName = %q; want =%q", u.DisplayName, "")
	}
	if u.PhotoURL != "" {
		t.Errorf("UpdateUser().PhotoURL = %q; want = %q", u.PhotoURL, "")
	}
}

func testRemoveCustomClaims(t *testing.T) {
	u, err := client.GetUser(context.Background(), testFixtures.sampleUserBlank.UID)
	if err != nil {
		t.Fatal(err)
	}
	want := map[string]interface{}{"custom": "claims"}
	if !reflect.DeepEqual(u.CustomClaims, want) {
		t.Errorf("CustomClaims = %v; want = %v", u.CustomClaims, want)
	}

	err = client.SetCustomUserClaims(context.Background(), u.UID, nil)
	if err != nil {
		t.Fatal(err)
	}
	u, err = client.GetUser(context.Background(), testFixtures.sampleUserBlank.UID)
	if u.CustomClaims != nil {
		t.Errorf("CustomClaims() = %#v; want = nil", u.CustomClaims)
	}
}

func testAddCustomClaims(t *testing.T) {
	u, err := client.GetUser(context.Background(), testFixtures.sampleUserBlank.UID)
	if err != nil {
		t.Fatal(err)
	}
	if u.CustomClaims != nil {
		t.Errorf("GetUser().CustomClaims = %v; want = nil", u.CustomClaims)
	}

	want := map[string]interface{}{"2custom": "2claims"}
	params := (&auth.UserToUpdate{}).CustomClaims(want)
	u, err = client.UpdateUser(context.Background(), u.UID, params)
	if err != nil {
		t.Fatal(err)
	}
	if !reflect.DeepEqual(u.CustomClaims, want) {
		t.Errorf("CustomClaims = %v; want = %v", u.CustomClaims, want)
	}
}

func testDeleteUsers(t *testing.T) {
	for _, id := range testFixtures.uidList {
		err := client.DeleteUser(context.Background(), id)
		if err != nil {
			t.Errorf("DeleteUser(%q) = %v; want = nil", id, err)
		}

		u, err := client.GetUser(context.Background(), id)
		if u != nil || err == nil {
			t.Errorf("GetUser(non-existing) = (%v, %v); want = (nil, error)", u, err)
		}
	}
}
<|endoftext|>


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

package source

import (
	"context"
	"fmt"
	"go/ast"
	"go/constant"
	"go/token"
	"go/types"
	"strings"

	"golang.org/x/tools/internal/event"
	"golang.org/x/tools/internal/lsp/lsppos"
	"golang.org/x/tools/internal/lsp/protocol"
	"golang.org/x/tools/internal/typeparams"
)

const (
	maxLabelLength = 28
)

type InlayHintFunc func(node ast.Node, tmap *lsppos.TokenMapper, info *types.Info, q *types.Qualifier) []protocol.InlayHint

type Hint struct {
	Name string
	Doc  string
	Run  InlayHintFunc
}

const (
	ParameterNames             = "parameterNames"
	AssignVariableTypes        = "assignVariableTypes"
	ConstantValues             = "constantValues"
	RangeVariableTypes         = "rangeVariableTypes"
	CompositeLiteralTypes      = "compositeLiteralTypes"
	CompositeLiteralFieldNames = "compositeLiteralFields"
	FunctionTypeParameters     = "functionTypeParameters"
)

var AllInlayHints = map[string]*Hint{
	AssignVariableTypes: {
		Name: AssignVariableTypes,
		Doc: `Enable/disable inlay hints for variable types in assign statements:

	i/* int/*, j/* int/* := 0, len(r)-1`,
		Run: assignVariableTypes,
	},
	ParameterNames: {
		Name: ParameterNames,
		Doc: `Enable/disable inlay hints for parameter names:

	parseInt(/* str: */ "123", /* radix: */ 8)`,
		Run: parameterNames,
	},
	ConstantValues: {
		Name: ConstantValues,
		Doc: `Enable/disable inlay hints for constant values:

	const (
		KindNone   Kind = iota/* = 0*/
		KindPrint/*  = 1*/
		KindPrintf/* = 2*/
		KindErrorf/* = 3*/
	)`,
		Run: constantValues,
	},
	RangeVariableTypes: {
		Name: RangeVariableTypes,
		Doc: `Enable/disable inlay hints for variable types in range statements:

	for k/* int*/, v/* string/* := range []string{} {
		fmt.Println(k, v)
	}`,
		Run: rangeVariableTypes,
	},
	CompositeLiteralTypes: {
		Name: CompositeLiteralTypes,
		Doc: `Enable/disable inlay hints for composite literal types:

	for _, c := range []struct {
		in, want string
	}{
		/*struct{ in string; want string }*/{"Hello, world", "dlrow ,olleH"},
	}`,
		Run: compositeLiteralTypes,
	},
	CompositeLiteralFieldNames: {
		Name: CompositeLiteralFieldNames,
		Doc: `Enable/disable inlay hints for composite literal field names:

	{in: "Hello, world", want: "dlrow ,olleH"}`,
		Run: compositeLiteralFields,
	},
	FunctionTypeParameters: {
		Name: FunctionTypeParameters,
		Doc: `Enable/disable inlay hints for implicit type parameters on generic functions:

	myFoo/*[int, string]*/(1, "hello")`,
		Run: funcTypeParams,
	},
}

func InlayHint(ctx context.Context, snapshot Snapshot, fh FileHandle, _ protocol.Range) ([]protocol.InlayHint, error) {
	ctx, done := event.Start(ctx, "source.InlayHint")
	defer done()

	pkg, pgf, err := GetParsedFile(ctx, snapshot, fh, NarrowestPackage)
	if err != nil {
		return nil, fmt.Errorf("getting file for InlayHint: %w", err)
	}

	// Collect a list of the inlay hints that are enabled.
	inlayHintOptions := snapshot.View().Options().InlayHintOptions
	var enabledHints []InlayHintFunc
	for hint, enabled := range inlayHintOptions.Hints {
		if !enabled {
			continue
		}
		if h, ok := AllInlayHints[hint]; ok {
			enabledHints = append(enabledHints, h.Run)
		}
	}
	if len(enabledHints) == 0 {
		return nil, nil
	}

	tmap := lsppos.NewTokenMapper(pgf.Src, pgf.Tok)
	info := pkg.GetTypesInfo()
	q := Qualifier(pgf.File, pkg.GetTypes(), info)

	var hints []protocol.InlayHint
	ast.Inspect(pgf.File, func(node ast.Node) bool {
		for _, fn := range enabledHints {
			hints = append(hints, fn(node, tmap, info, &q)...)
		}
		return true
	})
	return hints, nil
}

func parameterNames(node ast.Node, tmap *lsppos.TokenMapper, info *types.Info, _ *types.Qualifier) []protocol.InlayHint {
	callExpr, ok := node.(*ast.CallExpr)
	if !ok {
		return nil
	}
	signature, ok := info.TypeOf(callExpr.Fun).(*types.Signature)
	if !ok {
		return nil
	}

	var hints []protocol.InlayHint
	for i, v := range callExpr.Args {
		start, ok := tmap.Position(v.Pos())
		if !ok {
			continue
		}
		params := signature.Params()
		// When a function has variadic params, we skip args after
		// params.Len().
		if i > params.Len()-1 {
			break
		}
		value := params.At(i).Name()
		// param.Name is empty for built-ins like append
		if value == "" {
			continue
		}
		if signature.Variadic() && i == params.Len()-1 {
			value = value + "..."
		}
		hints = append(hints, protocol.InlayHint{
			Position:     &start,
			Label:        buildLabel(value + ":"),
			Kind:         protocol.Parameter,
			PaddingRight: true,
		})
	}
	return hints
}

func funcTypeParams(node ast.Node, tmap *lsppos.TokenMapper, info *types.Info, _ *types.Qualifier) []protocol.InlayHint {
	ce, ok := node.(*ast.CallExpr)
	if !ok {
		return nil
	}
	id, ok := ce.Fun.(*ast.Ident)
	if !ok {
		return nil
	}
	inst := typeparams.GetInstances(info)[id]
	if inst.TypeArgs == nil {
		return nil
	}
	start, ok := tmap.Position(id.End())
	if !ok {
		return nil
	}
	var args []string
	for i := 0; i < inst.TypeArgs.Len(); i++ {
		args = append(args, inst.TypeArgs.At(i).String())
	}
	if len(args) == 0 {
		return nil
	}
	return []protocol.InlayHint{{
		Position: &start,
		Label:    buildLabel("[" + strings.Join(args, ", ") + "]"),
		Kind:     protocol.Type,
	}}
}

func assignVariableTypes(node ast.Node, tmap *lsppos.TokenMapper, info *types.Info, q *types.Qualifier) []protocol.InlayHint {
	stmt, ok := node.(*ast.AssignStmt)
	if !ok || stmt.Tok != token.DEFINE {
		return nil
	}

	var hints []protocol.InlayHint
	for _, v := range stmt.Lhs {
		if h := variableType(v, tmap, info, q); h != nil {
			hints = append(hints, *h)
		}
	}
	return hints
}

func rangeVariableTypes(node ast.Node, tmap *lsppos.TokenMapper, info *types.Info, q *types.Qualifier) []protocol.InlayHint {
	rStmt, ok := node.(*ast.RangeStmt)
	if !ok {
		return nil
	}
	var hints []protocol.InlayHint
	if h := variableType(rStmt.Key, tmap, info, q); h != nil {
		hints = append(hints, *h)
	}
	if h := variableType(rStmt.Value, tmap, info, q); h != nil {
		hints = append(hints, *h)
	}
	return hints
}

func variableType(e ast.Expr, tmap *lsppos.TokenMapper, info *types.Info, q *types.Qualifier) *protocol.InlayHint {
	typ := info.TypeOf(e)
	if typ == nil {
		return nil
	}
	end, ok := tmap.Position(e.End())
	if !ok {
		return nil
	}
	return &protocol.InlayHint{
		Position:    &end,
		Label:       buildLabel(types.TypeString(typ, *q)),
		Kind:        protocol.Type,
		PaddingLeft: true,
	}
}

func constantValues(node ast.Node, tmap *lsppos.TokenMapper, info *types.Info, _ *types.Qualifier) []protocol.InlayHint {
	genDecl, ok := node.(*ast.GenDecl)
	if !ok || genDecl.Tok != token.CONST {
		return nil
	}

	var hints []protocol.InlayHint
	for _, v := range genDecl.Specs {
		spec, ok := v.(*ast.ValueSpec)
		if !ok {
			continue
		}
		end, ok := tmap.Position(v.End())
		if !ok {
			continue
		}
		// Show hints when values are missing or at least one value is not
		// a basic literal.
		showHints := len(spec.Values) == 0
		checkValues := len(spec.Names) == len(spec.Values)
		var values []string
		for i, w := range spec.Names {
			obj, ok := info.ObjectOf(w).(*types.Const)
			if !ok || obj.Val().Kind() == constant.Unknown {
				return nil
			}
			if checkValues {
				switch spec.Values[i].(type) {
				case *ast.BadExpr:
					return nil
				case *ast.BasicLit:
				default:
					if obj.Val().Kind() != constant.Bool {
						showHints = true
					}
				}
			}
			values = append(values, fmt.Sprintf("%v", obj.Val()))
		}
		if !showHints || len(values) == 0 {
			continue
		}
		hints = append(hints, protocol.InlayHint{
			Position:    &end,
			Label:       buildLabel("= " + strings.Join(values, ", ")),
			PaddingLeft: true,
		})
	}
	return hints
}

func compositeLiteralFields(node ast.Node, tmap *lsppos.TokenMapper, info *types.Info, q *types.Qualifier) []protocol.InlayHint {
	compLit, ok := node.(*ast.CompositeLit)
	if !ok {
		return nil
	}
	typ := info.TypeOf(compLit)
	if typ == nil {
		return nil
	}
	if t, ok := typ.(*types.Pointer); ok {
		typ = t.Elem()
	}
	strct, ok := typ.Underlying().(*types.Struct)
	if !ok {
		return nil
	}

	var hints []protocol.InlayHint

	for i, v := range compLit.Elts {
		if _, ok := v.(*ast.KeyValueExpr); !ok {
			start, ok := tmap.Position(v.Pos())
			if !ok {
				continue
			}
			if i > strct.NumFields()-1 {
				break
			}
			hints = append(hints, protocol.InlayHint{
				Position:     &start,
				Label:        buildLabel(strct.Field(i).Name() + ":"),
				Kind:         protocol.Parameter,
				PaddingRight: true,
			})
		}
	}
	return hints
}

func compositeLiteralTypes(node ast.Node, tmap *lsppos.TokenMapper, info *types.Info, q *types.Qualifier) []protocol.InlayHint {
	compLit, ok := node.(*ast.CompositeLit)
	if !ok {
		return nil
	}
	typ := info.TypeOf(compLit)
	if typ == nil {
		return nil
	}
	if compLit.Type != nil {
		return nil
	}
	prefix := ""
	if t, ok := typ.(*types.Pointer); ok {
		typ = t.Elem()
		prefix = "&"
	}
	// The type for this composite literal is implicit, add an inlay hint.
	start, ok := tmap.Position(compLit.Lbrace)
	if !ok {
		return nil
	}
	return []protocol.InlayHint{{
		Position: &start,
		Label:    buildLabel(fmt.Sprintf("%s%s", prefix, types.TypeString(typ, *q))),
		Kind:     protocol.Type,
	}}
}

func buildLabel(s string) []protocol.InlayHintLabelPart {
	label := protocol.InlayHintLabelPart{
		Value: s,
	}
	if len(s) > maxLabelLength+len("...") {
		label.Value = s[:maxLabelLength] + "..."
	}
	return []protocol.InlayHintLabelPart{label}
}
<commit_msg>internal/lsp: respect range for inlay hints<commit_after>// Copyright 2022 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package source

import (
	"context"
	"fmt"
	"go/ast"
	"go/constant"
	"go/token"
	"go/types"
	"strings"

	"golang.org/x/tools/internal/event"
	"golang.org/x/tools/internal/lsp/lsppos"
	"golang.org/x/tools/internal/lsp/protocol"
	"golang.org/x/tools/internal/typeparams"
)

const (
	maxLabelLength = 28
)

type InlayHintFunc func(node ast.Node, tmap *lsppos.TokenMapper, info *types.Info, q *types.Qualifier) []protocol.InlayHint

type Hint struct {
	Name string
	Doc  string
	Run  InlayHintFunc
}

const (
	ParameterNames             = "parameterNames"
	AssignVariableTypes        = "assignVariableTypes"
	ConstantValues             = "constantValues"
	RangeVariableTypes         = "rangeVariableTypes"
	CompositeLiteralTypes      = "compositeLiteralTypes"
	CompositeLiteralFieldNames = "compositeLiteralFields"
	FunctionTypeParameters     = "functionTypeParameters"
)

var AllInlayHints = map[string]*Hint{
	AssignVariableTypes: {
		Name: AssignVariableTypes,
		Doc: `Enable/disable inlay hints for variable types in assign statements:

	i/* int/*, j/* int/* := 0, len(r)-1`,
		Run: assignVariableTypes,
	},
	ParameterNames: {
		Name: ParameterNames,
		Doc: `Enable/disable inlay hints for parameter names:

	parseInt(/* str: */ "123", /* radix: */ 8)`,
		Run: parameterNames,
	},
	ConstantValues: {
		Name: ConstantValues,
		Doc: `Enable/disable inlay hints for constant values:

	const (
		KindNone   Kind = iota/* = 0*/
		KindPrint/*  = 1*/
		KindPrintf/* = 2*/
		KindErrorf/* = 3*/
	)`,
		Run: constantValues,
	},
	RangeVariableTypes: {
		Name: RangeVariableTypes,
		Doc: `Enable/disable inlay hints for variable types in range statements:

	for k/* int*/, v/* string/* := range []string{} {
		fmt.Println(k, v)
	}`,
		Run: rangeVariableTypes,
	},
	CompositeLiteralTypes: {
		Name: CompositeLiteralTypes,
		Doc: `Enable/disable inlay hints for composite literal types:

	for _, c := range []struct {
		in, want string
	}{
		/*struct{ in string; want string }*/{"Hello, world", "dlrow ,olleH"},
	}`,
		Run: compositeLiteralTypes,
	},
	CompositeLiteralFieldNames: {
		Name: CompositeLiteralFieldNames,
		Doc: `Enable/disable inlay hints for composite literal field names:

	{in: "Hello, world", want: "dlrow ,olleH"}`,
		Run: compositeLiteralFields,
	},
	FunctionTypeParameters: {
		Name: FunctionTypeParameters,
		Doc: `Enable/disable inlay hints for implicit type parameters on generic functions:

	myFoo/*[int, string]*/(1, "hello")`,
		Run: funcTypeParams,
	},
}

func InlayHint(ctx context.Context, snapshot Snapshot, fh FileHandle, pRng protocol.Range) ([]protocol.InlayHint, error) {
	ctx, done := event.Start(ctx, "source.InlayHint")
	defer done()

	pkg, pgf, err := GetParsedFile(ctx, snapshot, fh, NarrowestPackage)
	if err != nil {
		return nil, fmt.Errorf("getting file for InlayHint: %w", err)
	}

	// Collect a list of the inlay hints that are enabled.
	inlayHintOptions := snapshot.View().Options().InlayHintOptions
	var enabledHints []InlayHintFunc
	for hint, enabled := range inlayHintOptions.Hints {
		if !enabled {
			continue
		}
		if h, ok := AllInlayHints[hint]; ok {
			enabledHints = append(enabledHints, h.Run)
		}
	}
	if len(enabledHints) == 0 {
		return nil, nil
	}

	tmap := lsppos.NewTokenMapper(pgf.Src, pgf.Tok)
	info := pkg.GetTypesInfo()
	q := Qualifier(pgf.File, pkg.GetTypes(), info)

	// Set the range to the full file if the range is not valid.
	start, end := pgf.File.Pos(), pgf.File.End()
	if pRng.Start.Line < pRng.End.Line || pRng.Start.Character < pRng.End.Character {
		// Adjust start and end for the specified range.
		rng, err := pgf.Mapper.RangeToSpanRange(pRng)
		if err != nil {
			return nil, err
		}
		start, end = rng.Start, rng.End
	}

	var hints []protocol.InlayHint
	ast.Inspect(pgf.File, func(node ast.Node) bool {
		// If not in range, we can stop looking.
		if node == nil || node.End() < start || node.Pos() > end {
			return false
		}
		for _, fn := range enabledHints {
			hints = append(hints, fn(node, tmap, info, &q)...)
		}
		return true
	})
	return hints, nil
}

func parameterNames(node ast.Node, tmap *lsppos.TokenMapper, info *types.Info, _ *types.Qualifier) []protocol.InlayHint {
	callExpr, ok := node.(*ast.CallExpr)
	if !ok {
		return nil
	}
	signature, ok := info.TypeOf(callExpr.Fun).(*types.Signature)
	if !ok {
		return nil
	}

	var hints []protocol.InlayHint
	for i, v := range callExpr.Args {
		start, ok := tmap.Position(v.Pos())
		if !ok {
			continue
		}
		params := signature.Params()
		// When a function has variadic params, we skip args after
		// params.Len().
		if i > params.Len()-1 {
			break
		}
		value := params.At(i).Name()
		// param.Name is empty for built-ins like append
		if value == "" {
			continue
		}
		if signature.Variadic() && i == params.Len()-1 {
			value = value + "..."
		}
		hints = append(hints, protocol.InlayHint{
			Position:     &start,
			Label:        buildLabel(value + ":"),
			Kind:         protocol.Parameter,
			PaddingRight: true,
		})
	}
	return hints
}

func funcTypeParams(node ast.Node, tmap *lsppos.TokenMapper, info *types.Info, _ *types.Qualifier) []protocol.InlayHint {
	ce, ok := node.(*ast.CallExpr)
	if !ok {
		return nil
	}
	id, ok := ce.Fun.(*ast.Ident)
	if !ok {
		return nil
	}
	inst := typeparams.GetInstances(info)[id]
	if inst.TypeArgs == nil {
		return nil
	}
	start, ok := tmap.Position(id.End())
	if !ok {
		return nil
	}
	var args []string
	for i := 0; i < inst.TypeArgs.Len(); i++ {
		args = append(args, inst.TypeArgs.At(i).String())
	}
	if len(args) == 0 {
		return nil
	}
	return []protocol.InlayHint{{
		Position: &start,
		Label:    buildLabel("[" + strings.Join(args, ", ") + "]"),
		Kind:     protocol.Type,
	}}
}

func assignVariableTypes(node ast.Node, tmap *lsppos.TokenMapper, info *types.Info, q *types.Qualifier) []protocol.InlayHint {
	stmt, ok := node.(*ast.AssignStmt)
	if !ok || stmt.Tok != token.DEFINE {
		return nil
	}

	var hints []protocol.InlayHint
	for _, v := range stmt.Lhs {
		if h := variableType(v, tmap, info, q); h != nil {
			hints = append(hints, *h)
		}
	}
	return hints
}

func rangeVariableTypes(node ast.Node, tmap *lsppos.TokenMapper, info *types.Info, q *types.Qualifier) []protocol.InlayHint {
	rStmt, ok := node.(*ast.RangeStmt)
	if !ok {
		return nil
	}
	var hints []protocol.InlayHint
	if h := variableType(rStmt.Key, tmap, info, q); h != nil {
		hints = append(hints, *h)
	}
	if h := variableType(rStmt.Value, tmap, info, q); h != nil {
		hints = append(hints, *h)
	}
	return hints
}

func variableType(e ast.Expr, tmap *lsppos.TokenMapper, info *types.Info, q *types.Qualifier) *protocol.InlayHint {
	typ := info.TypeOf(e)
	if typ == nil {
		return nil
	}
	end, ok := tmap.Position(e.End())
	if !ok {
		return nil
	}
	return &protocol.InlayHint{
		Position:    &end,
		Label:       buildLabel(types.TypeString(typ, *q)),
		Kind:        protocol.Type,
		PaddingLeft: true,
	}
}

func constantValues(node ast.Node, tmap *lsppos.TokenMapper, info *types.Info, _ *types.Qualifier) []protocol.InlayHint {
	genDecl, ok := node.(*ast.GenDecl)
	if !ok || genDecl.Tok != token.CONST {
		return nil
	}

	var hints []protocol.InlayHint
	for _, v := range genDecl.Specs {
		spec, ok := v.(*ast.ValueSpec)
		if !ok {
			continue
		}
		end, ok := tmap.Position(v.End())
		if !ok {
			continue
		}
		// Show hints when values are missing or at least one value is not
		// a basic literal.
		showHints := len(spec.Values) == 0
		checkValues := len(spec.Names) == len(spec.Values)
		var values []string
		for i, w := range spec.Names {
			obj, ok := info.ObjectOf(w).(*types.Const)
			if !ok || obj.Val().Kind() == constant.Unknown {
				return nil
			}
			if checkValues {
				switch spec.Values[i].(type) {
				case *ast.BadExpr:
					return nil
				case *ast.BasicLit:
				default:
					if obj.Val().Kind() != constant.Bool {
						showHints = true
					}
				}
			}
			values = append(values, fmt.Sprintf("%v", obj.Val()))
		}
		if !showHints || len(values) == 0 {
			continue
		}
		hints = append(hints, protocol.InlayHint{
			Position:    &end,
			Label:       buildLabel("= " + strings.Join(values, ", ")),
			PaddingLeft: true,
		})
	}
	return hints
}

func compositeLiteralFields(node ast.Node, tmap *lsppos.TokenMapper, info *types.Info, q *types.Qualifier) []protocol.InlayHint {
	compLit, ok := node.(*ast.CompositeLit)
	if !ok {
		return nil
	}
	typ := info.TypeOf(compLit)
	if typ == nil {
		return nil
	}
	if t, ok := typ.(*types.Pointer); ok {
		typ = t.Elem()
	}
	strct, ok := typ.Underlying().(*types.Struct)
	if !ok {
		return nil
	}

	var hints []protocol.InlayHint

	for i, v := range compLit.Elts {
		if _, ok := v.(*ast.KeyValueExpr); !ok {
			start, ok := tmap.Position(v.Pos())
			if !ok {
				continue
			}
			if i > strct.NumFields()-1 {
				break
			}
			hints = append(hints, protocol.InlayHint{
				Position:     &start,
				Label:        buildLabel(strct.Field(i).Name() + ":"),
				Kind:         protocol.Parameter,
				PaddingRight: true,
			})
		}
	}
	return hints
}

func compositeLiteralTypes(node ast.Node, tmap *lsppos.TokenMapper, info *types.Info, q *types.Qualifier) []protocol.InlayHint {
	compLit, ok := node.(*ast.CompositeLit)
	if !ok {
		return nil
	}
	typ := info.TypeOf(compLit)
	if typ == nil {
		return nil
	}
	if compLit.Type != nil {
		return nil
	}
	prefix := ""
	if t, ok := typ.(*types.Pointer); ok {
		typ = t.Elem()
		prefix = "&"
	}
	// The type for this composite literal is implicit, add an inlay hint.
	start, ok := tmap.Position(compLit.Lbrace)
	if !ok {
		return nil
	}
	return []protocol.InlayHint{{
		Position: &start,
		Label:    buildLabel(fmt.Sprintf("%s%s", prefix, types.TypeString(typ, *q))),
		Kind:     protocol.Type,
	}}
}

func buildLabel(s string) []protocol.InlayHintLabelPart {
	label := protocol.InlayHintLabelPart{
		Value: s,
	}
	if len(s) > maxLabelLength+len("...") {
		label.Value = s[:maxLabelLength] + "..."
	}
	return []protocol.InlayHintLabelPart{label}
}
<|endoftext|>


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

package source

import (
	"context"
	"go/ast"
	"go/token"
	"go/types"

	"golang.org/x/tools/internal/lsp/protocol"
	"golang.org/x/tools/internal/telemetry/trace"
)

// ReferenceInfo holds information about reference to an identifier in Go source.
type ReferenceInfo struct {
	Name string
	mappedRange
	ident         *ast.Ident
	obj           types.Object
	pkg           Package
	isDeclaration bool
}

// References returns a list of references for a given identifier within the packages
// containing i.File. Declarations appear first in the result.
func References(ctx context.Context, s Snapshot, f FileHandle, pp protocol.Position, includeDeclaration bool) ([]*ReferenceInfo, error) {
	ctx, done := trace.StartSpan(ctx, "source.References")
	defer done()

	qualifiedObjs, err := qualifiedObjsAtProtocolPos(ctx, s, f, pp)
	if err != nil {
		return nil, err
	}

	var (
		references []*ReferenceInfo
		seen       = make(map[token.Position]bool)
		fset       = s.View().Session().Cache().FileSet()
	)

	for _, qo := range qualifiedObjs {
		var searchPkgs []Package

		// Only search dependents if the object is exported.
		if qo.obj.Exported() {
			reverseDeps, err := s.GetReverseDependencies(ctx, qo.pkg.ID())
			if err != nil {
				return nil, err
			}

			for _, ph := range reverseDeps {
				pkg, err := ph.Check(ctx)
				if err != nil {
					return nil, err
				}
				searchPkgs = append(searchPkgs, pkg)
			}
		}

		// Add the package in which the identifier is declared.
		searchPkgs = append(searchPkgs, qo.pkg)

		for _, pkg := range searchPkgs {
			for ident, obj := range pkg.GetTypesInfo().Uses {
				if obj != qo.obj {
					continue
				}

				pos := fset.Position(ident.Pos())
				if seen[pos] {
					continue
				}
				seen[pos] = true

				rng, err := posToMappedRange(s.View(), pkg, ident.Pos(), ident.End())
				if err != nil {
					return nil, err
				}
				references = append(references, &ReferenceInfo{
					Name:        ident.Name,
					ident:       ident,
					pkg:         pkg,
					obj:         obj,
					mappedRange: rng,
				})
			}
		}
	}

	if includeDeclaration {
		rng, err := objToMappedRange(s.View(), qualifiedObjs[0].pkg, qualifiedObjs[0].obj)
		if err != nil {
			return nil, err
		}

		ident, _ := qualifiedObjs[0].node.(*ast.Ident)
		references = append(references, &ReferenceInfo{
			mappedRange:   rng,
			Name:          qualifiedObjs[0].obj.Name(),
			ident:         ident,
			obj:           qualifiedObjs[0].obj,
			pkg:           qualifiedObjs[0].pkg,
			isDeclaration: true,
		})
	}

	return references, nil
}
<commit_msg>internal/lsp/source: return obj decl first in find-references<commit_after>// Copyright 2019 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package source

import (
	"context"
	"go/ast"
	"go/token"
	"go/types"

	"golang.org/x/tools/internal/lsp/protocol"
	"golang.org/x/tools/internal/telemetry/trace"
)

// ReferenceInfo holds information about reference to an identifier in Go source.
type ReferenceInfo struct {
	Name string
	mappedRange
	ident         *ast.Ident
	obj           types.Object
	pkg           Package
	isDeclaration bool
}

// References returns a list of references for a given identifier within the packages
// containing i.File. Declarations appear first in the result.
func References(ctx context.Context, s Snapshot, f FileHandle, pp protocol.Position, includeDeclaration bool) ([]*ReferenceInfo, error) {
	ctx, done := trace.StartSpan(ctx, "source.References")
	defer done()

	qualifiedObjs, err := qualifiedObjsAtProtocolPos(ctx, s, f, pp)
	if err != nil {
		return nil, err
	}

	var (
		references []*ReferenceInfo
		seen       = make(map[token.Position]bool)
		fset       = s.View().Session().Cache().FileSet()
	)

	// Make sure declaration is the first item in the response.
	if includeDeclaration {
		rng, err := objToMappedRange(s.View(), qualifiedObjs[0].pkg, qualifiedObjs[0].obj)
		if err != nil {
			return nil, err
		}

		ident, _ := qualifiedObjs[0].node.(*ast.Ident)
		references = append(references, &ReferenceInfo{
			mappedRange:   rng,
			Name:          qualifiedObjs[0].obj.Name(),
			ident:         ident,
			obj:           qualifiedObjs[0].obj,
			pkg:           qualifiedObjs[0].pkg,
			isDeclaration: true,
		})
	}

	for _, qo := range qualifiedObjs {
		var searchPkgs []Package

		// Only search dependents if the object is exported.
		if qo.obj.Exported() {
			reverseDeps, err := s.GetReverseDependencies(ctx, qo.pkg.ID())
			if err != nil {
				return nil, err
			}

			for _, ph := range reverseDeps {
				pkg, err := ph.Check(ctx)
				if err != nil {
					return nil, err
				}
				searchPkgs = append(searchPkgs, pkg)
			}
		}

		// Add the package in which the identifier is declared.
		searchPkgs = append(searchPkgs, qo.pkg)

		for _, pkg := range searchPkgs {
			for ident, obj := range pkg.GetTypesInfo().Uses {
				if obj != qo.obj {
					continue
				}

				pos := fset.Position(ident.Pos())
				if seen[pos] {
					continue
				}
				seen[pos] = true

				rng, err := posToMappedRange(s.View(), pkg, ident.Pos(), ident.End())
				if err != nil {
					return nil, err
				}
				references = append(references, &ReferenceInfo{
					Name:        ident.Name,
					ident:       ident,
					pkg:         pkg,
					obj:         obj,
					mappedRange: rng,
				})
			}
		}
	}

	return references, nil
}
<|endoftext|>