diff --git a/doc b/doc index 5849f875ea..79748981d8 160000 --- a/doc +++ b/doc @@ -1 +1 @@ -Subproject commit 5849f875ea6cae038d4881eba326256202e711be +Subproject commit 79748981d8abf9d76fe045b0151698f98d43c05c diff --git a/src/Stmt.cc b/src/Stmt.cc index 5e6ac3fb39..7e7ba23a18 100644 --- a/src/Stmt.cc +++ b/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); diff --git a/src/Stmt.h b/src/Stmt.h index a6676d678d..a9bf7cddf8 100644 --- a/src/Stmt.h +++ b/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 { diff --git a/src/parse.y b/src/parse.y index db7e0f846f..c0980ce8de 100644 --- a/src/parse.y +++ b/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 diff --git a/testing/btest/Baseline/language.for/out b/testing/btest/Baseline/language.for/out index dccc00ce3e..d7e75c46e1 100644 --- a/testing/btest/Baseline/language.for/out +++ b/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) diff --git a/testing/btest/Baseline/language.key-value-for/out b/testing/btest/Baseline/language.key-value-for/out new file mode 100644 index 0000000000..43a5609374 --- /dev/null +++ b/testing/btest/Baseline/language.key-value-for/out @@ -0,0 +1,4 @@ +1, hello +55, goodbye +goodbye, world, 55 +hello, world, 1 diff --git a/testing/btest/core/leaks/kv-iteration.bro b/testing/btest/core/leaks/kv-iteration.bro new file mode 100644 index 0000000000..5c7a9f1f62 --- /dev/null +++ b/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; + } diff --git a/testing/btest/language/for.bro b/testing/btest/language/for.bro index eb99a2705d..5f0c211597 100644 --- a/testing/btest/language/for.bro +++ b/testing/btest/language/for.bro @@ -13,7 +13,7 @@ event bro_init() local vv: vector of string = vector( "a", "b", "c" ); local ct: count = 0; - # Test a "for" loop without "break" or "next" + # Test a "for" loop without "break" or "next" ct = 0; for ( i in vv ) ++ct; @@ -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 +} diff --git a/testing/btest/language/key-value-for.bro b/testing/btest/language/key-value-for.bro new file mode 100644 index 0000000000..97591dcacf --- /dev/null +++ b/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; + }