PowerShell · iSazonov · Sep 19, 2018 · Jun 27, 2018 · Jun 27, 2018 · Jun 27, 2018
diff --git a/src/System.Management.Automation/engine/SessionState.cs b/src/System.Management.Automation/engine/SessionState.cs
@@ -65,8 +65,8 @@ internal SessionStateInternal(SessionStateInternal parent, bool linkToGlobal, Ex
             _workingLocationStack = new Dictionary<String, Stack<PathInfo>>(StringComparer.OrdinalIgnoreCase);
 
             // Conservative choice to limit the Set-Location history in order to limit memory impact in case of a regression.
-            const int locationHistoryLimit = 20;
-            _SetLocationHistory = new BoundedStack<PathInfo>(locationHistoryLimit);
+            const uint locationHistoryLimit = 20;
+            _setLocationHistory = new HistoryStack<PathInfo>(locationHistoryLimit);
 
             GlobalScope = new SessionStateScope(null);
             ModuleScope = GlobalScope;

diff --git a/src/System.Management.Automation/engine/SessionStateLocationAPIs.cs b/src/System.Management.Automation/engine/SessionStateLocationAPIs.cs
@@ -206,23 +206,26 @@ internal PathInfo SetLocation(string path, CmdletProviderContext context)
             ProviderInfo provider = null;
             string providerId = null;
 
-            // Replace path with last working directory when '-' was passed.
-            bool pushNextLocation = true;
-            if (originalPath.Equals("-", StringComparison.OrdinalIgnoreCase))
+            switch (originalPath)
             {
-                if (_SetLocationHistory.Count <= 0)
-                {
-                    throw new InvalidOperationException(SessionStateStrings.SetContentToLastLocationWhenHistoryIsEmpty);
-                }
-                var previousLocation =  _SetLocationHistory.Pop();
-                path = previousLocation.Path;
-                pushNextLocation = false;
-            }
-
-            if (pushNextLocation)
-            {
-                var newPushPathInfo = GetNewPushPathInfo();
-                _SetLocationHistory.Push(newPushPathInfo);
+                case string originalPathSwitch when originalPathSwitch.Equals("-", StringComparison.OrdinalIgnoreCase):
+                    if (_setLocationHistory.UndoCount <= 0)
+                    {
+                        throw new InvalidOperationException(SessionStateStrings.LocationUndoStackIsEmpty);
+                    }
+                    path = _setLocationHistory.Undo(this.CurrentLocation).Path;
+                    break;
+                case string originalPathSwitch when originalPathSwitch.Equals("+", StringComparison.OrdinalIgnoreCase):
+                    if (_setLocationHistory.RedoCount <= 0)
+                    {
+                        throw new InvalidOperationException(SessionStateStrings.LocationRedoStackIsEmpty);
+                    }
+                    path = _setLocationHistory.Redo(this.CurrentLocation).Path;
+                    break;
+                default:
+                    var pushPathInfo = GetNewPushPathInfo();
+                    _setLocationHistory.Push(pushPathInfo);
+                    break;
             }
 
             PSDriveInfo previousWorkingDrive = CurrentDrive;
@@ -776,9 +779,9 @@ CmdletProviderContext normalizePathContext
         #region push-Pop current working directory
 
         /// <summary>
-        /// A bounded stack for the location history of Set-Location
+        /// Location history for Set-Location that supports Undo/Redo using bounded stacks.
         /// </summary>
-        private BoundedStack<PathInfo> _SetLocationHistory;
+        private readonly HistoryStack<PathInfo> _setLocationHistory;
 
         /// <summary>
         /// A stack of the most recently pushed locations
@@ -816,13 +819,13 @@ internal void PushCurrentLocation(string stackName)
             }
 
             // Push the directory/drive pair onto the stack
-            var newPushPathInfo = GetNewPushPathInfo();
-            locationStack.Push(newPushPathInfo);
+            var pushPathInfo = GetNewPushPathInfo();
+            locationStack.Push(pushPathInfo);
         }
 
         private PathInfo GetNewPushPathInfo()
         {
-             // Create a new instance of the directory/drive pair
+            // Create a new instance of the directory/drive pair
             ProviderInfo provider = CurrentDrive.Provider;
             string mshQualifiedPath =
                 LocationGlobber.GetMshQualifiedPath(CurrentDrive.CurrentLocation, CurrentDrive);

diff --git a/src/System.Management.Automation/engine/Utils.cs b/src/System.Management.Automation/engine/Utils.cs
@@ -1407,17 +1407,65 @@ public static void SetTestHook(string property, object value)
     }
 
     /// <summary>
-    /// An bounded stack based on a linked list.
+    /// Provides undo/redo functionality by using 2 instances of <seealso cref="BoundedStack{T}"/>.
+    /// </summary>
+    internal class HistoryStack<T>
+    {
+        private readonly BoundedStack<T> _boundedUndoStack;
+        private readonly BoundedStack<T> _boundedRedoStack;
+
+        internal HistoryStack(uint capacity)
+        {
+            _boundedUndoStack = new BoundedStack<T>(capacity);
+            _boundedRedoStack = new BoundedStack<T>(capacity);
+        }
+
+        internal void Push(T item)
+        {
+            _boundedUndoStack.Push(item);
+            if (RedoCount >= 0)
+            {
+                _boundedRedoStack.Clear();
+            }
+        }
+
+        /// <summary>
+        /// Handles bounded history stacks by pushing the current item to the redoStack and returning the item from the popped undoStack.
+        /// </summary>
+        internal T Undo(T currentItem)
+        {
+            T previousItem = _boundedUndoStack.Pop();
+            _boundedRedoStack.Push(currentItem);
+            return previousItem;
+        }
+
+        /// <summary>
+        /// Handles bounded history stacks by pushing the current item to the undoStack and returning the item from the popped redoStack.
+        /// </summary>
+        internal T Redo(T currentItem)
+        {
+            var nextItem = _boundedRedoStack.Pop();
+            _boundedUndoStack.Push(currentItem);
+            return nextItem;
+        }
+
+        internal int UndoCount => _boundedUndoStack.Count;
+
+        internal int RedoCount => _boundedRedoStack.Count;
+    }
+
+    /// <summary>
+    /// A bounded stack based on a linked list.
     /// </summary>
     internal class BoundedStack<T> : LinkedList<T>
     {
-        private readonly int _capacity;
+        private readonly uint _capacity;
 
         /// <summary>
         /// Lazy initialisation, i.e. it sets only its limit but does not allocate the memory for the given capacity.
         /// </summary>
         /// <param name="capacity"></param>
-        internal BoundedStack(int capacity)
+        internal BoundedStack(uint capacity)
         {
             _capacity = capacity;
         }
@@ -1430,7 +1478,7 @@ internal void Push(T item)
         {
             this.AddFirst(item);
 
-            if(this.Count > _capacity)
+            if (this.Count > _capacity)
             {
                 this.RemoveLast();
             }
@@ -1442,6 +1490,11 @@ internal void Push(T item)
         /// <returns></returns>
         internal T Pop()
         {
+            if (this.First == null)
+            {
+                throw new InvalidOperationException(SessionStateStrings.BoundedStackIsEmpty);
+            }
+
             var item = this.First.Value;
             try
             {

diff --git a/src/System.Management.Automation/resources/SessionStateStrings.resx b/src/System.Management.Automation/resources/SessionStateStrings.resx
@@ -279,9 +279,12 @@
   <data name="GetContentWriterDynamicParametersProviderException" xml:space="preserve">
     <value>The dynamic parameters for the GetContentWriter operation cannot be retrieved for the '{0}' provider for path '{1}'. {2}</value>
   </data>
-  <data name="SetContentToLastLocationWhenHistoryIsEmpty" xml:space="preserve">
+  <data name="LocationUndoStackIsEmpty" xml:space="preserve">
     <value>There is no location history left to navigate backwards.</value>
   </data>
+  <data name="LocationRedoStackIsEmpty" xml:space="preserve">
+    <value>There is no location history left to navigate forwards.</value>
+  </data>
   <data name="BoundedStackIsEmpty" xml:space="preserve">
     <value>The BoundedStack is empty.</value>
   </data>

diff --git a/test/csharp/test_Utils.cs b/test/csharp/test_Utils.cs
@@ -3,6 +3,7 @@
 using Xunit;
 using System;
 using System.Management.Automation;
+using System.Management.Automation.Internal;
 using System.Reflection;
 
 namespace PSTests.Parallel
@@ -15,5 +16,60 @@ public static void TestIsWinPEHost()
             Skip.IfNot(Platform.IsWindows);
             Assert.False(Utils.IsWinPEHost());
         }
+
+        [Fact]
+        public static void TestHistoryStack()
+        {
+            var historyStack = new HistoryStack<string>(20);
+            Assert.Equal(0, historyStack.UndoCount);
+            Assert.Equal(0, historyStack.RedoCount);
+
+            historyStack.Push("first item");
+            historyStack.Push("second item");
+            Assert.Equal(2, historyStack.UndoCount);
+            Assert.Equal(0, historyStack.RedoCount);
+
+            Assert.Equal("second item", historyStack.Undo("second item"));
+            Assert.Equal("first item", historyStack.Undo("first item"));
+            Assert.Equal(0, historyStack.UndoCount);
+            Assert.Equal(2, historyStack.RedoCount);
+
+            Assert.Equal("first item", historyStack.Redo("first item"));
+            Assert.Equal(1, historyStack.UndoCount);
+            Assert.Equal(1, historyStack.RedoCount);
+
+            // Pushing a new item should invalidate the RedoCount
+            historyStack.Push("third item");
+            Assert.Equal(2, historyStack.UndoCount);
+            Assert.Equal(0, historyStack.RedoCount);
+
+            // Check for the correct exception when the Redo/Undo stack is empty.
+            Assert.Throws<InvalidOperationException>(() => historyStack.Redo("bar"));
+            historyStack.Undo("third item");
+            historyStack.Undo("first item");
+            Assert.Equal(0, historyStack.UndoCount);
+            Assert.Throws<InvalidOperationException>(() => historyStack.Undo("foo"));
+        }
+
+        [Fact]
+        public static void TestBoundedStack()
+        {
+            uint capacity = 20;
+            var boundedStack = new BoundedStack<string>(capacity);
+            Assert.Throws<InvalidOperationException>(() => boundedStack.Pop());
+
+            for (int i = 0; i < capacity; i++)
+            {
+                boundedStack.Push($"{i}");
+            }
+
+            for (int i = 0; i < capacity; i++)
+            {
+                var poppedItem = boundedStack.Pop();
+                Assert.Equal($"{20 - 1 - i}", poppedItem);
+            }
+
+            Assert.Throws<InvalidOperationException>(() => boundedStack.Pop());
+        }
     }
 }
diff --git a/test/powershell/Modules/Microsoft.PowerShell.Management/Set-Location.Tests.ps1 b/test/powershell/Modules/Microsoft.PowerShell.Management/Set-Location.Tests.ps1
@@ -120,6 +120,18 @@ Describe "Set-Location" -Tags "CI" {
             (Get-Location).Path | Should -Be ($initialLocation).Path
         }
 
+        It 'Should go to last location back, forth and back again when specifying minus, plus and minus as a path' {
+            $initialLocation = (Get-Location).Path
+            Set-Location ([System.IO.Path]::GetTempPath())
+            $tempPath = (Get-Location).Path
+            Set-Location -
+            (Get-Location).Path | Should -Be $initialLocation
+            Set-Location +
+            (Get-Location).Path | Should -Be $tempPath
+            Set-Location -
+            (Get-Location).Path | Should -Be $initialLocation
+        }
+
         It 'Should go back to previous locations when specifying minus twice' {
             $initialLocation = (Get-Location).Path
             Set-Location ([System.IO.Path]::GetTempPath())
@@ -137,11 +149,19 @@ Describe "Set-Location" -Tags "CI" {
             foreach ($i in 1..$maximumLocationHistory) {
                 Set-Location ([System.IO.Path]::GetTempPath())
             }
+            $tempPath = (Get-Location).Path
+            # Go back up to the maximum
             foreach ($i in 1..$maximumLocationHistory) {
                 Set-Location -
             }
             (Get-Location).Path | Should Be $initialLocation
             { Set-Location - } | Should -Throw -ErrorId 'System.InvalidOperationException,Microsoft.PowerShell.Commands.SetLocationCommand'
+            # Go forwards up to the maximum
+            foreach ($i in 1..($maximumLocationHistory)) {
+                Set-Location +
+            }
+            (Get-Location).Path | Should -Be $tempPath
+            { Set-Location + } | Should -Throw -ErrorId 'System.InvalidOperationException,Microsoft.PowerShell.Commands.SetLocationCommand'
         }
     }