AST profiling enhnacements in support of script optimization for lambdas/whens

src/script_opt/CPP/Stmts.cc

@@ -420,9 +420,9 @@ void CPPCompile::GenWhenStmt(const WhenStmt* w)
 	NL();
 
 	Emit("std::vector<ValPtr> CPP__local_aggrs;");
-	for ( auto l : wi->WhenExprLocals() )
+	for ( auto& l : wi->WhenExprLocals() )
 		if ( IsAggr(l->GetType()) )
-			Emit("CPP__local_aggrs.emplace_back(%s);", IDNameStr(l));
+			Emit("CPP__local_aggrs.emplace_back(%s);", IDNameStr(l.get()));
 
 	Emit("CPP__wi->Instantiate(%s);", GenExpr(wi->Lambda(), GEN_NATIVE));
 

src/script_opt/ProfileFunc.cc

@@ -28,7 +28,23 @@ ProfileFunc::ProfileFunc(const Func* func, const StmtPtr& body, bool _abs_rec_fi
 	profiled_func = func;
 	profiled_body = body.get();
 	abs_rec_fields = _abs_rec_fields;
-	Profile(func->GetType().get(), body);
+
+	auto ft = func->GetType()->AsFuncType();
+	auto& fcaps = ft->GetCaptures();
+
+	if ( fcaps )
+		{
+		int offset = 0;
+
+		for ( auto& c : *fcaps )
+			{
+			auto cid = c.Id().get();
+			captures.insert(cid);
+			captures_offsets[cid] = offset++;
+			}
+		}
+
+	Profile(ft, body);
 	}
 
 ProfileFunc::ProfileFunc(const Stmt* s, bool _abs_rec_fields)
@@ -48,10 +64,15 @@ ProfileFunc::ProfileFunc(const Expr* e, bool _abs_rec_fields)
 		{
 		auto func = e->AsLambdaExpr();
 
-		for ( auto oid : func->OuterIDs() )
+		int offset = 0;
-			captures.insert(oid);
 
-		Profile(func->GetType()->AsFuncType(), func->Ingredients().Body());
+		for ( auto oid : func->OuterIDs() )
+			{
+			captures.insert(oid);
+			captures_offsets[oid] = offset++;
+			}
+
+		Profile(func->GetType()->AsFuncType(), func->Ingredients()->Body());
 		}
 
 	else
@@ -91,9 +112,9 @@ TraversalCode ProfileFunc::PreStmt(const Stmt* s)
 
 			auto w = s->AsWhenStmt();
 			auto wi = w->Info();
-			auto wl = wi ? wi->Lambda() : nullptr;
+
-			if ( wl )
+			for ( auto wl : wi->WhenNewLocals() )
-				lambdas.push_back(wl.get());
+				when_locals.insert(wl);
 			}
 			break;
 
@@ -171,6 +192,11 @@ TraversalCode ProfileFunc::PreExpr(const Expr* e)
 			auto n = e->AsNameExpr();
 			auto id = n->Id();
 
+			// Turns out that NameExpr's can be constructed using a
+			// different Type* than that of the identifier itself,
+			// so be sure we track the latter too.
+			TrackType(id->GetType());
+
 			if ( id->IsGlobal() )
 				{
 				globals.insert(id);
@@ -179,10 +205,10 @@ TraversalCode ProfileFunc::PreExpr(const Expr* e)
 				const auto& t = id->GetType();
 				if ( t->Tag() == TYPE_FUNC && t->AsFuncType()->Flavor() == FUNC_FLAVOR_EVENT )
 					events.insert(id->Name());
+
+				break;
 				}
 
-			else
-				{
 			// This is a tad ugly.  Unfortunately due to the
 			// weird way that Zeek function *declarations* work,
 			// there's no reliable way to get the list of
@@ -197,12 +223,6 @@ TraversalCode ProfileFunc::PreExpr(const Expr* e)
 				params.insert(id);
 
 			locals.insert(id);
-				}
-
-			// Turns out that NameExpr's can be constructed using a
-			// different Type* than that of the identifier itself,
-			// so be sure we track the latter too.
-			TrackType(id->GetType());
 
 			break;
 			}
@@ -350,7 +370,12 @@ TraversalCode ProfileFunc::PreExpr(const Expr* e)
 					params.insert(i);
 				}
 
-			// Avoid recursing into the body.
+			// In general, we don't want to recurse into the body.
+			// However, we still want to *profile* it so we can
+			// identify calls within it.
+			ProfileFunc body_pf(l->Ingredients()->Body().get(), false);
+			script_calls.insert(body_pf.ScriptCalls().begin(), body_pf.ScriptCalls().end());
+
 			return TC_ABORTSTMT;
 			}
 

src/script_opt/ProfileFunc.h

@@ -100,6 +100,8 @@ public:
 	const IDSet& Globals() const { return globals; }
 	const IDSet& AllGlobals() const { return all_globals; }
 	const IDSet& Locals() const { return locals; }
+	const IDSet& Captures() const { return captures; }
+	const auto& CapturesOffsets() const { return captures_offsets; }
 	const IDSet& WhenLocals() const { return when_locals; }
 	const IDSet& Params() const { return params; }
 	const std::unordered_map<const ID*, int>& Assignees() const { return assignees; }
@@ -208,6 +210,9 @@ protected:
 	// If we're profiling a lambda function, this holds the captures.
 	IDSet captures;
 
+	// This maps capture identifiers to their offsets.
+	std::map<const ID*, int> captures_offsets;
+
 	// Constants seen in the function.
 	std::vector<const ConstExpr*> constants;
 
@@ -224,7 +229,8 @@ protected:
 	// The same, but in a deterministic order, with duplicates removed.
 	std::vector<const Type*> ordered_types;
 
-	// Script functions that this script calls.
+	// Script functions that this script calls.  Includes calls made
+	// by lambdas and when bodies, as the goal is to identify recursion.
 	std::unordered_set<ScriptFunc*> script_calls;
 
 	// Same for BiF's, though for them we record the corresponding global