From c2ac39cd4636f493bb147b6ba5eb50fece8fc24d Mon Sep 17 00:00:00 2001
From: Ben Lesh <ben@benlesh.com>
Date: Tue, 8 Jan 2019 23:49:54 -0800
Subject: [PATCH] fix(types): support union type inference for merge operators

---
 spec-dtslint/observables/concat-spec.ts    | 11 +++++++++++
 spec-dtslint/operators/concatMap-spec.ts   |  4 ++++
 spec-dtslint/operators/concatMapTo-spec.ts |  5 +++++
 spec-dtslint/operators/mergeMap-spec.ts    |  4 ++++
 spec-dtslint/operators/mergeMapTo-spec.ts  |  5 +++++
 src/internal/observable/concat.ts          | 23 ++++++++++------------
 src/internal/operators/concat.ts           |  2 +-
 src/internal/operators/concatMap.ts        | 16 +++++++--------
 src/internal/operators/concatMapTo.ts      | 16 +++++++--------
 src/internal/operators/mergeAll.ts         |  8 +++-----
 src/internal/operators/mergeMap.ts         | 18 ++++++++---------
 src/internal/operators/mergeMapTo.ts       | 14 ++++++-------
 12 files changed, 75 insertions(+), 51 deletions(-)

diff --git a/spec-dtslint/observables/concat-spec.ts b/spec-dtslint/observables/concat-spec.ts
index c3397910eb..3cfdf8af28 100644
--- a/spec-dtslint/observables/concat-spec.ts
+++ b/spec-dtslint/observables/concat-spec.ts
@@ -52,3 +52,14 @@ it('should enforce types', () => {
   const o = concat(5); // $ExpectError
   const p = concat(of(5), 6); // $ExpectError
 });
+
+it('should support union types', () => {
+  const u = Math.random() > 0.5 ? of(123) : of('abc');
+  const o = concat(u, u, u); // $ExpectType Observable<string | number>
+});
+
+it('should support different union types', () => {
+  const u1 = Math.random() > 0.5 ? of(123) : of('abc');
+  const u2 = Math.random() > 0.5 ? of(true) : of([1, 2, 3]);
+  const o = concat(u1, u2); // $ExpectType Observable<string | number | boolean | number[]>
+});
diff --git a/spec-dtslint/operators/concatMap-spec.ts b/spec-dtslint/operators/concatMap-spec.ts
index 852ed5ecce..1724c484ca 100644
--- a/spec-dtslint/operators/concatMap-spec.ts
+++ b/spec-dtslint/operators/concatMap-spec.ts
@@ -29,6 +29,10 @@ it('should support an undefined resultSelector', () => {
   const o = of(1, 2, 3).pipe(concatMap(p => of(Boolean(p)), undefined)); // $ExpectType Observable<boolean>
 });
 
+it('should support union-type projections', () => {
+  const o = of(Math.random()).pipe(concatMap(n => n > 0.5 ? of('life') : of(42))); // $ExpectType Observable<string | number>
+});
+
 it('should enforce types', () => {
   const o = of(1, 2, 3).pipe(concatMap()); // $ExpectError
 });
diff --git a/spec-dtslint/operators/concatMapTo-spec.ts b/spec-dtslint/operators/concatMapTo-spec.ts
index 517f863e80..29b5b39613 100644
--- a/spec-dtslint/operators/concatMapTo-spec.ts
+++ b/spec-dtslint/operators/concatMapTo-spec.ts
@@ -37,6 +37,11 @@ it('should support an undefined resultSelector', () => {
   const o = of(1, 2, 3).pipe(concatMapTo(of('foo'), undefined)); // $ExpectType Observable<string>
 });
 
+it('should support union types', () => {
+  const s = Math.random() > 0.5 ? of(123) : of('abc');
+  const r = of(1, 2, 3).pipe(concatMapTo(s)); // $ExpectType<string | number>
+});
+
 it('should enforce types', () => {
   const o = of(1, 2, 3).pipe(concatMapTo()); // $ExpectError
 });
diff --git a/spec-dtslint/operators/mergeMap-spec.ts b/spec-dtslint/operators/mergeMap-spec.ts
index 056a81f7d7..5eeb83199d 100644
--- a/spec-dtslint/operators/mergeMap-spec.ts
+++ b/spec-dtslint/operators/mergeMap-spec.ts
@@ -41,6 +41,10 @@ it('should support a undefined resultSelector and concurrent parameter', () => {
   const o = of(1, 2, 3).pipe(mergeMap(p => of(Boolean(p)), undefined, 4)); // $ExpectType Observable<boolean>
 });
 
+it('should support union-type projections', () => {
+  const o = of(Math.random()).pipe(mergeMap(n => n > 0.5 ? of('life') : of(42))); // $ExpectType Observable<string | number>
+});
+
 it('should enforce types', () => {
   const o = of(1, 2, 3).pipe(mergeMap()); // $ExpectError
 });
diff --git a/spec-dtslint/operators/mergeMapTo-spec.ts b/spec-dtslint/operators/mergeMapTo-spec.ts
index eaf1f18804..b3a4863d7a 100644
--- a/spec-dtslint/operators/mergeMapTo-spec.ts
+++ b/spec-dtslint/operators/mergeMapTo-spec.ts
@@ -41,6 +41,11 @@ it('should support a resultSelector and concurrent parameter', () => {
   const o = of(1, 2, 3).pipe(mergeMapTo(of('foo'), (a, b) => b, 4)); // $ExpectType Observable<string>
 });
 
+it('should support union types', () => {
+  const s = Math.random() > 0.5 ? of(123) : of('abc');
+  const r = of(1, 2, 3).pipe(mergeMapTo(s)); // $ExpectType<string | number>
+});
+
 it('should enforce types', () => {
   const o = of(1, 2, 3).pipe(mergeMapTo()); // $ExpectError
 });
diff --git a/src/internal/observable/concat.ts b/src/internal/observable/concat.ts
index d98725dfa8..3fd86f9377 100644
--- a/src/internal/observable/concat.ts
+++ b/src/internal/observable/concat.ts
@@ -1,19 +1,19 @@
 import { Observable } from '../Observable';
-import { ObservableInput, SchedulerLike } from '../types';
+import { ObservableInput, SchedulerLike, ObservedValueOf } from '../types';
 import { isScheduler } from '../util/isScheduler';
 import { of } from './of';
 import { from } from './from';
 import { concatAll } from '../operators/concatAll';
 
 /* tslint:disable:max-line-length */
-export function concat<T>(v1: ObservableInput<T>, scheduler?: SchedulerLike): Observable<T>;
-export function concat<T, T2>(v1: ObservableInput<T>, v2: ObservableInput<T2>, scheduler?: SchedulerLike): Observable<T | T2>;
-export function concat<T, T2, T3>(v1: ObservableInput<T>, v2: ObservableInput<T2>, v3: ObservableInput<T3>, scheduler?: SchedulerLike): Observable<T | T2 | T3>;
-export function concat<T, T2, T3, T4>(v1: ObservableInput<T>, v2: ObservableInput<T2>, v3: ObservableInput<T3>, v4: ObservableInput<T4>, scheduler?: SchedulerLike): Observable<T | T2 | T3 | T4>;
-export function concat<T, T2, T3, T4, T5>(v1: ObservableInput<T>, v2: ObservableInput<T2>, v3: ObservableInput<T3>, v4: ObservableInput<T4>, v5: ObservableInput<T5>, scheduler?: SchedulerLike): Observable<T | T2 | T3 | T4 | T5>;
-export function concat<T, T2, T3, T4, T5, T6>(v1: ObservableInput<T>, v2: ObservableInput<T2>, v3: ObservableInput<T3>, v4: ObservableInput<T4>, v5: ObservableInput<T5>, v6: ObservableInput<T6>, scheduler?: SchedulerLike): Observable<T | T2 | T3 | T4 | T5 | T6>;
-export function concat<T>(...observables: (ObservableInput<T> | SchedulerLike)[]): Observable<T>;
-export function concat<T, R>(...observables: (ObservableInput<any> | SchedulerLike)[]): Observable<R>;
+export function concat<O1 extends ObservableInput<any>>(v1: O1, scheduler?: SchedulerLike): Observable<ObservedValueOf<O1>>;
+export function concat<O1 extends ObservableInput<any>, O2 extends ObservableInput<any>>(v1: O1, v2: O2, scheduler?: SchedulerLike): Observable<ObservedValueOf<O1> | ObservedValueOf<O2>>;
+export function concat<O1 extends ObservableInput<any>, O2 extends ObservableInput<any>, O3 extends ObservableInput<any>>(v1: O1, v2: O2, v3: O3, scheduler?: SchedulerLike): Observable<ObservedValueOf<O1> | ObservedValueOf<O2> | ObservedValueOf<O3>>;
+export function concat<O1 extends ObservableInput<any>, O2 extends ObservableInput<any>, O3 extends ObservableInput<any>, O4 extends ObservableInput<any>>(v1: O1, v2: O2, v3: O3, v4: O4, scheduler?: SchedulerLike): Observable<ObservedValueOf<O1> | ObservedValueOf<O2> | ObservedValueOf<O3> | ObservedValueOf<O4>>;
+export function concat<O1 extends ObservableInput<any>, O2 extends ObservableInput<any>, O3 extends ObservableInput<any>, O4 extends ObservableInput<any>, O5 extends ObservableInput<any>>(v1: O1, v2: O2, v3: O3, v4: O4, v5: O5, scheduler?: SchedulerLike): Observable<ObservedValueOf<O1> | ObservedValueOf<O2> | ObservedValueOf<O3> | ObservedValueOf<O4> | ObservedValueOf<O5>>;
+export function concat<O1 extends ObservableInput<any>, O2 extends ObservableInput<any>, O3 extends ObservableInput<any>, O4 extends ObservableInput<any>, O5 extends ObservableInput<any>, O6 extends ObservableInput<any>>(v1: O1, v2: O2, v3: O3, v4: O4, v5: O5, v6: O6, scheduler?: SchedulerLike): Observable<ObservedValueOf<O1> | ObservedValueOf<O2> | ObservedValueOf<O3> | ObservedValueOf<O4> | ObservedValueOf<O5> | ObservedValueOf<O6>>;
+export function concat<O extends ObservableInput<any>>(...observables: (O | SchedulerLike)[]): Observable<ObservedValueOf<O>>;
+export function concat<O1 extends ObservableInput<any>, R>(...observables: (ObservableInput<any> | SchedulerLike)[]): Observable<R>;
 /* tslint:enable:max-line-length */
 /**
  * Creates an output Observable which sequentially emits all values from given
@@ -113,9 +113,6 @@ export function concat<T, R>(...observables: (ObservableInput<any> | SchedulerLi
  * @name concat
  * @owner Observable
  */
-export function concat<T, R>(...observables: Array<ObservableInput<any> | SchedulerLike>): Observable<R> {
-  if (observables.length === 1 || (observables.length === 2 && isScheduler(observables[1]))) {
-    return from(<any>observables[0]);
-  }
+export function concat<O extends ObservableInput<any>, R>(...observables: Array<O | SchedulerLike>): Observable<ObservedValueOf<O> | R> {
   return concatAll<R>()(of(...observables));
 }
diff --git a/src/internal/operators/concat.ts b/src/internal/operators/concat.ts
index 0c1179c69c..aee64666f9 100644
--- a/src/internal/operators/concat.ts
+++ b/src/internal/operators/concat.ts
@@ -25,5 +25,5 @@ export function concat<T, R>(...observables: Array<ObservableInput<any> | Schedu
  * @deprecated Deprecated in favor of static {@link concat}.
  */
 export function concat<T, R>(...observables: Array<ObservableInput<any> | SchedulerLike>): OperatorFunction<T, R> {
-  return (source: Observable<T>) => source.lift.call(concatStatic<T, R>(source, ...observables));
+  return (source: Observable<T>) => source.lift.call(concatStatic(source, ...observables));
 }
diff --git a/src/internal/operators/concatMap.ts b/src/internal/operators/concatMap.ts
index cc397052aa..c40afb3889 100644
--- a/src/internal/operators/concatMap.ts
+++ b/src/internal/operators/concatMap.ts
@@ -1,12 +1,12 @@
 import { mergeMap } from './mergeMap';
-import { ObservableInput, OperatorFunction } from '../types';
+import { ObservableInput, OperatorFunction, ObservedValueOf } from '../types';
 
 /* tslint:disable:max-line-length */
-export function concatMap<T, R>(project: (value: T, index: number) =>  ObservableInput<R>): OperatorFunction<T, R>;
+export function concatMap<T, O extends ObservableInput<any>>(project: (value: T, index: number) =>  O): OperatorFunction<T, ObservedValueOf<O>>;
 /** @deprecated resultSelector no longer supported, use inner map instead */
-export function concatMap<T, R>(project: (value: T, index: number) => ObservableInput<R>, resultSelector: undefined): OperatorFunction<T, R>;
+export function concatMap<T, O extends ObservableInput<any>>(project: (value: T, index: number) => O, resultSelector: undefined): OperatorFunction<T, ObservedValueOf<O>>;
 /** @deprecated resultSelector no longer supported, use inner map instead */
-export function concatMap<T, I, R>(project: (value: T, index: number) =>  ObservableInput<I>, resultSelector: (outerValue: T, innerValue: I, outerIndex: number, innerIndex: number) => R): OperatorFunction<T, R>;
+export function concatMap<T, R, O extends ObservableInput<any>>(project: (value: T, index: number) =>  O, resultSelector: (outerValue: T, innerValue: ObservedValueOf<O>, outerIndex: number, innerIndex: number) => R): OperatorFunction<T, R>;
 /* tslint:enable:max-line-length */
 
 /**
@@ -66,9 +66,9 @@ export function concatMap<T, I, R>(project: (value: T, index: number) =>  Observ
  * @method concatMap
  * @owner Observable
  */
-export function concatMap<T, I, R>(
-  project: (value: T, index: number) =>  ObservableInput<I>,
-  resultSelector?: (outerValue: T, innerValue: I, outerIndex: number, innerIndex: number) => R
-): OperatorFunction<T, I|R> {
+export function concatMap<T, R, O extends ObservableInput<any>>(
+  project: (value: T, index: number) => O,
+  resultSelector?: (outerValue: T, innerValue: ObservedValueOf<O>, outerIndex: number, innerIndex: number) => R
+): OperatorFunction<T, ObservedValueOf<O>|R> {
   return mergeMap(project, resultSelector, 1);
 }
diff --git a/src/internal/operators/concatMapTo.ts b/src/internal/operators/concatMapTo.ts
index 487f91b22b..627d4ec1cc 100644
--- a/src/internal/operators/concatMapTo.ts
+++ b/src/internal/operators/concatMapTo.ts
@@ -1,12 +1,12 @@
 import { concatMap } from './concatMap';
-import { ObservableInput, OperatorFunction } from '../types';
+import { ObservableInput, OperatorFunction, ObservedValueOf } from '../types';
 
 /* tslint:disable:max-line-length */
-export function concatMapTo<T>(observable: ObservableInput<T>): OperatorFunction<any, T>;
+export function concatMapTo<T, O extends ObservableInput<any>>(observable: O): OperatorFunction<T, ObservedValueOf<O>>;
 /** @deprecated */
-export function concatMapTo<T>(observable: ObservableInput<T>, resultSelector: undefined): OperatorFunction<any, T>;
+export function concatMapTo<T, O extends ObservableInput<any>>(observable: O, resultSelector: undefined): OperatorFunction<T, ObservedValueOf<O>>;
 /** @deprecated */
-export function concatMapTo<T, I, R>(observable: ObservableInput<I>, resultSelector: (outerValue: T, innerValue: I, outerIndex: number, innerIndex: number) => R): OperatorFunction<T, R>;
+export function concatMapTo<T, R, O extends ObservableInput<any>>(observable: O, resultSelector: (outerValue: T, innerValue: ObservedValueOf<O>, outerIndex: number, innerIndex: number) => R): OperatorFunction<T, R>;
 /* tslint:enable:max-line-length */
 
 /**
@@ -62,9 +62,9 @@ export function concatMapTo<T, I, R>(observable: ObservableInput<I>, resultSelec
  * @method concatMapTo
  * @owner Observable
  */
-export function concatMapTo<T, I, R>(
-  innerObservable: ObservableInput<I>,
-  resultSelector?: (outerValue: T, innerValue: I, outerIndex: number, innerIndex: number) => R
-): OperatorFunction<T, R> {
+export function concatMapTo<T, R, O extends ObservableInput<any>>(
+  innerObservable: O,
+  resultSelector?: (outerValue: T, innerValue: ObservedValueOf<O>, outerIndex: number, innerIndex: number) => R
+): OperatorFunction<T, ObservedValueOf<O>|R> {
   return concatMap(() => innerObservable, resultSelector);
 }
diff --git a/src/internal/operators/mergeAll.ts b/src/internal/operators/mergeAll.ts
index bf8ea5f7bd..4aab502652 100644
--- a/src/internal/operators/mergeAll.ts
+++ b/src/internal/operators/mergeAll.ts
@@ -1,9 +1,7 @@
 
 import { mergeMap } from './mergeMap';
 import { identity } from '../util/identity';
-import { MonoTypeOperatorFunction, OperatorFunction, ObservableInput } from '../types';
-
-export function mergeAll<T>(concurrent?: number): OperatorFunction<ObservableInput<T>, T>;
+import { OperatorFunction, ObservableInput } from '../types';
 
 /**
  * Converts a higher-order Observable into a first-order Observable which
@@ -57,6 +55,6 @@ export function mergeAll<T>(concurrent?: number): OperatorFunction<ObservableInp
  * @method mergeAll
  * @owner Observable
  */
-export function mergeAll<T>(concurrent: number = Number.POSITIVE_INFINITY): MonoTypeOperatorFunction<T> {
-  return mergeMap<T, T>(identity as (value: T, index: number) => ObservableInput<T>, concurrent);
+export function mergeAll<T>(concurrent: number = Number.POSITIVE_INFINITY): OperatorFunction<ObservableInput<T>, T> {
+  return mergeMap(identity, concurrent);
 }
diff --git a/src/internal/operators/mergeMap.ts b/src/internal/operators/mergeMap.ts
index d843dcdc97..a5b963bcc2 100644
--- a/src/internal/operators/mergeMap.ts
+++ b/src/internal/operators/mergeMap.ts
@@ -5,16 +5,16 @@ import { Subscription } from '../Subscription';
 import { subscribeToResult } from '../util/subscribeToResult';
 import { OuterSubscriber } from '../OuterSubscriber';
 import { InnerSubscriber } from '../InnerSubscriber';
-import { ObservableInput, OperatorFunction } from '../types';
+import { ObservableInput, OperatorFunction, ObservedValueOf } from '../types';
 import { map } from './map';
 import { from } from '../observable/from';
 
 /* tslint:disable:max-line-length */
-export function mergeMap<T, R>(project: (value: T, index: number) => ObservableInput<R>, concurrent?: number): OperatorFunction<T, R>;
+export function mergeMap<T, O extends ObservableInput<any>>(project: (value: T, index: number) => O, concurrent?: number): OperatorFunction<T, ObservedValueOf<O>>;
 /** @deprecated resultSelector no longer supported, use inner map instead */
-export function mergeMap<T, R>(project: (value: T, index: number) => ObservableInput<R>, resultSelector: undefined, concurrent?: number): OperatorFunction<T, R>;
+export function mergeMap<T, O extends ObservableInput<any>>(project: (value: T, index: number) => O, resultSelector: undefined, concurrent?: number): OperatorFunction<T, ObservedValueOf<O>>;
 /** @deprecated resultSelector no longer supported, use inner map instead */
-export function mergeMap<T, I, R>(project: (value: T, index: number) => ObservableInput<I>, resultSelector: (outerValue: T, innerValue: I, outerIndex: number, innerIndex: number) => R, concurrent?: number): OperatorFunction<T, R>;
+export function mergeMap<T, R, O extends ObservableInput<any>>(project: (value: T, index: number) => O, resultSelector: (outerValue: T, innerValue: ObservedValueOf<O>, outerIndex: number, innerIndex: number) => R, concurrent?: number): OperatorFunction<T, R>;
 /* tslint:enable:max-line-length */
 
 /**
@@ -70,16 +70,16 @@ export function mergeMap<T, I, R>(project: (value: T, index: number) => Observab
  * @method mergeMap
  * @owner Observable
  */
-export function mergeMap<T, I, R>(
-  project: (value: T, index: number) => ObservableInput<I>,
-  resultSelector?: ((outerValue: T, innerValue: I, outerIndex: number, innerIndex: number) => R) | number,
+export function mergeMap<T, R, O extends ObservableInput<any>>(
+  project: (value: T, index: number) => O,
+  resultSelector?: ((outerValue: T, innerValue: ObservedValueOf<O>, outerIndex: number, innerIndex: number) => R) | number,
   concurrent: number = Number.POSITIVE_INFINITY
-): OperatorFunction<T, I|R> {
+): OperatorFunction<T, ObservedValueOf<O>|R> {
   if (typeof resultSelector === 'function') {
     // DEPRECATED PATH
     return (source: Observable<T>) => source.pipe(
       mergeMap((a, i) => from(project(a, i)).pipe(
-        map((b, ii) => resultSelector(a, b, i, ii)),
+        map((b: any, ii: number) => resultSelector(a, b, i, ii)),
       ), concurrent)
     );
   } else if (typeof resultSelector === 'number') {
diff --git a/src/internal/operators/mergeMapTo.ts b/src/internal/operators/mergeMapTo.ts
index 57186e253a..65e34b0ef4 100644
--- a/src/internal/operators/mergeMapTo.ts
+++ b/src/internal/operators/mergeMapTo.ts
@@ -1,12 +1,12 @@
 import { Observable } from '../Observable';
-import { OperatorFunction } from '../../internal/types';
+import { OperatorFunction, ObservedValueOf } from '../../internal/types';
 import { mergeMap } from './mergeMap';
 import { ObservableInput } from '../types';
 
 /* tslint:disable:max-line-length */
-export function mergeMapTo<T>(innerObservable: ObservableInput<T>, concurrent?: number): OperatorFunction<any, T>;
+export function mergeMapTo<T, O extends ObservableInput<any>>(innerObservable: O, concurrent?: number): OperatorFunction<any, ObservedValueOf<O>>;
 /** @deprecated */
-export function mergeMapTo<T, I, R>(innerObservable: ObservableInput<I>, resultSelector: (outerValue: T, innerValue: I, outerIndex: number, innerIndex: number) => R, concurrent?: number): OperatorFunction<T, R>;
+export function mergeMapTo<T, R, O extends ObservableInput<any>>(innerObservable: O, resultSelector: (outerValue: T, innerValue: ObservedValueOf<O>, outerIndex: number, innerIndex: number) => R, concurrent?: number): OperatorFunction<T, R>;
 /* tslint:enable:max-line-length */
 
 /**
@@ -46,11 +46,11 @@ export function mergeMapTo<T, I, R>(innerObservable: ObservableInput<I>, resultS
  * @method mergeMapTo
  * @owner Observable
  */
-export function mergeMapTo<T, I, R>(
-  innerObservable: ObservableInput<I>,
-  resultSelector?: ((outerValue: T, innerValue: I, outerIndex: number, innerIndex: number) => R) | number,
+export function mergeMapTo<T, R, O extends ObservableInput<any>>(
+  innerObservable: O,
+  resultSelector?: ((outerValue: T, innerValue: ObservedValueOf<O>, outerIndex: number, innerIndex: number) => R) | number,
   concurrent: number = Number.POSITIVE_INFINITY
-): OperatorFunction<T, I|R> {
+): OperatorFunction<T, ObservedValueOf<O>|R> {
   if (typeof resultSelector === 'function') {
     return mergeMap(() => innerObservable, resultSelector, concurrent);
   }