Add key-value for loop

src/Stmt.cc

@@ -1421,12 +1421,48 @@ 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;
+		// Valdate that key-value for loop is being used on a table
+		if ( e->Type()->Tag() == TYPE_TABLE )
+		  {
+			// Type of values table holds
+			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;
 
+	if (value_var)
+	  {
+		Unref(value_var);
+	  }
+
 	Unref(body);
 	}
 
@@ -1443,12 +1479,18 @@ 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,7 +1606,60 @@ 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:

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/language/for.bro

@@ -40,5 +40,16 @@ 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,23 @@
+# @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;
+	}
+}