de.fraunhofer.aisec.cpg.helpers.IdentitySet.kt Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of cpg-core Show documentation
A simple library to extract a code property graph out of source code. It has support for multiple passes that can extend the analysis after the graph is constructed.
There is a newer version: 8.3.0
Show newest version
/*
 * Copyright (c) 2022, Fraunhofer AISEC. All rights reserved.
 *
 * 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
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * 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.
 *
 *                    $$$$$$\  $$$$$$$\   $$$$$$\
 *                   $$  __$$\ $$  __$$\ $$  __$$\
 *                   $$ /  \__|$$ |  $$ |$$ /  \__|
 *                   $$ |      $$$$$$$  |$$ |$$$$\
 *                   $$ |      $$  ____/ $$ |\_$$ |
 *                   $$ |  $$\ $$ |      $$ |  $$ |
 *                   \$$$$$   |$$ |      \$$$$$   |
 *                    \______/ \__|       \______/
 *
 */
package de.fraunhofer.aisec.cpg.helpers

import de.fraunhofer.aisec.cpg.graph.Node
import java.lang.UnsupportedOperationException
import java.util.*
import java.util.concurrent.atomic.AtomicInteger

/**
 * This class implements the [MutableSet] interface with an underlying map and reference-equality
 * instead of object-equality. That means, objects are only considered equal, if they are the *same*
 * object. This logic is primarily implemented by the underlying [IdentityHashMap].
 *
 * The use case of this [MutableSet] is quite simple: In order to avoid loops while traversing in
 * the CPG AST we often need to store [Node] objects in a work-list (usually a set), in order to
 * filter out nodes that were already visited or processed (for example, see
 * [SubgraphWalker.flattenAST]. However, using a normal set triggers object-equality functions, such
 * as [Node.hashCode] or even worse [Node.equals], if the hashcode is the same. This can potentially
 * be very resource-intensive if nodes are very similar but not the *same*, in a work-list however
 * we only want just to avoid to place the exact node twice.
 */
class IdentitySet : MutableSet {
    /**
     * The backing hashmap for our set. The [IdentityHashMap] offers reference-equality for keys and
     * values. In this case we use it to determine, if a node is already in our set or not. The
     * value of the map is not used and is always true. A [Boolean] is used because it seems to be
     * the smallest data type possible.
     */
    private val map: IdentityHashMap = IdentityHashMap()
    private val counter = AtomicInteger()

    override operator fun contains(element: T): Boolean {
        // We are using the backing reference-equality based map to check, if the element is already
        // in the set.
        return map.containsKey(element)
    }

    override fun equals(other: Any?): Boolean {
        if (other !is IdentitySet<*>) return false
        val otherSet = other as? IdentitySet<*>
        return otherSet != null && this.containsAll(otherSet) && otherSet.containsAll(this)
    }

    override fun add(element: T): Boolean {
        // Since we are a Set, we only want to add elements that are not already there
        if (!contains(element)) {
            map[element] = counter.addAndGet(1)
            return true
        }

        return false
    }

    override fun containsAll(elements: Collection): Boolean {
        return elements.all { map.containsKey(it) }
    }

    override fun isEmpty(): Boolean {
        return map.isEmpty()
    }

    override fun iterator(): MutableIterator {
        return map.keys.iterator()
    }

    /**
     * Returns the contents of this [IdentitySet] as a sorted [List] according to order the nodes
     * were inserted to. This is particularly useful, if you need to look up values in the list
     * according to their "closeness" to the root AST node.
     */
    fun toSortedList(): List {
        return map.entries.sortedBy { it.value }.map { it.key }
    }

    override fun addAll(elements: Collection): Boolean {
        // We need to keep track, whether we modified the set
        var modified = false

        elements.forEach {
            if (add(it)) {
                modified = true
            }
        }

        return modified
    }

    override fun clear() {
        map.clear()
    }

    override fun remove(element: T): Boolean {
        return map.remove(element) != null
    }

    override fun removeAll(elements: Collection): Boolean {
        // We need to keep track, whether we modified the set
        var modified = false

        elements.forEach {
            if (remove(it)) {
                modified = true
            }
        }

        return modified
    }

    override fun retainAll(elements: Collection): Boolean {
        throw UnsupportedOperationException()
    }

    override fun hashCode() = map.hashCode()

    override val size: Int
        get() = map.size
}

fun  identitySetOf(vararg elements: T): IdentitySet {
    val set = IdentitySet()
    for (element in elements) set.add(element)

    return set
}

infix fun  IdentitySet.union(other: Iterable): IdentitySet {
    val set = identitySetOf()
    set += this
    set += other
    return set
}

fun  Collection.toIdentitySet(): IdentitySet {
    val set = identitySetOf()
    set += this
    return set
}