feature completeness for ZAM

src/script_opt/CPP/Driver.cc

@@ -77,6 +77,12 @@ void CPPCompile::Compile(bool report_uncompilable)
 			continue;
 			}
 
+		if ( is_when_lambda(f) )
+			{
+			func.SetSkip(true);
+			continue;
+			}
+
 		const char* reason;
 		if ( IsCompilable(func, &reason) )
 			{

src/script_opt/Inline.cc

@@ -115,16 +115,21 @@ void Inliner::Analyze()
 		const auto& func = func_ptr.get();
 		const auto& body = f.Body();
 
+		if ( ! should_analyze(func_ptr, body) )
+			continue;
+
 		// Candidates are non-event, non-hook, non-recursive,
-		// non-compiled functions ... that don't use lambdas or when's,
-		// since we don't currently compute the closures/frame
-		// sizes for them correctly, and more fundamentally since
-		// we don't compile them and hence inlining them will
-		// make the parent non-compilable.
-		if ( should_analyze(func_ptr, body) && func->Flavor() == FUNC_FLAVOR_FUNCTION &&
-		     non_recursive_funcs.count(func) > 0 && f.Profile()->NumLambdas() == 0 &&
-		     f.Profile()->NumWhenStmts() == 0 && body->Tag() != STMT_CPP )
-			inline_ables.insert(func);
+		// non-compiled functions ...
+		if ( func->Flavor() != FUNC_FLAVOR_FUNCTION )
+			continue;
+
+		if ( non_recursive_funcs.count(func) == 0 )
+			continue;
+
+		if ( body->Tag() == STMT_CPP )
+			continue;
+
+		inline_ables.insert(func);
 		}
 
 	for ( auto& f : funcs )
@@ -165,6 +170,11 @@ ExprPtr Inliner::CheckForInlining(CallExprPtr c)
 		// We don't inline indirect calls.
 		return c;
 
+	if ( c->IsInWhen() )
+		// Don't inline these, as doing so requires propagating
+		// the in-when attribute to the inlined function body.
+		return c;
+
 	auto n = f->AsNameExpr();
 	auto func = n->Id();
 

src/script_opt/ScriptOpt.cc

@@ -36,15 +36,35 @@ static ZAMCompiler* ZAM = nullptr;
 static bool generating_CPP = false;
 static std::string CPP_dir; // where to generate C++ code
 
+static std::unordered_set<const ScriptFunc*> lambdas;
+static std::unordered_set<const ScriptFunc*> when_lambdas;
 static ScriptFuncPtr global_stmts;
 
 void analyze_func(ScriptFuncPtr f)
 	{
 	// Even if we're analyzing only a subset of the scripts, we still
 	// track all functions here because the inliner will need the full list.
+	ASSERT(f->GetScope());
 	funcs.emplace_back(f, f->GetScope(), f->CurrentBody(), f->CurrentPriority());
 	}
 
+void analyze_lambda(LambdaExpr* l)
+	{
+	auto& mf = l->MasterFunc();
+	analyze_func(mf);
+	lambdas.insert(mf.get());
+	}
+
+void analyze_when_lambda(LambdaExpr* l)
+	{
+	when_lambdas.insert(l->MasterFunc().get());
+	}
+
+bool is_when_lambda(const ScriptFunc* f)
+	{
+	return when_lambdas.count(f) > 0;
+	}
+
 const FuncInfo* analyze_global_stmts(Stmt* stmts)
 	{
 	// We ignore analysis_options.only_{files,funcs} - if they're in use, later
@@ -208,7 +228,8 @@ static void optimize_func(ScriptFunc* f, std::shared_ptr<ProfileFunc> pf, ScopeP
 
 	rc->SetReadyToOptimize();
 
-	auto ud = std::make_shared<UseDefs>(body, rc);
+	auto ft = cast_intrusive<FuncType>(f->GetType());
+	auto ud = std::make_shared<UseDefs>(body, rc, ft);
 	ud->Analyze();
 
 	if ( analysis_options.dump_uds )
@@ -445,10 +466,10 @@ static void analyze_scripts_for_ZAM(std::unique_ptr<ProfileFuncs>& pfs)
 		}
 
 	// Re-profile the functions, now without worrying about compatibility
-	// with compilation to C++.  Note that the first profiling pass earlier
-	// may have marked some of the functions as to-skip, so first clear
-	// those markings.  Once we have full compile-to-C++ and ZAM support
-	// for all Zeek language features, we can remove the re-profiling here.
+	// with compilation to C++.
+
+	// The first profiling pass earlier may have marked some of the
+	// functions as to-skip, so clear those markings.
 	for ( auto& f : funcs )
 		f.SetSkip(false);
 
@@ -499,6 +520,7 @@ static void analyze_scripts_for_ZAM(std::unique_ptr<ProfileFuncs>& pfs)
 	for ( auto& f : funcs )
 		{
 		auto func = f.Func();
+		bool is_lambda = lambdas.count(func) > 0;
 
 		if ( ! analysis_options.only_funcs.empty() || ! analysis_options.only_files.empty() )
 			{
@@ -506,14 +528,17 @@ static void analyze_scripts_for_ZAM(std::unique_ptr<ProfileFuncs>& pfs)
 				continue;
 			}
 
-		else if ( ! analysis_options.compile_all && inl && inl->WasInlined(func) &&
+		else if ( ! analysis_options.compile_all && ! is_lambda && inl && inl->WasInlined(func) &&
 		          func_used_indirectly.count(func) == 0 )
+			{
 			// No need to compile as it won't be called directly.
 			continue;
+			}
 
 		auto new_body = f.Body();
 		optimize_func(func, f.ProfilePtr(), f.Scope(), new_body);
 		f.SetBody(new_body);
+
 		did_one = true;
 		}
 

src/script_opt/ScriptOpt.h

@@ -168,6 +168,16 @@ extern std::unordered_set<const Func*> non_recursive_funcs;
 // Analyze a given function for optimization.
 extern void analyze_func(ScriptFuncPtr f);
 
+// Same, for lambdas.
+extern void analyze_lambda(LambdaExpr* f);
+
+// Same, for lambdas used in "when" statements.  For these, analyze_lambda()
+// has already been called.
+extern void analyze_when_lambda(LambdaExpr* f);
+
+// Whether a given script function is a "when" lambda.
+extern bool is_when_lambda(const ScriptFunc* f);
+
 // Analyze the given top-level statement(s) for optimization.  Returns
 // a pointer to a FuncInfo for an argument-less quasi-function that can
 // be Invoked, or its body executed directly, to execute the statements.

src/script_opt/ZAM/Support.cc

@@ -20,20 +20,6 @@ bool ZAM_error = false;
 
 bool is_ZAM_compilable(const ProfileFunc* pf, const char** reason)
 	{
-	if ( pf->NumLambdas() > 0 )
-		{
-		if ( reason )
-			*reason = "use of lambda";
-		return false;
-		}
-
-	if ( pf->NumWhenStmts() > 0 )
-		{
-		if ( reason )
-			*reason = "use of \"when\"";
-		return false;
-		}
-
 	auto b = pf->ProfiledBody();
 	auto is_hook = pf->ProfiledFunc()->Flavor() == FUNC_FLAVOR_HOOK;
 	if ( b && ! script_is_valid(b, is_hook) )