// Encodes datum at the end of key, using direction `dir` for the encoding.
// The key is a span end key, which is exclusive, but `val` needs to
// be inclusive. So if datum is the last end constraint, we transform it accordingly.
func encodeInclusiveEndValue(
	key roachpb.Key, datum parser.Datum, dir encoding.Direction,
	isLastEndConstraint bool) roachpb.Key {
	// Since the end of a span is exclusive, if the last constraint is an
	// inclusive one, we might need to make the key exclusive by applying a
	// PrefixEnd().  We normally avoid doing this by transforming "a = x" to
	// "a = x±1" for the last end constraint, depending on the encoding direction
	// (since this keeps the key nice and pretty-printable).
	// However, we might not be able to do the ±1.
	needExclusiveKey := false
	if isLastEndConstraint {
		if dir == encoding.Ascending {
			if datum.IsMax() || !datum.HasNext() {
				needExclusiveKey = true
			} else {
				datum = datum.Next()
			}
		} else {
			if datum.IsMin() || !datum.HasPrev() {
				needExclusiveKey = true
			} else {
				datum = datum.Prev()
			}
		}
	}
	key, pErr := sqlbase.EncodeTableKey(key, datum, dir)
	if pErr != nil {
		panic(pErr)
	}
	if needExclusiveKey {
		key = key.PrefixEnd()
	}
	return key
}

func encodeEndConstraintDescending(
	spans []sqlbase.Span, c *parser.ComparisonExpr, isLastEndConstraint bool,
) {
	switch c.Operator {
	case parser.IsNot:
		// An IS NOT NULL expressions allows us to constrain the end of the range
		// to stop at NULL.
		if c.Right != parser.DNull {
			panic(fmt.Sprintf("expected NULL operand for IS NOT operator, found %v", c.Right))
		}
		for i := range spans {
			spans[i].End = encoding.EncodeNotNullDescending(spans[i].End)
		}
	default:
		datum := c.Right.(parser.Datum)
		if c.Operator != parser.GT {
			for i := range spans {
				spans[i].End = encodeInclusiveEndValue(
					spans[i].End, datum, encoding.Descending, isLastEndConstraint)
			}
			break
		}
		if !isLastEndConstraint {
			panic(fmt.Sprintf("can't have other end constraints after a '>' constraint, found %v", c.Operator))
		}
		key, err := sqlbase.EncodeTableKey(nil, datum, encoding.Descending)
		if err != nil {
			panic(err)
		}
		// Append the constraint to all of the existing spans.
		for i := range spans {
			spans[i].End = append(spans[i].End, key...)
		}
	}
}

func encodeStartConstraintDescending(
	spans []sqlbase.Span, c *parser.ComparisonExpr) {
	switch c.Operator {
	case parser.Is:
		// An IS NULL expressions allows us to constrain the start of the range
		// to begin at NULL.
		if c.Right != parser.DNull {
			panic(fmt.Sprintf("expected NULL operand for IS operator, found %v", c.Right))
		}
		for i := range spans {
			spans[i].Start = encoding.EncodeNullDescending(spans[i].Start)
		}
	case parser.NE:
		panic("'!=' operators should have been transformed to 'IS NOT NULL'")
	case parser.LE, parser.EQ:
		datum := c.Right.(parser.Datum)
		key, pErr := sqlbase.EncodeTableKey(nil, datum, encoding.Descending)
		if pErr != nil {
			panic(pErr)
		}
		// Append the constraint to all of the existing spans.
		for i := range spans {
			spans[i].Start = append(spans[i].Start, key...)
		}
	case parser.LT:
		// A "<" constraint is the last start constraint. Since the constraint
		// is exclusive and the start key is inclusive, we're going to apply
		// a .PrefixEnd(). Note that a "<" is usually transformed to a "<=".
		datum := c.Right.(parser.Datum)
		key, pErr := sqlbase.EncodeTableKey(nil, datum, encoding.Descending)
		if pErr != nil {
			panic(pErr)
		}
		// Append the constraint to all of the existing spans.
		for i := range spans {
			spans[i].Start = append(spans[i].Start, key...)
			spans[i].Start = spans[i].Start.PrefixEnd()
		}

	default:
		panic(fmt.Sprintf("unexpected operator: %s", c))
	}
}

// Splits spans according to a constraint like (...) in <tuple>.
// If the constraint is (a,b) IN ((1,2),(3,4)), each input span
// will be split into two: the first one will have "1/2" appended to
// the start and/or end, the second one will have "3/4" appended to
// the start and/or end.
//
// Returns the exploded spans.
func applyInConstraint(
	spans []sqlbase.Span,
	c indexConstraint,
	firstCol int,
	index *sqlbase.IndexDescriptor,
	isLastEndConstraint bool,
) []sqlbase.Span {
	var e *parser.ComparisonExpr
	// It might be that the IN constraint is a start constraint, an
	// end constraint, or both, depending on how whether we had
	// start and end constraints for all the previous index cols.
	if c.start != nil && c.start.Operator == parser.In {
		e = c.start
	} else {
		e = c.end
	}
	tuple := *e.Right.(*parser.DTuple)
	existingSpans := spans
	spans = make([]sqlbase.Span, 0, len(existingSpans)*len(tuple))
	for _, datum := range tuple {
		// start and end will accumulate the end constraint for
		// the current element of the tuple.
		var start, end []byte

		switch t := datum.(type) {
		case *parser.DTuple:
			// The constraint is a tuple of tuples, meaning something like
			// (...) IN ((1,2),(3,4)).
			for j, tupleIdx := range c.tupleMap {
				var err error
				var colDir encoding.Direction
				if colDir, err = index.ColumnDirections[firstCol+j].ToEncodingDirection(); err != nil {
					panic(err)
				}

				if start, err = sqlbase.EncodeTableKey(start, (*t)[tupleIdx], colDir); err != nil {
					panic(err)
				}
				end = encodeInclusiveEndValue(
					end, (*t)[tupleIdx], colDir, isLastEndConstraint && (j == len(c.tupleMap)-1))
			}
		default:
			// The constraint is a tuple of values, meaning something like
			// a IN (1,2).
			var colDir encoding.Direction
			var err error
			if colDir, err = index.ColumnDirections[firstCol].ToEncodingDirection(); err != nil {
				panic(err)
			}
			if start, err = sqlbase.EncodeTableKey(nil, datum, colDir); err != nil {
				panic(err)
			}

			end = encodeInclusiveEndValue(nil, datum, colDir, isLastEndConstraint)
			// TODO(andrei): assert here that we end is not \xff\xff...
			// encodeInclusiveEndValue sometimes calls key.PrefixEnd(),
			// which doesn't work if the input is \xff\xff... However,
			// that shouldn't happen: datum should not have that encoding.
		}
		for _, s := range existingSpans {
			if c.start != nil {
				s.Start = append(append(roachpb.Key(nil), s.Start...), start...)
			}
			if c.end != nil {
				s.End = append(append(roachpb.Key(nil), s.End...), end...)
			}
			spans = append(spans, s)
		}
	}
	return spans
}