refinements to LambdaExpr's to provide flexibility, support for ZVal captures

src/Expr.cc

@@ -4644,14 +4644,15 @@ void CallExpr::ExprDescribe(ODesc* d) const
 		args->Describe(d);
 	}
 
-LambdaExpr::LambdaExpr(std::unique_ptr<FunctionIngredients> arg_ing, IDPList arg_outer_ids,
-                       StmtPtr when_parent)
+LambdaExpr::LambdaExpr(std::shared_ptr<FunctionIngredients> arg_ing, IDPList arg_outer_ids,
+                       std::string name, StmtPtr when_parent)
 	: Expr(EXPR_LAMBDA)
 	{
 	ingredients = std::move(arg_ing);
 	outer_ids = std::move(arg_outer_ids);
 
-	SetType(ingredients->GetID()->GetType());
+	auto ingr_t = ingredients->GetID()->GetType<FuncType>();
+	SetType(ingr_t);
 
 	if ( ! CheckCaptures(when_parent) )
 		{
@@ -4659,47 +4660,61 @@ LambdaExpr::LambdaExpr(std::unique_ptr<FunctionIngredients> arg_ing, IDPList arg
 		return;
 		}
 
-	// Install a dummy version of the function globally for use only
-	// when broker provides a closure.
-	auto dummy_func = make_intrusive<ScriptFunc>(ingredients->GetID());
-	dummy_func->AddBody(ingredients->Body(), ingredients->Inits(), ingredients->FrameSize(),
-	                    ingredients->Priority());
+	// Install a master version of the function globally.  This is used
+	// by both broker (for transmitting closures) and script optimization
+	// (replacing its AST body with a compiled one).
+	master_func = make_intrusive<ScriptFunc>(ingredients->GetID());
+	master_func->SetOuterIDs(outer_ids);
 
-	dummy_func->SetOuterIDs(outer_ids);
+	// When we build the body, it will get updated with initialization
+	// statements.  Update the ingredients to reflect the new body,
+	// and no more need for initializers.
+	master_func->AddBody(*ingredients);
+	master_func->SetScope(ingredients->Scope());
+	ingredients->ClearInits();
 
-	// Get the body's "string" representation.
-	ODesc d;
-	dummy_func->Describe(&d);
-
-	for ( ;; )
-		{
-		hash128_t h;
-		KeyedHash::Hash128(d.Bytes(), d.Len(), &h);
-
-		my_name = "lambda_<" + std::to_string(h[0]) + ">";
-		auto fullname = make_full_var_name(current_module.data(), my_name.data());
-		const auto& id = global_scope()->Find(fullname);
-
-		if ( id )
-			// Just try again to make a unique lambda name.
-			// If two peer processes need to agree on the same
-			// lambda name, this assumes they're loading the same
-			// scripts and thus have the same hash collisions.
-			d.Add(" ");
-		else
-			break;
-		}
+	if ( name.empty() )
+		BuildName();
+	else
+		my_name = name;
 
 	// Install that in the global_scope
 	lambda_id = install_ID(my_name.c_str(), current_module.c_str(), true, false);
 
 	// Update lamb's name
-	dummy_func->SetName(my_name.c_str());
+	master_func->SetName(my_name.c_str());
 
-	auto v = make_intrusive<FuncVal>(std::move(dummy_func));
+	auto v = make_intrusive<FuncVal>(master_func);
 	lambda_id->SetVal(std::move(v));
-	lambda_id->SetType(ingredients->GetID()->GetType());
+	lambda_id->SetType(ingr_t);
 	lambda_id->SetConst();
+
+	captures = ingr_t->GetCaptures();
+
+	analyze_lambda(this);
+	}
+
+LambdaExpr::LambdaExpr(LambdaExpr* orig) : Expr(EXPR_LAMBDA)
+	{
+	master_func = orig->master_func;
+	ingredients = orig->ingredients;
+	lambda_id = orig->lambda_id;
+	my_name = orig->my_name;
+	private_captures = orig->private_captures;
+
+	// We need to have our own copies of the outer IDs and captures so
+	// we can rename them when inlined.
+	for ( auto i : orig->outer_ids )
+		outer_ids.append(i);
+
+	if ( orig->captures )
+		{
+		captures = std::vector<FuncType::Capture>{};
+		for ( auto& c : *orig->captures )
+			captures->push_back(c);
+		}
+
+	SetType(orig->GetType());
 	}
 
 bool LambdaExpr::CheckCaptures(StmtPtr when_parent)
@@ -4784,6 +4799,32 @@ bool LambdaExpr::CheckCaptures(StmtPtr when_parent)
 	return true;
 	}
 
+void LambdaExpr::BuildName()
+	{
+	// Get the body's "string" representation.
+	ODesc d;
+	master_func->Describe(&d);
+
+	for ( ;; )
+		{
+		hash128_t h;
+		KeyedHash::Hash128(d.Bytes(), d.Len(), &h);
+
+		my_name = "lambda_<" + std::to_string(h[0]) + ">";
+		auto fullname = make_full_var_name(current_module.data(), my_name.data());
+		const auto& id = global_scope()->Find(fullname);
+
+		if ( id )
+			// Just try again to make a unique lambda name.
+			// If two peer processes need to agree on the same
+			// lambda name, this assumes they're loading the same
+			// scripts and thus have the same hash collisions.
+			d.Add(" ");
+		else
+			break;
+		}
+	}
+
 ScopePtr LambdaExpr::GetScope() const
 	{
 	return ingredients->Scope();
@@ -4792,12 +4833,23 @@ ScopePtr LambdaExpr::GetScope() const
 ValPtr LambdaExpr::Eval(Frame* f) const
 	{
 	auto lamb = make_intrusive<ScriptFunc>(ingredients->GetID());
-	lamb->AddBody(ingredients->Body(), ingredients->Inits(), ingredients->FrameSize(),
-	              ingredients->Priority());
 
+	StmtPtr body = master_func->GetBodies()[0].stmts;
+
+	if ( run_state::is_parsing )
+		// We're evaluating this lambda at parse time, which happens
+		// for initializations.  If we're doing script optimization
+		// then the current version of the body might be left in an
+		// inconsistent state (e.g., if it's replaced with ZAM code)
+		// causing problems if we execute this lambda subsequently.
+		// To avoid that problem, we duplicate the AST so it's
+		// distinct.
+		body = body->Duplicate();
+
+	lamb->AddBody(*ingredients, body);
 	lamb->CreateCaptures(f);
 
-	// Set name to corresponding dummy func.
+	// Set name to corresponding master func.
 	// Allows for lookups by the receiver.
 	lamb->SetName(my_name.c_str());
 
@@ -4807,6 +4859,22 @@ ValPtr LambdaExpr::Eval(Frame* f) const
 void LambdaExpr::ExprDescribe(ODesc* d) const
 	{
 	d->Add(expr_name(Tag()));
+
+	if ( captures && d->IsReadable() )
+		{
+		d->Add("[");
+
+		for ( auto& c : *captures )
+			{
+			if ( &c != &(*captures)[0] )
+				d->AddSP(", ");
+
+			d->Add(c.Id()->Name());
+			}
+
+		d->Add("]");
+		}
+
 	ingredients->Body()->Describe(d);
 	}
 

src/Expr.h

@@ -28,8 +28,10 @@ namespace detail
 class Frame;
 class Scope;
 class FunctionIngredients;
+class WhenInfo;
 using IDPtr = IntrusivePtr<ID>;
 using ScopePtr = IntrusivePtr<Scope>;
+using ScriptFuncPtr = IntrusivePtr<ScriptFunc>;
 
 enum ExprTag : int
 	{
@@ -1452,12 +1454,17 @@ protected:
 class LambdaExpr final : public Expr
 	{
 public:
-	LambdaExpr(std::unique_ptr<FunctionIngredients> ingredients, IDPList outer_ids,
-	           StmtPtr when_parent = nullptr);
+	LambdaExpr(std::shared_ptr<FunctionIngredients> ingredients, IDPList outer_ids,
+	           std::string name = "", StmtPtr when_parent = nullptr);
 
 	const std::string& Name() const { return my_name; }
+
 	const IDPList& OuterIDs() const { return outer_ids; }
-	const FunctionIngredients& Ingredients() const { return *ingredients; }
+
+	// Lambda's potentially have their own private copy of captures,
+	// to enable updates to the set during script optimization.
+	using CaptureList = std::vector<FuncType::Capture>;
+	const std::optional<CaptureList>& GetCaptures() const { return captures; }
 
 	ValPtr Eval(Frame* f) const override;
 	TraversalCode Traverse(TraversalCallback* cb) const override;
@@ -1466,19 +1473,43 @@ public:
 
 	// Optimization-related:
 	ExprPtr Duplicate() override;
-	ExprPtr Inline(Inliner* inl) override;
 
+	const ScriptFuncPtr& MasterFunc() const { return master_func; }
+
+	const std::shared_ptr<FunctionIngredients>& Ingredients() const { return ingredients; }
+
+	bool IsReduced(Reducer* c) const override;
+	bool HasReducedOps(Reducer* c) const override;
 	ExprPtr Reduce(Reducer* c, StmtPtr& red_stmt) override;
+	StmtPtr ReduceToSingletons(Reducer* c) override;
 
 protected:
+	// Constructor used for script optimization.
+	LambdaExpr(LambdaExpr* orig);
+
 	void ExprDescribe(ODesc* d) const override;
 
 private:
-	bool CheckCaptures(StmtPtr when_parent);
+	friend class WhenInfo;
 
-	std::unique_ptr<FunctionIngredients> ingredients;
+	// "Private" captures are captures that correspond to "when"
+	// condition locals.  These aren't true captures in that they
+	// don't come from the outer frame when the lambda is constructed,
+	// but they otherwise behave like captures in that they persist
+	// across function invocations.
+	void SetPrivateCaptures(const IDSet& pcaps) { private_captures = pcaps; }
+
+	bool CheckCaptures(StmtPtr when_parent);
+	void BuildName();
+
+	void UpdateCaptures(Reducer* c);
+
+	std::shared_ptr<FunctionIngredients> ingredients;
+	ScriptFuncPtr master_func;
 	IDPtr lambda_id;
 	IDPList outer_ids;
+	std::optional<CaptureList> captures;
+	IDSet private_captures;
 
 	std::string my_name;
 	};

src/parse.y

@@ -1566,7 +1566,7 @@ lambda_body:
 
 			// Gather the ingredients for a Func from the
 			// current scope.
-			auto ingredients = std::make_unique<FunctionIngredients>(
+			auto ingredients = std::make_shared<FunctionIngredients>(
 				current_scope(), IntrusivePtr{AdoptRef{}, $3}, current_module.c_str());
 			auto outer_ids = gather_outer_ids(pop_scope(), ingredients->Body());
 

src/script_opt/CPP/DeclFunc.cc

@@ -29,8 +29,8 @@ void CPPCompile::DeclareLambda(const LambdaExpr* l, const ProfileFunc* pf)
 	ASSERT(is_CPP_compilable(pf));
 
 	auto lname = Canonicalize(l->Name().c_str()) + "_lb";
-	auto body = l->Ingredients().Body();
-	auto l_id = l->Ingredients().GetID();
+	auto body = l->Ingredients()->Body();
+	auto l_id = l->Ingredients()->GetID();
 	auto& ids = l->OuterIDs();
 
 	for ( auto id : ids )

src/script_opt/CPP/Driver.cc

@@ -150,7 +150,7 @@ void CPPCompile::Compile(bool report_uncompilable)
 	for ( const auto& l : pfs.Lambdas() )
 		{
 		const auto& n = l->Name();
-		const auto body = l->Ingredients().Body().get();
+		const auto body = l->Ingredients()->Body().get();
 		if ( lambda_ASTs.count(n) > 0 )
 			// Reuse previous body.
 			body_names[body] = body_names[lambda_ASTs[n]];
@@ -176,7 +176,7 @@ void CPPCompile::Compile(bool report_uncompilable)
 			continue;
 
 		CompileLambda(l, pfs.ExprProf(l).get());
-		lambda_ASTs[n] = l->Ingredients().Body().get();
+		lambda_ASTs[n] = l->Ingredients()->Body().get();
 		}
 
 	NL();

src/script_opt/CPP/GenFunc.cc

@@ -23,8 +23,8 @@ void CPPCompile::CompileFunc(const FuncInfo& func)
 void CPPCompile::CompileLambda(const LambdaExpr* l, const ProfileFunc* pf)
 	{
 	auto lname = Canonicalize(l->Name().c_str()) + "_lb";
-	auto body = l->Ingredients().Body();
-	auto l_id = l->Ingredients().GetID();
+	auto body = l->Ingredients()->Body();
+	auto l_id = l->Ingredients()->GetID();
 	auto& ids = l->OuterIDs();
 
 	DefineBody(l_id->GetType<FuncType>(), pf, lname, body, &ids, FUNC_FLAVOR_FUNCTION);

src/script_opt/Expr.cc

@@ -2425,23 +2425,60 @@ StmtPtr CallExpr::ReduceToSingletons(Reducer* c)
 
 ExprPtr LambdaExpr::Duplicate()
 	{
-	auto ingr = std::make_unique<FunctionIngredients>(*ingredients);
-	ingr->SetBody(ingr->Body()->Duplicate());
-	return SetSucc(new LambdaExpr(std::move(ingr), outer_ids));
+	return SetSucc(new LambdaExpr(this));
 	}
 
-ExprPtr LambdaExpr::Inline(Inliner* inl)
+bool LambdaExpr::IsReduced(Reducer* c) const
 	{
-	// Don't inline these, we currently don't get the closure right.
-	return ThisPtr();
+	if ( ! captures )
+		return true;
+
+	for ( auto& cp : *captures )
+		{
+		auto& cid = cp.Id();
+
+		if ( private_captures.count(cid.get()) == 0 && ! c->ID_IsReduced(cid) )
+			return NonReduced(this);
+		}
+
+	return true;
+	}
+
+bool LambdaExpr::HasReducedOps(Reducer* c) const
+	{
+	return IsReduced(c);
 	}
 
 ExprPtr LambdaExpr::Reduce(Reducer* c, StmtPtr& red_stmt)
 	{
 	if ( c->Optimizing() )
 		return ThisPtr();
-	else
-		return AssignToTemporary(c, red_stmt);
+
+	UpdateCaptures(c);
+
+	return AssignToTemporary(c, red_stmt);
+	}
+
+StmtPtr LambdaExpr::ReduceToSingletons(Reducer* c)
+	{
+	UpdateCaptures(c);
+	return nullptr;
+	}
+
+void LambdaExpr::UpdateCaptures(Reducer* c)
+	{
+	if ( captures )
+		{
+		for ( auto& cp : *captures )
+			{
+			auto& cid = cp.Id();
+
+			if ( private_captures.count(cid.get()) == 0 )
+				cp.SetID(c->UpdateID(cid));
+			}
+
+		c->UpdateIDs(&outer_ids);
+		}
 	}
 
 ExprPtr EventExpr::Duplicate()