Post review fixes, now building text in single pass.

9 files changed, 54 insertions, 63 deletions

diff --git a/core/src/main/kotlin/Formats/FormatDescriptor.kt b/core/src/main/kotlin/Formats/FormatDescriptor.kt
index 58c9472b..7cc50acb 100644
--- a/core/src/main/kotlin/Formats/FormatDescriptor.kt
+++ b/core/src/main/kotlin/Formats/FormatDescriptor.kt
@@ -1,7 +1,7 @@
 package org.jetbrains.dokka.Formats
 
-import Samples.SampleProcessingService
 import org.jetbrains.dokka.*
+import org.jetbrains.dokka.Samples.SampleProcessingService
 import kotlin.reflect.KClass
 
 interface FormatDescriptor {
diff --git a/core/src/main/kotlin/Formats/StandardFormats.kt b/core/src/main/kotlin/Formats/StandardFormats.kt
index 45a653b0..dabde881 100644
--- a/core/src/main/kotlin/Formats/StandardFormats.kt
+++ b/core/src/main/kotlin/Formats/StandardFormats.kt
@@ -1,9 +1,9 @@
 package org.jetbrains.dokka.Formats
 
-import Samples.DefaultSampleProcessingService
-import Samples.KotlinWebsiteSampleProcessingService
-import Samples.SampleProcessingService
 import org.jetbrains.dokka.*
+import org.jetbrains.dokka.Samples.DefaultSampleProcessingService
+import org.jetbrains.dokka.Samples.KotlinWebsiteSampleProcessingService
+import org.jetbrains.dokka.Samples.SampleProcessingService
 import kotlin.reflect.KClass
 
 abstract class KotlinFormatDescriptorBase : FormatDescriptor {
diff --git a/core/src/main/kotlin/Kotlin/DescriptorDocumentationParser.kt b/core/src/main/kotlin/Kotlin/DescriptorDocumentationParser.kt
index a7aa6452..d1f64eeb 100644
--- a/core/src/main/kotlin/Kotlin/DescriptorDocumentationParser.kt
+++ b/core/src/main/kotlin/Kotlin/DescriptorDocumentationParser.kt
@@ -1,14 +1,13 @@
 package org.jetbrains.dokka.Kotlin
 
-import Samples.SampleProcessingService
 import com.google.inject.Inject
 import com.intellij.psi.PsiDocCommentOwner
 import com.intellij.psi.PsiNamedElement
 import com.intellij.psi.util.PsiTreeUtil
 import org.jetbrains.dokka.*
+import org.jetbrains.dokka.Samples.SampleProcessingService
 import org.jetbrains.kotlin.descriptors.*
 import org.jetbrains.kotlin.idea.kdoc.findKDoc
-import org.jetbrains.kotlin.idea.kdoc.getKDocLinkResolutionScope
 import org.jetbrains.kotlin.incremental.components.NoLookupLocation
 import org.jetbrains.kotlin.kdoc.parser.KDocKnownTag
 import org.jetbrains.kotlin.kdoc.psi.impl.KDocSection
@@ -16,17 +15,10 @@ import org.jetbrains.kotlin.kdoc.psi.impl.KDocTag
 import org.jetbrains.kotlin.load.java.descriptors.JavaCallableMemberDescriptor
 import org.jetbrains.kotlin.load.java.descriptors.JavaClassDescriptor
 import org.jetbrains.kotlin.name.FqName
-import org.jetbrains.kotlin.name.Name
-import org.jetbrains.kotlin.psi.KtBlockExpression
-import org.jetbrains.kotlin.psi.KtCallExpression
-import org.jetbrains.kotlin.psi.KtDeclarationWithBody
-import org.jetbrains.kotlin.psi.KtFile
-import org.jetbrains.kotlin.resolve.DescriptorToSourceUtils
 import org.jetbrains.kotlin.resolve.DescriptorUtils
 import org.jetbrains.kotlin.resolve.annotations.argumentValue
 import org.jetbrains.kotlin.resolve.constants.StringValue
 import org.jetbrains.kotlin.resolve.scopes.DescriptorKindFilter
-import org.jetbrains.kotlin.resolve.scopes.ResolutionScope
 import org.jetbrains.kotlin.resolve.scopes.getDescriptorsFiltered
 import org.jetbrains.kotlin.resolve.source.PsiSourceElement
 
diff --git a/core/src/main/kotlin/Samples/DefaultSampleProcessingService.kt b/core/src/main/kotlin/Samples/DefaultSampleProcessingService.kt
index 26841ea6..e6539135 100644
--- a/core/src/main/kotlin/Samples/DefaultSampleProcessingService.kt
+++ b/core/src/main/kotlin/Samples/DefaultSampleProcessingService.kt
@@ -1,4 +1,4 @@
-package Samples
+package org.jetbrains.dokka.Samples
 
 import com.google.inject.Inject
 import com.intellij.psi.PsiElement
diff --git a/core/src/main/kotlin/Samples/KotlinWebsiteSampleProcessingService.kt b/core/src/main/kotlin/Samples/KotlinWebsiteSampleProcessingService.kt
index 864734ea..648a6b67 100644
--- a/core/src/main/kotlin/Samples/KotlinWebsiteSampleProcessingService.kt
+++ b/core/src/main/kotlin/Samples/KotlinWebsiteSampleProcessingService.kt
@@ -1,14 +1,11 @@
-package Samples
+package org.jetbrains.dokka.Samples
 
 import com.google.inject.Inject
-import com.intellij.openapi.util.TextRange
 import com.intellij.psi.PsiElement
 import org.jetbrains.dokka.DocumentationOptions
 import org.jetbrains.dokka.DokkaLogger
 import org.jetbrains.dokka.DokkaResolutionFacade
 import org.jetbrains.kotlin.psi.*
-import org.jetbrains.kotlin.psi.psiUtil.startOffset
-import java.util.*
 
 open class KotlinWebsiteSampleProcessingService
 @Inject constructor(options: DocumentationOptions,
@@ -16,66 +13,68 @@ open class KotlinWebsiteSampleProcessingService
                     resolutionFacade: DokkaResolutionFacade)
     : DefaultSampleProcessingService(options, logger, resolutionFacade) {
 
+    private class SampleBuilder() : KtVisitorVoid() {
+        val builder = StringBuilder()
+        val text: String
+            get() = builder.toString()
 
-    fun buildReplacementData(psiElement: PsiElement): Map<TextRange, String> {
-
-        val result = TreeMap<TextRange, String>({ o1, o2 -> o1.startOffset.compareTo(o2.startOffset) })
         fun convertAssertPrints(expression: KtCallExpression) {
             val (argument, commentArgument) = expression.valueArguments
             val comment = commentArgument.getArgumentExpression() as KtStringTemplateExpression
             val commentText = comment.entries.joinToString("") { it.text }
-            result[expression.textRange] = "println(${argument.text}) //$commentText"
+            builder.apply {
+                append("println(")
+                append(argument.text)
+                append(") // ")
+                append(commentText)
+            }
         }
 
         fun convertAssertTrue(expression: KtCallExpression) {
             val (argument) = expression.valueArguments
-            result[expression.textRange] = "println(\"${argument.text} is \${${argument.text}}\") //true "
+            builder.apply {
+                append("println(\"")
+                append(argument.text)
+                append(" is \${")
+                append(argument.text)
+                append("}\") // true")
+            }
         }
 
-        if (psiElement is KtElement) {
-            val visitor = object : KtTreeVisitor<Any>() {
-                override fun visitCallExpression(expression: KtCallExpression, data: Any?): Void? {
-                    when (expression.calleeExpression?.text) {
-                        "assertPrints" -> convertAssertPrints(expression)
-                        "assertTrue" -> convertAssertTrue(expression)
-                        else -> super.visitCallExpression(expression, data)
-                    }
-                    return null
-                }
+        override fun visitCallExpression(expression: KtCallExpression) {
+            when (expression.calleeExpression?.text) {
+                "assertPrints" -> convertAssertPrints(expression)
+                "assertTrue" -> convertAssertTrue(expression)
+                else -> super.visitCallExpression(expression)
             }
-            psiElement.acceptChildren(visitor)
         }
-        return result
-    }
 
-    private fun String.applyReplacements(baseOffset: Int, replacementData: Map<TextRange, String>): String {
-        val partsList = arrayListOf<String>()
-        var prevRange = TextRange(0, baseOffset)
-        for ((range, replacement) in replacementData) {
-            partsList.add(substring(prevRange.endOffset - baseOffset, range.startOffset - baseOffset))
-            partsList.add(replacement)
-            prevRange = range
+        override fun visitElement(element: PsiElement?) {
+            if (element != null) {
+                if (element.children.isEmpty())
+                    builder.append(element.text)
+                else
+                    element.acceptChildren(this)
+            }
         }
-        partsList.add(substring(prevRange.endOffset - baseOffset))
-        return partsList.joinToString(separator = "")
     }
 
+    private fun PsiElement.buildSampleText(): String {
+        val sampleBuilder = SampleBuilder()
+        this.accept(sampleBuilder)
+        return sampleBuilder.text
+    }
 
-    override fun processSampleBody(psiElement: PsiElement): String {
-
-        val replacementData = buildReplacementData(psiElement)
-
-        return when (psiElement) {
-            is KtDeclarationWithBody -> {
-                val bodyExpression = psiElement.bodyExpression
-                val bodyExpressionText = bodyExpression!!.text.applyReplacements(bodyExpression.startOffset, replacementData)
-                when (bodyExpression) {
-                    is KtBlockExpression -> bodyExpressionText.removeSurrounding("{", "}")
-                    else -> bodyExpressionText
-                }
+    override fun processSampleBody(psiElement: PsiElement) = when (psiElement) {
+        is KtDeclarationWithBody -> {
+            val bodyExpression = psiElement.bodyExpression
+            val bodyExpressionText = bodyExpression!!.buildSampleText()
+            when (bodyExpression) {
+                is KtBlockExpression -> bodyExpressionText.removeSurrounding("{", "}")
+                else -> bodyExpressionText
             }
-            else -> psiElement.text.applyReplacements(psiElement.startOffset, replacementData)
         }
+        else -> psiElement.buildSampleText()
     }
 }
 
diff --git a/core/src/main/kotlin/Samples/SampleProcessingService.kt b/core/src/main/kotlin/Samples/SampleProcessingService.kt
index 61abda6d..7f46299f 100644
--- a/core/src/main/kotlin/Samples/SampleProcessingService.kt
+++ b/core/src/main/kotlin/Samples/SampleProcessingService.kt
@@ -1,4 +1,4 @@
-package Samples
+package org.jetbrains.dokka.Samples
 
 import org.jetbrains.dokka.ContentNode
 import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
diff --git a/core/src/main/kotlin/Utilities/DokkaModule.kt b/core/src/main/kotlin/Utilities/DokkaModule.kt
index 4b2a99d3..e1ae829a 100644
--- a/core/src/main/kotlin/Utilities/DokkaModule.kt
+++ b/core/src/main/kotlin/Utilities/DokkaModule.kt
@@ -1,12 +1,12 @@
 package org.jetbrains.dokka.Utilities
 
-import Samples.SampleProcessingService
 import com.google.inject.Binder
 import com.google.inject.Module
 import com.google.inject.Provider
 import com.google.inject.name.Names
 import org.jetbrains.dokka.*
 import org.jetbrains.dokka.Formats.FormatDescriptor
+import org.jetbrains.dokka.Samples.SampleProcessingService
 import org.jetbrains.kotlin.cli.jvm.compiler.KotlinCoreEnvironment
 import java.io.File
 
diff --git a/core/src/main/kotlin/javadoc/dokka-adapters.kt b/core/src/main/kotlin/javadoc/dokka-adapters.kt
index e3a58331..920eca2d 100644
--- a/core/src/main/kotlin/javadoc/dokka-adapters.kt
+++ b/core/src/main/kotlin/javadoc/dokka-adapters.kt
@@ -1,10 +1,10 @@
 package org.jetbrains.dokka.javadoc
 
-import Samples.DefaultSampleProcessingService
 import com.google.inject.Inject
 import com.sun.tools.doclets.formats.html.HtmlDoclet
 import org.jetbrains.dokka.*
 import org.jetbrains.dokka.Formats.FormatDescriptor
+import org.jetbrains.dokka.Samples.DefaultSampleProcessingService
 
 class JavadocGenerator @Inject constructor (val options: DocumentationOptions, val logger: DokkaLogger) : Generator {
     override fun buildPages(nodes: Iterable<DocumentationNode>) {
diff --git a/core/testdata/format/website/sampleWithAsserts.md b/core/testdata/format/website/sampleWithAsserts.md
index 3353ce46..98d7df33 100644
--- a/core/testdata/format/website/sampleWithAsserts.md
+++ b/core/testdata/format/website/sampleWithAsserts.md
@@ -14,8 +14,8 @@ layout: api
 
 fun main(args: Array<String>) {
 //sampleStart
-println(a()) //Hello, Work
-println("a() == b() is ${a() == b()}") //true
+println(a()) // Hello, Work
+println("a() == b() is ${a() == b()}") // true
 //sampleEnd
 }
 ```