Fix up $using index and member lookups (#4)

Fixes the logic for extracting the values of a $using variable that has
an index and member lookup with constant values. The $using
implementation needs to extract the final value to add to the final
parameter value rather than just the variable object itself.

src/RemoteForge/Commands/InvokeRemoteCommand.cs

@@ -194,8 +194,6 @@ await RunWorker(
 
     private Hashtable GetUsingParameters(string script)
     {
-        ScriptBlock? indexAccesor = null;
-
         Hashtable usingParams = new();
 
         // ParseInput will not fail on errors, to allow providing anything as
@@ -233,30 +231,14 @@ private Hashtable GetUsingParameters(string script)
                     ThrowTerminatingError(err);
                 }
 
-                if (usingAst.SubExpression is IndexExpressionAst usingIndex)
+                try
                 {
-                    // ConstantExpressionAst, StringConstantExpressionAst, and
-                    // VariableExpressionAst with IsConstantVariable() all
-                    // should return a value here. Those 3 scenarios are the
-                    // only index scenarios that $using supports.
-                    // SafeGetValue() fails in other scenarios which we warn
-                    // about.
-                    object index;
-                    try
-                    {
-                        index = usingIndex.Index.SafeGetValue();
-                    }
-                    catch (InvalidOperationException e)
-                    {
-                        WriteWarning($"Failed to extract $using value: {e.Message}");
-                        continue;
-                    }
-
-                    // PowerShell doesn't have a nice API to dynamically
-                    // invoke the Item[] lookup so to avoid complex reflection
-                    // code we just fallback to doing it in pwsh itself.
-                    indexAccesor ??= System.Management.Automation.ScriptBlock.Create("$args[0][$args[1]]");
-                    value = indexAccesor.Invoke(value, index).FirstOrDefault();
+                    value = ExtractUsingExpressionValue(value, usingAst.SubExpression);
+                }
+                catch (Exception e)
+                {
+                    WriteWarning($"Failed to extract $using value: {e.Message}");
+                    continue;
                 }
 
                 usingParams.Add(key, value);
@@ -266,6 +248,72 @@ private Hashtable GetUsingParameters(string script)
         return usingParams;
     }
 
+    private object? ExtractUsingExpressionValue(object? value, ExpressionAst ast)
+    {
+        if (ast is not MemberExpressionAst && ast is not IndexExpressionAst)
+        {
+            // No need to extract the inner value for simple $using:var entries.
+            return value;
+        }
+
+        // We need to replace the inner VariableExpressionAst that the
+        // member/index expressions are pulling from with the constant value.
+        VariableExpressionAst usingVariable = (VariableExpressionAst)ast.Find(a => a is VariableExpressionAst, false);
+
+        // We go up the hierarchy replacing the index and member expressions
+        // with the new constant value/the newly wrapped AST expressions.
+        ExpressionAst lookupAst = new ConstantExpressionAst(ast.Extent, value);
+        ExpressionAst currentAst = usingVariable;
+        while (true)
+        {
+            if (currentAst.Parent is IndexExpressionAst indexAst)
+            {
+                lookupAst = new IndexExpressionAst(
+                    indexAst.Extent,
+                    lookupAst,
+                    (ExpressionAst)indexAst.Index.Copy(),
+                    indexAst.NullConditional);
+                currentAst = indexAst;
+            }
+            else if (currentAst.Parent is MemberExpressionAst memberAst)
+            {
+                lookupAst = new MemberExpressionAst(
+                    memberAst.Extent,
+                    lookupAst,
+                    (ExpressionAst)memberAst.Member.Copy(),
+                    memberAst.Static,
+                    memberAst.NullConditional);
+                currentAst = memberAst;
+            }
+            else
+            {
+                break;
+            }
+        }
+
+        // With the new ScriptBlock we can just run it to get the final value.
+        ScriptBlock extractionScriptBlock = new ScriptBlockAst(
+            ast.Extent,
+            null,
+            new StatementBlockAst(
+                ast.Extent,
+                new StatementAst[]
+                {
+                    new PipelineAst(
+                        ast.Extent,
+                        new CommandBaseAst[]
+                        {
+                            new CommandExpressionAst(
+                                ast.Extent,
+                                lookupAst,
+                                null)
+                        })
+                },
+                null),
+            false).GetScriptBlock();
+        return extractionScriptBlock.Invoke().FirstOrDefault();
+    }
+
     protected override void ProcessRecord()
     {
         foreach (PSObject? input in InputObject)

tests/InvokeRemote.tests.ps1

@@ -646,6 +646,23 @@ Describe 'Invoke-Remote $using: tests' {
         $actual | Should -BeNullOrEmpty
     }
 
+    It "Passes with index and member lookup" {
+        $var = @{
+            'Prop With Space' = @(
+                @{ Foo = 'value 1' }
+                @{ Foo = 'value 2' }
+            )
+        }
+        $actual = Invoke-Remote -ConnectionInfo PipeTest: -ScriptBlock {
+            $using:var.'Prop With Space'[0]['Foo']
+            $using:var["Prop With Space"][1].Foo
+        }
+
+        $actual.Count | Should -Be 2
+        $actual[0] | Should -Be 'value 1'
+        $actual[1] | Should -Be 'value 2'
+    }
+
     It "Fails if variable not defined locally" {
         {
             Invoke-Remote -ConnectionInfo PipeTest: -ScriptBlock { $using:UndefinedVar }
@@ -666,6 +683,6 @@ Describe 'Invoke-Remote $using: tests' {
         Invoke-Remote -ConnectionInfo PipeTest: -ScriptBlock '$idx = 0; $using:var[$idx]' -WarningAction SilentlyContinue -WarningVariable warn -ErrorAction Ignore
 
         $warn.Count | Should -Be 1
-        $warn[0].Message | Should -Be 'Failed to extract $using value: Cannot generate a PowerShell object for a ScriptBlock evaluating dynamic expressions. Dynamic expression: $idx.'
+        $warn[0].Message | Should -Be 'Failed to extract $using value: Index operation failed; the array index evaluated to null.'
     }
 }