Preliminary support for .kts files

Summary:
The formatting is inherited from non-kts files, so each statement is surrounded by newlines.

Follow-up: only force newlines between certain constrcuts, such as classes.

This unblocks https://github.com/facebookincubator/ktfmt/issues/35 .

Reviewed By: strulovich

Differential Revision: D25743859

fbshipit-source-id: 8dbda2866ea50eb4a73d94a72d2b7f10618a4cd0
diff --git a/core/src/main/java/com/facebook/ktfmt/KotlinInputAstVisitorBase.kt b/core/src/main/java/com/facebook/ktfmt/KotlinInputAstVisitorBase.kt
index f004559..84ad791 100644
--- a/core/src/main/java/com/facebook/ktfmt/KotlinInputAstVisitorBase.kt
+++ b/core/src/main/java/com/facebook/ktfmt/KotlinInputAstVisitorBase.kt
@@ -91,6 +91,7 @@
 import org.jetbrains.kotlin.psi.KtQualifiedExpression
 import org.jetbrains.kotlin.psi.KtReferenceExpression
 import org.jetbrains.kotlin.psi.KtReturnExpression
+import org.jetbrains.kotlin.psi.KtScript
 import org.jetbrains.kotlin.psi.KtSecondaryConstructor
 import org.jetbrains.kotlin.psi.KtSimpleNameExpression
 import org.jetbrains.kotlin.psi.KtStringTemplateExpression
@@ -2084,6 +2085,22 @@
       if (child.text.isBlank()) {
         continue
       }
+      if (child !is PsiComment && child !is KtScript) {
+        builder.blankLineWanted(OpsBuilder.BlankLineWanted.YES)
+      }
+      child.accept(this)
+    }
+    markForPartialFormat()
+  }
+
+  override fun visitScript(script: KtScript) {
+    markForPartialFormat()
+    for (child in script.blockExpression.children) {
+      if (child.text.isBlank()) {
+        continue
+      }
+      builder.forcedBreak()
+
       if (child !is PsiComment) {
         builder.blankLineWanted(OpsBuilder.BlankLineWanted.YES)
       }
diff --git a/core/src/main/java/com/facebook/ktfmt/Parser.kt b/core/src/main/java/com/facebook/ktfmt/Parser.kt
index cefb8e9..e491b9c 100644
--- a/core/src/main/java/com/facebook/ktfmt/Parser.kt
+++ b/core/src/main/java/com/facebook/ktfmt/Parser.kt
@@ -45,7 +45,7 @@
       val env =
           KotlinCoreEnvironment.createForProduction(
               disposable, configuration, EnvironmentConfigFiles.JVM_CONFIG_FILES)
-      val virtualFile = LightVirtualFile("temp.kt", KotlinFileType.INSTANCE, code)
+      val virtualFile = LightVirtualFile("temp.kts", KotlinFileType.INSTANCE, code)
       val ktFile = PsiManager.getInstance(env.project).findFile(virtualFile) as KtFile
       ktFile.collectDescendantsOfType<PsiErrorElement>().let {
         if (it.isNotEmpty()) throwParseError(code, it[0])
diff --git a/core/src/test/java/com/facebook/ktfmt/FormatterKtTest.kt b/core/src/test/java/com/facebook/ktfmt/FormatterKtTest.kt
index 0a6d9e2..80bbb93 100644
--- a/core/src/test/java/com/facebook/ktfmt/FormatterKtTest.kt
+++ b/core/src/test/java/com/facebook/ktfmt/FormatterKtTest.kt
@@ -28,6 +28,27 @@
 class FormatterKtTest {
 
   @Test
+  fun `support script (kts) files`() =
+      assertFormatted(
+          """
+        |package foo
+        |
+        |import java.io.File
+        |
+        |val one: String
+        |
+        |val two: String
+        |
+        |fun f() {
+        |  println("asd")
+        |}
+        |
+        |println("Called with args:")
+        |
+        |args.forEach { println(File + "-") }
+        |""".trimMargin())
+
+  @Test
   fun `call chains`() =
       assertFormatted(
           """
@@ -3362,14 +3383,15 @@
       |  //
       |}
       |
-      |fn () { }
+      |fn (
       |""".trimMargin()
     try {
       format(code)
       fail()
     } catch (e: ParseError) {
-      assertThat(e.lineColumn.line).isEqualTo(5)
+      assertThat(e.lineColumn.line).isEqualTo(6)
       assertThat(e.lineColumn.column).isEqualTo(0)
+      assertThat(e.errorDescription).contains("Expecting an expression")
     }
   }