Merge branch 'master' of https://github.com/ZekeMedley/zeek

* 'master' of https://github.com/ZekeMedley/zeek:
  Add key-value for loop

Fixes GH-154

doc

@@ -1 +1 @@
-Subproject commit 5849f875ea6cae038d4881eba326256202e711be
+Subproject commit 79748981d8abf9d76fe045b0151698f98d43c05c

src/Stmt.cc

@@ -1421,12 +1421,38 @@ ForStmt::ForStmt(id_list* arg_loop_vars, Expr* loop_expr)
 		e->Error("target to iterate over must be a table, set, vector, or string");
 	}
 
+ForStmt::ForStmt(id_list* arg_loop_vars, Expr* loop_expr, ID* val_var)
+	: ForStmt(arg_loop_vars, loop_expr)
+	{
+	value_var = val_var;
+
+	if ( e->Type()->IsTable() )
+		{
+		BroType* yield_type = e->Type()->AsTableType()->YieldType();
+
+		// Verify value_vars type if its already been defined
+		if ( value_var->Type() )
+			{
+			if ( ! same_type(value_var->Type(), yield_type) )
+				value_var->Type()->Error("type clash in iteration", yield_type);
+			}
+		else
+			{
+			delete add_local(value_var, yield_type->Ref(), INIT_NONE,
+			                 0, 0, VAR_REGULAR);
+			}
+		}
+	else
+		e->Error("key value for loops only support iteration over tables");
+	}
+
 ForStmt::~ForStmt()
 	{
 	loop_over_list(*loop_vars, i)
 		Unref((*loop_vars)[i]);
 	delete loop_vars;
 
+	Unref(value_var);
 	Unref(body);
 	}
 
@@ -1443,12 +1469,16 @@ Val* ForStmt::DoExec(Frame* f, Val* v, stmt_flow_type& flow) const
 			return 0;
 
 		HashKey* k;
+		TableEntryVal* current_tev;
 		IterCookie* c = loop_vals->InitForIteration();
-		while ( loop_vals->NextEntry(k, c) )
+		while ( (current_tev = loop_vals->NextEntry(k, c)) )
 			{
 			ListVal* ind_lv = tv->RecoverIndex(k);
 			delete k;
 
+			if ( value_var )
+				f->SetElement(value_var->Offset(), current_tev->Value()->Ref());
+
 			for ( int i = 0; i < ind_lv->Length(); i++ )
 				f->SetElement((*loop_vars)[i]->Offset(), ind_lv->Index(i)->Ref());
 			Unref(ind_lv);

src/Stmt.h

@@ -337,6 +337,8 @@ protected:
 class ForStmt : public ExprStmt {
 public:
 	ForStmt(id_list* loop_vars, Expr* loop_expr);
+	// Special constructor for key value for loop.
+	ForStmt(id_list* loop_vars, Expr* loop_expr, ID* val_var);
 	~ForStmt() override;
 
 	void AddBody(Stmt* arg_body)	{ body = arg_body; }
@@ -361,6 +363,9 @@ protected:
 
 	id_list* loop_vars;
 	Stmt* body;
+	// Stores the value variable being used for a key value for loop.
+	// Always set to nullptr unless special constructor is called.
+	ID* value_var = nullptr;
 };
 
 class NextStmt : public Stmt {

src/parse.y

@@ -1592,7 +1592,7 @@ for_head:
 			if ( loop_var )
 				{
 				if ( loop_var->IsGlobal() )
-					loop_var->Error("global used in for loop");
+					loop_var->Error("global variable used in for loop");
 				}
 
 			else
@@ -1606,8 +1606,62 @@ for_head:
 			}
 	|
 		TOK_FOR '(' '[' local_id_list ']' TOK_IN expr ')'
-			{ $$ = new ForStmt($4, $7); }
-		;
+			{
+			$$ = new ForStmt($4, $7);
+			}
+	|
+		TOK_FOR '(' TOK_ID ',' TOK_ID TOK_IN expr ')'
+			{
+			set_location(@1, @8);
+			const char* module = current_module.c_str();
+
+			// Check for previous definitions of key and
+			// value variables.
+			ID* key_var = lookup_ID($3, module);
+			ID* val_var = lookup_ID($5, module);
+
+			// Validate previous definitions as needed.
+			if ( key_var )
+				{
+				if ( key_var->IsGlobal() )
+					key_var->Error("global variable used in for loop");
+				}
+			else
+				key_var = install_ID($3, module, false, false);
+
+			if ( val_var )
+				{
+				if ( val_var->IsGlobal() )
+					val_var->Error("global variable used in for loop");
+				}
+			else
+				val_var = install_ID($5, module, false, false);
+
+			id_list* loop_vars = new id_list;
+			loop_vars->append(key_var);
+
+			$$ = new ForStmt(loop_vars, $7, val_var);
+			}
+	|
+		TOK_FOR '(' '[' local_id_list ']' ',' TOK_ID TOK_IN expr ')'
+			{
+			set_location(@1, @10);
+			const char* module = current_module.c_str();
+
+			// Validate value variable
+			ID* val_var = lookup_ID($7, module);
+
+			if ( val_var )
+				{
+				if ( val_var->IsGlobal() )
+					val_var->Error("global variable used in for loop");
+				}
+			else
+				val_var = install_ID($7, module, false, false);
+
+			$$ = new ForStmt($4, $9, val_var);
+			}
+	;
 
 local_id_list:
 		local_id_list ',' local_id

testing/btest/Baseline/language.for/out

@@ -1,3 +1,4 @@
 for loop (PASS)
 for loop with break (PASS)
 for loop with next (PASS)
+keys that are tuples (PASS)

testing/btest/Baseline/language.key-value-for/out

@@ -0,0 +1,4 @@
+1, hello
+55, goodbye
+goodbye, world, 55
+hello, world, 1

testing/btest/core/leaks/kv-iteration.bro

@@ -0,0 +1,22 @@
+# @TEST-GROUP: leaks
+# @TEST-REQUIRES: bro --help 2>&1 | grep -q mem-leaks
+
+# @TEST-EXEC: HEAP_CHECK_DUMP_DIRECTORY=. HEAPCHECK=local btest-bg-run bro bro -m -b -r $TRACES/http/get.trace %INPUT
+# @TEST-EXEC: btest-bg-wait 60
+
+event new_connection(c: connection)
+	{
+	local t: table[count] of string = table();
+	t[1] = "hello";
+	t[55] = "goodbye";
+
+	for (key, value in t)
+		print key, value;
+
+	local tkk: table[string, string] of count = table();
+	tkk["hello", "world"] = 1;
+	tkk["goodbye", "world"] = 55;
+
+	for ([k1, k2], val in tkk)
+		print k1, k2, val;
+	}

testing/btest/language/for.bro

@@ -40,5 +40,18 @@ event bro_init()
 		test_case("Error: this should not happen", F);
 	}
 	test_case("for loop with next", ct == 3 );
-}
 
+	# Test keys that are tuples
+
+	local t: table[count, count] of string = table();
+	t[1, 2] = "hi";
+
+	local s1: string = "";
+
+	for ( [i, j] in t )
+		s1 = fmt("%d %d %s", i, j, t[i,j]);
+
+	test_case("keys that are tuples", s1 == "1 2 hi");
+
+	# Tests for key value for loop are in key-value-for.bro
+}

testing/btest/language/key-value-for.bro

@@ -0,0 +1,22 @@
+# @TEST-EXEC: bro -b %INPUT >out
+# @TEST-EXEC: btest-diff out
+
+
+event bro_init()
+	{
+	# Test single keys
+
+	local t: table[count] of string = table();
+	t[1] = "hello";
+	t[55] = "goodbye";
+	for (key, value in t)
+		print key, value;
+
+	# Test multiple keys
+
+	local tkk: table[string, string] of count = table();
+	tkk["hello", "world"] = 1;
+	tkk["goodbye", "world"] = 55;
+	for ([k1, k2], val in tkk)
+		print k1, k2, val;
+	}