Fix recursion detection

Kotlin · Aug 24, 2021 · bd39684 · bd39684
1 parent 9de0755
commit bd39684
Show file tree

Hide file tree

Showing 2 changed files with 47 additions and 46 deletions.
diff --git a/jupyter-lib/api/src/main/kotlin/org/jetbrains/kotlinx/jupyter/api/VariableState.kt b/jupyter-lib/api/src/main/kotlin/org/jetbrains/kotlinx/jupyter/api/VariableState.kt
@@ -2,9 +2,6 @@ package org.jetbrains.kotlinx.jupyter.api
 
 import java.lang.reflect.Field
 import kotlin.reflect.KProperty
-import kotlin.reflect.KProperty1
-import kotlin.reflect.full.declaredMemberProperties
-import kotlin.reflect.jvm.isAccessible
 
 interface VariableState {
     val property: Field
@@ -32,6 +29,7 @@ data class VariableStateImpl(
 ) : VariableState {
     private var cachedValue: Result<Any?> = Result.success(null)
     override var isRecursive: Boolean = false
+    private var isLargeForString: Boolean = false
 
     // use of Java 9 required
     @SuppressWarnings("DEPRECATION")
@@ -58,10 +56,28 @@ data class VariableStateImpl(
         if (cachedValue.getOrNull() == null) {
             return@DependentLazyDelegate null
         }
-        handleIfRecursiveStructure()
+        try {
+            cachedValue.getOrNull().toString()
+            isRecursive = false
+            isLargeForString = false
+        } catch (e: VirtualMachineError) {
+            when (e) {
+                is StackOverflowError -> {
+                    isRecursive = true
+                }
+                is OutOfMemoryError -> {
+                    isLargeForString = true
+                }
+                else -> {
+                    return@DependentLazyDelegate null
+                }
+            }
+        }
 
         if (!isRecursive) {
             cachedValue.getOrNull().toString()
+        } else if (isLargeForString) {
+            "${cachedValue.getOrNull()!!::class.simpleName}: too large structure"
         } else {
             getRecursiveObjectName()
         }
@@ -72,46 +88,4 @@ data class VariableStateImpl(
 
     override val value: Result<Any?>
         get() = cachedValue
-
-    private fun handleIfRecursiveStructure() {
-        if (cachedValue.getOrNull() == null) return
-        traverseObject(cachedValue.getOrNull(), mutableSetOf())
-    }
-
-    private fun traverseObject(value: Any?, seenNodes: MutableSet<Any>) {
-        if (value == null) return
-        val membersProperties = try {
-            value::class.declaredMemberProperties
-        } catch (ex: Throwable) {
-            emptyList<Collection<KProperty1<Any, *>>>()
-        }
-
-        val receivedInstances: MutableList<Any?> = mutableListOf()
-        for (property in membersProperties) {
-            @Suppress("UNCHECKED_CAST")
-            property as KProperty1<Any, *>
-            try {
-                val wasAccessible = property.isAccessible
-                property.isAccessible = true
-                val callInstance = property.get(value)
-                property.isAccessible = wasAccessible
-                val wasSeen = callInstance != null && !seenNodes.add(callInstance)
-
-                if (wasSeen) {
-                    isRecursive = true
-                    return
-                }
-                receivedInstances.add(callInstance)
-            } catch (ex: Throwable) {
-                // there might be recursive elements inside the container
-                if (ex is StackOverflowError) {
-                    isRecursive = true
-                }
-                return
-            }
-        }
-        receivedInstances.forEach {
-            traverseObject(it, seenNodes)
-        }
-    }
 }
diff --git a/src/test/kotlin/org/jetbrains/kotlinx/jupyter/test/repl/ReplTests.kt b/src/test/kotlin/org/jetbrains/kotlinx/jupyter/test/repl/ReplTests.kt
@@ -771,6 +771,33 @@ class ReplVarsTest : AbstractSingleReplTest() {
         assertEquals(2, newData.fieldDescriptor.size)
     }
 
+    @Test
+    fun testProperBiRecursionHandling() {
+        eval(
+            """
+            val l = mutableListOf<Any>()
+            l.add(listOf(l))
+            
+            val m = mutableMapOf<Int, Any>(1 to l)
+            
+            val z = setOf(1, 2, 4)
+            """.trimIndent(),
+            jupyterId = 1
+        )
+        var state = repl.notebook.variablesState
+        assertEquals("ArrayList: recursive structure", state["l"]!!.stringValue)
+        assertEquals("LinkedHashMap: recursive structure", state["m"]!!.stringValue)
+        eval(
+            """
+            val m = mutableMapOf<Int, Any>(1 to "abc")
+            """.trimIndent(),
+            jupyterId = 2
+        )
+        state = repl.notebook.variablesState
+        assertEquals("ArrayList: recursive structure", state["l"]!!.stringValue)
+        assertNotEquals("LinkedHashMap: recursive structure", state["m"]!!.stringValue)
+    }
+
     @Test
     fun testUnchangedVars() {
         eval(