Fix JS source maps for optional arguments

Also, when function has expression body, use expression instead
of function declaration to mark return statement.

js/js.tests/test/org/jetbrains/kotlin/js/test/JsLineNumberTestGenerated.java

@@ -72,6 +72,12 @@ public class JsLineNumberTestGenerated extends AbstractJsLineNumberTest {
         doTest(fileName);
     }
 
+    @TestMetadata("complexExpressionAsDefaultArgument.kt")
+    public void testComplexExpressionAsDefaultArgument() throws Exception {
+        String fileName = KotlinTestUtils.navigationMetadata("js/js.translator/testData/lineNumbers/complexExpressionAsDefaultArgument.kt");
+        doTest(fileName);
+    }
+
     @TestMetadata("coroutine.kt")
     public void testCoroutine() throws Exception {
         String fileName = KotlinTestUtils.navigationMetadata("js/js.translator/testData/lineNumbers/coroutine.kt");

js/js.translator/src/org/jetbrains/kotlin/js/translate/utils/FunctionBodyTranslator.java

@@ -16,6 +16,7 @@
 
 package org.jetbrains.kotlin.js.translate.utils;
 
+import com.intellij.psi.PsiElement;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns;
 import org.jetbrains.kotlin.descriptors.DeclarationDescriptor;
@@ -84,7 +85,8 @@ public final class FunctionBodyTranslator extends AbstractTranslator {
             KtExpression defaultArgument = getDefaultArgument(valueParameter);
             JsBlock defaultArgBlock = new JsBlock();
             JsExpression defaultValue = Translation.translateAsExpression(defaultArgument, functionBodyContext, defaultArgBlock);
-            JsStatement assignStatement = assignment(jsNameRef, defaultValue).makeStmt();
+            PsiElement psi = KotlinSourceElementKt.getPsi(valueParameter.getSource());
+            JsStatement assignStatement = assignment(jsNameRef, defaultValue).source(psi).makeStmt();
             JsStatement thenStatement = JsAstUtils.mergeStatementInBlockIfNeeded(assignStatement, defaultArgBlock);
             JsBinaryOperation checkArgIsUndefined = equality(jsNameRef, Namer.getUndefinedExpression());
             checkArgIsUndefined.source(KotlinSourceElementKt.getPsi(valueParameter.getSource()));
@@ -150,7 +152,7 @@ public final class FunctionBodyTranslator extends AbstractTranslator {
             }
 
             JsReturn jsReturn = new JsReturn((JsExpression) node);
-            jsReturn.setSource(declaration);
+            jsReturn.setSource(declaration.getBodyExpression());
             MetadataProperties.setReturnTarget(jsReturn, descriptor);
             return jsReturn;
         });

js/js.translator/testData/lineNumbers/closure.kt

@@ -5,4 +5,4 @@ fun box(x: Int, y: Int): Int {
     return foo(y)
 }
 
-// LINES: 2 2 4 2 5
+// LINES: 2 4 2 5

js/js.translator/testData/lineNumbers/complexExpressionAsDefaultArgument.kt

@@ -0,0 +1,14 @@
+fun foo(
+        a: Int =
+            when (baz()) {
+                1 -> bar();
+                else -> 0
+            }
+): Int =
+        a + 1
+
+fun baz() = 1
+
+fun bar() = 2
+
+// LINES: 2 3 4 4 5 2 8 10 12

js/js.translator/testData/lineNumbers/coroutine.kt

@@ -14,4 +14,4 @@ suspend fun bar(): Unit {
     println(a + b)
 }
 
-// LINES: 4 5 4 6 4 9 9 9 9 9 9 9 9 * 10 11 11 11 11 11 * 11 12 13 13 13 13 13 13 14
+// LINES: 4 5 5 6 4 9 9 9 9 9 9 9 9 * 10 11 11 11 11 11 * 11 12 13 13 13 13 13 13 14

js/js.translator/testData/lineNumbers/expressionAsFunctionBody.kt

@@ -4,4 +4,4 @@ fun box() =
 fun foo() =
         23
 
-// LINES: 1 2 4 5
+// LINES: 2 5

js/js.translator/testData/lineNumbers/optionalArgs.kt

@@ -7,4 +7,4 @@ fun box(
     println(y)
 }
 
-// LINES: 2 2 3 4 6 7
+// LINES: 2 2 3 3 4 6 7