Add EntityCollector

jacodb-ets/src/main/kotlin/org/jacodb/ets/utils/Collector.kt

@@ -0,0 +1,54 @@
+/*
+ *  Copyright 2022 UnitTestBot contributors (utbot.org)
+ * <p>
+ *  Licensed under the Apache License, Version 2.0 (the "License");
+ *  you may not use this file except in compliance with the License.
+ *  You may obtain a copy of the License at
+ * <p>
+ *  http://www.apache.org/licenses/LICENSE-2.0
+ * <p>
+ *  Unless required by applicable law or agreed to in writing, software
+ *  distributed under the License is distributed on an "AS IS" BASIS,
+ *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *  See the License for the specific language governing permissions and
+ *  limitations under the License.
+ */
+
+@file:Suppress("PropertyName")
+
+package org.jacodb.ets.utils
+
+import org.jacodb.ets.base.EtsEntity
+import org.jacodb.ets.base.EtsStmt
+
+class EntityCollector<R : Any, C : MutableCollection<R>>(
+    val result: C,
+    val block: (EtsEntity) -> R?,
+) : AbstractHandler() {
+    override fun handle(value: EtsEntity) {
+        val item = block(value)
+        if (item != null) {
+            result += item
+        }
+    }
+
+    override fun handle(stmt: EtsStmt) {
+        // Do nothing.
+    }
+}
+
+fun <R : Any, C : MutableCollection<R>> EtsEntity.collectEntitiesTo(
+    destination: C,
+    block: (EtsEntity) -> R?,
+): C {
+    accept(EntityCollector(destination, block))
+    return destination
+}
+
+fun <R : Any, C : MutableCollection<R>> EtsStmt.collectEntitiesTo(
+    destination: C,
+    block: (EtsEntity) -> R?,
+): C {
+    accept(EntityCollector(destination, block))
+    return destination
+}

jacodb-ets/src/main/kotlin/org/jacodb/ets/utils/GetUses.kt → jacodb-ets/src/main/kotlin/org/jacodb/ets/utils/GetLocals.kt

@@ -16,11 +16,31 @@
 
 package org.jacodb.ets.utils
 
+import org.jacodb.ets.base.EtsAssignStmt
 import org.jacodb.ets.base.EtsEntity
+import org.jacodb.ets.base.EtsLocal
 import org.jacodb.ets.base.EtsStmt
+import org.jacodb.ets.model.EtsMethod
 
-fun EtsStmt.getUses(): Sequence<EtsEntity> =
-    getOperands().flatMap { sequenceOf(it) + it.getUses() }
+fun EtsMethod.getDeclaredLocals(): Set<EtsLocal> =
+    cfg.stmts.mapNotNullTo(mutableSetOf()) {
+        if (it is EtsAssignStmt && it.lhv is EtsLocal) {
+            it.lhv
+        } else {
+            null
+        }
+    }
 
-fun EtsEntity.getUses(): Sequence<EtsEntity> =
-    getOperands().flatMap { sequenceOf(it) + it.getUses() }
+fun EtsMethod.getLocals(): Set<EtsLocal> {
+    val result = mutableSetOf<EtsLocal>()
+    cfg.stmts.forEach { it.collectEntitiesTo(result) { it as? EtsLocal } }
+    return result
+}
+
+fun EtsStmt.getLocals(): Set<EtsLocal> {
+    return collectEntitiesTo(mutableSetOf()) { it as? EtsLocal }
+}
+
+fun EtsEntity.getLocals(): Set<EtsLocal> {
+    return collectEntitiesTo(mutableSetOf()) { it as? EtsLocal }
+}

jacodb-ets/src/main/kotlin/org/jacodb/ets/utils/GetValues.kt

@@ -20,8 +20,10 @@ import org.jacodb.ets.base.EtsEntity
 import org.jacodb.ets.base.EtsStmt
 import org.jacodb.ets.base.EtsValue
 
-fun EtsStmt.getValues(): Sequence<EtsValue> =
-    getOperands().filterIsInstance<EtsValue>()
+fun EtsStmt.getValues(): Set<EtsValue> {
+    return collectEntitiesTo(mutableSetOf()) { it as? EtsValue }
+}
 
-fun EtsEntity.getValues(): Sequence<EtsValue> =
-    getOperands().filterIsInstance<EtsValue>()
+fun EtsEntity.getValues(): Set<EtsValue> {
+    return collectEntitiesTo(mutableSetOf()) { it as? EtsValue }
+}