<commit_before>/*
 * Copyright 2020 Google LLC
 *
 * 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
 *
 *     https://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 com.google.idea.perf.methodtracer

import com.google.idea.perf.agent.MethodListener
import com.google.idea.perf.util.ConcurrentAppendOnlyList
import org.objectweb.asm.Type
import java.lang.reflect.Method
import java.util.concurrent.locks.ReentrantLock
import kotlin.concurrent.withLock

// Things to improve:
// - Somehow gc or recycle Tracepoints that are no longer used.
// - Pretty-print method descriptors for better UX.

/** Keeps track of which methods should be traced via bytecode instrumentation. */
object TracerConfig {
    private val tracepoints = ConcurrentAppendOnlyList<Tracepoint>()
    private val lock = ReentrantLock() // Protects the data structures below.
    private val classConfigs = mutableMapOf<String, ClassConfig>() // Keyed by 'JVM' class name.

    private class TracepointProperties(
        val flags: Int = TracepointFlags.TRACE_ALL,
        val parameters: Int = 0
    ) {
        companion object {
            val DEFAULT = TracepointProperties()
        }
    }

    /** Specifies which methods to instrument for a given class. */
    private class ClassConfig {
        /** Set of simple method names to instrumental and their associated properties. */
        val methodNames = mutableMapOf<String, TracepointProperties>()

        /** Map from method signature to method ID. */
        val methodIds = mutableMapOf<String, Int>()
    }

    fun traceMethods(
        classJvmName: String,
        methodName: String,
        flags: Int,
        parameters: Collection<Int>
    ) {
        var parameterBits = 0
        for (index in parameters) {
            parameterBits = parameterBits or (1 shl index)
        }

        lock.withLock {
            val classConfig = classConfigs.getOrPut(classJvmName) { ClassConfig() }
            classConfig.methodNames[methodName] = TracepointProperties(flags, parameterBits)

            // If the tracepoint already exists, set tracepoint properties.
            for ((signature, methodId) in classConfig.methodIds) {
                if (signature.substringBefore('(') == methodName) {
                    val tracepoint = getTracepoint(methodId)
                    tracepoint.setFlags(TracepointFlags.TRACE_ALL)
                    tracepoint.parameters.set(parameterBits)
                }
            }
        }
    }

    fun traceMethod(method: Method) {
        lock.withLock {
            val classJvmName = method.declaringClass.name.replace('.', '/')
            val methodDesc = Type.getMethodDescriptor(method)
            traceMethod(classJvmName, method.name, methodDesc)
        }
    }

    private fun traceMethod(classJvmName: String, methodName: String, methodDesc: String) {
        lock.withLock {
            val classConfig = classConfigs.getOrPut(classJvmName) { ClassConfig() }
            val methodSignature = "$methodName$methodDesc"
            val tracepoint = createTracepoint(
                classJvmName, methodName, methodDesc, TracepointProperties.DEFAULT
            )
            classConfig.methodIds.getOrPut(methodSignature) { tracepoints.append(tracepoint) }
        }
    }

    fun untraceMethods(classJvmName: String, methodName: String) {
        lock.withLock {
            val classConfig = classConfigs[classJvmName] ?: return

            for ((signature, _) in classConfig.methodIds) {
                if (signature.substringBefore('(') == methodName) {
                    val methodDesc = signature.substring(methodName.length)
                    val methodId = getMethodId(classJvmName, methodName, methodDesc)
                    if (methodId != null) {
                        val tracepoint = getTracepoint(methodId)
                        tracepoint.unsetFlags(TracepointFlags.TRACE_ALL)
                        tracepoint.parameters.set(0)
                    }
                }
            }
        }
    }

    fun untraceMethod(method: Method) {
        lock.withLock {
            val classJvmName = method.declaringClass.name.replace('.', '/')
            val methodDesc = Type.getMethodDescriptor(method)
            untraceMethod(classJvmName, method.name, methodDesc)
        }
    }<commit_msg>Fix "trace tracer" not working after untrace<commit_after>/*
 * Copyright 2020 Google LLC
 *
 * 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
 *
 *     https://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 com.google.idea.perf.methodtracer

import com.google.idea.perf.agent.MethodListener
import com.google.idea.perf.util.ConcurrentAppendOnlyList
import org.objectweb.asm.Type
import java.lang.reflect.Method
import java.util.concurrent.locks.ReentrantLock
import kotlin.concurrent.withLock

// Things to improve:
// - Somehow gc or recycle Tracepoints that are no longer used.
// - Pretty-print method descriptors for better UX.

/** Keeps track of which methods should be traced via bytecode instrumentation. */
object TracerConfig {
    private val tracepoints = ConcurrentAppendOnlyList<Tracepoint>()
    private val lock = ReentrantLock() // Protects the data structures below.
    private val classConfigs = mutableMapOf<String, ClassConfig>() // Keyed by 'JVM' class name.

    private class TracepointProperties(
        val flags: Int = TracepointFlags.TRACE_ALL,
        val parameters: Int = 0
    ) {
        companion object {
            val DEFAULT = TracepointProperties()
        }
    }

    /** Specifies which methods to instrument for a given class. */
    private class ClassConfig {
        /** Set of simple method names to instrumental and their associated properties. */
        val methodNames = mutableMapOf<String, TracepointProperties>()

        /** Map from method signature to method ID. */
        val methodIds = mutableMapOf<String, Int>()
    }

    fun traceMethods(
        classJvmName: String,
        methodName: String,
        flags: Int,
        parameters: Collection<Int>
    ) {
        var parameterBits = 0
        for (index in parameters) {
            parameterBits = parameterBits or (1 shl index)
        }

        lock.withLock {
            val classConfig = classConfigs.getOrPut(classJvmName) { ClassConfig() }
            classConfig.methodNames[methodName] = TracepointProperties(flags, parameterBits)

            // If the tracepoint already exists, set tracepoint properties.
            for ((signature, methodId) in classConfig.methodIds) {
                if (signature.substringBefore('(') == methodName) {
                    val tracepoint = getTracepoint(methodId)
                    tracepoint.setFlags(TracepointFlags.TRACE_ALL)
                    tracepoint.parameters.set(parameterBits)
                }
            }
        }
    }

    fun traceMethod(method: Method) {
        lock.withLock {
            val classJvmName = method.declaringClass.name.replace('.', '/')
            val methodDesc = Type.getMethodDescriptor(method)
            traceMethod(classJvmName, method.name, methodDesc)
        }
    }

    private fun traceMethod(classJvmName: String, methodName: String, methodDesc: String) {
        lock.withLock {
            val classConfig = classConfigs.getOrPut(classJvmName) { ClassConfig() }
            val methodSignature = "$methodName$methodDesc"
            val methodId = classConfig.methodIds.getOrPut(methodSignature) {
                val tracepoint = createTracepoint(
                    classJvmName, methodName, methodDesc, TracepointProperties.DEFAULT
                )
                tracepoints.append(tracepoint)
            }

            val tracepoint = getTracepoint(methodId)
            tracepoint.parameters.set(0)
            tracepoint.setFlags(TracepointFlags.TRACE_ALL)
        }
    }

    fun untraceMethods(classJvmName: String, methodName: String) {
        lock.withLock {
            val classConfig = classConfigs[classJvmName] ?: return

            for ((signature, _) in classConfig.methodIds) {
                if (signature.substringBefore('(') == methodName) {
                    val methodDesc = signature.substring(methodName.length)
                    val methodId = getMethodId(classJvmName, methodName, methodDesc)
                    if (methodId != null) {
                        val tracepoint = getTracepoint(methodId)
                        tracepoint.unsetFlags(TracepointFlags.TRACE_ALL)
                        tracepoint.parameters.set(0)
                    }
                }
            }
        }
    }

    fun untraceMethod(method: Method) {
        lock.withLock {
            val classJvmName = method.declaringClass.name.replace('.', '/')
            val methodDesc = Type.getMethodDescriptor(method)
            untraceMethod(classJvmName, method.name, methodDesc)
        }
    }
{'fraction_non_alphanumeric': 0.0821917808219178, 'fraction_numerical': 0.002555714577795952, 'mean_word_length': 2.330949948927477, 'pattern_counts': {'":': 0, '<': 13, '<?xml version=': 0, '>': 13, 'https://': 2, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4332279', 'n_tokens_mistral': 2525, 'n_tokens_neox': 2544, 'n_words': 697}