diff --git a/Sources/SwiftKotlinFramework/KotlinTokenizer.swift b/Sources/SwiftKotlinFramework/KotlinTokenizer.swift
index 1f1c5ac..1b8c42f 100644
--- a/Sources/SwiftKotlinFramework/KotlinTokenizer.swift
+++ b/Sources/SwiftKotlinFramework/KotlinTokenizer.swift
@@ -394,6 +394,21 @@ public class KotlinTokenizer: SwiftTokenizer {
 
     }
     
+    open override func tokenize(_ codeBlock: CodeBlock) -> [Token] {
+        guard codeBlock.statements.count == 1,
+            let returnStatement = codeBlock.statements.first as? ReturnStatement,
+            let parent = codeBlock.lexicalParent as? Declaration else {
+            return super.tokenize(codeBlock)
+        }
+        let sameLine = parent is VariableDeclaration
+        let separator = sameLine ? codeBlock.newToken(.space, " ") : codeBlock.newToken(.linebreak, "\n")
+        let tokens = Array(tokenize(returnStatement).dropFirst(2))
+        return [
+            [codeBlock.newToken(.symbol, "=")],
+            sameLine ? tokens : indent(tokens)
+        ].joined(token: separator)
+    }
+    
     // MARK: - Statements
 
     open override func tokenize(_ statement: GuardStatement) -> [Token] {
diff --git a/Tests/SwiftKotlinFrameworkTests/Tests/KotlinTokenizer/annotations.kt b/Tests/SwiftKotlinFrameworkTests/Tests/KotlinTokenizer/annotations.kt
index 6ec5b3c..25096fe 100644
--- a/Tests/SwiftKotlinFrameworkTests/Tests/KotlinTokenizer/annotations.kt
+++ b/Tests/SwiftKotlinFrameworkTests/Tests/KotlinTokenizer/annotations.kt
@@ -12,6 +12,5 @@ fun collectCustomerProviders(customerProvider: () -> String) {
     customerProviders.append(customerProvider)
 }
 
-fun foo(code: (() -> String)) : String {
-    return "foo ${bar(code)}"
-}
+fun foo(code: (() -> String)) : String =
+    "foo ${bar(code)}"
diff --git a/Tests/SwiftKotlinFrameworkTests/Tests/KotlinTokenizer/extensions.kt b/Tests/SwiftKotlinFrameworkTests/Tests/KotlinTokenizer/extensions.kt
index af8175c..e21c5ac 100644
--- a/Tests/SwiftKotlinFrameworkTests/Tests/KotlinTokenizer/extensions.kt
+++ b/Tests/SwiftKotlinFrameworkTests/Tests/KotlinTokenizer/extensions.kt
@@ -1,24 +1,15 @@
 val Double.km: Double
-    get() {
-        return this * 1000.0
-    }
+    get() = this * 1000.0
 val Double.m: Double
-    get() {
-        return this
-    }
+    get() = this
 
-open fun Double.toKm() : Double {
-    return this * 1000.0
-}
+open fun Double.toKm() : Double =
+    this * 1000.0
 
-fun Double.toMeter() : Double {
-    return this
-}
+fun Double.toMeter() : Double =
+    this
 
-public fun Double.Companion.toKm() : Double {
-    return this * 1000.0
-}
+public fun Double.Companion.toKm() : Double =
+    this * 1000.0
 public val Double.Companion.m: Double
-    get() {
-        return this
-    }
+    get() = this
diff --git a/Tests/SwiftKotlinFrameworkTests/Tests/KotlinTokenizer/functions.kt b/Tests/SwiftKotlinFrameworkTests/Tests/KotlinTokenizer/functions.kt
index 7c9ea3b..4cfcc94 100644
--- a/Tests/SwiftKotlinFrameworkTests/Tests/KotlinTokenizer/functions.kt
+++ b/Tests/SwiftKotlinFrameworkTests/Tests/KotlinTokenizer/functions.kt
@@ -7,7 +7,11 @@ fun method(param: String) : String {}
 
 fun method(param: (Int) -> Unit) {}
 
-fun findRestaurant(restaurantId: Int) : ServiceTask<Restaurant> {
-    return NetworkRequestServiceTask<Restaurant>(networkSession = networkSession, endpoint = "restaurants/")
-}
+fun findRestaurant(restaurantId: Int) : ServiceTask<Restaurant> =
+    NetworkRequestServiceTask<Restaurant>(networkSession = networkSession, endpoint = "restaurants/")
 restaurantService.findRestaurant(restaurantId = restaurant.id, param = param)
+
+fun tokenize(codeBlock: String?) : List<String> {
+    val statement = codeBlock ?: return listOf()
+    return someOtherMethod(statement = statement)
+}
diff --git a/Tests/SwiftKotlinFrameworkTests/Tests/KotlinTokenizer/functions.swift b/Tests/SwiftKotlinFrameworkTests/Tests/KotlinTokenizer/functions.swift
index 99289e5..d19f793 100644
--- a/Tests/SwiftKotlinFrameworkTests/Tests/KotlinTokenizer/functions.swift
+++ b/Tests/SwiftKotlinFrameworkTests/Tests/KotlinTokenizer/functions.swift
@@ -11,3 +11,10 @@ func findRestaurant(restaurantId: Int) -> ServiceTask<Restaurant> {
 }
 
 restaurantService.findRestaurant(restaurantId: restaurant.id, param: param)
+
+func tokenize(_ codeBlock: String?) -> [String] {
+    guard let statement = codeBlock else {
+        return []
+    }
+    return someOtherMethod(statement: statement)
+}
diff --git a/Tests/SwiftKotlinFrameworkTests/Tests/KotlinTokenizer/properties.kt b/Tests/SwiftKotlinFrameworkTests/Tests/KotlinTokenizer/properties.kt
index ba0bfbd..9349506 100644
--- a/Tests/SwiftKotlinFrameworkTests/Tests/KotlinTokenizer/properties.kt
+++ b/Tests/SwiftKotlinFrameworkTests/Tests/KotlinTokenizer/properties.kt
@@ -6,9 +6,7 @@ interface Hello {
 
 class A {
     val stateObservable1: Observable<RestaurantsListState>
-        get() {
-            return state.asObservable()
-        }
+        get() = state.asObservable()
     val stateObservable2: Observable<RestaurantsListState>
         get() {
             return state.asObservable()
diff --git a/Tests/SwiftKotlinFrameworkTests/Tests/KotlinTokenizer/statics.kt b/Tests/SwiftKotlinFrameworkTests/Tests/KotlinTokenizer/statics.kt
index c275b08..3e57a3a 100644
--- a/Tests/SwiftKotlinFrameworkTests/Tests/KotlinTokenizer/statics.kt
+++ b/Tests/SwiftKotlinFrameworkTests/Tests/KotlinTokenizer/statics.kt
@@ -25,13 +25,11 @@ class A {
         
         fun method() {}
         
-        fun create() : A? {
-            return null
-        }
+        fun create() : A? =
+            null
         
-        fun withParams(param: Int) : A? {
-            return null
-        }
+        fun withParams(param: Int) : A? =
+            null
     }
 }