From 986cb5be273676550e50dd79a66f526828cb71d2 Mon Sep 17 00:00:00 2001
From: Simon Ogorodnik <Simon.Ogorodnik@jetbrains.com>
Date: Thu, 31 May 2018 20:32:21 +0300
Subject: [PATCH] New J2K: Add more elements to code builder

---
 .../org/jetbrains/kotlin/j2k/NewCodeBuilder.kt | 18 ++++++++++++++++++
 1 file changed, 18 insertions(+)

diff --git a/j2k/newSrc/org/jetbrains/kotlin/j2k/NewCodeBuilder.kt b/j2k/newSrc/org/jetbrains/kotlin/j2k/NewCodeBuilder.kt
index 5178868a44d..6b6fb57480a 100644
--- a/j2k/newSrc/org/jetbrains/kotlin/j2k/NewCodeBuilder.kt
+++ b/j2k/newSrc/org/jetbrains/kotlin/j2k/NewCodeBuilder.kt
@@ -85,11 +85,29 @@ class NewCodeBuilder {
             ktFunction.returnType.accept(this, data)
             if (ktFunction.block !== JKBodyStub) {
                 printer.printlnWithNoIndent("{")
+                printer.pushIndent()
                 ktFunction.block.accept(this, data)
+                printer.popIndent()
                 printer.printWithNoIndent("}")
             }
         }
 
+        override fun visitBlock(block: JKBlock, data: Unit) {
+            block.acceptChildren(this, data)
+        }
+
+        override fun visitExpressionStatement(expressionStatement: JKExpressionStatement, data: Unit) {
+            printer.printIndent()
+            expressionStatement.expression.accept(this, data)
+            printer.printlnWithNoIndent()
+        }
+
+        override fun visitReturnStatement(returnStatement: JKReturnStatement, data: Unit) {
+            printer.print("return ")
+            returnStatement.expression.accept(this, data)
+            printer.printlnWithNoIndent()
+        }
+
         override fun visitKtProperty(ktProperty: JKKtProperty, data: Unit) {
             // TODO: Fix this
             if (ktProperty.modifierList.modifiers.any { (it as? JKJavaModifier)?.type == JKJavaModifier.JavaModifierType.FINAL }) {