Rewrite stages, fixes and refactoring

Author: tom-andersen

firebase-firestore/src/main/java/com/google/firebase/firestore/Pipeline.kt

@@ -16,8 +16,6 @@ package com.google.firebase.firestore
 
 import com.google.android.gms.tasks.Task
 import com.google.android.gms.tasks.TaskCompletionSource
-import com.google.common.collect.FluentIterable
-import com.google.common.collect.ImmutableList
 import com.google.firebase.Timestamp
 import com.google.firebase.firestore.model.DocumentKey
 import com.google.firebase.firestore.model.Values
@@ -62,7 +60,7 @@ open class AbstractPipeline
 internal constructor(
   internal val firestore: FirebaseFirestore,
   internal val userDataReader: UserDataReader,
-  internal val stages: FluentIterable<Stage<*>>
+  internal val stages: List<Stage<*>>
 ) {
   private fun toStructuredPipelineProto(options: InternalOptions?): StructuredPipeline {
     val builder = StructuredPipeline.newBuilder()
@@ -95,7 +93,7 @@ internal constructor(
     private val userDataWriter =
       UserDataWriter(firestore, DocumentSnapshot.ServerTimestampBehavior.DEFAULT)
     private val taskCompletionSource = TaskCompletionSource<PipelineSnapshot>()
-    private val results: ImmutableList.Builder<PipelineResult> = ImmutableList.builder()
+    private val results: MutableList<PipelineResult> = mutableListOf()
     override fun onDocument(
       key: DocumentKey?,
       data: Map<String, Value>,
@@ -115,7 +113,7 @@ internal constructor(
     }
 
     override fun onComplete(executionTime: Timestamp) {
-      taskCompletionSource.setResult(PipelineSnapshot(executionTime, results.build()))
+      taskCompletionSource.setResult(PipelineSnapshot(executionTime, results))
     }
 
     override fun onError(exception: FirebaseFirestoreException) {
@@ -131,16 +129,16 @@ class Pipeline
 private constructor(
   firestore: FirebaseFirestore,
   userDataReader: UserDataReader,
-  stages: FluentIterable<Stage<*>>
+  stages: List<Stage<*>>
 ) : AbstractPipeline(firestore, userDataReader, stages) {
   internal constructor(
     firestore: FirebaseFirestore,
     userDataReader: UserDataReader,
     stage: Stage<*>
-  ) : this(firestore, userDataReader, FluentIterable.of(stage))
+  ) : this(firestore, userDataReader, listOf(stage))
 
   private fun append(stage: Stage<*>): Pipeline {
-    return Pipeline(firestore, userDataReader, stages.append(stage))
+    return Pipeline(firestore, userDataReader, stages.plus(stage))
   }
 
   fun execute(): Task<PipelineSnapshot> = execute(null)
@@ -760,17 +758,18 @@ class RealtimePipeline
 internal constructor(
   firestore: FirebaseFirestore,
   userDataReader: UserDataReader,
-  stages: FluentIterable<Stage<*>>
+  stages: List<Stage<*>>
 ) : AbstractPipeline(firestore, userDataReader, stages) {
   internal constructor(
     firestore: FirebaseFirestore,
     userDataReader: UserDataReader,
     stage: Stage<*>
-  ) : this(firestore, userDataReader, FluentIterable.of(stage))
+  ) : this(firestore, userDataReader, listOf(stage))
 
-  private fun append(stage: Stage<*>): RealtimePipeline {
-    return RealtimePipeline(firestore, userDataReader, stages.append(stage))
-  }
+  private fun with(stages: List<Stage<*>>): RealtimePipeline =
+    RealtimePipeline(firestore, userDataReader, stages)
+
+  private fun append(stage: Stage<*>): RealtimePipeline = with(stages.plus(stage))
 
   fun execute(): Task<PipelineSnapshot> = execute(null)
 
@@ -790,6 +789,33 @@ internal constructor(
     append(SortStage(arrayOf(order, *additionalOrders)))
 
   fun where(condition: BooleanExpr): RealtimePipeline = append(WhereStage(condition))
+
+  internal fun rewriteStages(): RealtimePipeline {
+    var hasOrder = false
+    return with(
+      buildList {
+        for (stage in stages) when (stage) {
+          // Stages whose semantics depend on ordering
+          is LimitStage,
+          is OffsetStage -> {
+            if (!hasOrder) {
+              hasOrder = true
+              add(SortStage.BY_DOCUMENT_ID)
+            }
+            add(stage)
+          }
+          is SortStage -> {
+            hasOrder = true
+            add(stage.withStableOrdering())
+          }
+          else -> add(stage)
+        }
+        if (!hasOrder) {
+          add(SortStage.BY_DOCUMENT_ID)
+        }
+      }
+    )
+  }
 }
 
 /**

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

@@ -20,8 +20,7 @@ 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.FirebaseFirestore
-import com.google.firebase.firestore.UserDataReader
+import com.google.firebase.firestore.RealtimePipeline
 import com.google.firebase.firestore.model.MutableDocument
 import com.google.firebase.firestore.model.Values
 import com.google.firebase.firestore.model.Values.encodeValue
@@ -43,7 +42,7 @@ import kotlin.math.log10
 import kotlin.math.pow
 import kotlin.math.sqrt
 
-internal class EvaluationContext(val db: FirebaseFirestore, val userDataReader: UserDataReader)
+internal class EvaluationContext(val pipeline: RealtimePipeline)
 
 internal typealias EvaluateDocument = (input: MutableDocument) -> EvaluateResult
 

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

@@ -227,7 +227,7 @@ abstract class Expr internal constructor() {
         override fun evaluateContext(
           context: EvaluationContext
         ): (input: MutableDocument) -> EvaluateResult {
-          val result = EvaluateResultValue(toProto(context.userDataReader))
+          val result = EvaluateResultValue(toProto(context.pipeline.userDataReader))
           return { _ -> result }
         }
       }
@@ -4124,8 +4124,8 @@ abstract class Expr internal constructor() {
 
 /** Expressions that have an alias are [Selectable] */
 abstract class Selectable : Expr() {
-  internal abstract fun getAlias(): String
-  internal abstract fun getExpr(): Expr
+  internal abstract val alias: String
+  internal abstract val expr: Expr
 
   internal companion object {
     fun toSelectable(o: Any): Selectable {
@@ -4140,10 +4140,8 @@ abstract class Selectable : Expr() {
 }
 
 /** Represents an expression that will be given the alias in the output document. */
-class ExprWithAlias internal constructor(private val alias: String, private val expr: Expr) :
+class ExprWithAlias internal constructor(override val alias: String, override val expr: Expr) :
   Selectable() {
-  override fun getAlias() = alias
-  override fun getExpr() = expr
   override fun toProto(userDataReader: UserDataReader): Value = expr.toProto(userDataReader)
   override fun evaluateContext(context: EvaluationContext) = expr.evaluateContext(context)
 }
@@ -4164,12 +4162,16 @@ class Field internal constructor(private val fieldPath: ModelFieldPath) : Select
      *
      * @return An [Field] representing the document ID.
      */
-    @JvmField val DOCUMENT_ID: Field = field(FieldPath.documentId())
+    @JvmField val DOCUMENT_ID: Field = Field(KEY_PATH)
+
+    @JvmField internal val UPDATE_TIME: Field = Field(UPDATE_TIME_PATH)
+
+    @JvmField internal val CREATE_TIME: Field = Field(CREATE_TIME_PATH)
   }
 
-  override fun getAlias(): String = fieldPath.canonicalString()
+  override val alias: String = fieldPath.canonicalString()
 
-  override fun getExpr(): Expr = this
+  override val expr: Expr = this
 
   override fun toProto(userDataReader: UserDataReader) = toProto()
 
@@ -4180,7 +4182,8 @@ class Field internal constructor(private val fieldPath: ModelFieldPath) : Select
     block@{ input: MutableDocument ->
       EvaluateResultValue(
         when (fieldPath) {
-          KEY_PATH -> encodeValue(DocumentReference.forPath(input.key.path, context.db))
+          KEY_PATH ->
+            encodeValue(DocumentReference.forPath(input.key.path, context.pipeline.firestore))
           CREATE_TIME_PATH -> encodeValue(input.createTime.timestamp)
           UPDATE_TIME_PATH -> encodeValue(input.version.timestamp)
           else -> input.getField(fieldPath) ?: return@block EvaluateResultUnset

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

@@ -18,6 +18,7 @@ import com.google.firebase.firestore.CollectionReference
 import com.google.firebase.firestore.FirebaseFirestore
 import com.google.firebase.firestore.UserDataReader
 import com.google.firebase.firestore.VectorValue
+import com.google.firebase.firestore.model.DocumentKey.KEY_FIELD_NAME
 import com.google.firebase.firestore.model.MutableDocument
 import com.google.firebase.firestore.model.ResourcePath
 import com.google.firebase.firestore.model.Values
@@ -292,9 +293,17 @@ internal constructor(
   private val fields: Array<out Selectable>,
   options: InternalOptions = InternalOptions.EMPTY
 ) : Stage<AddFieldsStage>("add_fields", options) {
+  init {
+    for (field in fields) {
+      val alias = field.alias
+      require(alias != Field.DOCUMENT_ID.alias, { "Alias ${Field.DOCUMENT_ID.alias} is reserved" })
+      require(alias != Field.CREATE_TIME.alias, { "Alias ${Field.CREATE_TIME.alias} is reserved" })
+      require(alias != Field.UPDATE_TIME.alias, { "Alias ${Field.UPDATE_TIME.alias} is reserved" })
+    }
+  }
   override fun self(options: InternalOptions) = AddFieldsStage(fields, options)
   override fun args(userDataReader: UserDataReader): Sequence<Value> =
-    sequenceOf(encodeValue(fields.associate { it.getAlias() to it.toProto(userDataReader) }))
+    sequenceOf(encodeValue(fields.associate { it.alias to it.toProto(userDataReader) }))
 }
 
 /**
@@ -368,8 +377,8 @@ internal constructor(
   fun withGroups(group: Selectable, vararg additionalGroups: Any) =
     AggregateStage(
       accumulators,
-      mapOf(group.getAlias() to group.getExpr())
-        .plus(additionalGroups.map(Selectable::toSelectable).associateBy(Selectable::getAlias))
+      mapOf(group.alias to group.expr)
+        .plus(additionalGroups.map(Selectable::toSelectable).associateBy(Selectable::alias))
     )
 
   override fun args(userDataReader: UserDataReader): Sequence<Value> =
@@ -539,7 +548,7 @@ internal constructor(private val offset: Int, options: InternalOptions = Interna
 }
 
 internal class SelectStage
-private constructor(private val fields: Array<out Selectable>, options: InternalOptions) :
+private constructor(internal val fields: Array<out Selectable>, options: InternalOptions) :
   Stage<SelectStage>("select", options) {
   companion object {
     @JvmStatic
@@ -555,14 +564,18 @@ private constructor(private val fields: Array<out Selectable>, options: Internal
   }
   override fun self(options: InternalOptions) = SelectStage(fields, options)
   override fun args(userDataReader: UserDataReader): Sequence<Value> =
-    sequenceOf(encodeValue(fields.associate { it.getAlias() to it.toProto(userDataReader) }))
+    sequenceOf(encodeValue(fields.associate { it.alias to it.toProto(userDataReader) }))
 }
 
 internal class SortStage
 internal constructor(
   private val orders: Array<out Ordering>,
   options: InternalOptions = InternalOptions.EMPTY
 ) : Stage<SortStage>("sort", options) {
+  companion object {
+    internal val BY_DOCUMENT_ID = SortStage(arrayOf(Field.DOCUMENT_ID.ascending()))
+  }
+
   override fun self(options: InternalOptions) = SortStage(orders, options)
   override fun args(userDataReader: UserDataReader): Sequence<Value> =
     orders.asSequence().map { it.toProto(userDataReader) }
@@ -603,6 +616,16 @@ internal constructor(
         .forEach { p -> emit(p.first) }
     }
   }
+
+  internal fun withStableOrdering(): SortStage {
+    val position = orders.indexOfFirst { (it.expr as? Field)?.alias == KEY_FIELD_NAME }
+    return if (position < 0) {
+      // Append the DocumentId to orders to make ordering stable.
+      SortStage(orders.asList().plus(Field.DOCUMENT_ID.ascending()).toTypedArray(), options)
+    } else {
+      this
+    }
+  }
 }
 
 internal class DistinctStage
@@ -612,14 +635,22 @@ internal constructor(
 ) : Stage<DistinctStage>("distinct", options) {
   override fun self(options: InternalOptions) = DistinctStage(groups, options)
   override fun args(userDataReader: UserDataReader): Sequence<Value> =
-    sequenceOf(encodeValue(groups.associate { it.getAlias() to it.toProto(userDataReader) }))
+    sequenceOf(encodeValue(groups.associate { it.alias to it.toProto(userDataReader) }))
 }
 
 internal class RemoveFieldsStage
 internal constructor(
   private val fields: Array<out Field>,
   options: InternalOptions = InternalOptions.EMPTY
 ) : Stage<RemoveFieldsStage>("remove_fields", options) {
+  init {
+    for (field in fields) {
+      val alias = field.alias
+      require(alias != Field.DOCUMENT_ID.alias, { "Alias ${Field.DOCUMENT_ID.alias} is required" })
+      require(alias != Field.CREATE_TIME.alias, { "Alias ${Field.CREATE_TIME.alias} is required" })
+      require(alias != Field.UPDATE_TIME.alias, { "Alias ${Field.UPDATE_TIME.alias} is required" })
+    }
+  }
   override fun self(options: InternalOptions) = RemoveFieldsStage(fields, options)
   override fun args(userDataReader: UserDataReader): Sequence<Value> =
     fields.asSequence().map(Field::toProto)
@@ -746,11 +777,11 @@ internal constructor(
      */
     @JvmStatic
     fun withField(arrayField: String, alias: String): UnnestStage =
-      UnnestStage(Expr.field(arrayField).alias(alias))
+      UnnestStage(Expr.Companion.field(arrayField).alias(alias))
   }
   override fun self(options: InternalOptions) = UnnestStage(selectable, options)
   override fun args(userDataReader: UserDataReader): Sequence<Value> =
-    sequenceOf(encodeValue(selectable.getAlias()), selectable.toProto(userDataReader))
+    sequenceOf(encodeValue(selectable.alias), selectable.toProto(userDataReader))
 
   /**
    * Adds index field to emitted documents