# Extract some data from profile flamegraphs produced by cargo flamegraph

import re, sys

FUNCTIONS = {
	"rustls::record_layer::RecordLayer::decrypt_incoming": "decrypt",

	"<[a-zA-Z0-9_:]+ as rustls::crypto::cipher::MessageEncrypter>::encrypt": "encrypt",

	"<[a-zA-Z0-9_:]+ as rustls::crypto::tls13::Hkdf>::expander_for_okm": "hkdf",

	"<[a-zA-Z0-9_:]+ as rustls::crypto::SecureRandom>::fill": "rand",

	"<[a-zA-Z0-9_:]+ as rustls::crypto::SupportedKxGroup>::start": "kx",

	"<[a-zA-Z0-9_:]+ as rustls::crypto::hash::Hash>::start": "hash",
	"<[a-zA-Z0-9_:]+ as rustls::crypto::hash::Context>::finish": "hash",
	"<[a-zA-Z0-9_:]+ as rustls::crypto::hash::Context>::update": "hash",
	"<[a-zA-Z0-9_:]+ as rustls::crypto::hash::Context>::fork_finish": "hash",

	"<[a-zA-Z0-9_:]+ as rustls::crypto::tls13::Hkdf>::extract_from_secret": "hkdf",

	"<[a-zA-Z0-9_:]+ as rustls::crypto::ActiveKeyExchange>::complete": "kx",

	"<[a-zA-Z0-9_:]+ as rustls::crypto::tls13::HkdfExpander>::hash_len": "hkdf",
	"<[a-zA-Z0-9_:]+ as rustls::crypto::tls13::HkdfExpander>::expand_slice": "hkdf",
	"<[a-zA-Z0-9_:]+ as rustls::crypto::tls13::Hkdf>::extract_from_secret": "hkdf",
	"<[a-zA-Z0-9_:]+ as rustls::crypto::tls13::Hkdf>::hmac_sign": "hkdf",

	"ring::hkdf::fill_okm": "hkdf",
	"aws_lc_0_32_2_HKDF": "hkdf",

	"rustls_openssl::tls13::<impl rustls::crypto::cipher::Tls13AeadAlgorithm for rustls_openssl::aead::Algorithm>::encrypter": "encrypt",
	"rustls::crypto::aws_lc_rs::tls13::AeadAlgorithm::encrypter": "encrypt",

	"rustls::crypto::aws_lc_rs::tls13::AeadAlgorithm::decrypter": "decrypt",
}

def extract_function(data, name):
	r = re.compile(name.replace("<", "&lt;").replace(">", "&gt;") + " \\(([0-9,]+) samples, ([0-9.]+)%\\)")
	samples = 0
	percents = 0.0
	for match in r.finditer(data):
		samples += int(match.group(1).replace(",", ""))
		percents += float(match.group(2))
	return (samples, percents)

if __name__ == "__main__":
	filename = sys.argv[1]
	f = open(filename, "r")
	c = f.read()
	results = {}
	for function in FUNCTIONS:
		samples, percents = extract_function(c, function)
		title = FUNCTIONS[function]
		if title not in results:
			results[title] = [0, 0.0]
		results[title][0] += samples
		results[title][1] += percents
	for title in results:
		print(title, results[title])