Fix return of lambda within BlockExpression

Author: zspitz

Shared/CSharpCodeWriter.cs

@@ -213,8 +213,16 @@ protected override void WriteLambda(LambdaExpression expr) {
             Write("{");
             Indent();
             WriteEOL();
-            if (expr.Body.Type != typeof(void)) { Write("return "); }
-            WriteNode("Body", expr.Body, CreateMetadata(CSharpMultilineBlockTypes.Block));
+
+            CSharpMultilineBlockTypes blockType = CSharpMultilineBlockTypes.Block;
+            if (expr.Body.Type != typeof(void)) {
+                if (expr.Body is BlockExpression bexpr && bexpr.HasMultipleLines()) {
+                    blockType = CSharpMultilineBlockTypes.Return;
+                } else {
+                    Write("return ");
+                }
+            }
+            WriteNode("Body", expr.Body, CreateMetadata(blockType));
             WriteStatementEnd(expr.Body);
             WriteEOL(true);
             Write("}");
@@ -535,7 +543,7 @@ protected override void WriteIndex(IndexExpression expr) =>
         protected override void WriteBlock(BlockExpression expr, object metadata) {
             var (blockType, parentIsBlock) = (CSharpBlockMetadata)metadata ?? CreateMetadata(Inline, false);
             bool introduceNewBlock;
-            if (blockType == CSharpMultilineBlockTypes.Block) {
+            if (blockType.In(CSharpMultilineBlockTypes.Block, CSharpMultilineBlockTypes.Return)) {
                 introduceNewBlock = expr.Variables.Any() && parentIsBlock;
                 if (introduceNewBlock) {
                     Write("{");
@@ -550,7 +558,16 @@ protected override void WriteBlock(BlockExpression expr, object metadata) {
                 expr.Expressions.ForEach((subexpr, index) => {
                     if (index > 0) { WriteEOL(); }
                     if (subexpr is LabelExpression) { TrimEnd(); }
-                    WriteNode($"Expressions[{index}]", subexpr, CreateMetadata(blockType, true));
+                    if (blockType == CSharpMultilineBlockTypes.Return &&  index == expr.Expressions.Count-1) {
+                       if (subexpr is BlockExpression bexpr && bexpr.HasMultipleLines()) {
+                            WriteNode($"Expressions[{index}]", subexpr, CreateMetadata(CSharpMultilineBlockTypes.Return, true));
+                        } else {
+                            Write("return ");
+                            WriteNode($"Expressions[{index}]", subexpr, CreateMetadata(CSharpMultilineBlockTypes.Block, true));
+                        }
+                    } else {
+                        WriteNode($"Expressions[{index}]", subexpr, CreateMetadata(CSharpMultilineBlockTypes.Block, true));
+                    }
                     WriteStatementEnd(subexpr);
                 });
                 if (introduceNewBlock) {
@@ -712,7 +729,7 @@ protected override void WriteGoto(GotoExpression expr) {
                 case GotoExpressionKind.Goto:
                     gotoKeyword = "goto";
                     break;
-                case Return:
+                case GotoExpressionKind.Return:
                     gotoKeyword = "return";
                     break;
                 default:

Shared/CSharpMultilineBlockTypes.cs

@@ -8,6 +8,7 @@ namespace ExpressionToString {
     public enum CSharpMultilineBlockTypes {
         Inline,
         Test,
-        Block
+        Block,
+        Return
     }
 }

Shared/VBCodeWriter.cs

@@ -270,8 +270,16 @@ protected override void WriteLambda(LambdaExpression expr) {
 
             Indent();
             WriteEOL();
-            if (expr.Body.Type != typeof(void)) { Write("Return "); }
-            WriteNode("Body", expr.Body, CreateMetadata(true, Lambda));
+
+            bool returnBlock = false;
+            if (expr.Body.Type != typeof(void)) {
+                if (expr.Body is BlockExpression bexpr && bexpr.HasMultipleLines()) {
+                    returnBlock = true;
+                } else {
+                    Write("Return ");
+                }
+            }
+            WriteNode("Body", expr.Body, CreateMetadata(true, Lambda, returnBlock));
             WriteEOL(true);
             Write($"End {lambdaKeyword}");
         }
@@ -649,7 +657,7 @@ protected override void WriteIndex(IndexExpression expr) =>
             WriteIndexerAccess("Object", expr.Object, "Arguments", expr.Arguments);
 
         protected override void WriteBlock(BlockExpression expr, object metadata) {
-            var (isInMultiline, parentType) = (VBExpressionMetadata)metadata ?? CreateMetadata(false, null);
+            var (isInMultiline, parentType, returnBlock) = (VBExpressionMetadata)metadata ?? CreateMetadata(false, null, false);
             var useBlockConstruct = !isInMultiline ||
                 (expr.Variables.Any() && parentType == Block);
             if (useBlockConstruct) {
@@ -665,7 +673,16 @@ protected override void WriteBlock(BlockExpression expr, object metadata) {
             expr.Expressions.ForEach((subexpr, index) => {
                 if (index > 0 || expr.Variables.Count > 0) { WriteEOL(); }
                 if (subexpr is LabelExpression) { TrimEnd(); }
-                WriteNode($"Expressions[{index}]", subexpr, CreateMetadata(true, Block));
+
+                var outgoingMetadata = CreateMetadata(true, Block);
+                if (returnBlock && index == expr.Expressions.Count - 1) {
+                    if (subexpr is BlockExpression bexpr && bexpr.HasMultipleLines()) {
+                        outgoingMetadata = CreateMetadata(true, Block, true);
+                    } else {
+                        Write("Return ");
+                    }
+                }
+                WriteNode($"Expressions[{index}]", subexpr, outgoingMetadata);
             });
             if (useBlockConstruct) {
                 WriteEOL(true);

Shared/VBExpressionMetadata.cs

@@ -9,13 +9,16 @@ namespace ExpressionToString {
     internal class VBExpressionMetadata {
         internal bool IsInMutiline { get; private set; } = false;
         internal ExpressionType? ExpressionType { get; private set; } = null;
-        internal static VBExpressionMetadata CreateMetadata(bool isInMultiline = false, ExpressionType? expressionType = null) => new VBExpressionMetadata {
+        internal bool ReturnBkock { get; private set; } = false;
+        internal static VBExpressionMetadata CreateMetadata(bool isInMultiline = false, ExpressionType? expressionType = null, bool returnBlock = false) => new VBExpressionMetadata {
             IsInMutiline = isInMultiline,
-            ExpressionType = expressionType
+            ExpressionType = expressionType,
+            ReturnBkock = returnBlock
         };
-        internal void Deconstruct(out bool isInMultiline, out ExpressionType? expressionType) {
+        internal void Deconstruct(out bool isInMultiline, out ExpressionType? expressionType, out bool returnBlock) {
             isInMultiline = IsInMutiline;
             expressionType = ExpressionType;
+            returnBlock = ReturnBkock;
         }
     }
 }

Tests.Common/Constructed/MakeLambda.cs

@@ -178,5 +178,73 @@ Return Function(s1 As String, s2 As String) s1 + s2
     )
 )"
         );
+
+        [Fact]
+        [Trait("Category", Lambdas)]
+        public void LambdaMultilineBlockNonvoidReturn() => RunTest(
+            Lambda(
+                Block(
+                    Constant(true),
+                    Constant(true)
+                )
+            ),
+            @"() => {
+    true;
+    return true;
+}",
+            @"Function()
+    True
+    Return True
+End Function",
+            @"Lambda(
+    Block(
+        Constant(true),
+        Constant(true)
+    )
+)"
+        );
+
+        [Fact]
+        [Trait("Category", Lambdas)]
+        public void LambdaMultilineNestedBlockNonvoidReturn() => RunTest(
+            Lambda(
+                Block(
+                    Constant(true),
+                    Block(
+                        new[] { s1, s2 },
+                        Constant(true),
+                        Constant(true)
+                    )
+                )
+            ),
+            @"() => {
+    true;
+    {
+        string s1;
+        string s2;
+        true;
+        return true;
+    }
+}",
+            @"Function()
+    True
+    Block
+        Dim s1 As String
+        Dim s2 As String
+        True
+        Return True
+    End Block
+End Function",
+            @"Lambda(
+    Block(
+        Constant(true),
+        Block(new[] { s1, s2 },
+            Constant(true),
+            Constant(true)
+        )
+    )
+)"
+        );
+
     }
 }