fixes for standalone C++ scripts making types & variables/functions available

src/script_opt/CPP/CPP-load.bif

@@ -22,7 +22,8 @@ function load_CPP%(h: count%): bool
 	%{
 	auto cb = detail::standalone_callbacks.find(h);
 
-	if ( cb == detail::standalone_callbacks.end() )
+	if ( cb == detail::standalone_callbacks.end() ||
+	     ! detail::CPP_init_hook )
 		{
 		reporter->Error("load of non-existing C++ code (%" PRIu64 ")", h);
 		return zeek::val_mgr->False();
@@ -38,5 +39,8 @@ function load_CPP%(h: count%): bool
 	// compiled scripts.
 	detail::standalone_activations.push_back(cb->second);
 
+	// Proceed with activation.
+	(*detail::CPP_init_hook)();
+
 	return zeek::val_mgr->True();
 	%}

src/script_opt/CPP/Driver.cc

@@ -303,6 +303,9 @@ void CPPCompile::GenEpilog()
 	CheckInitConsistency(to_do);
 	auto nc = GenDependentInits(to_do);
 
+	if ( standalone )
+		GenStandaloneActivation();
+
 	NL();
 	Emit("void init__CPP()");
 
@@ -319,6 +322,9 @@ void CPPCompile::GenEpilog()
 	NL();
 	InitializeFieldMappings();
 
+	if ( standalone )
+		Emit("standalone_init__CPP();");
+
 	EndBlock(true);
 
 	GenInitHook();

src/script_opt/CPP/Exprs.cc

@@ -246,15 +246,18 @@ string CPPCompile::GenCallExpr(const CallExpr* c, GenType gt)
 		auto f_id = f->AsNameExpr()->Id();
 		const auto& params = f_id->GetType()->AsFuncType()->Params();
 		auto id_name = f_id->Name();
-		auto fname = Canonicalize(id_name) + "_zf";
 
-		bool is_compiled = compiled_funcs.count(fname) > 0;
+		bool is_compiled = compiled_simple_funcs.count(id_name) > 0;
 		bool was_compiled = hashed_funcs.count(id_name) > 0;
 
 		if ( is_compiled || was_compiled )
 			{
+			string fname;
+
 			if ( was_compiled )
 				fname = hashed_funcs[id_name];
+			else
+				fname = compiled_simple_funcs[id_name];
 
 			if ( args_l->Exprs().length() > 0 )
 				gen = fname + "(" + GenArgs(params, args_l) +

src/script_opt/CPP/Func.cc

@@ -12,6 +12,7 @@ namespace zeek::detail {
 using namespace std;
 
 unordered_map<p_hash_type, CompiledScript> compiled_scripts;
+unordered_map<string, unordered_set<p_hash_type>> added_bodies;
 unordered_map<p_hash_type, void (*)()> standalone_callbacks;
 vector<void (*)()> standalone_activations;
 

src/script_opt/CPP/Func.h

@@ -108,6 +108,14 @@ struct CompiledScript {
 // Maps hashes to compiled information.
 extern std::unordered_map<p_hash_type, CompiledScript> compiled_scripts;
 
+// When using standalone-code, tracks which function bodies have had
+// compiled versions added to them.  Needed so that we don't replace
+// the body twice, leading to two copies.  Indexed first by the name
+// of the function, and then via the hash of the body that has been
+// added to it.
+extern std::unordered_map<std::string, std::unordered_set<p_hash_type>>
+ added_bodies;
+
 // Maps hashes to standalone script initialization callbacks.
 extern std::unordered_map<p_hash_type, void (*)()> standalone_callbacks;
 

src/script_opt/CPP/Inits.cc

@@ -460,9 +460,6 @@ void CPPCompile::GenInitHook()
 	{
 	NL();
 
-	if ( standalone )
-		GenStandaloneActivation();
-
 	Emit("int hook_in_init()");
 
 	StartBlock();
@@ -482,6 +479,15 @@ void CPPCompile::GenInitHook()
 
 void CPPCompile::GenStandaloneActivation()
 	{
+	NL();
+
+	Emit("void standalone_activation__CPP()");
+	StartBlock();
+	for ( auto& a : activations )
+		Emit(a);
+	EndBlock();
+
+	NL();
 	Emit("void standalone_init__CPP()");
 	StartBlock();
 
@@ -497,11 +503,6 @@ void CPPCompile::GenStandaloneActivation()
 	for ( const auto& func : funcs )
 		{
 		auto f = func.Func();
-
-		if ( f->Flavor() == FUNC_FLAVOR_FUNCTION )
-			// No need to explicitly add bodies.
-			continue;
-
 		auto fname = BodyName(func);
 		auto bname = Canonicalize(fname.c_str()) + "_zf";
 
@@ -534,8 +535,11 @@ void CPPCompile::GenStandaloneActivation()
 		     fn, GenTypeName(ft), hashes);
 		}
 
-	EndBlock();
 	NL();
+	Emit("CPP_activation_funcs.push_back(standalone_activation__CPP);");
+	Emit("CPP_activation_hook = activate__CPPs;");
+
+	EndBlock();
 	}
 
 void CPPCompile::GenLoad()

src/script_opt/CPP/RuntimeInit.cc

@@ -9,12 +9,30 @@ namespace zeek::detail {
 using namespace std;
 
 vector<CPP_init_func> CPP_init_funcs;
+vector<CPP_init_func> CPP_activation_funcs;
 
 // Calls all of the initialization hooks, in the order they were added.
 void init_CPPs()
 	{
-	for ( auto f : CPP_init_funcs )
+	static bool need_init = true;
-		f();
+
+	if ( need_init )
+		for ( auto f : CPP_init_funcs )
+			f();
+
+	need_init = false;
+	}
+
+// Calls all of the registered activation hooks for standalone code.
+void activate__CPPs()
+	{
+	static bool need_init = true;
+
+	if ( need_init )
+		for ( auto f : CPP_activation_funcs )
+			f();
+
+	need_init = false;
 	}
 
 // This is a trick used to register the presence of compiled code.
@@ -30,6 +48,19 @@ static int flag_init_CPP()
 static int dummy = flag_init_CPP();
 
 
+void register_type__CPP(TypePtr t, const std::string& name)
+	{
+	if ( t->GetName().size() > 0 )
+		// Already registered.
+		return;
+
+	t->SetName(name);
+
+	auto id = install_ID(name.c_str(), GLOBAL_MODULE_NAME, true, false);
+	id->SetType(t);
+	id->MakeType();
+	}
+
 void register_body__CPP(CPPStmtPtr body, int priority, p_hash_type hash,
                         vector<string> events)
 	{
@@ -78,7 +109,19 @@ void activate_bodies__CPP(const char* fn, TypePtr t, vector<p_hash_type> hashes)
 		fg->SetType(ft);
 		}
 
-	auto f = fg->GetVal()->AsFunc();
+	auto v = fg->GetVal();
+	if ( ! v )
+		{ // Create it.
+		std::vector<StmtPtr> no_bodies;
+		std::vector<int> no_priorities;
+		auto sf = make_intrusive<ScriptFunc>(fn, ft, no_bodies,
+		                                     no_priorities);
+
+		v = make_intrusive<FuncVal>(move(sf));
+		fg->SetVal(v);
+		}
+
+	auto f = v->AsFunc();
 	const auto& bodies = f->GetBodies();
 
 	// Track hashes of compiled bodies already associated with f.
@@ -115,6 +158,7 @@ void activate_bodies__CPP(const char* fn, TypePtr t, vector<p_hash_type> hashes)
 		auto cs = compiled_scripts[h];
 
 		f->AddBody(cs.body, no_inits, num_params, cs.priority);
+		added_bodies[fn].insert(h);
 
 		events.insert(cs.events.begin(), cs.events.end());
 		}

src/script_opt/CPP/RuntimeInit.h

@@ -20,6 +20,15 @@ typedef void (*CPP_init_func)();
 // Tracks the initialization hooks for different compilation runs.
 extern std::vector<CPP_init_func> CPP_init_funcs;
 
+// Tracks the activation hooks for different "standalone" compilations.
+extern std::vector<CPP_init_func> CPP_activation_funcs;
+
+// Activates all previously registered standalone code.
+extern void activate__CPPs();
+
+// Registers the given global type, if not already present.
+extern void register_type__CPP(TypePtr t, const std::string& name);
+
 // Registers the given compiled function body as associated with the
 // given priority and hash.  "events" is a list of event handlers
 // relevant for the function body, which should be registered if the

src/script_opt/CPP/Types.cc

@@ -134,7 +134,8 @@ void CPPCompile::ExpandTypeVar(const TypePtr& t)
 
 	auto& script_type_name = t->GetName();
 	if ( script_type_name.size() > 0 )
-		AddInit(t, tn + "->SetName(\"" + script_type_name + "\");");
+		AddInit(t, "register_type__CPP(" + tn + ", \"" +
+		           script_type_name + "\");");
 
 	AddInit(t);
 	}

src/script_opt/CPP/Vars.cc

@@ -168,7 +168,12 @@ void CPPCompile::AddBiF(const ID* b, bool is_var)
 	if ( AddGlobal(n, "bif", true) )
 		Emit("Func* %s;", globals[n]);
 
-	AddInit(b, globals[n], string("lookup_bif__CPP(\"") + bn + "\")");
+	auto lookup = string("lookup_bif__CPP(\"") + bn + "\")";
+
+	if ( standalone )
+		AddActivation(globals[n] + " = " + lookup + ";");
+	else
+		AddInit(b, globals[n], lookup);
 	}
 
 bool CPPCompile::AddGlobal(const string& g, const char* suffix, bool track)

src/script_opt/ScriptOpt.cc

@@ -24,6 +24,7 @@ AnalyOpt analysis_options;
 std::unordered_set<const Func*> non_recursive_funcs;
 
 void (*CPP_init_hook)() = nullptr;
+void (*CPP_activation_hook)() = nullptr;
 
 // Tracks all of the loaded functions (including event handlers and hooks).
 static std::vector<FuncInfo> funcs;
@@ -354,8 +355,21 @@ void analyze_scripts()
 				{
 				auto b = s->second.body;
 				b->SetHash(hash);
-				f.Func()->ReplaceBody(f.Body(), b);
+
-				f.SetBody(b);
+				// We may have already updated the body if
+				// we're using code compiled for standalone.
+				if ( f.Body()->Tag() != STMT_CPP )
+					{
+					auto func = f.Func();
+					if ( added_bodies[func->Name()].count(hash) > 0 )
+						// We've already added the
+						// replacement.  Delete orig.
+						func->ReplaceBody(f.Body(), nullptr);
+					else
+						func->ReplaceBody(f.Body(), b);
+
+					f.SetBody(b);
+					}
 
 				for ( auto& e : s->second.events )
 					{

src/script_opt/ScriptOpt.h

@@ -147,5 +147,9 @@ extern void analyze_scripts();
 // to a non-empty value.
 extern void (*CPP_init_hook)();
 
+// Used for "standalone" C++-compiled scripts to complete their activation;
+// called after parsing and BiF initialization, but before zeek_init.
+extern void (*CPP_activation_hook)();
+
 
 } // namespace zeek::detail

src/zeek-setup.cc

@@ -849,6 +849,9 @@ SetupResult setup(int argc, char** argv, Options* zopts)
 		// we don't have any other source for it.
 		run_state::detail::update_network_time(util::current_time());
 
+	if ( CPP_activation_hook )
+		(*CPP_activation_hook)();
+
 	if ( zeek_init )
 		event_mgr.Enqueue(zeek_init, Args{});