From aa9d9c904fa1af40a843aeca701326354cb7523a Mon Sep 17 00:00:00 2001
From: Jon Siwek <jsiwek@corelight.com>
Date: Mon, 18 May 2020 15:55:30 -0700
Subject: [PATCH] Switch some TopkVal methods to use IntrusivePtr

---
 src/probabilistic/Topk.cc   | 19 +++++++------------
 src/probabilistic/Topk.h    | 10 +++++-----
 src/probabilistic/top-k.bif |  4 ++--
 3 files changed, 14 insertions(+), 19 deletions(-)

diff --git a/src/probabilistic/Topk.cc b/src/probabilistic/Topk.cc
index 13c5c9adc8..33857416f4 100644
--- a/src/probabilistic/Topk.cc
+++ b/src/probabilistic/Topk.cc
@@ -18,11 +18,6 @@ static void topk_element_hash_delete_func(void* val)
 	delete e;
 	}
 
-Element::~Element()
-	{
-	Unref(value);
-	}
-
 void TopkVal::Typify(IntrusivePtr<BroType> t)
 	{
 	assert(!hash && !type);
@@ -116,7 +111,7 @@ void TopkVal::Merge(const TopkVal* value, bool doPrune)
 				{
 				olde = new Element();
 				olde->epsilon = 0;
-				olde->value = e->value->Ref();
+				olde->value = e->value;
 				// insert at bucket position 0
 				if ( buckets.size() > 0 )
 					{
@@ -188,7 +183,7 @@ IntrusivePtr<Val> TopkVal::DoClone(CloneState* state)
 	return state->NewClone(this, std::move(clone));
 	}
 
-VectorVal* TopkVal::GetTopK(int k) const // returns vector
+IntrusivePtr<VectorVal> TopkVal::GetTopK(int k) const // returns vector
 	{
 	if ( numElements == 0 )
 		{
@@ -225,7 +220,7 @@ VectorVal* TopkVal::GetTopK(int k) const // returns vector
 		it--;
 		}
 
-	return t.release();
+	return t;
 	}
 
 uint64_t TopkVal::GetCount(Val* value) const
@@ -276,7 +271,7 @@ uint64_t TopkVal::GetSum() const
 	return sum;
 	}
 
-void TopkVal::Encountered(Val* encountered)
+void TopkVal::Encountered(IntrusivePtr<Val> encountered)
 	{
 	// ok, let's see if we already know this one.
 
@@ -297,7 +292,7 @@ void TopkVal::Encountered(Val* encountered)
 		{
 		e = new Element();
 		e->epsilon = 0;
-		e->value = encountered->Ref(); // or no ref?
+		e->value = std::move(encountered);
 
 		// well, we do not know this one yet...
 		if ( numElements < size )
@@ -437,7 +432,7 @@ broker::expected<broker::data> TopkVal::DoSerialize() const
 			{
 			Element* element = *eit;
 			d.emplace_back(element->epsilon);
-			auto v = bro_broker::val_to_data(element->value);
+			auto v = bro_broker::val_to_data(element->value.get());
 			if ( ! v )
 				return broker::ec::invalid_data;
 
@@ -509,7 +504,7 @@ bool TopkVal::DoUnserialize(const broker::data& data)
 
 			Element* e = new Element();
 			e->epsilon = *epsilon;
-			e->value = val.release();
+			e->value = std::move(val);
 			e->parent = b;
 
 			b->elements.insert(b->elements.end(), e);
diff --git a/src/probabilistic/Topk.h b/src/probabilistic/Topk.h
index e6e162f2d6..c71dbfe769 100644
--- a/src/probabilistic/Topk.h
+++ b/src/probabilistic/Topk.h
@@ -27,10 +27,8 @@ struct Bucket {
 
 struct Element {
 	uint64_t epsilon;
-	Val* value;
+	IntrusivePtr<Val> value;
 	Bucket* parent;
-
-	~Element();
 };
 
 class TopkVal : public OpaqueVal {
@@ -57,7 +55,7 @@ public:
 	 *
 	 * @param value The encountered element
 	 */
-	void Encountered(Val* value);
+	void Encountered(IntrusivePtr<Val> value);
 
 	/**
 	 * Get the first *k* elements of the result vector. At the moment,
@@ -68,7 +66,7 @@ public:
 	 *
 	 * @returns The top-k encountered elements
 	 */
-	VectorVal* GetTopK(int k) const;
+	IntrusivePtr<VectorVal> GetTopK(int k) const;
 
 	/**
 	 * Get the current count tracked in the top-k data structure for a
@@ -155,6 +153,8 @@ private:
 	 * @returns HashKey for value
 	 */
 	HashKey* GetHash(Val* v) const; // this probably should go somewhere else.
+	HashKey* GetHash(const IntrusivePtr<Val>& v) const
+		{ return GetHash(v.get()); }
 
 	/**
 	 * Set the type that this TopK instance tracks
diff --git a/src/probabilistic/top-k.bif b/src/probabilistic/top-k.bif
index d771df332f..947b514c08 100644
--- a/src/probabilistic/top-k.bif
+++ b/src/probabilistic/top-k.bif
@@ -34,7 +34,7 @@ function topk_add%(handle: opaque of topk, value: any%): any
 	%{
 	assert(handle);
 	probabilistic::TopkVal* h = (probabilistic::TopkVal*) handle;
-	h->Encountered(value);
+	h->Encountered({NewRef{}, value});
 
 	return nullptr;
 	%}
@@ -53,7 +53,7 @@ function topk_get_top%(handle: opaque of topk, k: count%): any_vec
 	%{
 	assert(handle);
 	probabilistic::TopkVal* h = (probabilistic::TopkVal*) handle;
-	return IntrusivePtr{AdoptRef{}, h->GetTopK(k)};
+	return h->GetTopK(k);
 	%}
 
 ## Get an overestimated count of how often a value has been encountered.