From 1f7924754eadb250a96056808faa8a1e6b5bc5df Mon Sep 17 00:00:00 2001 From: ZekeMedley Date: Thu, 14 Mar 2019 09:46:16 -0700 Subject: [PATCH] Add key-value for loop --- src/Stmt.cc | 44 +++++++++++++- src/Stmt.h | 5 ++ src/parse.y | 59 ++++++++++++++++++- testing/btest/Baseline/language.for/out | 1 + .../btest/Baseline/language.key-value-for/out | 4 ++ testing/btest/language/for.bro | 15 ++++- testing/btest/language/key-value-for.bro | 23 ++++++++ 7 files changed, 145 insertions(+), 6 deletions(-) create mode 100644 testing/btest/Baseline/language.key-value-for/out create mode 100644 testing/btest/language/key-value-for.bro diff --git a/src/Stmt.cc b/src/Stmt.cc index 5e6ac3fb39..036019d9fd 100644 --- a/src/Stmt.cc +++ b/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); 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..8b09620a24 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,61 @@ 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/language/for.bro b/testing/btest/language/for.bro index eb99a2705d..37d7d2e7c2 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,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 +} diff --git a/testing/btest/language/key-value-for.bro b/testing/btest/language/key-value-for.bro new file mode 100644 index 0000000000..3d69a97f06 --- /dev/null +++ b/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; + } +}