diff --git a/java/kotlin-extractor2/src/main/kotlin/KotlinFileExtractor.kt b/java/kotlin-extractor2/src/main/kotlin/KotlinFileExtractor.kt
index 8c55294217a..efa8b441142 100644
--- a/java/kotlin-extractor2/src/main/kotlin/KotlinFileExtractor.kt
+++ b/java/kotlin-extractor2/src/main/kotlin/KotlinFileExtractor.kt
@@ -5447,46 +5447,48 @@ OLD: KE1
 
     abstract inner class StmtExprParent {
         abstract fun stmt(e: KtExpression, callable: Label<out DbCallable>): StmtParent
-/*
-OLD: KE1
-        abstract fun expr(e: IrExpression, callable: Label<out DbCallable>): ExprParent
-*/
+        abstract fun expr(e: KtExpression, callable: Label<out DbCallable>): ExprParent
     }
 
     inner class StmtParent(val parent: Label<out DbStmtparent>, val idx: Int) : StmtExprParent() {
         override fun stmt(e: KtExpression, callable: Label<out DbCallable>) = this
 
-/*
-OLD: KE1
-        override fun expr(e: IrExpression, callable: Label<out DbCallable>) =
+        override fun expr(e: KtExpression, callable: Label<out DbCallable>) =
             extractExpressionStmt(tw.getLocation(e), parent, idx, callable).let { id ->
                 ExprParent(id, 0, id)
             }
-*/
     }
 
-/*
-OLD: KE1
     inner class ExprParent(
         val parent: Label<out DbExprparent>,
         val idx: Int,
         val enclosingStmt: Label<out DbStmt>
     ) : StmtExprParent() {
-        override fun stmt(e: IrExpression, callable: Label<out DbCallable>): StmtParent {
+        override fun stmt(e: KtExpression, callable: Label<out DbCallable>): StmtParent {
             val id = tw.getFreshIdLabel<DbStmtexpr>()
-            val type = useType(e.type)
-            val locId = tw.getLocation(e)
-            tw.writeExprs_stmtexpr(id, type.javaResult.id, parent, idx)
+            val et = e.expressionType
+            if (et != null) {
+                val type = useType(et)
+                val locId = tw.getLocation(e)
+                tw.writeExprs_stmtexpr(id, type.javaResult.id, parent, idx)
+            } else {
+                logger.errorElement("Unexpected null type: ${e.javaClass}", e)
+            }
+/*
+OLD: KE1
             tw.writeExprsKotlinType(id, type.kotlinResult.id)
             extractExprContext(id, locId, callable, enclosingStmt)
+*/
             return StmtParent(id, 0)
         }
 
-        override fun expr(e: IrExpression, callable: Label<out DbCallable>): ExprParent {
+        override fun expr(e: KtExpression, callable: Label<out DbCallable>): ExprParent {
             return this
         }
     }
 
+/*
+OLD: KE1
     private fun getStatementOriginOperator(origin: IrStatementOrigin?) =
         when (origin) {
             IrStatementOrigin.PLUSEQ -> "plus"
@@ -5816,6 +5818,7 @@ OLD: KE1
 
         return false
     }
+*/
 
     private fun extractExpressionStmt(
         locId: Label<DbLocation>,
@@ -5828,6 +5831,8 @@ OLD: KE1
             tw.writeHasLocation(it, locId)
         }
 
+/*
+OLD: KE1
     private fun extractExpressionStmt(
         e: IrExpression,
         callable: Label<out DbCallable>,
@@ -6029,10 +6034,11 @@ OLD: KE1
                     val locId = tw.getLocation(e)
                     tw.writeStmts_returnstmt(id, stmtParent.parent, stmtParent.idx, callable)
                     tw.writeHasLocation(id, locId)
-/*
-OLD: KE1
-                    extractExpressionExpr(e.value, callable, id, 0, id)
-*/
+                    val returned = e.getReturnedExpression()
+                    if (returned != null) {
+                        extractExpression(returned, callable, ExprParent(id, 0, id))
+                    }
+                    // TODO: e.getLabeledExpression()
                 }
 /*
 OLD: KE1
diff --git a/java/kotlin-extractor2/src/main/kotlin/TrapWriter.kt b/java/kotlin-extractor2/src/main/kotlin/TrapWriter.kt
index b0c7b147cd4..72c245d1c8b 100644
--- a/java/kotlin-extractor2/src/main/kotlin/TrapWriter.kt
+++ b/java/kotlin-extractor2/src/main/kotlin/TrapWriter.kt
@@ -367,8 +367,9 @@ OLD: KE1
     }
 */
 
-    /** Gets a label for the location of `e`. */
-    fun getLocation(e: PsiElement): Label<DbLocation> {
+    private data class Location(val startLine: Int, val startColumn: Int, val endLine: Int, val endColumn: Int)
+
+    private fun getLocationInfo(e: PsiElement): Location {
         val range = e.getTextRange()
         val document = e.getContainingFile().getViewProvider().getDocument()
         val start = range.getStartOffset()
@@ -377,7 +378,14 @@ OLD: KE1
         val end = range.getEndOffset()
         val endLine0 = document.getLineNumber(end)
         val endCol1 = end - document.getLineStartOffset(endLine0)
-        return getLocation(fileId, startLine0 + 1, startCol0 + 1, endLine0 + 1, endCol1)
+        // TODO: unknown/synthetic locations?
+        return Location(startLine0 + 1, startCol0 + 1, endLine0 + 1, endCol1)
+    }
+
+    /** Gets a label for the location of `e`. */
+    fun getLocation(e: PsiElement): Label<DbLocation> {
+        val loc = getLocationInfo(e)
+        return getLocation(fileId, loc.startLine, loc.startColumn, loc.endLine, loc.endColumn)
     }
 
 /*
@@ -395,19 +403,27 @@ OLD: KE1
         // we can do is return a whole-file location.
         return getWholeFileLocation()
     }
+*/
+
     /**
      * Gets the location of `e` as a human-readable string. Only used in log messages and exception
      * messages.
      */
-    open fun getLocationString(e: IrElement): String {
+    /* TODO open */ fun getLocationString(e: PsiElement): String {
+        val loc = getLocationInfo(e)
+        return "file://$filePath:${loc.startLine}:${loc.startColumn}:${loc.endLine}:${loc.endColumn}"
+
+/*
+OLD: KE1
         // We don't have a FileEntry to look up the offsets in, so all
         // we can do is return a whole-file location. We omit the
         // `:0:0:0:0` so that it is easy to distinguish from a location
         // where we have actually determined the start/end lines/columns
         // to be 0.
         return "file://$filePath"
-    }
 */
+    }
+
     /** Gets a label for the location representing the whole of this file. */
     fun getWholeFileLocation(): Label<DbLocation> {
         return getWholeFileLocation(fileId)
diff --git a/java/kotlin-extractor2/src/main/kotlin/utils/Logger.kt b/java/kotlin-extractor2/src/main/kotlin/utils/Logger.kt
index d9522a2d5b9..7b76dd0e7f2 100644
--- a/java/kotlin-extractor2/src/main/kotlin/utils/Logger.kt
+++ b/java/kotlin-extractor2/src/main/kotlin/utils/Logger.kt
@@ -377,7 +377,7 @@ OLD: KE1
 */
 
     fun errorElement(msg: String, element: PsiElement /* TODO , exn: Throwable? = null */) {
-        val locationString = "TODO" // OLD: KE1: ftw.getLocationString(element)
+        val locationString = ftw.getLocationString(element)
         val mkLocationId = { ftw.getLocation(element) }
         loggerBase.diagnostic(
             ftw.getDiagnosticTrapWriter(),