From 783c23cbfbf60597cec1c42643b386a16b587206 Mon Sep 17 00:00:00 2001
From: John-David Dalton <john.david.dalton@gmail.com>
Date: Sat, 14 Nov 2015 14:15:11 -0800
Subject: [PATCH] Make `_.at` support shortcut fusion in a chain sequence when
 operating on an array and selecting a single index. [closes #1636]

---
 lodash.js    |  7 +++++++
 test/test.js | 59 +++++++++++++++++++++++++++++++++++++++++++++++++---
 2 files changed, 63 insertions(+), 3 deletions(-)

diff --git a/lodash.js b/lodash.js
index 832182d7f..e13b7e8cc 100644
--- a/lodash.js
+++ b/lodash.js
@@ -13893,6 +13893,13 @@
       };
     });
 
+    LazyWrapper.prototype.at = rest(function(paths) {
+      paths = baseFlatten(paths);
+      return (paths.length == 1 && isIndex(paths[0]))
+        ? this.slice(paths[0], paths[0] + 1)
+        : new LazyWrapper(this);
+    });
+
     LazyWrapper.prototype.compact = function() {
       return this.filter(identity);
     };
diff --git a/test/test.js b/test/test.js
index 2fc7e1d68..6456bffaf 100644
--- a/test/test.js
+++ b/test/test.js
@@ -1154,7 +1154,8 @@
 
   (function() {
     var args = arguments,
-        array = ['a', 'b', 'c'];
+        array = ['a', 'b', 'c'],
+        object = { 'a': [{ 'b': { 'c': 3 } }, 4] };
 
     QUnit.test('should return the elements corresponding to the specified keys', function(assert) {
       assert.expect(1);
@@ -1232,8 +1233,8 @@
     QUnit.test('should work with an object for `object`', function(assert) {
       assert.expect(1);
 
-      var actual = _.at({ 'a': 1, 'b': 2, 'c': 3 }, ['c', 'a']);
-      assert.deepEqual(actual, [3, 1]);
+      var actual = _.at(object, ['a[0].b.c', 'a[1]']);
+      assert.deepEqual(actual, [3, 4]);
     });
 
     QUnit.test('should pluck inherited property values', function(assert) {
@@ -1245,6 +1246,58 @@
       var actual = _.at(new Foo, 'b');
       assert.deepEqual(actual, [2]);
     });
+
+    QUnit.test('should work in a lazy chain sequence', function(assert) {
+      assert.expect(4);
+
+      if (!isNpm) {
+        var largeArray = lodashStable.range(LARGE_ARRAY_SIZE),
+            smallArray = array;
+
+        lodashStable.each([[2], [2, 1]], function(paths) {
+          lodashStable.times(2, function(index) {
+            var array = index ? largeArray : smallArray,
+                wrapped = _(array).map(identity).at(paths);
+
+            assert.deepEqual(wrapped.value(), _.at(_.map(array, identity), paths));
+          });
+        });
+      }
+      else {
+        skipTest(assert, 4);
+      }
+    });
+
+    QUnit.test('should support shortcut fusion', function(assert) {
+      assert.expect(2);
+
+      if (!isNpm) {
+        var count = 0,
+            array = lodashStable.range(LARGE_ARRAY_SIZE),
+            iteratee = function(value) { count++; return square(value); },
+            actual = _(array).map(iteratee).at(LARGE_ARRAY_SIZE - 1).value();
+
+        assert.strictEqual(count, 1);
+        assert.deepEqual(actual, [square(LARGE_ARRAY_SIZE -1)]);
+      }
+      else {
+        skipTest(assert, 2);
+      }
+    });
+
+    QUnit.test('work with an object for `object` when chaining', function(assert) {
+      assert.expect(1);
+
+      if (!isNpm) {
+        var paths = ['a[0].b.c', 'a[1]'],
+            wrapped = _(object).map(identity).at(paths);
+
+        assert.deepEqual(wrapped.value(), _.at(_.map(object, identity), paths));
+      }
+      else {
+        skipTest(assert);
+      }
+    });
   }(1, 2, 3));
 
   /*--------------------------------------------------------------------------*/