ZAM internals have a notion of "captures" as global-like variables

src/script_opt/ZAM/Vars.cc

@@ -24,6 +24,17 @@ bool ZAMCompiler::IsUnused(const IDPtr& id, const Stmt* where) const
 	return ! usage || ! usage->HasID(id.get());
 	}
 
+bool ZAMCompiler::IsCapture(const ID* id) const
+	{
+	auto c = pf->CapturesOffsets();
+	return c.find(id) != c.end();
+	}
+
+int ZAMCompiler::CaptureOffset(const ID* id) const
+	{
+	return pf->CapturesOffsets().find(id)->second;
+	}
+
 void ZAMCompiler::LoadParam(const ID* id)
 	{
 	if ( id->IsType() )
@@ -69,6 +80,23 @@ const ZAMStmt ZAMCompiler::LoadGlobal(const ID* id)
 	return AddInst(z, true);
 	}
 
+const ZAMStmt ZAMCompiler::LoadCapture(const ID* id)
+	{
+	ZOp op;
+
+	if ( ZVal::IsManagedType(id->GetType()) )
+		op = OP_LOAD_MANAGED_CAPTURE_VV;
+	else
+		op = OP_LOAD_CAPTURE_VV;
+
+	auto slot = RawSlot(id);
+
+	ZInstI z(op, slot, CaptureOffset(id));
+	z.op_type = OP_VV_I2;
+
+	return AddInst(z, true);
+	}
+
 int ZAMCompiler::AddToFrame(const ID* id)
 	{
 	frame_layout1[id] = frame_sizeI;
@@ -83,6 +111,9 @@ int ZAMCompiler::FrameSlot(const ID* id)
 	if ( id->IsGlobal() )
 		(void)LoadGlobal(id);
 
+	else if ( IsCapture(id) )
+		(void)LoadCapture(id);
+
 	return slot;
 	}
 
@@ -103,6 +134,13 @@ int ZAMCompiler::Frame1Slot(const ID* id, ZAMOp1Flavor fl)
 	if ( id->IsGlobal() )
 		pending_global_store = global_id_to_info[id];
 
+	else if ( IsCapture(id) )
+		pending_capture_store = CaptureOffset(id);
+
+	// Make sure we don't think we're storing to both a global and
+	// a capture.
+	ASSERT(pending_global_store == -1 || pending_capture_store == -1);
+
 	return slot;
 	}
 

src/script_opt/ZAM/ZInst.cc

@@ -383,16 +383,20 @@ bool ZInstI::IsDirectAssignment() const
 
 	switch ( op )
 		{
-		case OP_ASSIGN_VV_N:
 		case OP_ASSIGN_VV_A:
+		case OP_ASSIGN_VV_D:
+		case OP_ASSIGN_VV_F:
+		case OP_ASSIGN_VV_I:
+		case OP_ASSIGN_VV_L:
+		case OP_ASSIGN_VV_N:
 		case OP_ASSIGN_VV_O:
 		case OP_ASSIGN_VV_P:
 		case OP_ASSIGN_VV_R:
 		case OP_ASSIGN_VV_S:
-		case OP_ASSIGN_VV_F:
 		case OP_ASSIGN_VV_T:
+		case OP_ASSIGN_VV_U:
 		case OP_ASSIGN_VV_V:
-		case OP_ASSIGN_VV_L:
+		case OP_ASSIGN_VV_a:
 		case OP_ASSIGN_VV_f:
 		case OP_ASSIGN_VV_t:
 		case OP_ASSIGN_VV:
@@ -403,6 +407,21 @@ bool ZInstI::IsDirectAssignment() const
 		}
 	}
 
+bool ZInstI::HasCaptures() const
+	{
+	switch ( op )
+		{
+		case OP_LAMBDA_VV:
+		case OP_WHEN_V:
+		case OP_WHEN_TIMEOUT_VV:
+		case OP_WHEN_TIMEOUT_VC:
+			return true;
+
+		default:
+			return false;
+		}
+	}
+
 bool ZInstI::HasSideEffects() const
 	{
 	return op_side_effects[op];
@@ -647,6 +666,11 @@ bool ZInstI::IsGlobalLoad() const
 	return global_ops.count(op) > 0;
 	}
 
+bool ZInstI::IsCaptureLoad() const
+	{
+	return op == OP_LOAD_CAPTURE_VV || op == OP_LOAD_MANAGED_CAPTURE_VV;
+	}
+
 void ZInstI::InitConst(const ConstExpr* ce)
 	{
 	auto v = ce->ValuePtr();

src/script_opt/ZAM/ZInst.h

@@ -5,6 +5,7 @@
 #pragma once
 
 #include "zeek/Desc.h"
+#include "zeek/Func.h"
 #include "zeek/script_opt/ZAM/BuiltInSupport.h"
 #include "zeek/script_opt/ZAM/Support.h"
 #include "zeek/script_opt/ZAM/ZOp.h"
@@ -220,6 +221,9 @@ public:
 	// True if this instruction is of the form "v1 = v2".
 	bool IsDirectAssignment() const;
 
+	// True if this instruction includes captures in its aux slots.
+	bool HasCaptures() const;
+
 	// True if this instruction has side effects when executed, so
 	// should not be pruned even if it has a dead assignment.
 	bool HasSideEffects() const;
@@ -247,9 +251,17 @@ public:
 	// the ZAM frame.
 	bool IsGlobalLoad() const;
 
+	// True if the instruction corresponds to loading a capture into
+	// the ZAM frame.
+	bool IsCaptureLoad() const;
+
+	// True if the instruction does not correspond to a load from the
+	// ZAM frame.
+	bool IsNonLocalLoad() const { return IsGlobalLoad() || IsCaptureLoad(); }
+
 	// True if the instruction corresponds to some sort of load,
-	// either from the interpreter frame or of a global.
-	bool IsLoad() const { return op_type == OP_VV_FRAME || IsGlobalLoad(); }
+	// either from the interpreter frame or of a global/capture.
+	bool IsLoad() const { return op_type == OP_VV_FRAME || IsNonLocalLoad(); }
 
 	// True if the instruction corresponds to storing a global.
 	bool IsGlobalStore() const { return op == OP_STORE_GLOBAL_V; }
@@ -317,6 +329,7 @@ public:
 			slots = ints = new int[n];
 			constants = new ValPtr[n];
 			types = new TypePtr[n];
+			is_managed = new bool[n];
 			}
 		}
 
@@ -325,6 +338,7 @@ public:
 		delete[] ints;
 		delete[] constants;
 		delete[] types;
+		delete[] is_managed;
 		delete[] cat_args;
 		}
 
@@ -384,6 +398,7 @@ public:
 		ints[i] = slot;
 		constants[i] = nullptr;
 		types[i] = t;
+		is_managed[i] = t ? ZVal::IsManagedType(t) : false;
 		}
 
 	// Same but for constants.
@@ -392,6 +407,7 @@ public:
 		ints[i] = -1;
 		constants[i] = c;
 		types[i] = nullptr;
+		is_managed[i] = false;
 		}
 
 	// Member variables.  We could add accessors for manipulating
@@ -404,13 +420,26 @@ public:
 	// if not, it's nil).
 	//
 	// We track associated types, too, enabling us to use
-	// ZVal::ToVal to convert frame slots or constants to ValPtr's.
+	// ZVal::ToVal to convert frame slots or constants to ValPtr's;
+	// and, as a performance optimization, whether those types
+	// indicate the slot needs to be managed.
 
 	int n; // size of arrays
 	int* slots = nullptr; // either nil or points to ints
 	int* ints = nullptr;
 	ValPtr* constants = nullptr;
 	TypePtr* types = nullptr;
+	bool* is_managed = nullptr;
+
+	// Ingredients associated with lambdas ...
+	ScriptFuncPtr master_func;
+
+	// ... and its name.
+	std::string lambda_name;
+
+	// For "when" statements.  Needs to be non-const so we can
+	// Instantiate() it as needed.
+	WhenInfo* wi;
 
 	// A parallel array for the cat() built-in replacement.
 	std::unique_ptr<CatArg>* cat_args = nullptr;
@@ -418,10 +447,10 @@ public:
 	// Used for accessing function names.
 	const ID* id_val = nullptr;
 
-	// Whether the instruction can lead to globals changing.
-	// Currently only needed by the optimizer, but convenient
-	// to store here.
-	bool can_change_globals = false;
+	// Whether the instruction can lead to globals/captures changing.
+	// Currently only needed by the optimizer, but convenient to
+	// store here.
+	bool can_change_non_locals = false;
 
 	// The following is only used for OP_CONSTRUCT_KNOWN_RECORD_V,
 	// to map elements in slots/constants/types to record field offsets.