diff --git a/core/src/main/java/com/facebook/ktfmt/Formatter.kt b/core/src/main/java/com/facebook/ktfmt/Formatter.kt
index d1520e4..dc8bd73 100644
--- a/core/src/main/java/com/facebook/ktfmt/Formatter.kt
+++ b/core/src/main/java/com/facebook/ktfmt/Formatter.kt
@@ -15,6 +15,8 @@
 package com.facebook.ktfmt
 
 import com.facebook.ktfmt.kdoc.KDocCommentsHelper
+import com.google.common.collect.ImmutableList
+import com.google.common.collect.Range
 import com.google.googlejavaformat.Doc
 import com.google.googlejavaformat.DocBuilder
 import com.google.googlejavaformat.OpsBuilder
@@ -35,22 +37,19 @@
 fun format(code: String, maxWidth: Int): String {
   val file = Parser.parse(code)
 
-  val javaInput = KotlinInput(code, file)
+  val kotlinInput = KotlinInput(code, file)
   val options = JavaFormatterOptions.defaultOptions()
-  val javaOutput = JavaOutput("\n", javaInput, KDocCommentsHelper("\n"))
-  val builder = OpsBuilder(javaInput, javaOutput)
+  val javaOutput = JavaOutput("\n", kotlinInput, KDocCommentsHelper("\n"))
+  val builder = OpsBuilder(kotlinInput, javaOutput)
   file.accept(KotlinInputAstVisitor(builder))
-  builder.sync(javaInput.text.length)
+  builder.sync(kotlinInput.text.length)
   builder.drain()
   val doc = DocBuilder().withOps(builder.build()).build()
   doc.computeBreaks(javaOutput.commentsHelper, maxWidth, Doc.State(+0, 0))
   doc.write(javaOutput)
   javaOutput.flush()
 
-  val stringBuilder = StringBuilder()
-  (0 until javaOutput.lineCount).forEach {
-    stringBuilder.append(javaOutput.getLine(it))
-    stringBuilder.append('\n')
-  }
-  return stringBuilder.toString()
+  val tokenRangeSet =
+      kotlinInput.characterRangesToTokenRanges(ImmutableList.of(Range.closedOpen(0, code.length)))
+  return JavaOutput.applyReplacements(code, javaOutput.getFormatReplacements(tokenRangeSet))
 }
diff --git a/core/src/main/java/com/facebook/ktfmt/KotlinInput.kt b/core/src/main/java/com/facebook/ktfmt/KotlinInput.kt
index ca8ebcf..802f056 100644
--- a/core/src/main/java/com/facebook/ktfmt/KotlinInput.kt
+++ b/core/src/main/java/com/facebook/ktfmt/KotlinInput.kt
@@ -15,12 +15,17 @@
 package com.facebook.ktfmt
 
 import com.google.common.base.MoreObjects
+import com.google.common.collect.DiscreteDomain
 import com.google.common.collect.ImmutableList
 import com.google.common.collect.ImmutableMap
 import com.google.common.collect.ImmutableRangeMap
+import com.google.common.collect.Iterables.getLast
 import com.google.common.collect.Range
+import com.google.common.collect.RangeSet
+import com.google.common.collect.TreeRangeSet
 import com.google.googlejavaformat.Input
 import com.google.googlejavaformat.Newlines
+import com.google.googlejavaformat.java.FormatterException
 import com.google.googlejavaformat.java.JavaOutput
 import com.intellij.openapi.util.text.StringUtil
 import com.intellij.psi.PsiComment
@@ -86,6 +91,52 @@
     }
   }
 
+  @Throws(FormatterException::class)
+  fun characterRangesToTokenRanges(characterRanges: Collection<Range<Int>>): RangeSet<Int> {
+    val tokenRangeSet = TreeRangeSet.create<Int>()
+    for (characterRange0 in characterRanges) {
+      val characterRange = characterRange0.canonical(DiscreteDomain.integers())
+      tokenRangeSet.add(
+          characterRangeToTokenRange(
+              characterRange.lowerEndpoint(),
+              characterRange.upperEndpoint() - characterRange.lowerEndpoint()))
+    }
+    return tokenRangeSet
+  }
+
+  /**
+   * Convert from an offset and length flag pair to a token range.
+   *
+   * @param offset the `0`-based offset in characters
+   * @param length the length in characters
+   * @return the `0`-based [Range] of tokens
+   * @throws FormatterException
+   */
+  @Throws(FormatterException::class)
+  internal fun characterRangeToTokenRange(offset: Int, length: Int): Range<Int> {
+    var length = length
+    val requiredLength = offset + length
+    if (requiredLength > text.length) {
+      throw FormatterException(
+          String.format(
+              "error: invalid length %d, offset + length (%d) is outside the file",
+              length, requiredLength))
+    }
+    when {
+      length < 0 -> return EMPTY_RANGE
+      length == 0 -> // 0 stands for "format the line under the cursor"
+        length = 1
+    }
+    val enclosed = getPositionTokenMap()
+        .subRangeMap(Range.closedOpen(offset, offset + length))
+        .asMapOfRanges()
+        .values
+    return if (enclosed.isEmpty()) {
+      EMPTY_RANGE
+    } else Range.closedOpen(
+        enclosed.iterator().next().tok.index, getLast(enclosed).getTok().getIndex() + 1)
+  }
+
   private fun makePositionToColumnMap(toks: List<KotlinTok>): ImmutableMap<Int, Int> {
     val builder = ImmutableMap.builder<Int, Int>()
     for (tok in toks) {
diff --git a/core/src/main/java/com/facebook/ktfmt/KotlinInputAstVisitor.kt b/core/src/main/java/com/facebook/ktfmt/KotlinInputAstVisitor.kt
index 68a6cfe..2a04360 100644
--- a/core/src/main/java/com/facebook/ktfmt/KotlinInputAstVisitor.kt
+++ b/core/src/main/java/com/facebook/ktfmt/KotlinInputAstVisitor.kt
@@ -14,6 +14,8 @@
 
 package com.facebook.ktfmt
 
+import com.google.common.base.Throwables
+import com.google.common.collect.ImmutableList
 import com.google.googlejavaformat.Doc
 import com.google.googlejavaformat.FormattingError
 import com.google.googlejavaformat.Indent
@@ -123,6 +125,9 @@
   /** Standard indentation for a long expression or function call, it is different than block indentation on purpose */
   private val expressionBreakIndent: Indent.Const = Indent.Const.make(+4, 1)
 
+  /** A record of whether we have visited into an expression.  */
+  private val inExpression = ArrayDeque(ImmutableList.of(false))
+
   /** Example: `fun foo(n: Int) { println(n) }` */
   override fun visitNamedFunction(function: KtNamedFunction) {
     builder.sync(function)
@@ -1591,6 +1596,47 @@
   }
 
   /**
+   * visitElement is called for almost all types of AST nodes.
+   * We use it to keep track of whether we're currently inside an expression or not.
+   */
+  override fun visitElement(element: PsiElement) {
+    inExpression.addLast(element is KtExpression || inExpression.peekLast())
+    val previous = builder.depth()
+    try {
+      super.visitElement(element)
+    } catch (e: FormattingError) {
+      throw e
+    } catch (t: Throwable) {
+      throw FormattingError(builder.diagnostic(Throwables.getStackTraceAsString(t)))
+    } finally {
+      inExpression.removeLast()
+    }
+    builder.checkClosed(previous)
+  }
+
+  override fun visitKtFile(file: KtFile) {
+    markForPartialFormat()
+    super.visitKtFile(file)
+    markForPartialFormat()
+  }
+
+  private fun inExpression(): Boolean {
+    return inExpression.peekLast()
+  }
+
+  /**
+   * markForPartialFormat is used to delineate the smallest areas of code that must be formatted together.
+   *
+   * When only parts of the code are being formatted, the requested area is expanded until it's
+   * covered by an area marked by this method.
+   */
+  private fun markForPartialFormat() {
+    if (!inExpression()) {
+      builder.markForPartialFormat()
+    }
+  }
+
+  /**
    * Emit a [Doc.Token].
    *
    * @param token the [String] to wrap in a [Doc.Token]