From 98ed21b014e4e01281efb5c8d26864175c80e40b Mon Sep 17 00:00:00 2001
From: Juan Romero <romero126@gmail.com>
Date: Tue, 18 Feb 2020 10:12:42 -0800
Subject: [PATCH] Introduce Get-EditorServicesParserAst to expand on
 ExpandAlias

---
 .../Commands/GetEditorServicesParserAst.cs    | 66 +++++++++++++++++++
 .../Handlers/ExpandAliasHandler.cs            | 31 ++-------
 2 files changed, 70 insertions(+), 27 deletions(-)
 create mode 100644 src/PowerShellEditorServices.Hosting/Commands/GetEditorServicesParserAst.cs

diff --git a/src/PowerShellEditorServices.Hosting/Commands/GetEditorServicesParserAst.cs b/src/PowerShellEditorServices.Hosting/Commands/GetEditorServicesParserAst.cs
new file mode 100644
index 000000000..eba68ce91
--- /dev/null
+++ b/src/PowerShellEditorServices.Hosting/Commands/GetEditorServicesParserAst.cs
@@ -0,0 +1,66 @@
+//
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+//
+
+using System;
+using System.Collections.ObjectModel;
+using System.IO;
+using System.Linq;
+using System.Management.Automation;
+using Microsoft.PowerShell.Commands;
+
+namespace Microsoft.PowerShell.EditorServices.Commands
+{
+    
+    /// <summary>
+    /// The Get-EditorServicesParserAst command will parse and expand out data parsed by ast.
+    /// </summary>
+    [Cmdlet(VerbsCommon.Get, "EditorServicesParserAst")]
+    public sealed class GetEditorServicesParserAst : PSCmdlet
+    {
+
+        /// <summary>
+        /// The Scriptblock or block of code that gets parsed by ast.
+        /// </summary>
+        [Parameter(Mandatory = true)]
+        public string ScriptBlock { get; set; }
+
+        /// <summary>
+        /// Specify a specific command type
+        /// [System.Management.Automation.CommandTypes]
+        /// </summary>
+        [Parameter(Mandatory = true)]        
+        public CommandTypes CommandType { get; set; }
+
+        /// <summary>
+        /// Specify a specific token type
+        /// [System.Management.Automation.PSTokenType]
+        /// </summary>
+        [Parameter(Mandatory = true)]
+        public PSTokenType PSTokenType { get; set; }
+
+        protected override void EndProcessing()
+        {
+            var errors = new Collection<PSParseError>();
+
+            var tokens =
+                System.Management.Automation.PSParser.Tokenize(ScriptBlock, out errors)
+                .Where(token => token.Type == this.PSTokenType)
+                .OrderBy(token => token.Content);
+
+            foreach (PSToken token in tokens)
+            {
+                if (PSTokenType == PSTokenType.Command)
+                {
+                    var result = SessionState.InvokeCommand.GetCommand(token.Content, CommandType);
+                    WriteObject(result);
+                }
+                else {
+                    WriteObject(token);
+                }
+
+            }
+        }
+    }
+}
\ No newline at end of file
diff --git a/src/PowerShellEditorServices/Services/PowerShellContext/Handlers/ExpandAliasHandler.cs b/src/PowerShellEditorServices/Services/PowerShellContext/Handlers/ExpandAliasHandler.cs
index 662b91098..0f937d0fd 100644
--- a/src/PowerShellEditorServices/Services/PowerShellContext/Handlers/ExpandAliasHandler.cs
+++ b/src/PowerShellEditorServices/Services/PowerShellContext/Handlers/ExpandAliasHandler.cs
@@ -40,37 +40,14 @@ public ExpandAliasHandler(ILoggerFactory factory, PowerShellContextService power
 
         public async Task<ExpandAliasResult> Handle(ExpandAliasParams request, CancellationToken cancellationToken)
         {
-            const string script = @"
-function __Expand-Alias {
 
-    param($targetScript)
-
-    [ref]$errors=$null
-
-    $tokens = [System.Management.Automation.PsParser]::Tokenize($targetScript, $errors).Where({$_.type -eq 'command'}) |
-                    Sort-Object Start -Descending
-
-    foreach ($token in  $tokens) {
-        $definition=(Get-Command ('`'+$token.Content) -CommandType Alias -ErrorAction SilentlyContinue).Definition
-
-        if($definition) {
-            $lhs=$targetScript.Substring(0, $token.Start)
-            $rhs=$targetScript.Substring($token.Start + $token.Length)
-
-            $targetScript=$lhs + $definition + $rhs
-       }
-    }
-
-    $targetScript
-}";
-
-            // TODO: Refactor to not rerun the function definition every time.
             var psCommand = new PSCommand();
             psCommand
-                .AddScript(script)
                 .AddStatement()
-                .AddCommand("__Expand-Alias")
-                .AddArgument(request.Text);
+                .AddCommand("Get-EditorServicesParserAst")
+                .AddParameter("ScriptBlock", request.Text)
+                .AddParameter("CommandType", CommandTypes.Alias)
+                .AddParameter("PSTokenType", PSTokenType.Command);
             var result = await _powerShellContextService.ExecuteCommandAsync<string>(psCommand).ConfigureAwait(false);
 
             return new ExpandAliasResult