datetime: implement interval math

Signed-off-by: Vicent Marti <vmg@strn.cat>

go/mysql/datetime/types.go → go/mysql/datetime/datetime.go

@@ -94,7 +94,7 @@ func (t Time) FormatDecimal() decimal.Decimal {
 }
 
 func (t Time) ToDateTime() (out DateTime) {
-	return FromStdTime(t.ToStdTime(time.Local))
+	return NewDateTimeFromStd(t.ToStdTime(time.Local))
 }
 
 func (t Time) IsZero() bool {
@@ -413,6 +413,20 @@ func (t Time) ToStdTime(loc *time.Location) (out time.Time) {
 	return t.toStdTime(year, month, day, loc)
 }
 
+func (t Time) AddInterval(itv *Interval, stradd bool) (Time, uint8, bool) {
+	dt := DateTime{Time: t}
+	ok := dt.addInterval(itv)
+	return dt.Time, itv.precision(stradd), ok
+}
+
+func (t Time) toSeconds() int {
+	tsecs := t.Hour()*secondsPerHour + t.Minute()*secondsPerMinute + t.Second()
+	if t.Neg() {
+		return -tsecs
+	}
+	return tsecs
+}
+
 func (d Date) ToStdTime(loc *time.Location) (out time.Time) {
 	return time.Date(d.Year(), time.Month(d.Month()), d.Day(), 0, 0, 0, 0, loc)
 }
@@ -471,6 +485,12 @@ func (d Date) Compare(d2 Date) int {
 	return 0
 }
 
+func (d Date) AddInterval(itv *Interval) (Date, bool) {
+	dt := DateTime{Date: d}
+	ok := dt.addInterval(itv)
+	return dt.Date, ok
+}
+
 func (dt DateTime) FormatInt64() int64 {
 	d := dt.Round(0)
 	return d.Date.FormatInt64()*1000000 + d.Time.FormatInt64()
@@ -493,7 +513,7 @@ func (dt DateTime) Compare(dt2 DateTime) int {
 	case zerodate1 || zerodate2:
 		// if we're comparing a time to a datetime, we need to normalize them
 		// both into datetimes; this normalization is not trivial because negative
-		// times result in a date change, to let the standard library handle this
+		// times result in a date change, so let the standard library handle this
 		return dt.ToStdTime(time.Local).Compare(dt2.ToStdTime(time.Local))
 	}
 	if cmp := dt.Date.Compare(dt2.Date); cmp != 0 {
@@ -502,6 +522,11 @@ func (dt DateTime) Compare(dt2 DateTime) int {
 	return dt.Time.Compare(dt2.Time)
 }
 
+func (dt DateTime) AddInterval(itv *Interval, stradd bool) (DateTime, uint8, bool) {
+	ok := dt.addInterval(itv)
+	return dt, itv.precision(stradd), ok
+}
+
 func (dt DateTime) Round(p int) (r DateTime) {
 	if dt.Time.nanosecond == 0 {
 		return dt
@@ -521,28 +546,124 @@ func (dt DateTime) Round(p int) (r DateTime) {
 	r = dt
 	if n == 1e9 {
 		r.Time.nanosecond = 0
-		return FromStdTime(r.ToStdTime(time.Local).Add(time.Second))
+		return NewDateTimeFromStd(r.ToStdTime(time.Local).Add(time.Second))
 	}
 	r.Time.nanosecond = uint32(n)
 	return r
 }
 
-func FromStdTime(t time.Time) DateTime {
+func (dt DateTime) toSeconds() int {
+	return (dt.Date.Day()-1)*secondsPerDay + dt.Time.toSeconds()
+}
+
+func (dt *DateTime) addInterval(itv *Interval) bool {
+	switch {
+	case itv.unit.HasTimeParts():
+		if !itv.inRange() {
+			return false
+		}
+
+		nsec := dt.Time.Nanosecond() + itv.nsec
+		sec := dt.toSeconds() + itv.toSeconds() + (nsec / int(time.Second))
+		nsec = nsec % int(time.Second)
+
+		if nsec < 0 {
+			nsec += int(time.Second)
+			sec--
+		}
+
+		days := sec / secondsPerDay
+		sec -= days * secondsPerDay
+
+		if sec < 0 {
+			sec += secondsPerDay
+			days--
+		}
+
+		dt.Time.nanosecond = uint32(nsec)
+		dt.Time.second = uint8(sec % secondsPerMinute)
+		dt.Time.minute = uint8((sec / secondsPerMinute) % secondsPerMinute)
+		dt.Time.hour = uint16(sec / secondsPerHour)
+
+		daynum := mysqlDayNumber(dt.Date.Year(), dt.Date.Month(), 1) + days
+		if daynum < 0 || daynum > maxDay {
+			return false
+		}
+
+		dt.Date.year, dt.Date.month, dt.Date.day = mysqlDateFromDayNumber(daynum)
+		return true
+
+	case itv.unit.HasDayParts():
+		daynum := mysqlDayNumber(dt.Date.Year(), dt.Date.Month(), dt.Date.Day())
+		daynum += itv.day
+		dt.Date.year, dt.Date.month, dt.Date.day = mysqlDateFromDayNumber(daynum)
+		return true
+
+	case itv.unit.HasMonthParts():
+		months := dt.Date.Year()*12 + itv.year*12 + (dt.Date.Month() - 1) + itv.month
+		if months < 0 || months >= 120000 {
+			return false
+		}
+
+		year := months / 12
+		month := (months % 12) + 1
+
+		dt.Date.year = uint16(year)
+		dt.Date.month = uint8(month)
+
+		// MySQL quirk: if the original date was in a day that the new month
+		// doesn't have, the date is offset backwards to the last day of
+		// the new month. This is the opposite to normal date handling where
+		// we'd offset days into the next month.
+		if dim := daysIn(time.Month(month), year); dt.Date.Day() > dim {
+			dt.Date.day = uint8(dim)
+		}
+		return true
+
+	case itv.unit == IntervalYear:
+		if itv.year > 10000 {
+			return false
+		}
+
+		year := dt.Date.Year() + itv.year
+		dt.Date.year = uint16(year)
+
+		// MySQL quirk: if the original date was Feb 29th on a leap year, and
+		// the resulting year is not a leap year, the date is offset backwards.
+		// This is the opposite to what normal date handling does.
+		if dt.Date.Month() == 2 && dt.Date.Day() == 29 && !isLeap(year) {
+			dt.Date.day = 28
+		}
+		return true
+
+	default:
+		panic("unexpected IntervalType")
+	}
+}
+
+func NewDateFromStd(t time.Time) Date {
 	year, month, day := t.Date()
+	return Date{
+		year:  uint16(year),
+		month: uint8(month),
+		day:   uint8(day),
+	}
+}
+
+func NewTimeFromStd(t time.Time) Time {
 	hour, min, sec := t.Clock()
 	nsec := t.Nanosecond()
+	return Time{
+		hour:       uint16(hour),
+		minute:     uint8(min),
+		second:     uint8(sec),
+		nanosecond: uint32(nsec),
+	}
+}
 
+func NewDateTimeFromStd(t time.Time) DateTime {
 	return DateTime{
-		Date: Date{
-			year:  uint16(year),
-			month: uint8(month),
-			day:   uint8(day),
-		},
-		Time: Time{
-			hour:       uint16(hour),
-			minute:     uint8(min),
-			second:     uint8(sec),
-			nanosecond: uint32(nsec),
-		},
+		Date: NewDateFromStd(t),
+		Time: NewTimeFromStd(t),
 	}
 }

go/mysql/datetime/helpers.go

@@ -198,7 +198,7 @@ func getnuml(s string, l int) (int, string, bool) {
 }
 
 func getnumn(s string) (int, string, bool) {
-	if !isDigit(s, 0) {
+	if len(s) == 0 || !('0' <= s[0] && s[0] <= '9') {
 		return 0, s, false
 	}
 
@@ -229,6 +229,14 @@ var daysBefore = [...]int32{
 	31 + 28 + 31 + 30 + 31 + 30 + 31 + 31 + 30 + 31 + 30 + 31,
 }
 
+var daysInMonth = [...]int{
+	31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31,
+}
+
+var daysInMonthLeap = [...]int{
+	31, 29, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31,
+}
+
 func daysIn(m time.Month, year int) int {
 	if m == time.February && isLeap(year) {
 		return 29
@@ -240,6 +248,13 @@ func isLeap(year int) bool {
 	return year%4 == 0 && (year%100 != 0 || year%400 == 0)
 }
 
+func daysInYear(year int) int {
+	if isLeap(year) {
+		return 366
+	}
+	return 365
+}
+
 func parseNanoseconds[bytes []byte | string](value bytes, nbytes int) (ns int, l int, ok bool) {
 	if value[0] != '.' {
 		return 0, 0, false
@@ -268,3 +283,9 @@ func parseNanoseconds[bytes []byte | string](value bytes, nbytes int) (ns int, l
 
 	return
 }
+
+const (
+	secondsPerMinute = 60
+	secondsPerHour   = 60 * secondsPerMinute
+	secondsPerDay    = 24 * secondsPerHour
+)

go/mysql/datetime/mydate.go

@@ -0,0 +1,83 @@
+/*
+Copyright 2023 The Vitess 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 datetime
+
+// mysqlDayNumber converts a date into an absolute day number.
+// This is an algorithm that has been reverse engineered from MySQL;
+// the tables used as a reference can be found in `testdata/year_to_daynr.json`.
+// It is worth noting that this absolute day number does not match the
+// day count that traditional datetime systems use (e.g. the daycount
+// algorithm in the Go standard library). It is often off by one, possibly
+// because of incorrect leap year handling, but the inverse of the algorithm
+// in mysqlDateFromDayNumber takes this into account. Hence, the results
+// of this function can ONLY be passed to mysqlDateFromDayNumber; using
+// a day number with one of Go's datetime APIs will return incorrect results.
+// This API should only be used when performing datetime calculations (addition
+// and subtraction), so that the results match MySQL's. All other date handling
+// operations must use our helpers based on Go's standard library.
+func mysqlDayNumber(year, month, day int) int {
+	if year == 0 && month == 0 {
+		return 0
+	}
+
+	days := 365*year + 31*(month-1) + day
+	switch month {
+	case 1, 2:
+		year = year - 1
+	default:
+		days = days - (month*4+23)/10
+	}
+
+	leapAdjust := ((year/100 + 1) * 3) / 4
+	return days + year/4 - leapAdjust
+}
+
+// mysqlDateFromDayNumber converts an absolute day number into a date (a year, month, day triplet).
+// This is an algorithm that has been reverse engineered from MySQL;
+// the tables used as a reference can be found in `testdata/daynr_to_date.json`.
+// See the warning from mysqlDayNumber: the day number used as an argument to
+// this function must come from mysqlDayNumber or the results won't be correct.
+// This API should only be used when performing datetime calculations (addition
+// and subtraction), so that the results match MySQL's. All other date handling
+// operations must use our helpers based on Go's standard library.
+func mysqlDateFromDayNumber(daynr int) (uint16, uint8, uint8) {
+	if daynr <= 365 || daynr >= 3652500 {
+		return 0, 0, 0
+	}
+
+	year := daynr * 100 / 36525
+	leapAdjust := (((year-1)/100 + 1) * 3) / 4
+	yday := (daynr - year*365) - (year-1)/4 + leapAdjust
+
+	if diy := daysInYear(year); yday > diy {
+		yday -= diy
+		year++
+	}
+
+	daycount := daysInMonth
+	if isLeap(year) {
+		daycount = daysInMonthLeap
+	}
+	for month, dim := range daycount {
+		if yday <= dim {
+			return uint16(year), uint8(month + 1), uint8(yday)
+		}
+		yday -= dim
+	}
+
+	panic("unreachable: yday is too large?")
+}

go/mysql/datetime/mydate_test.go

@@ -0,0 +1,59 @@
+/*
+Copyright 2023 The Vitess 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 datetime
+
+import (
+	"encoding/json"
+	"os"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestDayNumber(t *testing.T) {
+	td, err := os.Open("testdata/year_to_daynr.json")
+	require.NoError(t, err)
+	defer td.Close()
+
+	var expected []int
+	err = json.NewDecoder(td).Decode(&expected)
+	require.NoError(t, err)
+
+	for year, daynr := range expected {
+		assert.Equal(t, daynr, mysqlDayNumber(year, 1, 1))
+	}
+}
+
+func TestDayNumberFields(t *testing.T) {
+	td, err := os.Open("testdata/daynr_to_date.json")
+	require.NoError(t, err)
+	defer td.Close()
+
+	var expected [][4]int
+	err = json.NewDecoder(td).Decode(&expected)
+	require.NoError(t, err)
+
+	for _, tc := range expected {
+		y, m, d := mysqlDateFromDayNumber(tc[0])
+		assert.Equal(t, tc[1], int(y))
+		assert.Equal(t, tc[2], int(m))
+		assert.Equal(t, tc[3], int(d))
+
+		assert.Equalf(t, tc[0], mysqlDayNumber(tc[1], tc[2], tc[3]), "date %d-%d-%d", tc[1], tc[2], tc[3])
+	}
+}