ui: EventSet: Add some optimisations

Change-Id: Ib5b99513c65be586d2972f3527cd0d747415bf52
diff --git a/ui/src/base/array_utils.ts b/ui/src/base/array_utils.ts
index 7236a43..fb845e0 100644
--- a/ui/src/base/array_utils.ts
+++ b/ui/src/base/array_utils.ts
@@ -40,3 +40,8 @@
   }
   return true;
 }
+
+export function isArrayOf<P, Q>(
+    predicate: (x: P|Q) => x is P, xs: (P|Q)[]): xs is P[] {
+  return xs.every(predicate);
+}
diff --git a/ui/src/common/event_set.ts b/ui/src/common/event_set.ts
index 2eb65e5..4b365de 100644
--- a/ui/src/common/event_set.ts
+++ b/ui/src/common/event_set.ts
@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+import {arrayEquals, isArrayOf} from '../base/array_utils';
 import {intersect} from '../base/set_utils';
 
 // Contents:
@@ -104,15 +105,15 @@
 }
 
 interface UnionEventSet<T extends KeySet> extends EventSet<T> {
-  readonly parents: EventSet<KeySet>[];
+  readonly parents: EventSet<T>[];
   readonly isUnion: true;
-  create(...selections: EventSet<KeySet>[]): UnionEventSet<T>;
+  create(...events: EventSet<KeySet>[]): UnionEventSet<T>;
 }
 
 interface IntersectionEventSet<T extends KeySet> extends EventSet<T> {
-  readonly parents: EventSet<KeySet>[];
+  readonly parents: EventSet<T>[];
   readonly isIntersection: true;
-  create(...selections: EventSet<KeySet>[]): IntersectionEventSet<T>;
+  create(...events: EventSet<KeySet>[]): IntersectionEventSet<T>;
 }
 
 interface FilterEventSet<T extends KeySet> extends EventSet<T> {
@@ -510,7 +511,7 @@
           id: e.id,
         };
         for (const [k, v] of Object.entries(keys)) {
-          // While the static typing prevents folks from doing hitting
+          // While the static typing prevents folks from hitting
           // this in the common case people can still on purpose pass
           // keysets and lie about the types.
           result[k] = (e as UntypedEvent)[k] ?? getKeyDefault(k, v);
@@ -555,9 +556,96 @@
 // A critical pre-condition of this function is that EventSets are
 // immutable - this allows us to reuse parts of the input event set tree
 // in the output.
-export function optimise<T extends KeySet>(events: EventSet<T>): EventSet<T> {
+export function optimise<T extends KeySet>(eventSet: EventSet<T>): EventSet<T> {
+  // Empty EventSet can't be futher optimised.
+  if (isEmptyEventSet(eventSet)) {
+    return eventSet;
+  }
+
+  if (isConcreteEventSet(eventSet)) {
+    // A concrete events with zero elements is the empty events.
+    if (eventSet.events.length === 0) {
+      return new EmptyEventSet(eventSet.keys);
+    }
+    // ...but otherwise can not be optimised further.
+    return eventSet;
+  }
+
+  if (isUnionEventSet(eventSet)) {
+    const keys = eventSet.keys;
+
+    let newParents: EventSet<T>[] = eventSet.parents.slice();
+
+    // Empty sets don't contribute to the union.
+    newParents = newParents.filter((p) => !isEmptyEventSet(p));
+
+    // union([]) is empty.
+    if (newParents.length === 0) {
+      return new EmptyEventSet(keys);
+    }
+
+    if (newParents.length === 1) {
+      return newParents[0];
+    }
+
+    // The union of concrete EventSets is a concrete EventSets with all
+    // the events in.
+    if (isArrayOf<ConcreteEventSet<T>, EventSet<T>>(
+            isConcreteEventSet, newParents)) {
+      const seen = new Set<string>();
+      const events = [];
+      for (const p of newParents) {
+        for (const e of p.events) {
+          if (!seen.has(e.id)) {
+            events.push(e);
+            seen.add(e.id);
+          }
+        }
+      }
+      return ConcreteEventSet.from(eventSet.keys, events);
+    }
+
+    if (arrayEquals(newParents, eventSet.parents)) {
+      return eventSet;
+    } else {
+      return eventSet.create(...newParents);
+    }
+  }
+
+  if (isIntersectionEventSet(eventSet)) {
+    // For any x: intersect([x, 0]) is 0
+    for (const parent of eventSet.parents) {
+      if (isEmptyEventSet(parent)) {
+        return parent;
+      }
+    }
+    return eventSet;
+  }
+
+  if (isFilterEventSet(eventSet)) {
+    const parent = eventSet.parent;
+
+    if (isEmptyEventSet(parent)) {
+      return parent;
+    }
+
+    return eventSet;
+  }
+
+  if (isSortEventSet(eventSet)) {
+    const parent = eventSet.parent;
+
+    if (isEmptyEventSet(parent)) {
+      return parent;
+    }
+
+    return eventSet;
+  }
+
   // TODO(hjd): Re-add the optimisations from the prototype.
-  return events;
+  // TODO(hjd): Union([a, a]) === a but maybe not worth optimising.
+
+  return eventSet;
 }
 
 // EXPR ===============================================================
diff --git a/ui/src/common/event_set_unittest.ts b/ui/src/common/event_set_unittest.ts
index 6fb901d..2242e87 100644
--- a/ui/src/common/event_set_unittest.ts
+++ b/ui/src/common/event_set_unittest.ts
@@ -98,8 +98,7 @@
       const events: EventSet<KeySet> = EmptyEventSet.get();
       const filtered = await events.filter(c(true));
 
-      // TODO(hjd): This should be true after the optimisation() lands.
-      // expect(filtered).toBe(events);
+      expect(filtered).toBe(events);
       expect(await filtered.isEmpty()).toEqual(true);
       expect(await filtered.count()).toEqual(0);
     });
@@ -111,8 +110,7 @@
         expression: c(0),
       });
 
-      // TODO(hjd): This should be true after the optimisations land.
-      // expect(sorted).toBe(events);
+      expect(sorted).toBe(events);
       expect(await sorted.isEmpty()).toEqual(true);
       expect(await sorted.count()).toEqual(0);
     });