selectMany now accepts arrays in projection.

BSick7 · Apr 21, 2014 · 91338ca · 91338ca
1 parent beb9d61
commit 91338ca
Show file tree

Hide file tree

Showing 5 changed files with 32 additions and 7 deletions.
diff --git a/dist/ex.d.ts b/dist/ex.d.ts
@@ -25,6 +25,7 @@ declare module exjs {
         reverse(): IEnumerable<T>;
         select<TResult>(selector: (t: T, index?: number) => TResult): IEnumerable<TResult>;
         selectMany<TResult>(selector: (t: T) => IEnumerable<TResult>): IEnumerable<TResult>;
+        selectMany<TResult>(selector: (t: T) => TResult[]): IEnumerable<TResult>;
         skip(count: number): IEnumerable<T>;
         skipWhile(predicate: (t: T, index?: number) => boolean): IEnumerable<T>;
         sum(selector?: (t: T) => number): number;
@@ -75,6 +76,7 @@ declare module exjs {
         public reverse(): IEnumerable<T>;
         public select<TResult>(selector: (t: T, index?: number) => TResult): IEnumerable<TResult>;
         public selectMany<TResult>(selector: (t: T) => IEnumerable<TResult>): IEnumerable<TResult>;
+        public selectMany<TResult>(selector: (t: T) => TResult[]): IEnumerable<TResult>;
         public skip(count: number): IEnumerable<T>;
         public skipWhile(predicate: (t: T, index?: number) => boolean): IEnumerable<T>;
         public sum(selector?: (t: T) => number): number;

diff --git a/dist/ex.js b/dist/ex.js
@@ -156,9 +156,11 @@ var exjs;
         Enumerable.prototype.select = function (selector) {
             throw new Error("Not implemented");
         };
+
         Enumerable.prototype.selectMany = function (selector) {
             throw new Error("Not implemented");
         };
+
         Enumerable.prototype.skip = function (count) {
             throw new Error("Not implemented");
         };
@@ -695,7 +697,6 @@ var exjs;
 
     function selectManyEnumerator(prev, selector) {
         var t;
-        var j = 0;
         var active;
         var e = {
             current: undefined,
@@ -706,7 +707,11 @@ var exjs;
                 while (!active || !active.moveNext()) {
                     if (!t.moveNext())
                         return false;
-                    active = selector(t.current).getEnumerator();
+                    var selected = selector(t.current);
+                    var en = selected;
+                    if (en instanceof Array)
+                        en = en.en();
+                    active = en.getEnumerator();
                 }
                 e.current = active.current;
                 return true;

diff --git a/src/enumerable.ts b/src/enumerable.ts
@@ -25,6 +25,7 @@ module exjs {
         reverse(): IEnumerable<T>;
         select<TResult>(selector: (t: T, index?: number) => TResult): IEnumerable<TResult>;
         selectMany<TResult>(selector: (t: T) => IEnumerable<TResult>): IEnumerable<TResult>;
+        selectMany<TResult>(selector: (t: T) => TResult[]): IEnumerable<TResult>;
         skip(count: number): IEnumerable<T>;
         skipWhile(predicate: (t: T, index?: number) => boolean): IEnumerable<T>;
         sum(selector?: (t: T) => number): number;
@@ -190,7 +191,11 @@ module exjs {
         orderByDescending<TKey>(keySelector: (t: T) => TKey, comparer?: (f: TKey, s: TKey) => number): IOrderedEnumerable<T> { throw new Error("Not implemented"); }
         reverse(): IEnumerable<T> { throw new Error("Not implemented"); }
         select<TResult>(selector: (t: T, index?: number) => TResult): IEnumerable<TResult> { throw new Error("Not implemented"); }
-        selectMany<TResult>(selector: (t: T) => IEnumerable<TResult>): IEnumerable<TResult> { throw new Error("Not implemented"); }
+
+        selectMany<TResult>(selector: (t: T) => IEnumerable<TResult>): IEnumerable<TResult>;
+        selectMany<TResult>(selector: (t: T) => TResult[]): IEnumerable<TResult>;
+        selectMany<TResult>(selector: (t: T) => any): IEnumerable<TResult> { throw new Error("Not implemented"); }
+
         skip(count: number): IEnumerable<T> { throw new Error("Not implemented"); }
         skipWhile(predicate: (t: T, index?: number) => boolean): IEnumerable<T> { throw new Error("Not implemented"); }
         sum(selector?: (t: T) => number): number {

diff --git a/src/select.ts b/src/select.ts
@@ -1,4 +1,5 @@
 /// <reference path="enumerable.ts" />
+/// <reference path="array.ts" />
 
 module exjs {
     function selectEnumerator<T, TResult>(prev: IEnumerable<T>, selector: (t: T, index?: number) => TResult): IEnumerator<TResult> {
@@ -17,9 +18,8 @@ module exjs {
         return e;
     }
 
-    function selectManyEnumerator<T, TResult>(prev: IEnumerable<T>, selector: (t: T) => IEnumerable<TResult>): IEnumerator<TResult> {
+    function selectManyEnumerator<T, TResult>(prev: IEnumerable<T>, selector: (t: T) => any): IEnumerator<TResult> {
         var t: IEnumerator<T>;
-        var j = 0;
         var active: IEnumerator<TResult>;
         var e = {
             current: undefined,
@@ -29,7 +29,11 @@ module exjs {
                 while (!active || !active.moveNext()) {
                     if (!t.moveNext())
                         return false;
-                    active = selector(t.current).getEnumerator();
+                    var selected = selector(t.current);
+                    var en:IEnumerable<TResult> = selected;
+                    if (en instanceof Array)
+                        en = (<TResult[]><any>en).en();
+                    active = en.getEnumerator();
                 }
                 e.current = active.current;
                 return true;
@@ -43,7 +47,7 @@ module exjs {
         e.getEnumerator = () => selectEnumerator(<IEnumerable<T>>this, selector);
         return e;
     };
-    Enumerable.prototype.selectMany = function<T,TResult>(selector: (t: T) => IEnumerable<TResult>): IEnumerable<TResult> {
+    (<any>Enumerable.prototype).selectMany = function<T,TResult>(selector: (t: T) => any): IEnumerable<TResult> {
         var e = new Enumerable<TResult>();
         e.getEnumerator = () => selectManyEnumerator<T, TResult>(<IEnumerable<T>>this, selector);
         return e;

diff --git a/test/enum.ts b/test/enum.ts
@@ -350,6 +350,15 @@ test("selectMany", () => {
     strictEqual(res2[3], 4);
     strictEqual(res2[4], 5);
     strictEqual(res2[5], 6);
+
+    res2  = arr2.en().selectMany(t => t.i).toArray();
+    strictEqual(res2.length, 6);
+    strictEqual(res2[0], 1);
+    strictEqual(res2[1], 2);
+    strictEqual(res2[2], 3);
+    strictEqual(res2[3], 4);
+    strictEqual(res2[4], 5);
+    strictEqual(res2[5], 6);
 });
 
 test("skip", () => {