Timestamp expressions WIP

Author: tom-andersen

firebase-firestore/src/main/java/com/google/firebase/firestore/model/Values.kt

@@ -731,4 +731,8 @@ internal object Values {
       is VectorValue -> encodeValue(value)
       else -> throw IllegalArgumentException("Unexpected type: $value")
     }
+
+  @JvmStatic
+  fun timestamp(seconds: Long, nanos: Int): Timestamp =
+    Timestamp.newBuilder().setSeconds(seconds).setNanos(nanos).build()
 }

firebase-firestore/src/main/java/com/google/firebase/firestore/pipeline/EvaluateResult.kt

@@ -17,12 +17,11 @@ internal sealed class EvaluateResult(val value: Value?) {
     fun long(long: Long) = EvaluateResultValue(encodeValue(long))
     fun long(int: Int) = EvaluateResultValue(encodeValue(int.toLong()))
     fun string(string: String) = EvaluateResultValue(encodeValue(string))
+    fun timestamp(timestamp: Timestamp): EvaluateResult =
+      EvaluateResultValue(encodeValue(timestamp))
     fun timestamp(seconds: Long, nanos: Int): EvaluateResult =
       if (seconds !in -62_135_596_800 until 253_402_300_800) EvaluateResultError
-      else
-        EvaluateResultValue(
-          encodeValue(Timestamp.newBuilder().setSeconds(seconds).setNanos(nanos).build())
-        )
+      else timestamp(Values.timestamp(seconds, nanos))
   }
   internal inline fun evaluateNonNull(f: (Value) -> EvaluateResult): EvaluateResult =
     if (value?.hasNullValue() == true) f(value) else this

firebase-firestore/src/main/java/com/google/firebase/firestore/pipeline/evaluation.kt

@@ -3,6 +3,7 @@ package com.google.firebase.firestore.pipeline
 import com.google.common.math.LongMath
 import com.google.common.math.LongMath.checkedAdd
 import com.google.common.math.LongMath.checkedMultiply
+import com.google.common.math.LongMath.checkedSubtract
 import com.google.firebase.firestore.UserDataReader
 import com.google.firebase.firestore.model.MutableDocument
 import com.google.firebase.firestore.model.Values
@@ -273,9 +274,71 @@ internal val evaluateStrJoin = notImplemented // TODO: Does not exist in express
 
 // === Date / Timestamp Functions ===
 
-internal val evaluateTimestampAdd = notImplemented
+private const val L_NANOS_PER_SECOND: Long = 1000_000_000
+private const val I_NANOS_PER_SECOND: Int = 1000_000_000
 
-internal val evaluateTimestampSub = notImplemented
+private const val L_MICROS_PER_SECOND: Long = 1000_000
+private const val I_MICROS_PER_SECOND: Int = 1000_000
+
+private const val L_MILLIS_PER_SECOND: Long = 1000
+private const val I_MILLIS_PER_SECOND: Int = 1000
+
+internal fun plus(t: Timestamp, seconds: Long, nanos: Long): Timestamp =
+  if (nanos == 0L) {
+    plus(t, seconds)
+  } else {
+    val nanoSum = t.nanos + nanos // Overflow not possible since nanos is 0 to 1 000 000.
+    val secondsSum: Long = checkedAdd(checkedAdd(t.seconds, seconds), nanoSum / L_NANOS_PER_SECOND)
+    Values.timestamp(secondsSum, (nanoSum % I_NANOS_PER_SECOND).toInt())
+  }
+
+private fun plus(t: Timestamp, seconds: Long): Timestamp =
+  if (seconds == 0L) t
+  else Values.timestamp(checkedAdd(t.seconds, seconds), t.nanos)
+
+internal fun minus(t: Timestamp, seconds: Long, nanos: Long): Timestamp =
+  if (nanos == 0L) {
+    minus(t, seconds)
+  } else {
+    val nanoSum = t.nanos - nanos // Overflow not possible since nanos is 0 to 1 000 000.
+    val secondsSum: Long = checkedSubtract(t.seconds, checkedSubtract(seconds, nanoSum / L_NANOS_PER_SECOND))
+    Values.timestamp(secondsSum, (nanoSum % I_NANOS_PER_SECOND).toInt())
+  }
+
+private fun minus(t: Timestamp, seconds: Long): Timestamp =
+  if (seconds == 0L) t
+  else Values.timestamp(checkedSubtract(t.seconds, seconds), t.nanos)
+
+
+internal val evaluateTimestampAdd =
+  ternaryTimestampFunction { t: Timestamp, u: String, n: Long ->
+    EvaluateResult.timestamp(
+      when (u) {
+        "microsecond" -> plus(t, n / L_MICROS_PER_SECOND, (n % L_MICROS_PER_SECOND) * 1000)
+        "millisecond" -> plus(t, n / L_MILLIS_PER_SECOND, (n % L_MILLIS_PER_SECOND) * 1000_000)
+        "second" -> plus(t, n)
+        "minute" -> plus(t, checkedMultiply(n, 60))
+        "hour" -> plus(t, checkedMultiply(n, 3600))
+        "day" -> plus(t, checkedMultiply(n, 86400))
+        else -> return@ternaryTimestampFunction EvaluateResultError
+      }
+    )
+  }
+
+internal val evaluateTimestampSub =
+  ternaryTimestampFunction { t: Timestamp, u: String, n: Long ->
+    EvaluateResult.timestamp(
+      when (u) {
+        "microsecond" -> minus(t, n / L_MICROS_PER_SECOND, (n % L_MICROS_PER_SECOND) * 1000)
+        "millisecond" -> minus(t, n / L_MILLIS_PER_SECOND, (n % L_MILLIS_PER_SECOND) * 1000_000)
+        "second" -> minus(t, n)
+        "minute" -> minus(t, checkedMultiply(n, 60))
+        "hour" -> minus(t, checkedMultiply(n, 3600))
+        "day" -> minus(t, checkedMultiply(n, 86400))
+        else -> return@ternaryTimestampFunction EvaluateResultError
+      }
+    )
+  }
 
 internal val evaluateTimestampTrunc = notImplemented // TODO: Does not exist in expressions.kt yet.
 
@@ -284,39 +347,46 @@ internal val evaluateTimestampToUnixMicros = unaryFunction { t: Timestamp ->
     if (t.seconds < Long.MIN_VALUE / 1_000_000) {
       // To avoid overflow when very close to Long.MIN_VALUE, add 1 second, multiply, then subtract
       // again.
-      val micros = checkedMultiply(t.seconds + 1, 1_000_000)
-      val adjustment = t.nanos.toLong() / 1_000 - 1_000_000
+      val micros = checkedMultiply(t.seconds + 1, L_MICROS_PER_SECOND)
+      val adjustment = t.nanos.toLong() / L_MILLIS_PER_SECOND - L_MICROS_PER_SECOND
       checkedAdd(micros, adjustment)
     } else {
-      val micros = checkedMultiply(t.seconds, 1_000_000)
-      checkedAdd(micros, t.nanos.toLong() / 1_000)
+      val micros = checkedMultiply(t.seconds, L_MICROS_PER_SECOND)
+      checkedAdd(micros, t.nanos.toLong() / L_MILLIS_PER_SECOND)
     }
   )
 }
 
 internal val evaluateTimestampToUnixMillis = unaryFunction { t: Timestamp ->
   EvaluateResult.long(
     if (t.seconds < 0 && t.nanos > 0) {
-      val millis = checkedMultiply(t.seconds + 1, 1000)
-      val adjustment = t.nanos.toLong() / 1000_000 - 1000
+      val millis = checkedMultiply(t.seconds + 1, L_MILLIS_PER_SECOND)
+      val adjustment = t.nanos.toLong() / L_MICROS_PER_SECOND - L_MILLIS_PER_SECOND
       checkedAdd(millis, adjustment)
     } else {
-      val millis = checkedMultiply(t.seconds, 1000)
-      checkedAdd(millis, t.nanos.toLong() / 1000_000)
+      val millis = checkedMultiply(t.seconds, L_MILLIS_PER_SECOND)
+      checkedAdd(millis, t.nanos.toLong() / L_MICROS_PER_SECOND)
     }
   )
 }
 
 internal val evaluateTimestampToUnixSeconds = unaryFunction { t: Timestamp ->
-  if (t.nanos !in 0 until 1_000_000_000) EvaluateResultError else EvaluateResult.long(t.seconds)
+  if (t.nanos !in 0 until L_NANOS_PER_SECOND) EvaluateResultError
+  else EvaluateResult.long(t.seconds)
 }
 
 internal val evaluateUnixMicrosToTimestamp = unaryFunction { micros: Long ->
-  EvaluateResult.timestamp(Math.floorDiv(micros, 1000_000), Math.floorMod(micros, 1000_000))
+  EvaluateResult.timestamp(
+    Math.floorDiv(micros, L_MICROS_PER_SECOND),
+    Math.floorMod(micros, I_MICROS_PER_SECOND)
+  )
 }
 
 internal val evaluateUnixMillisToTimestamp = unaryFunction { millis: Long ->
-  EvaluateResult.timestamp(Math.floorDiv(millis, 1000), Math.floorMod(millis, 1000))
+  EvaluateResult.timestamp(
+    Math.floorDiv(millis, L_MILLIS_PER_SECOND),
+    Math.floorMod(millis, I_MILLIS_PER_SECOND)
+  )
 }
 
 internal val evaluateUnixSecondsToTimestamp = unaryFunction { seconds: Long ->
@@ -457,6 +527,43 @@ private inline fun binaryFunction(crossinline function: (String, String) -> Eval
     function
   )
 
+private inline fun ternaryTimestampFunction(
+  crossinline function: (Timestamp, String, Long) -> EvaluateResult
+): EvaluateFunction = ternaryNullableValueFunction { timestamp: Value, unit: Value, number: Value ->
+  val t: Timestamp =
+    when (timestamp.valueTypeCase) {
+      Value.ValueTypeCase.NULL_VALUE -> return@ternaryNullableValueFunction EvaluateResult.NULL
+      Value.ValueTypeCase.TIMESTAMP_VALUE -> timestamp.timestampValue
+      else -> return@ternaryNullableValueFunction EvaluateResultError
+    }
+  val u: String =
+    if (unit.hasStringValue()) unit.stringValue
+    else return@ternaryNullableValueFunction EvaluateResultError
+  val n: Long =
+    when (number.valueTypeCase) {
+      Value.ValueTypeCase.NULL_VALUE -> return@ternaryNullableValueFunction EvaluateResult.NULL
+      Value.ValueTypeCase.INTEGER_VALUE -> number.integerValue
+      else -> return@ternaryNullableValueFunction EvaluateResultError
+    }
+  function(t, u, n)
+}
+
+private inline fun ternaryNullableValueFunction(
+  crossinline function: (Value, Value, Value) -> EvaluateResult
+): EvaluateFunction = { params ->
+  if (params.size != 3)
+    throw Assert.fail("Function should have exactly 3 params, but %d were given.", params.size)
+  val p1 = params[0]
+  val p2 = params[1]
+  val p3 = params[2]
+  block@{ input: MutableDocument ->
+    val v1 = p1(input).value ?: return@block EvaluateResultError
+    val v2 = p2(input).value ?: return@block EvaluateResultError
+    val v3 = p3(input).value ?: return@block EvaluateResultError
+    catch { function(v1, v2, v3) }
+  }
+}
+
 private inline fun <T1, T2> binaryFunctionType(
   valueTypeCase1: Value.ValueTypeCase,
   crossinline valueExtractor1: (Value) -> T1,

firebase-firestore/src/main/java/com/google/firebase/firestore/pipeline/expressions.kt

@@ -780,8 +780,7 @@ abstract class Expr internal constructor() {
      * @return A new [Expr] representing the addition operation.
      */
     @JvmStatic
-    fun add(first: Expr, second: Expr): Expr =
-      FunctionExpr("add", evaluateAdd, first, second)
+    fun add(first: Expr, second: Expr): Expr = FunctionExpr("add", evaluateAdd, first, second)
 
     /**
      * Creates an expression that adds numeric expressions with a constant.
@@ -791,8 +790,7 @@ abstract class Expr internal constructor() {
      * @return A new [Expr] representing the addition operation.
      */
     @JvmStatic
-    fun add(first: Expr, second: Number): Expr =
-      FunctionExpr("add", evaluateAdd, first, second)
+    fun add(first: Expr, second: Number): Expr = FunctionExpr("add", evaluateAdd, first, second)
 
     /**
      * Creates an expression that adds a numeric field with a numeric expression.

firebase-firestore/src/test/java/com/google/firebase/firestore/core/PipelineTests.kt

@@ -4,8 +4,12 @@ import com.google.common.truth.Truth.assertThat
 import com.google.firebase.firestore.RealtimePipelineSource
 import com.google.firebase.firestore.TestUtil
 import com.google.firebase.firestore.model.MutableDocument
+import com.google.firebase.firestore.model.Values
 import com.google.firebase.firestore.pipeline.Expr.Companion.field
+import com.google.firebase.firestore.pipeline.minus
+import com.google.firebase.firestore.pipeline.plus
 import com.google.firebase.firestore.testutil.TestUtilKtx.doc
+import com.google.protobuf.Timestamp
 import kotlinx.coroutines.flow.flowOf
 import kotlinx.coroutines.flow.toList
 import kotlinx.coroutines.runBlocking
@@ -26,4 +30,33 @@ internal class PipelineTests {
 
     assertThat(list).hasSize(1)
   }
+
+  @Test
+  fun xxx(): Unit = runBlocking {
+    val zero: Timestamp = Values.timestamp(0, 0)
+    
+    assertThat(plus(zero, 0, 0))
+      .isEqualTo(zero)
+
+    assertThat(plus(Values.timestamp(1, 1), 1, 1))
+      .isEqualTo(Values.timestamp(2, 2))
+
+    assertThat(plus(Values.timestamp(1, 1), 0, 1))
+      .isEqualTo(Values.timestamp(1, 2))
+
+    assertThat(plus(Values.timestamp(1, 1), 1, 0))
+      .isEqualTo(Values.timestamp(2, 1))
+
+    assertThat(minus(zero, 0, 0))
+      .isEqualTo(zero)
+
+    assertThat(minus(Values.timestamp(1, 1), 1, 1))
+      .isEqualTo(zero)
+
+    assertThat(minus(Values.timestamp(1, 1), 0, 1))
+      .isEqualTo(Values.timestamp(1, 0))
+
+    assertThat(minus(Values.timestamp(1, 1), 1, 0))
+      .isEqualTo(Values.timestamp(0, 1))
+  }
 }