Fix reference counting issues related to lambdas/closures

For example, circular references between a lambda function the frame
it's stored within and/or its closure could cause memory leaks.

This also fixes other various reference-count ownership issues that
could lead to memory errors.

There may still be some potential/undiscovered issues because the "outer
ID" finding logic doesn't look quite right as the AST traversal descends
within nested lambdas and considers their locals as "outer", but
possibly the other logic for locating values in closures or cloning
closures just works around that behavior.

src/Brofiler.cc

@@ -15,6 +15,17 @@ Brofiler::Brofiler()
 
 Brofiler::~Brofiler()
 	{
+	for ( auto& s : stmts )
+		Unref(s);
+	}
+
+void Brofiler::AddStmt(Stmt* s)
+	{
+	if ( ignoring != 0 )
+		return;
+
+	::Ref(s);
+	stmts.push_back(s);
 	}
 
 bool Brofiler::ReadStats()
@@ -109,7 +120,7 @@ bool Brofiler::WriteStats()
 		return false;
 		}
 
-	for ( list<const Stmt*>::const_iterator it = stmts.begin();
+	for ( list<Stmt*>::const_iterator it = stmts.begin();
 	      it != stmts.end(); ++it )
 		{
 		ODesc location_info;

src/Brofiler.h

@@ -38,13 +38,13 @@ class Brofiler {
 	void IncIgnoreDepth() { ignoring++; }
 	void DecIgnoreDepth() { ignoring--; }
 
-	void AddStmt(const Stmt* s) { if ( ignoring == 0 ) stmts.push_back(s); }
+	void AddStmt(Stmt* s);
 
 private:
 	/**
 	 * The current, global Brofiler instance creates this list at parse-time.
 	 */
-	list<const Stmt*> stmts;
+	list<Stmt*> stmts;
 
 	/**
 	 * Indicates whether new statments will not be considered as part of

src/Expr.cc

@@ -4339,6 +4339,7 @@ LambdaExpr::LambdaExpr(std::unique_ptr<function_ingredients> arg_ing,
 
 	// Install a dummy version of the function globally for use only
 	// when broker provides a closure.
+	::Ref(ingredients->body);
 	BroFunc* dummy_func = new BroFunc(
 		ingredients->id,
 		ingredients->body,
@@ -4378,13 +4379,15 @@ LambdaExpr::LambdaExpr(std::unique_ptr<function_ingredients> arg_ing,
 	dummy_func->SetName(my_name.c_str());
 
 	Val* v = new Val(dummy_func);
+	Unref(dummy_func);
 	id->SetVal(v); // id will unref v when its done.
 	id->SetType(ingredients->id->Type()->Ref());
 	id->SetConst();
 	}
 
 Val* LambdaExpr::Eval(Frame* f) const
 	{
+	::Ref(ingredients->body);
 	BroFunc* lamb = new BroFunc(
 		ingredients->id,
 		ingredients->body,
@@ -4398,7 +4401,9 @@ Val* LambdaExpr::Eval(Frame* f) const
 	// Allows for lookups by the receiver.
 	lamb->SetName(my_name.c_str());
 
-	return new Val(lamb);
+	auto rval = new Val(lamb);
+	Unref(lamb);
+	return rval;
 	}
 
 void LambdaExpr::ExprDescribe(ODesc* d) const

src/Frame.cc

@@ -31,20 +31,49 @@ Frame::Frame(int arg_size, const BroFunc* func, const val_list* fn_args)
 
 Frame::~Frame()
 	{
+	for ( auto& func : functions_with_closure_frame_reference )
+		{
+		func->StrengthenClosureReference(this);
+		Unref(func);
+		}
+
 	// Deleting a Frame that is a view is a no-op.
 	Unref(trigger);
-	Unref(closure);
+
+	if ( ! weak_closure_ref )
+		Unref(closure);
 
 	for ( auto& i : outer_ids )
 		Unref(i);
 
 	Release();
+
+	delete [] weak_refs;
+	}
+
+void Frame::AddFunctionWithClosureRef(BroFunc* func)
+	{
+	::Ref(func);
+	functions_with_closure_frame_reference.emplace_back(func);
 	}
 
-void Frame::SetElement(int n, Val* v)
+void Frame::SetElement(int n, Val* v, bool weak_ref)
 	{
-	Unref(frame[n]);
+	UnrefElement(n);
 	frame[n] = v;
+
+	if ( weak_ref )
+		{
+		if ( ! weak_refs )
+			{
+			weak_refs = new bool[size];
+
+			for ( auto i = 0; i < size; ++i )
+				weak_refs[i] = false;
+			}
+
+		weak_refs[n] = true;
+		}
 	}
 
 void Frame::SetElement(const ID* id, Val* v)
@@ -62,8 +91,15 @@ void Frame::SetElement(const ID* id, Val* v)
 	if ( offset_map.size() )
 		{
 		auto where = offset_map.find(std::string(id->Name()));
+
 		if ( where != offset_map.end() )
-			SetElement(where->second, v);
+			{
+			// Need to add a Ref to 'v' since the SetElement() for
+			// id->Offset() below is otherwise responsible for keeping track
+			// of the implied reference count of the passed-in 'v' argument.
+			// i.e. if we end up storing it twice, we need an addition Ref.
+			SetElement(where->second, v->Ref());
+			}
 		}
 
 	SetElement(id->Offset(), v);
@@ -92,15 +128,15 @@ void Frame::Reset(int startIdx)
 	{
 	for ( int i = startIdx; i < size; ++i )
 		{
-		Unref(frame[i]);
+		UnrefElement(i);
 		frame[i] = 0;
 		}
 	}
 
 void Frame::Release()
 	{
 	for ( int i = 0; i < size; ++i )
-		Unref(frame[i]);
+		UnrefElement(i);
 
 	delete [] frame;
 	}
@@ -145,7 +181,34 @@ Frame* Frame::Clone() const
 	return other;
 	}
 
-Frame* Frame::SelectiveClone(const id_list& selection) const
+static bool val_is_func(Val* v, BroFunc* func)
+	{
+	if ( v->Type()->Tag() != TYPE_FUNC )
+		return false;
+
+	return v->AsFunc() == func;
+	}
+
+static Val* clone_if_not_func(Val** frame, int offset, BroFunc* func,
+                              Frame* other)
+	{
+	auto v = frame[offset];
+
+	if ( ! v )
+		return nullptr;
+
+	if ( val_is_func(v, func) )
+		{
+		other->SetElement(offset, v, true);
+		return v;
+		}
+
+	auto rval = v->Clone();
+	other->SetElement(offset, rval);
+	return rval;
+	}
+
+Frame* Frame::SelectiveClone(const id_list& selection, BroFunc* func) const
 	{
 	if ( selection.length() == 0 )
 		return nullptr;
@@ -171,15 +234,15 @@ Frame* Frame::SelectiveClone(const id_list& selection) const
 			auto where = offset_map.find(std::string(id->Name()));
 			if ( where != offset_map.end() )
 				{
-				other->frame[where->second] = frame[where->second]->Clone();
+				clone_if_not_func(frame, where->second, func, other);
 				continue;
 				}
 			}
 
 		if ( ! frame[id->Offset()] )
 			reporter->InternalError("Attempted to clone an id ('%s') with no associated value.", id->Name());
 
-		other->frame[id->Offset()] = frame[id->Offset()]->Clone();
+		clone_if_not_func(frame, id->Offset(), func, other);
 		}
 
 	/**
@@ -379,6 +442,7 @@ std::pair<bool, Frame*> Frame::Unserialize(const broker::vector& data)
 	// Frame takes ownership of unref'ing elements in outer_ids
 	rf->outer_ids = std::move(outer_ids);
 	rf->closure = closure;
+	rf->weak_closure_ref = false;
 
 	for ( int i = 0; i < frame_size; ++i )
 		{
@@ -437,7 +501,7 @@ void Frame::CaptureClosure(Frame* c, id_list arg_outer_ids)
 
 	closure = c;
 	if ( closure )
-		Ref(closure);
+		weak_closure_ref = true;
 
 	/**
 	 * Want to capture closures by copy?

src/Frame.h

@@ -44,8 +44,11 @@ class Frame :  public BroObj {
 	 *
 	 * @param n the index to set
 	 * @param v the value to set it to
+	 * @param weak_ref whether the frame owns the value and should unref
+	 * it upon destruction.  Used to break circular references between
+	 * lambda functions and closure frames.
 	 */
-	void SetElement(int n, Val* v);
+	void SetElement(int n, Val* v, bool weak_ref = false);
 
 	/**
 	 * Associates *id* and *v* in the frame. Future lookups of
@@ -149,7 +152,7 @@ class Frame :  public BroObj {
 	 * *selection* have been cloned. All other values are made to be
 	 * null.
 	 */
-	Frame* SelectiveClone(const id_list& selection) const;
+	Frame* SelectiveClone(const id_list& selection, BroFunc* func) const;
 
 	/**
 	 * Serializes the Frame into a Broker representation.
@@ -215,8 +218,28 @@ class Frame :  public BroObj {
 	void SetDelayed()	{ delayed = true; }
 	bool HasDelayed() const	{ return delayed; }
 
+	/**
+	 * Track a new function that refers to this frame for use as a closure.
+	 * This frame's destructor will then upgrade that functions reference
+	 * from weak to strong (by making a copy).  The initial use of
+	 * weak references prevents unbreakable circular references that
+	 * otherwise cause memory leaks.
+	 */
+	void AddFunctionWithClosureRef(BroFunc* func);
 
 private:
+
+	/**
+	 * Unrefs the value at offset 'n' frame unless it's a weak reference.
+	 */
+	void UnrefElement(int n)
+		{
+		if ( weak_refs && weak_refs[n] )
+			return;
+
+		Unref(frame[n]);
+		}
+
 	/** Have we captured this id? */
 	bool IsOuterID(const ID* in) const;
 
@@ -242,8 +265,13 @@ class Frame :  public BroObj {
 	/** Associates ID's offsets with values. */
 	Val** frame;
 
+	/** Values that are weakly referenced by the frame.  Used to
+	 * prevent circular reference memory leaks in lambda/closures */
+	bool* weak_refs = nullptr;
+
 	/** The enclosing frame of this frame. */
 	Frame* closure;
+	bool weak_closure_ref = false;
 
 	/** ID's used in this frame from the enclosing frame. */
 	id_list outer_ids;
@@ -268,6 +296,8 @@ class Frame :  public BroObj {
 	Trigger* trigger;
 	const CallExpr* call;
 	bool delayed;
+
+	std::vector<BroFunc*> functions_with_closure_frame_reference;
 };
 
 /**